Pivot-SD：训练扩散语言模型，只需挑出那几个"命悬一线"的词

你有没有想过这样一个问题：一篇作文里，哪些字句是真正决定它好坏的？

不是通篇都重要。开头那句定调的话，中间那个转折的论点，往往比后面大段的铺陈更关键。写作文的老师批改时,其实心里也有一套"轻重缓急"的标尺,只不过他可能自己都没意识到。

现在把这个问题丢给AI训练,事情就变得棘手起来。大语言模型在学习怎么写得更好的时候,到底该在哪些词上下功夫?这篇来自KAIST和多伦多大学团队的论文,就是在死磕这个问题,只不过它瞄准的是一种比较特殊的AI模型:掩码扩散语言模型。

核心问题:扩散模型的生成顺序,根本不按套路来

先说说传统的大语言模型是怎么写东西的。

GPT这类模型写文章,是从左到右一个字一个字往下蹦的,写完"今天"才能写"天气",这个顺序是死的,改不了。

掩码扩散语言模型*（简称dLM）：一种和GPT不同的文本生成模型,它不是从左到右逐字生成,而是先把整段回答全部"打码"（变成掩码符号），然后像拼图一样,按照模型自己感觉"有把握"的顺序,逐步把每个位置的真实文字填回去。

这意味着什么?意味着dLM可能先把结尾的句号填上,再回头填开头的主语,它的生成顺序完全取决于模型在哪个位置更"有自信"。这种打破线性顺序的生成方式,听起来很酷,但也带来一个新麻烦:你没法再用"第几个字"这种位置信息,去判断这个字到底重不重要了。

论文里提到一个关键观察:在dLM填词的过程中,有那么几步,一旦某个词被确定下来,剩下所有还蒙着的位置,瞬间就变得清晰了很多。而在另外一些步骤,填上一个词之后,剩下的位置该怎么写几乎毫无变化。

论文作者管这种"一子落下、全局皆明"的关键步骤叫做"pivot"（枢轴词）。

打个比方。你和朋友玩猜谜游戏,谜底是一道数学题的答案。你问了十个问题,其中九个得到的回答都是"是"或"不是",信息量平淡无奇;但有一个问题你问出了"这个数是不是三位数",对方说"是",瞬间你把搜索范围从几千种可能压缩到了几百种。如果你事后复盘这场游戏,只去反思那九个平淡的问题,却忽略了那个让范围骤降的关键提问,你根本抓不住这场游戏真正的胜负手在哪里。dLM的填词过程也是如此,大部分填词步骤只是"确认"而已,只有少数几步才是真正"打开局面"的那一下。

现有的训练方法在干什么

现在的问题是,目前给dLM做后续训练（也就是让模型在已经学会说话的基础上,进一步学会更好地推理、更准确地做数学题和写代码）的方法,基本上没有用上这个"枢轴词"的信号。

论文梳理了两类常见做法。一类是把dLM当成普通的自回归模型来训练,直接在最终生成的完整文本上做监督微调(SFT)*。

监督微调（SFT）：拿一批标准答案（或模型自己生成的好答案）去训练模型,让它更倾向于模仿这些答案的写法。

这种做法的问题是,它把填词顺序完全抹掉了,一个在填词过程中起到关键决定作用的词,和一个在最后随手补上的无关紧要的词,在训练时拿到的权重是完全一样的。更糟的是,如果一条生成轨迹最后答案是错的,SFT式的训练往往会把整条轨迹里原本正确的中间步骤也一并"骂"错,这不是很冤吗?一道数学题你前面七步都对,只是最后抄错了一个数字,结果整条解题过程都被当成反面教材,这显然不公平。

另一类做法是在线强化学习(RL)*,比如diffu-GRPO、wd1++这些方法。

在线强化学习：模型一边生成答案,一边根据答案对错获得"奖励"或"惩罚",并实时调整自己的参数,这个过程需要反复采样、反复训练,计算成本很高。

这类方法确实关注了dLM的中间去噪步骤,但它们主要是把这些步骤当作强化学习里的"动作"去打分,整个训练过程仍然需要模型不断地重新生成、不断地在线采样,计算开销相当惹人头疼。

Pivot-SD的解法:只盯着那几个最关键的词下手

论文提出的方法叫Pivot-SD,核心思路可以用一句话概括:先把模型生成答案的全过程录下来,找出其中真正决定走向的那几步,只针对这几步做训练,其他地方一个字都不碰。

具体怎么找这几步?论文设计了一个叫"信息增益"的打分方式。

信息增益*：衡量某一步填词之后,剩下还没填的位置的"不确定性"下降了多少。模型对一个位置的预测越模糊,这个位置的"熵"（混乱程度）就越高;一旦某步填词让周围一大片位置瞬间变得清晰,这一步的信息增益就很高,会被标记为枢轴步。

论文里的图1展示了一个真实的数学题生成轨迹,熵值用颜色画成了一张热力图,你能清楚看到有那么几条红色虚线标出的"分水岭"时刻,熵值在这几步猛然下跌,其余大部分步骤的颜色几乎没什么变化。

那么,找到这些枢轴词之后,怎么用它们来训练模型?

论文的做法很干脆:先让冻结不动的原始模型(也就是没有经过任何微调的基础模型)自己生成一批答案,每道题生成四条不同的轨迹,然后用一个外部的"裁判"(数学题用精确匹配答案判断对错,代码题用跑单元测试判断对错)给每条轨迹打上"成功"或"失败"的标签。

接下来是最巧妙的一步:对于成功轨迹里挑出的枢轴词,用交叉熵损失去"奖励"它们,让模型更倾向于在类似情境下也这么写;对于失败轨迹里挑出的枢轴词,用一种叫"反似然"(unlikelihood)*的损失去"惩罚"它们,让模型降低在类似情境下写出这个词的概率。

反似然损失：和常规的交叉熵损失刚好反过来,交叉熵是让模型"更想"写出某个词,反似然则是让模型"更不想"写出某个词,专门用来给错误答案里的关键决策点"降权"。

而失败轨迹里除了那几个枢轴词之外的其他部分,一个字都不动。

这里有个很关键的细节:训练的时候,模型看到的不是最终完成的文本,而是那个枢轴词被填上时,周围依然大片蒙着码的"半成品状态"。换句话说,模型是在原汁原味地重演当时做决策的那个场景,而不是事后看着写完的作文反省。

这就好比复盘一场棋局。如果你只看棋局的最终结果,输了就是输了,你很难说清楚到底是哪一步走错了。但如果你能把棋盘倒回到那个真正扭转局势的关键落子时刻,盯着那一手棋单独复盘,效果显然会好很多,因为你是在那个"信息还没完全揭开"的紧张时刻去纠正判断,而不是在信息都摆明之后马后炮。如果不这么做,直接对整盘棋复盘打分,那些无关紧要的平淡走法会稀释掉关键手的权重,你最后学到的东西会很模糊。

数字有多夸张

论文给出的数字相当惊人:整个训练过程只用了200道训练题,每道题生成4条轨迹,每条轨迹只从中挑出10个枢轴词进行训练。10个词是什么概念?相对于256个词的回答长度,占比不到4%。

就是靠这不到4%的监督量,Pivot-SD在LLaDA-8B-Instruct这个80亿参数的扩散语言模型上,刷新了数学题和代码生成任务的成绩。下面这张表格摘录了几个关键对比:

在256长度的生成设定下,Pivot-SD在MATH数据集上拿到37.47分,比表现最好的对照方法SFT-GT(35.47分)高出2分;在GSM8K数学题上拿到79.51分,比训练了5000步、用了10倍题量的diffu-GRPO(77.63分)还高出近2分;在HumanEval+代码生成任务上拿到40.43分,比用同样数据训练的SFT-SD(35.78分)高出整整4.7分。

更让人意外的是训练效率。Pivot-SD整个流程(包括生成轨迹和训练)只需要2.8小时,而diffu-GRPO跑5000步要花22.9小时,wd1++要花17.3小时。也就是说,Pivot-SD用了不到对手八分之一的时间,反而拿到了更高的分数。

论文还专门做了一组消融实验,来验证到底是"挑枢轴词"这件事本身起了作用,还是随便挑几个词训练也行。结果显示,如果把惩罚机制随机施加到一个随机的去噪步骤上(而不是用信息增益挑出来的枢轴步),MATH分数会从37.47掉到29.93,甚至比完全不训练的基础模型(31.40)还低。这说明选对训练目标真的至关重要,盲目地"做点什么"甚至可能比"什么都不做"更糟。

这也呼应了开头那个作文批改的比方。如果老师批改作文的时候,随机圈出几个字就说"这里要改",而不是真正找准那个决定全文走向的句子,学生可能根本摸不着头脑该往哪个方向改进,甚至会被误导去改一些本来写得挺好的地方。

Pivot-SD换了个模型会不会失效

论文还检验了这个方法是不是只对LLaDA这一个模型管用。他们把同样的流程搬到另一个扩散语言模型Dream-7B上,结果Pivot-SD在数学、GSM8K、代码生成全部四项任务上都超过了基础模型和两个SFT对照组,说明这不是针对某个特定模型的"偏方",而是一个具有一定普适性的训练思路。

另外还有一项值得一提的发现:Pivot-SD训练出来的模型,在没见过的任务领域上表现也更稳。比如用MATH数据训练出来的模型,拿去测代码生成和GSM8K这些完全没训练过的任务,Pivot-SD的综合平均分(48.61)依然是所有对照方法里最高的。这说明这种"只抓关键决策点"的训练方式,学到的可能不只是某个任务的套路,而是更底层、更通用的推理习惯。

论文里还有一个很实在的发现,顺带提一下。在用代码数据训练一个在线强化学习对照组(wd1++)的时候,研究者发现这个模型一度彻底停止学习了,原因是基础模型生成的代码经常全部答案都通不过单元测试,导致整个批次的奖励信号全是零,梯度直接抵消,模型卡在原地动弹不得。这其实暴露了在线强化学习在数据稀缺、任务难度较高的场景下一个挺现实的脆弱点:如果运气不好连续几批样本全错,训练就可能彻底卡死。而Pivot-SD因为是离线训练,所有数据提前采样好、存好,不会有这种"训练中途卡壳"的风险。

写在后面

读到这篇论文的时候,我被一个细节打动了:研究者没有去发明一个全新的模型架构,也没有堆更多的参数,他们做的事情本质上是重新定义了"训练数据里谁重要"这个问题的答案。这件事听起来简单,但做起来需要先理解扩散模型那种"乱序填词"的特殊性,才能想到用熵的变化去衡量重要性。

还有一点挺有意思:论文作者专门强调,他们计算这个信息增益的分数,完全是复用生成过程中本来就要做的前向计算,没有额外增加一次模型调用。这个设计选择其实很克制,说明作者是真的在乎效率这件事,不是嘴上说说。

另一个值得琢磨的地方是失败轨迹的处理方式。大部分做法面对一条错误答案,会倾向于直接丢掉或者整体打低分,但Pivot-SD选择在失败轨迹里"挑着用",只惩罚那几个真正导致走偏的决策点,放过中间可能是正确的推理步骤。这种"精准打击"而不是"一棍子打死"的态度,某种程度上也是一种更接近人类复盘方式的处理逻辑,我们批评一个人做错事的时候,通常也不会说他做的所有事情都错了。

这篇论文留下的一个开放问题是,它目前只在一个固定长度的生成区块内做信息增益的计算和打分,如果未来要处理更长的推理链条,跨越多个生成区块去做这种"枢轴识别",会不会遇到新的麻烦?这个问题论文自己也承认还没有答案。

Q&A

Q1：Pivot-SD是什么？

A：Pivot-SD是一种针对掩码扩散语言模型的离线自蒸馏训练方法,它只挑选生成过程中信息增益最高的"枢轴词"进行训练,成功轨迹用交叉熵奖励,失败轨迹用反似然惩罚,只需200道题就能超过全量监督微调和在线强化学习的效果。

Q2：Pivot-SD和普通的监督微调有什么区别？

A：普通监督微调对所有生成的词一视同仁地训练,而Pivot-SD只训练那不到4%真正决定答案走向的关键词,并且对失败答案只惩罚导致出错的关键决策点,不会连带惩罚其中原本正确的部分。

Q3：Pivot-SD训练效率有多高？

A：Pivot-SD整个训练流程只需2.8小时,而对照的在线强化学习方法diffu-GRPO跑5000步需要22.9小时,wd1++需要17.3小时,Pivot-SD用更少时间拿到了更高的准确率。