E-MoE：当扩散语言模型学会"抽牌"而不是"独立猜测"

你有没有玩过这样一个文字接龙游戏：两个人同时被要求补全一句话"今天我去了___，买了___"，但彼此看不到对方的答案。结果很可能是，一个人写"北京"，另一个人写"苹果手机"，拼起来变成"今天我去了北京，买了苹果手机"，读起来好像没错，但其实这两个词根本不是在同一个语境下商量出来的，纯属巧合地凑在一起没露馅。

这正是当前主流扩散语言模型遇到的尴尬处境。

扩散语言模型*：一种不同于GPT式从左到右生成文字的模型，它先把整句话"打码"成一堆遮盖符号，再通过多轮"猜词"把遮盖的地方逐步填回正确的词，整个过程可以同时填很多个位置，而不是一个字一个字往下写。

这种方式理论上比自回归模型快得多，因为自回归模型写一句100个字的话，必须老老实实做100次前向计算，一个字接一个字往下吐；而扩散模型可以几步之内把整句话填满，天然适合并行加速。但快是快了，质量却常常打折扣，尤其是在"步数很少"的情况下，问题会被放大到肉眼可见的程度。

核心问题出在哪？独立猜词的代价

扩散语言模型的标准做法叫作"分解式"生成。

分解式（factorized）*：模型虽然能看到整个句子的上下文，但它对每个被遮盖的位置分别独立地给出一个概率分布，然后各个位置各自采样，互不商量。

这就是开头那个文字接龙游戏的真实写照。假设训练数据里只有两种合法组合："纽约"和"旧金山"，从来不会出现"纽约金山"或者"旧约市"这种混搭。但如果模型对"城市名"的第一个字和第二个字分别独立采样，哪怕两个位置单独看都给出了很合理的概率，拼在一起却完全可能产生数据里从未出现过的怪词。论文里给出的例子正是这样：模型同时猜测两个位置，一半概率给"纽"一半给"旧"，一半概率给"约"一半给"山"，于是"纽约"、"旧金山"、"纽金山"、"旧约"四种组合各有25%的概率出现，而后两种根本不是真实世界存在的地名。

这篇论文把这个现象起名叫"分解误差"（factorization error），并且指出，一次扩散步骤填的词越多，这种误差积累得越厉害，而步数少恰恰意味着每一步要填更多词，于是想要通过减少步数换取生成速度，反而会让这个问题雪上加霜，速度和质量成了一对矛盾。

在这篇论文之前，已经有研究者尝试过修补这个漏洞。VADD（Variational Autoencoding Discrete Diffusion）引入了一个连续的隐变量，让模型在生成文字前先"抽一个隐藏的主题"，这个主题会同时影响所有被遮盖的位置，相当于两个人在动笔之前先悄悄商量了一下今天要写什么话题。这个思路是对的，但VADD的隐变量是一个连续的高维向量，服从标准正态分布这种人为设定的先验，而且还需要额外训练一个独立的"识别网络"（recognition network，就像VAE里的编码器）去学习怎么从数据推断出这个隐变量。更麻烦的是，论文提到这类方法容易出现"后验坍缩"（posterior collapse），也就是说，辛辛苦苦训练出来的隐变量，最后模型根本没拿它当回事，所有信息还是从主干网络里走，隐变量成了摆设。

E-MoE的核心想法：别造新轮子，混合专家模型自带隐变量

这篇论文的作者们换了个思路。他们注意到现在很多扩散语言模型为了做大规模，已经普遍采用了混合专家架构。

混合专家模型（Mixture-of-Experts，MoE）*：不是让一个巨大的神经网络处理所有输入，而是准备很多个较小的"专家"网络，每次输入进来，由一个叫"路由器"（router）的小模块决定把这次计算交给哪几个专家处理，这样可以在参数量很大的情况下，实际计算量却保持可控。

既然路由器每次都要为每个词、每一层决定"这个词该交给哪个专家"，这个决定本身不就是一个天然存在的、离散的、每个位置都有的隐变量吗？根本不需要额外造一个连续向量出来，也不需要专门训练一个识别网络，路由器自己就既能当生成时的"先验"，也能当训练时的"后验"。

这就是E-MoE（Enhanced Mixture-of-Experts）这篇论文的核心贡献。它把混合专家模型里原本只是为了省计算量而设计的路由机制，重新利用起来，当作协调多个位置联合生成的隐藏纽带。

想象一个班级里有两位助教分别负责批改不同题型的作业，老师原本设计这套分工纯粹是为了减轻自己的工作量。但现在老师突然发现，只要记录下某份试卷被分给了哪位助教，其实就能反推出这份试卷大概率是哪种题型，这个"分配记录"本身包含着信息，完全可以拿来做别的用途，而且不需要额外雇佣第三个人专门记录。E-MoE做的正是这件事，它把路由决策这个"分配记录"，直接升级成了协调生成的隐变量。如果不这样做，而是像VADD那样另起炉灰训练一个识别网络，不仅多了一倍的训练开销，还要为这个新的连续隐空间设计合适的先验分布，这件事本身就不简单，自然语言的语义结构未必能用一个标准正态分布描述清楚。

具体怎么做：两次前向传播，一个路由器两种身份

E-MoE的训练过程需要两次前向传播。

第一次传播，模型看到的是完整无遮盖的真实句子，路由器根据这份"完整答案"给出的专家分配，叫作"后验"（posterior），相当于考试结束后老师看着标准答案给试卷分类。

第二次传播，模型看到的是带遮盖符号的残缺句子，路由器只能根据残缺信息给出专家分配，这叫"先验"（prior），相当于老师还没看到标准答案，只凭残缺的卷面内容猜测这应该归到哪一类。

训练的目标，就是让"只看残缺信息做出的猜测"尽量逼近"看到完整答案后的真实分类"，这中间的差距用KL散度（Kullback-Leibler divergence，一种衡量两个概率分布差异程度的指标）来衡量，并作为一项损失函数加入训练。

这里有个精妙的地方值得多说一句：VADD的先验和后验是两个完全不同的网络，一个用于生成，一个专门用于训练时的"作弊式"推断，两者参数互不共享；而E-MoE里先验和后验其实是同一个路由器，只是喂给它的输入不同，一次是残缺句子，一次是完整句子。这意味着E-MoE不需要多养一个网络，参数量和推理时的计算开销和普通分解式模型几乎一样。论文里的对比表格写得很清楚，VADD的总参数是256.5M，但实际推理时每个词只用到139.9M；E-MoE的总参数虽然高达536.1M（因为专家数量多），但每个词实际用到的激活参数只有139.4M，跟最朴素的分解式基线MDLM的139.3M几乎分毫不差。

到了真正生成文字的推理阶段，流程就简化了，只需要用先验路由器根据当前残缺句子选出专家，专家给出预测，一次前向传播搞定，速度和原始的分解式模型完全一样，没有任何额外开销。

训练时怎么让"选专家"这件事能被梯度更新呢？专家的选择本质上是一个离散的、不可微分的操作，直接用arg max选最大概率的专家会导致梯度传不过去。论文采用了Gumbel-Softmax技巧配合直通估计器。

Gumbel-Softmax*：一种让离散采样过程变得近似可微分的技巧，通过往概率分布的对数值上加一种特殊设计的随机噪声（Gumbel噪声），再用一个带温度参数的softmax函数做平滑近似，温度越低，这个近似越接近真正的离散选择，训练初期温度较高让梯度更容易流动，后期逐渐降温让选择更接近真实的硬性决策。

直通估计器*：前向传播时仍然做真正的离散选择（比如选出概率最高的专家），但反向传播时假装这个选择过程是用前面那个可微分的平滑近似完成的，用这种"障眼法"让离散操作也能参与梯度下降。

理论基础：一个站得住脚的证据下界

光有直觉上的设计还不够，论文给出了严格的数学推导，证明了这套混合专家式的生成过程确实有一个合理的训练目标，也就是负对数似然的上界（ELBO，Evidence Lower Bound的反面，常用于训练生成模型）。

证据下界（ELBO）*：生成模型训练中常用的一种技巧，因为真实数据的概率往往没法直接算出来，于是构造一个可以计算的"下界"函数，通过最大化这个下界，间接提升真实概率。

推导过程里有个细节特别值得一提，不同于以往多数工作把隐变量设置成"整句话共用一个"，E-MoE把隐变量设计成"每个位置、每一层各有一个"，并且假设同一位置上不同层的专家选择之间存在依赖关系，用链式法则把这种依赖关系拆解成逐层的条件概率。这样设计的好处是，隐变量的粒度足够细，能够捕捉到位置级别的关联，而不是只能表达整句话级别的粗糙信息。

三个实验场域：从玩具数据到真实文本

论文在三类任务上做了验证，分别是二维合成数据、二值化MNIST图像，以及LM1B文本数据集。

先看最直观的二维合成实验。研究者构造了两种分布，一种是"8个高斯团块"（8-modes），数据点被安排成8个彼此独立、不相连的小簇；另一种是"瑞士卷"（Swiss-roll），数据沿着一条细长的螺旋曲线分布。这两种数据都有一个共同特点，它们的"有效区域"占整个空间的比例很小，一旦模型犯了分解误差，生成的点很容易落到数据从未出现过的空白地带。

实验结果印证了前面的推理。在只用一步生成（NFE=1，NFE即Number of Function Evaluations，表示模型实际做了多少次前向推理）的极端条件下，分解式的MDLM模型生成的点几乎铺满了整个网格，完全看不出8个团块的轮廓，而在瑞士卷任务上更是直接填满了整个圆盘，螺旋形状荡然无存。E-MoE在同样只用一步的情况下，依然能准确地把点聚集在8个真实团块附近，螺旋的形状也清晰可辨。论文给出的"有效性"（validity，即生成的点落在数据支撑集附近的比例）指标也证实了这一点：在8-modes任务上，一步生成时MDLM只有37.7%的点算合格，E-MoE却达到了94.4%，VADD也有91.8%，但依然略逊于E-MoE。

再来看MNIST手写数字图像生成。论文把每张28×28（填充后变成32×32）的图像看作1024个二值像素组成的序列，相当于把每个像素当成一个只有"黑"和"白"两种状态的词。在极少步数下，分解式MDLM生成的图像基本是一团没有结构的噪点，而E-MoE从仅仅一步生成开始，就能画出能辨认出数字形状的图案。最终的测试集比特每维度误差（BPD，Bits-Per-Dimension，数值越低说明模型对数据分布的刻画越精确）上，E-MoE取得了0.062，优于VADD的0.064和MDLM的0.077，而且参数量跟VADD几乎一样（2.49M对2.50M）。

最后也是最重要的一关，真实世界文本生成，数据集是LM1B（One Billion Word Benchmark，一个常用于语言建模基准测试的大规模英文语料库）。

下面这张表格摘录了论文中最关键的几组对比数据：

| NFE（采样步数） | MDLM Gen-PPL | VADD Gen-PPL | E-MoE Gen-PPL |

|---|---|---|---|

| 1 | 1433.8 | 1270.8 | **643.8** |

| 2 | 997.5 | 763.2 | **383.7** |

| 4 | 477.8 | 370.5 | **235.4** |

| 8 | 260.5 | 220.7 | **174.9** |

| 16 | 179.4 | 168.4 | **146.4** |

（表中Gen-PPL即生成文本的困惑度，数值越低代表生成文本在GPT-2-large这个判别模型看来越通顺自然，越贴近真实人类语言；所有数值都在相近的样本熵H约4.34-4.37下比较，确保不是靠牺牲多样性来降低困惑度）

一步生成时，E-MoE把困惑度从MDLM的1433.8压到了643.8，降幅超过一半。另一个更直观的指标MAUVE（一种衡量生成文本整体分布与真实文本分布接近程度的指标，数值越高越好）差距拉得更大，两步生成时，E-MoE达到24.58%，而MDLM只有2.95%，VADD也只有3.76%，相当于八倍以上的差距。随着采样步数增加到32步以上，各个模型的差距逐渐收窄，这也验证了前面讲过的逻辑：步数越多，每一步需要同时填的词越少，分解误差本身就不那么致命了，E-MoE的优势自然会变得不那么明显。

这就好比一个团队赶工，如果deadline只剩下1小时，必须5个人同时各写一部分报告再拼接，这时候提前沟通好分工和口径（E-MoE的做法）和各自闷头写（分解式做法）的最终效果差距会非常大；但如果给足了24小时，可以让5个人轮流写、互相传阅、逐步修改，那么哪怕一开始没沟通好，最后也能改得七七八八，两种做法的差距自然就不那么突出了。如果不给E-MoE这个"提前商量"的机制，直接退化成贪婪选择专家（论文附录里专门做了这个消融实验），一步生成的困惑度会从626.1飙升到1079.5，几乎翻倍，证明了"混合多个专家路线"这件事本身，而不是单纯的参数量增加，才是性能提升的真正来源。

论文还特别对比了隐变量在训练过程中是否"被真正用起来"。VADD的隐变量KL值训练到后期稳定在大约0.14 nats/token（nats是用自然对数衡量信息量的单位），数值很低，意味着后验分布跟先验分布已经非常接近，隐变量传递的信息量很有限，这正是"后验坍缩"的典型表现。相比之下，E-MoE的路由KL值稳定在0.65 nats/token左右，是VADD的四倍多，没有用任何预热策略或者额外的辅助损失函数，说明这个离散隐变量在整个训练过程中始终在被认真使用，没有被模型"悄悄抛弃"。

这篇论文把时间拉回到了两种技术路线交汇的节点上。一边是扩散语言模型为了追求速度而放弃了传统自回归模型逐词生成的严谨性，用并行换取效率；另一边是混合专家架构本来只是为了在不增加计算量的前提下扩大模型容量而设计的工程技巧。这两条原本各自独立发展的技术路线，在E-MoE这里被巧妙地拧到了一起，路由决策顺手解决了并行生成的协调问题，没有花一分额外的推理成本。

写在后面

读完这篇论文，最触动我的其实不是实验数字，而是这个"废物利用"式的设计思路本身。

很多论文解决问题的方式是"发明一个新组件"，而E-MoE走了一条相反的路，它问的问题是"模型里已经存在的某个机制，能不能被赋予一个新的身份"。混合专家模型里的路由器，本来的职责仅仅是"分配计算资源"，没人指望它承担"协调多个位置联合生成"这种语义层面的任务。但论文作者观察到，路由决策本质上也是一种对输入信息的压缩和归类，这种压缩天然带有某种结构信息，稍加改造就能派上新用场。

这让我联想到很多工程领域里类似的"意外发现"，一个为了别的目的设计的日志系统，后来被用来做异常检测；一个为了省电设计的传感器，后来被用来做手势识别。好的设计往往自带多种可能性，只是最初的设计者没想到而已。

论文里还有个细节我觉得值得单独拎出来说，就是VADD的后验坍缩问题和E-MoE刻意回避这个问题的方式。VADD为了防止隐变量被忽略，专门设计了KL权重从很小的值逐渐退火到1的策略，这其实是一种"打补丁"式的解决方案，承认了问题存在但没有从根源上消除它。而E-MoE因为先验和后验共享同一个网络，本身就不太容易出现这种"后验学到的东西先验完全学不到"的割裂情况，这是一种更结构性的解决方式，而不是训练技巧上的缝缝补补。

当然论文自己也承认，这套方法的优势会随着采样步数增多而逐渐消失，本质上它解决的是"步数少、并行度高"这个特定场景下的问题。如果未来的硬件和算法让扩散模型可以用更多步数而不牺牲速度，这种方法的必要性会不会反而下降？这是个值得继续追问的问题，但至少现在，在"快"和"准"这对矛盾里，E-MoE给出了一个不需要额外代价的折中答案。

Q&A

Q1：E-MoE是什么？

A：E-MoE（Enhanced Mixture-of-Experts）是一种改进的掩码扩散语言模型，它把混合专家架构里本来用于分配计算资源的路由决策，重新用作协调多个位置联合生成文字的离散隐变量，不增加额外推理成本的情况下提升了少步数生成的质量。

Q2：E-MoE和VADD有什么区别？

A：VADD用一个连续的高斯隐变量，需要额外训练一个独立的识别网络，容易出现隐变量被忽略的"后验坍缩"问题。E-MoE用离散的专家路由决策做隐变量，先验和后验共用同一个路由器，不需要额外网络，也不容易坍缩。

Q3：E-MoE为什么在生成步数少时效果更好？

A：因为扩散模型步数越少，每一步要同时填的词就越多，各位置独立采样导致的"分解误差"就越严重。E-MoE通过共享的专家路由让多个位置的生成互相协调，专门缓解这种少步数场景下的协调失灵问题。