1. 扩散模型做文本生成到底卡在哪扩散模型在图像生成领域已经杀疯了从早期的DDPM到后来的Latent Diffusion画质和多样性都碾压了此前的GAN方案。但同样一套思路搬到文本生成上效果一直不太能打。我前阵子花了不少时间跟踪这个方向发现核心矛盾就一个文本是离散的而扩散过程天然是连续的。图像像素本身就是连续值加噪去噪都在同一个空间里做很自然。但文本是token序列你没法直接对一个词做“加一点点噪声”这种操作。所以早期做扩散语言模型的人基本都走了一条路把离散token映射到一个连续的嵌入空间在这个空间里做扩散最后再解码回文本。这个思路本身没问题但实际跑起来生成质量总是差口气。具体差在哪我总结下来主要是两个瓶颈。第一个是嵌入空间的尺度问题。文本嵌入向量的数值范围通常很小比如很多预训练模型输出的embedding每个维度的值集中在-1到1之间方差很小。扩散模型在前向过程里要逐步加高斯噪声如果原始信号的尺度太小信噪比会迅速恶化模型很难学到有意义的去噪方向。第二个是采样效率问题。扩散模型天生需要几十甚至上百步去噪迭代文本生成本来对序列长度就敏感步数一多误差累积非常严重生成出来的句子经常前言不搭后语。这两个问题不解决扩散语言模型就只能停留在论文里的BLEU分数上实际用起来根本没法看。而这次要聊的这个工作恰恰就是从这两个点切入的——缩放文本嵌入加上蒸馏思路不算复杂但效果提升很明显。下面我按自己的理解把这个方案的来龙去脉、核心细节和实操中要注意的地方完整拆一遍。2. 核心思路拆解为什么是缩放加蒸馏2.1 先搞清楚扩散语言模型的基本流程在讲这个方案之前得先把扩散语言模型的标准流程捋一遍不然很多细节会看不懂。整个流程大致分三个阶段编码阶段把文本序列通过一个预训练的语言模型比如BERT类的编码器转成连续嵌入向量。假设序列长度是L嵌入维度是d那你就得到一个L×d的矩阵。扩散阶段在这个嵌入矩阵上做标准的高斯扩散。前向过程逐步加噪反向过程从纯噪声开始一步步去噪最终得到一个“干净”的嵌入矩阵。解码阶段把去噪后的嵌入矩阵映射回离散token。这一步通常用一个训练好的解码器或者直接找嵌入空间里最近的词向量。这个流程听起来很顺但实际跑的时候编码阶段出来的嵌入矩阵有个致命特点数值尺度太小。我实测过几个常见的文本编码器输出的嵌入向量标准差普遍在0.1到0.3之间有些维度甚至接近0。这意味着什么呢扩散模型前向加噪的公式是x_t sqrt(alpha_t) * x_0 sqrt(1 - alpha_t) * epsilon其中x_0是原始嵌入epsilon是标准高斯噪声。如果x_0的尺度只有0.1量级而epsilon的尺度是1那在t稍微大一点的时候sqrt(alpha_t)*x_0这一项基本就被噪声淹没了。模型要学的去噪方向信噪比极低训练难度陡增。2.2 缩放操作为什么能解决问题这个方案的第一步就是在编码之后、扩散之前对嵌入矩阵做一个全局缩放。具体来说就是给每个嵌入向量乘上一个缩放因子s。这个s不是随便定的而是根据嵌入矩阵的统计特性来算的。我自己的做法是先统计整个训练集嵌入矩阵的标准差sigma然后令s 1 / sigma。这样缩放后的嵌入标准差就归一化到1左右和扩散过程加的噪声尺度匹配上了。你也可以用更精细的方案比如按维度做标准化但全局缩放实现简单效果已经够用。注意缩放因子一定要在训练集上统计然后固定住验证集和测试集用同一个s。我见过有人每个batch重新算s结果训练loss震荡得厉害模型根本收敛不了。缩放之后信噪比的问题就缓解了。x_0的尺度和epsilon在一个量级sqrt(alpha_t)*x_0在t较大的时候也不会被完全淹没模型能学到更稳定的去噪方向。这个操作看起来简单但效果提升是实打实的。论文里的消融实验显示不加缩放的话生成文本的困惑度要高出30%以上。2.3 蒸馏又是为了解决什么缩放解决了训练稳定性问题但采样效率问题还在。标准扩散模型采样要跑几百步文本生成里每一步都要过一遍解码器开销太大。而且步数多了之后去噪误差会累积生成质量反而下降。蒸馏的思路是先训练一个教师模型步数多、质量好然后让一个学生模型去模仿教师模型的去噪轨迹但学生模型只需要很少的步数比如4步、8步就能完成生成。具体做法是教师模型跑完整的T步去噪记录下每一步的输入和输出学生模型在训练时直接学习从某个中间状态一步跳到教师模型多步之后的结果。这里有个关键细节蒸馏的目标不是让学生模型完全复现教师模型的每一步而是让学生模型在少步数的约束下尽可能逼近教师模型的最终输出。所以损失函数通常设计成两部分一部分是学生输出和教师最终输出的距离另一部分是学生中间状态和教师对应状态的距离。两部分加权求和权重需要调。我试过的一个配置是教师模型用100步学生模型用8步蒸馏损失里最终输出的权重设为0.7中间状态的权重设为0.3。这个比例不是固定的要根据具体任务调。文本生成任务里最终输出的权重可以再高一点因为文本的局部连贯性比中间过程的精确性更重要。3. 实操细节从编码到采样的完整链路3.1 编码器的选择与嵌入提取编码器这块我建议直接用预训练好的语言模型不要从头训。原因很简单文本嵌入的质量直接决定扩散模型的上限预训练模型在大规模语料上已经学到了很好的语义表示你从头训不仅费资源效果还未必好。具体选哪个编码器要看你的任务。如果是通用文本生成BERT类的双向编码器就够用如果是对话生成可以考虑用对话预训练模型如果是特定领域比如医疗、法律最好在领域语料上再微调一下。我实测下来编码器微调和不微调生成质量差距能有20%左右。提取嵌入的时候有个坑不要用最后一层的输出。最后一层的输出往往过拟合了预训练任务语义信息反而不如中间层丰富。我通常取倒数第二层或倒数第三层的输出具体取哪层要做实验。有个简单的判断方法把嵌入拿去做相似度检索看语义相近的句子嵌入距离是不是更近哪层效果好就用哪层。提取出来的嵌入矩阵记得做缩放。缩放因子按前面说的在训练集上统计标准差然后取倒数。缩放后的嵌入保存下来后面扩散训练直接用。3.2 扩散过程的参数配置扩散过程的参数主要有三个总步数T、噪声调度策略、去噪网络结构。总步数T训练的时候可以设大一点比如1000步这样前向加噪更平滑模型学到的去噪方向更精细。但采样的时候不需要这么多步用蒸馏后的学生模型8步就够了。这里有个经验T设得太大训练时间会线性增长但效果提升会饱和。我试过T500和T1000最终生成质量差距不到5%但训练时间差了一倍。所以T500是个比较划算的选择。噪声调度策略常见的有线性调度和余弦调度。线性调度简单但在t接近0和T的时候信噪比变化太快模型学起来吃力。余弦调度更平滑我推荐用余弦调度。具体公式是alpha_t cos((t/T s) / (1 s) * pi/2)^2其中s是一个小常数通常取0.008。这个调度策略在图像扩散里已经被验证过搬到文本上同样有效。去噪网络结构通常用Transformer。输入是带噪嵌入和步数t的嵌入输出是预测的噪声。步数t的嵌入可以用正弦位置编码和Transformer里的位置编码类似。网络层数不用太深6到8层就够了太深容易过拟合。注意力头数4到8个隐藏维度跟编码器输出维度对齐。3.3 蒸馏训练的具体步骤蒸馏训练分两步走。第一步先训练教师模型。教师模型就是标准的扩散模型用完整的T步去噪。训练目标就是标准的去噪损失L_teacher E[||epsilon - epsilon_theta(x_t, t)||^2]其中epsilon是真实加的噪声epsilon_theta是模型预测的噪声。这个损失很成熟直接优化就行。教师模型训练到收敛生成质量达到你的要求就可以停了。第二步训练学生模型。学生模型的网络结构和教师模型一样但去噪步数少。训练的时候从训练集里采样一个嵌入x_0然后随机选一个步数t做前向加噪得到x_t。教师模型从x_t开始跑完整的T-t步去噪得到最终输出x_0_teacher。学生模型从x_t开始只跑K步K远小于T-t得到x_0_student。损失函数就是L_student w1 * ||x_0_student - x_0_teacher||^2 w2 * ||x_t_student - x_t_teacher||^2其中x_t_student和x_t_teacher分别是学生和教师在中间某一步的状态。w1和w2是权重我前面说了w1取0.7w2取0.3是个不错的起点。提示蒸馏训练的时候教师模型的参数要冻结只更新学生模型。另外学生模型的初始化可以用教师模型的参数这样收敛更快。3.4 解码与后处理去噪完成后你得到一个干净的嵌入矩阵接下来要把它解码回文本。解码方法有两种一种是最近邻搜索在词嵌入表里找和每个位置嵌入最接近的词另一种是训练一个解码器把嵌入矩阵映射回token序列。最近邻搜索实现简单但有个问题嵌入空间里距离近的词语义上未必合适。比如“猫”和“狗”的嵌入距离可能很近但在某个上下文里用“猫”还是“狗”差别很大。所以最近邻搜索生成的文本经常出现语义漂移。训练解码器效果更好但需要额外数据。我的做法是用编码器的逆过程训练一个从嵌入到token的解码器。训练数据就是原始文本和它的嵌入损失用交叉熵。解码器结构可以用Transformer层数和编码器对齐。训练好之后解码质量比最近邻搜索高不少。后处理这块主要是处理重复和截断。扩散模型生成的文本有时候会出现连续重复的词或者生成到一半突然截断。重复问题可以用n-gram惩罚来缓解截断问题可以在解码时加一个结束符预测。这些技巧在自回归生成里很成熟搬过来稍作调整就能用。4. 常见问题与排查技巧实录4.1 训练loss不下降或震荡这是最常见的问题。我踩过的坑里排第一的就是缩放因子没算对。如果你忘了做缩放或者缩放因子算错了训练loss会卡在一个很高的值下不去。排查方法很简单打印一下嵌入矩阵的标准差看看是不是在1左右。如果不是重新算缩放因子。第二个原因是噪声调度策略选得不好。线性调度在t接近0和T的时候信噪比变化太剧烈模型学不动。换成余弦调度loss曲线会平滑很多。第三个原因是学习率太大。扩散模型的训练对学习率很敏感我通常用1e-4到5e-5之间的值配合warmup。如果loss震荡先把学习率降一半试试。4.2 生成文本重复或不通顺重复问题通常和采样步数太少有关。学生模型步数少去噪不充分模型容易陷入局部最优生成重复内容。解决办法有两个一是增加学生模型的步数比如从4步加到8步二是在损失函数里加一个重复惩罚项对连续相同的token做惩罚。不通顺的问题多半是解码器没训好。检查一下解码器的训练loss如果loss很高说明解码器没学到嵌入到token的映射关系。可以增加解码器的训练数据或者调大解码器的容量。还有一个容易被忽略的原因编码器提取嵌入的时候用了最后一层的输出。前面说了最后一层过拟合预训练任务语义信息不丰富。换成倒数第二层试试生成质量会有明显提升。4.3 蒸馏后质量下降明显蒸馏本质上是一种压缩质量下降是正常的但下降太多就不对了。如果学生模型生成质量比教师模型差很多先检查蒸馏损失的权重。w1太小的话学生模型只关注中间状态忽略了最终输出生成质量肯定差。把w1调大比如调到0.8或0.9再训一轮看看。另一个原因是学生模型的容量太小。如果学生模型的层数或隐藏维度比教师模型少很多它根本学不到教师模型的去噪轨迹。这种情况下要么增大学生的容量要么减少蒸馏的压缩比。还有个细节蒸馏训练的时候教师模型的输出最好做一下平滑。因为教师模型本身也有噪声直接拿它的输出当目标学生模型会学到一些噪声模式。可以对教师输出做移动平均或者用多个教师模型的集成输出效果会更稳。4.4 采样速度还是不够快蒸馏之后采样步数已经从几百步降到几步了但如果你的应用对延迟极其敏感可能还是觉得慢。这时候可以考虑量化和剪枝。量化就是把模型参数从FP32降到FP16甚至INT8推理速度能提升2到4倍质量损失很小。剪枝就是去掉一些不重要的注意力头或神经元也能提速。另一个思路是缓存。扩散模型采样的时候相邻步的嵌入变化不大可以把一些中间结果缓存起来避免重复计算。这个技巧在图像扩散里用得比较多文本扩散同样适用。5. 这个方案还能怎么扩展缩放加蒸馏这个组合本质上解决的是扩散语言模型的两个基础问题训练稳定性和采样效率。把这两个问题解决之后扩散语言模型的应用空间就打开了。我目前看到几个比较有意思的扩展方向。一个是条件生成在去噪网络里加入条件信息比如类别标签、关键词、甚至另一段文本实现可控文本生成。这个在图像扩散里已经很成熟了搬到文本上主要解决的是条件信息的编码方式。另一个是多模态生成把文本嵌入和图像嵌入映射到同一个空间做联合扩散。这样可以用一个模型同时生成文本和图像或者用文本引导图像生成、用图像引导文本生成。这个方向目前还在早期但潜力很大。还有一个是少样本适配用少量领域数据微调扩散语言模型让它适应特定领域。因为扩散模型的训练目标比较稳定少样本微调比自回归模型更容易收敛。我试过用几百条数据微调生成质量已经能看了。最后说一个我自己的体会扩散语言模型现在还在快速迭代期很多技巧都是从图像扩散那边搬过来的但文本和图像的本质差异不能忽略。离散性带来的问题不是简单套用连续扩散的方案就能解决的。缩放和蒸馏只是第一步后面肯定还会有更多针对文本特性的改进。如果你也在做这个方向建议多关注嵌入空间的性质很多问题的根源都在那里。