首先语言扩散分为离散扩散和连续扩散,离散扩散是说比如4个词,刚开始是[MASK] [MASK] [MASK] [MASK] → [MASK] [MASK] [MASK] [France] → [The] [MASK] [MASK] [France] → [The] [capital] [MASK] [France] → [The] [capital] [of] [France],当然也可以一次确定两个词,这样就只需要unmasking两次。然后它指出这种方式为了加速推理,实现“少步”,一次应该unmasking多个词,但是这样生成质量就会下降(因为一次生成的多个词是没有交流的)。然后一些加速的改进还是在问“这个位置到底是 token A,还是 token B?”
于是作者采用了连续扩散的概念,不要急着决定这个位置究竟是什么 token,比如某个位置有0.5的概率是Paris,0.3的概率是London,0.2的概率是Berlin(注意实际维护的不是概率,而是连续向量,多权重相加),然后所有位置在embedding space里不断变化,直到最后一步才变成离散的token。
接着就是说难点,embedding geometry 怎么设计;noise schedule 怎么设计;training objective 怎么设计。没有从头开始训练,而是使用离散扩散(LLaDA)作为基模。
接下来来看method,还挺神奇的,对每个真实 token 先在一个低维 embedding space 中加入不同强度的 Gaussian noise,得到连续 noisy state;再用 softmax converter 将它转换成 LLaDA 原始 embedding 的 soft mixture,最后仍然让 LLaDA 用普通 cross-entropy 恢复原 token。
DSL-LLaDA 的核心思想非常直观:
传统 Masked Diffusion LM 中,一个 token 位置基本只有两种状态:
clean token:这个位置的信息完全可见;
[MASK]:这个位置的信息完全未知。
也就是说,中间不存在“部分确定”的状态。
这在高步数生成时问题不大,因为每一步只需要确定少量 token;但在 low-step 推理中,一次需要同时确定大量 token,过早进行 hard token commitment 很容易导致重复、提前终止等问题。
DSL-LLaDA 的核心改动,就是把这种 0/1 式的 clean / mask 状态扩展为连续的不确定状态 :
完全未知
↓
有一些可能性
↓
越来越确定
↓
最终确定为某个 token
作者额外引入一个较小的 token embedding:
W token ∈ R V × d W_{\text{token}}\in\mathbb{R}^{V\times d} W token ∈ R V × d
其中论文中使用:
对于真实 token x i x_i x i ,首先在这个小 embedding space 中加入不同强度的 Gaussian noise:
z i = γ i e x i + γ i ϵ i z_i=\gamma_i e_{x_i}+\sqrt{\gamma_i}\epsilon_i z i = γ i e x i + γ i ϵ i
其中 γ i \gamma_i γ i 表示 SNR。
γ i \gamma_i γ i 很小时,几乎不知道这个位置是什么;
γ i \gamma_i γ i 增大时,这个位置包含的真实 token 信息越来越多。
因此,一个位置不再只有 clean / mask 两种状态,而可以具有连续变化的不确定程度。
问题是,预训练 LLaDA 并不能直接处理这个新的低维连续向量 z i z_i z i 。
作者因此设计了一个非常简单的 converter:
首先将 z i z_i z i 与所有 token 的小 embedding 计算相似度,再经过 Softmax,得到它属于各个 token 的权重:
π i = Softmax ( β z i ⊤ W token ⊤ + b ) \pi_i=\operatorname{Softmax} \left( \beta z_i^\top W_{\text{token}}^\top+b \right) π i = Softmax ( β z i ⊤ W token ⊤ + b )
例如某个位置可能得到:
Paris 0.5
London 0.3
Berlin 0.2
这时并不立即选择 Paris,而是利用这些概率,对 LLaDA 原本的 token embedding 做加权:
h i = π i W backbone h_i=\pi_iW_{\text{backbone}} h i = π i W backbone
例如:
h i = 0.5 E Paris + 0.3 E London + 0.2 E Berlin h_i= 0.5E_{\text{Paris}} +0.3E_{\text{London}} +0.2E_{\text{Berlin}} h i = 0.5 E Paris + 0.3 E London + 0.2 E Berlin
最终得到的仍然是 LLaDA 原本能够接受的 d model d_{\text{model}} d model 维 embedding。
因此整个转换过程可以写成:
continuous noisy state
↓
小 token embedding + Softmax
↓
token probability
↓
LLaDA embedding 加权求和
↓
soft token embedding
↓
LLaDA
这样几乎不需要修改原始 LLaDA 的网络结构,就可以让它学会处理连续不确定的 token 状态。
训练过程其实非常简单。
DSL 只改变了 输入的加噪方式 ,并没有改变训练目标。
模型输入的是不同噪声程度的 soft token embedding:
但目标仍然是恢复原始 token:
L = − ∑ i log p θ , i ( x i ∣ h 1 : L ) \mathcal{L} = -\sum_i \log p_{\theta,i}(x_i|h_{1:L}) L = − i ∑ log p θ , i ( x i ∣ h 1 : L )
因此,可以简单理解为:
以前 LLaDA 只学会了处理“完全正确”和“完全 MASK”的输入,现在通过 DSL continuation training,让它进一步学会处理不同可信程度的 token 状态。
作者只在 LLaDA-8B-Instruct 上继续训练了 1000 steps。
推理时不再采用传统的 iterative unmasking,而是在 continuous embedding space 中不断去噪。
每一步 LLaDA 都会预测当前位置的 token probability,然后取概率最高的 Top-M token,对小 embedding e v e_v e v 做加权平均:
e ^ i = ∑ v ∈ S i p ~ θ , i ( v ) e v \hat e_i = \sum_{v\in S_i} \tilde p_{\theta,i}(v)e_v e ^ i = v ∈ S i ∑ p ~ θ , i ( v ) e v
得到模型当前预测的 clean embedding 。
接下来用这个 clean embedding 去更新当前 continuous state:
y i ← y i + Δ γ e ^ i − y i γ y_i \leftarrow y_i + \Delta\gamma \frac{\hat e_i-y_i}{\gamma} y i ← y i + Δ γ γ e ^ i − y i
直观上就是:
当前 noisy state
↓
LLaDA 预测 clean embedding
↓
朝 clean embedding 移动一点
↓
再次预测
↓
再次移动
↓
……
这个过程和连续生成模型中的 flow / denoising 思路非常相似:状态在连续空间中逐渐向更干净的状态演化。
最重要的是,中间过程始终不真正决定这个位置是哪一个 token 。
只有在最后一步,才执行:
x ^ i = arg max v p θ , i ( v ) \hat x_i = \arg\max_v p_{\theta,i}(v) x ^ i = arg v max p θ , i ( v )
将 continuous state 最终转换成离散 token。
因此 DSL-LLaDA 最核心的区别可以概括为:
传统 Masked Diffusion:
MASK → token
MASK → token
MASK → token
每一步不断进行 hard token commitment
DSL-LLaDA:
continuous noisy state
↓
continuous noisy state
↓
continuous noisy state
↓
clean state
↓
最终才转换为 token
其本质就是:
将原本只有 clean / mask 两种状态的离散扩散,扩展为具有连续不确定度的 embedding-space denoising,并利用“小 embedding + Softmax + 原始 LLaDA embedding 加权”这一简单结构,将连续状态无缝接入已有的 LLaDA。
其中“小 embedding → Softmax → 原 LLaDA embedding 的 soft mixture”确实是这篇方法里非常关键的桥接设计:论文用它把连续 DSL 状态映射回预训练 backbone 能直接处理的输入空间。
推理阶段则使用 Top-M token 得到 clean embedding estimate,再持续更新连续状态,最后才通过 argmax 做 hard decoding。