离散扩散采样不再每步掷骰子:指数残差流让轨迹在初始化后确定,并支撑一致性蒸馏
核心概要
离散扩散模型的采样此前每一步都要重新抽取随机数,轨迹无法固定。指数残差流为每个位置和词元配一个连续残差变量,初始化后按分段常微分方程确定演化,残差归零即触发一次词元替换,并证明该构造保持原扩散模型各时刻的边缘分布。在匹配模型评估次数下,其生成质量与祖先采样相当,确定性轨迹还可用于一致性蒸馏。
深度剖析
离散扩散模型现在可以有一种只依赖初始随机性的采样轨迹,而不必在每次跳变时重新掷骰子。 此前的标准做法是Gillespie算法或离散时间反向链的祖先采样,每次跳变都要抽取新的随机数,因此轨迹本身是随机的。 该构造被证明在每一时刻都保持原连续时间马尔可夫链的边缘分布,且残差在该时刻仍是独立标准指数变量并与离散状态独立。
做法是给每个位置和每个词元配一个非负的残差变量,让离散状态只在某个残差降到零时改变一个词元。 指数残差流沿用Gillespie第一反应法中的指数变量,但不再在跳变后重新抽样,而是让所有残差按由前向与反向速率矩阵决定的分段常微分方程继续确定演化。 为用于学习到的模型,采样在每个区间内冻结留一法去噪器而保留速率已知的时间依赖,并给出掩码扩散与均匀扩散的显式形式;当去噪器精确且数据分布可分解时该格式是精确的。
确定性轨迹使一致性蒸馏在离散扩散中有了可对齐的路径目标。 此前缺少离散版的概率流常微分方程,常被当作反对基于轨迹蒸馏的理由。 在TinyGSM上,均匀扩散的蒸馏在2与4次模型评估时把最强基线的准确率提高一倍,在8与16次时也更高,32次时接近教师;在OpenWebText上,均匀扩散的蒸馏在匹配熵下困惑度低于其他基线,掩码扩散的蒸馏则大致与第二好的基线持平。
启示与展望
演示范围是OpenWebText与TinyGSM上的掩码扩散和均匀扩散、Sudoku补全上的蒸馏实验,以及DiffusionGemma在MMLU-Pro的140道题上的对比;采样质量结论限定在匹配模型评估次数的条件下。
确定性轨迹的分布一般与原反向马尔可夫链不同,且不必是马尔可夫过程;蒸馏的最大增益出现在均匀扩散,掩码扩散的结果只是与所比较的基线相当。
