扩散应进入语言模型的何处？几何引导的隐藏状态替换

# 扩散应进入语言模型的哪个位置？几何引导的隐藏状态替换

来源：https://arxiv.org/html/2605.14368

Injin Kong¹   Hyoungjoon Lee²   Yohan Jo¹,†  
¹数据科学研究生院，首尔大学  
²生物系统与生物材料科学与工程系，首尔大学  
mtkong77@snu\.ac\.kr, hjoon721@snu\.ac\.kr, yohan\.jo@snu\.ac\.kr  

###### 摘要  
连续扩散语言模型落后于自回归Transformer，部分原因是扩散被应用于不适合语言去噪和令牌恢复的空间。我们提出DiHAL，一种几何引导的扩散–Transformer混合架构，它询问扩散应该进入预训练Transformer的哪个位置。DiHAL使用基于几何的代理对层进行评分，选择扩散友好的隐藏状态接口，并用扩散桥替换下层Transformer前缀，同时保留上层和原始LM头。通过重构所选层的隐藏状态而非令牌，DiHAL避免了直接的连续到离散恢复。在8B规模骨干模型上的实验表明，在固定的桥训练协议下，几何评分能够预测有效的浅层插入位置，并且在匹配扩散/恢复训练预算的诊断比较中，隐藏状态恢复优于连续扩散基线。这些结果表明，隐藏状态几何有助于识别在预训练语言模型中何处可以进行基于扩散的替换。  

22通讯作者。  

## 1 引言  

大型语言模型在各种语言生成任务中取得了显著进展，但这种进展伴随着规模和计算成本的增加(Brown et al., 2020 (https://arxiv.org/html/2605.14368#bib.bib1); Hoffmann et al., 2022 (https://arxiv.org/html/2605.14368#bib.bib2); Yang et al., 2025 (https://arxiv.org/html/2605.14368#bib.bib3))。扩散模型提供了一种基于迭代去噪的不同生成范式，并已成为图像生成的主流方法(Song et al., 2021 (https://arxiv.org/html/2605.14368#bib.bib4))。它们的成功激发了人们对基于扩散的语言生成的兴趣(Li et al., 2022 (https://arxiv.org/html/2605.14368#bib.bib8); Strudel et al., 2023 (https://arxiv.org/html/2605.14368#bib.bib9); Nie et al., 2025 (https://arxiv.org/html/2605.14368#bib.bib15))。然而，将扩散从图像迁移到文本是困难的，因为文本生成最终必须处理离散令牌。一个自然的应对措施是使扩散适应文本的离散性。先前的工作探索了离散令牌破坏、掩码扩散、连续到离散恢复，以及基于令牌嵌入、自条件嵌入或学习到的文本潜在变量的连续扩散(Li et al., 2022 (https://arxiv.org/html/2605.14368#bib.bib8); Strudel et al., 2023 (https://arxiv.org/html/2605.14368#bib.bib9); Lovelace et al., 2023 (https://arxiv.org/html/2605.14368#bib.bib21); Gong et al., 2023 (https://arxiv.org/html/2605.14368#bib.bib12); Zhang et al., 2025 (https://arxiv.org/html/2605.14368#bib.bib23))。尽管如此，基于扩散的语言模型仍然落后于自回归Transformer，尤其是在连续扩散设置中(Jo and Hwang, 2026 (https://arxiv.org/html/2605.14368#bib.bib42))。一个常见的解释是，去噪后的连续向量最终必须映射回离散令牌，因此表示空间中的微小误差就可能改变恢复出的令牌(Li et al., 2022 (https://arxiv.org/html/2605.14368#bib.bib8))。  

为什么这种差距仍然存在？我们从另一个假设出发：离散性很重要，但可能无法完全解释这种差距。Transformer语言模型也使用离散令牌，然而大部分计算发生在连续隐藏状态中，这些状态随后被映射到词汇表logits(Vaswani et al., 2017 (https://arxiv.org/html/2605.14368#bib.bib20))。这表明困难可能并非源于连续性本身，而是源于在不具合适几何特性的连续空间中应用扩散。如果连续空间的选择至关重要，那么核心问题就变为：什么特性使得一个表示空间适合扩散？我们将这样的空间称为*扩散友好*空间：一个容易去噪、在分数估计不完美时保持稳定、并且足够简单以便扩散学习的空间。稍后我们将利用Langevin动力学和集中理论(Villani, 2009 (https://arxiv.org/html/2605.14368#bib.bib44); Bakry et al., 2014 (https://arxiv.org/html/2605.14368#bib.bib45); Ledoux, 2001 (https://arxiv.org/html/2605.14368#bib.bib46))来阐述这些要求。  

在语言模型中，这样的空间能在哪里找到？一个预训练Transformer在令牌嵌入层和LM头之间已经包含了许多连续的隐藏空间。这些隐藏状态不会直接解码为令牌；它们被后续的Transformer层处理，然后由LM头产生最终的令牌分布(Vaswani et al., 2017 (https://arxiv.org/html/2605.14368#bib.bib20))。因此，在内部层进行扩散可以针对隐藏状态恢复，而非直接的令牌恢复(Lovelace et al., 2023 (https://arxiv.org/html/2605.14368#bib.bib21); Rombach et al., 2022 (https://arxiv.org/html/2605.14368#bib.bib43))。由于隐藏状态几何随深度变化，我们提出疑问：哪个Transformer层提供了最扩散友好的表示空间？  

为了回答这个问题，我们提出DiHAL（Diffusion-Transformer Hybrid Architecture for Language Generation，面向语言生成的扩散-Transformer混合架构），一种基于*定位与替换*策略的混合架构。如图1 (https://arxiv.org/html/2605.14368#S1.F1)所示，DiHAL使用基于几何的标准定位扩散友好的层，然后用一个*扩散桥*替换下层Transformer层，该桥重构所选层的隐藏状态，同时保留上层和原始LM头用于令牌预测。这减少了连续到离散的恢复误差，并将连续扩散应用于语言重新定义为在内部为去噪选择合适的表示空间。  
我们的贡献有三方面。  

- ·我们将预训练Transformer语言模型中的扩散插入形式化为一个*几何引导的接口选择问题*，并提出了实用的逐层代理——局部紧致性、全局刚度和有效秩——用于识别扩散友好的隐藏空间。  
- ·我们引入了一个固定的几何评分，该评分在无需穷举逐层桥训练的情况下缩小了有效插入层的搜索范围，并且在单周期桥训练协议下，与8B规模骨干模型上的隐藏状态重构质量强相关。  
- ·我们提出了DiHAL，一种定位与替换混合架构，用条件扩散桥替换下层Transformer层，并重用上层和LM头。在诊断性的扩散/恢复预算下，DiHAL表明隐藏状态恢复可以在生成困惑度和多样性上优于嵌入层、潜在变量和连续到离散接口。

参见图注

图1：定位与替换框架。逐层几何代理对Transformer层进行评分，选择插入点，并引导用扩散模块进行替换。

## 2 背景  

Transformer语言模型接收离散令牌作为输入，但大部分计算发生在连续的隐藏空间中。给定x1:Tx_{1:T}，一个自回归模型将p(x1:T)=∏t=1Tp(xt∣x<t)p(x_{1:T}) = \prod_{t=1}^T p(x_t \mid x_{<t})分解，其中每个条件分布由一组Transformer层参数化。这些层的输出（隐藏状态）并不直接对应令牌；它们通过LM头映射到词汇表上的logits。因此，只要后续层和LM头能将隐藏状态解码为令牌，我们就可以将其视为连续空间的元素。

扩散模型(Li et al., 2022 (https://arxiv.org/html/2605.14368#bib.bib8); Austin et al., 2021 (https://arxiv.org/html/2605.14368#bib.bib24))学习通过模拟将数据分布扰动为噪声的正向过程并训练一个神经网络来逆转该过程，从而生成数据。对于连续扩散，正向过程通常由随机微分方程(SDE)定义：
dXt=f(Xt,t)dt+g(t)dWt, X0∼pdata,
dX_t = f(X_t, t) dt + g(t) dW_t, \quad X_0 \sim p_{\text{data}},
其中WtW_t是布朗运动。反向过程由时间相关的分数函数∇logpt(Xt)\nabla \log p_t(X_t)给出，该函数通过去噪分数匹配进行估计：E[∥sθ(Xt,t)−∇logpt(Xt)∥2]\mathbb{E}[\|s_\theta(X_t, t) - \nabla \log p_t(X_t)\|^2]。生成时，从先验pT≈N(0,σ2I)p_T \approx \mathcal{N}(0, \sigma^2 I)采样，并模拟反向SDE。

先前的工作探索了不同的连续空间：令牌嵌入、自条件特征或潜在变量(Li et al., 2022; Lovelace et al., 2023; Gong et al., 2023; Zhang et al., 2025)（同上）。通常，这些方法将连续变量解码回令牌，无论是通过直接投影到词汇表（连续到离散）还是通过稍后映射到logits的隐藏状态。相反，我们考虑选择一个Transformer内部隐藏层，并用一个扩散桥替换其下方的所有层，该桥以该层的隐藏状态为目标。然后，剩余的上层和LM头将这些重构的隐藏状态解码为令牌。我们将此表述为接口选择问题：给定预训练模型中的L层，找出第ℓ∗\ell^*层，使得在第ℓ∗\ell^*层的隐藏状态空间上运行扩散，然后通过上层和LM头进行解码，能够产生最佳结果。我们通过几何代理来引导此选择，而不对所有层进行密集的桥训练。

## 3 扩散友好的隐藏状态  

我们通过几何度量来形式化扩散友好性：强对数凹性（用于快速收敛和稳定性）、低维流形结构（用于表示紧致性），以及它们之间的相互作用。这些抽象概念通过谱代理（局部协方差浓度、全局刚度、有效秩）进行近似，从而提供计算上可行的层评分。我们将第ℓ\ell层在位置tt处的隐藏状态定义为hℓ,t∈Rdℓh_{\ell, t} \in \mathbb{R}^{d_\ell}。为简洁起见，除非需要明确位置，否则我们将省略tt。

### 3.1 理论动机  

尽管Transformer隐藏状态并非来自真正的对数凹性(3)密度，但标准的扩散收敛论证提供了一个有用的类比。考虑一个分布μ∈P2(Rd)\mu \in \mathcal{P}_2(\mathbb{R}^d)，其中P2(Rd)\mathcal{P}_2(\mathbb{R}^d)是在ℝd\mathbb{R}^d上具有有限二阶矩的概率测度集合。设其密度p=e−U/Zp = e^{-U}/Z，其中U∈C2(Rd)U \in C^2(\mathbb{R}^d)是mm-强凸的，且∇U\nabla U是全局Lipschitz的。令XtX_t为正则化Langevin扩散过程，该过程将μ\mu作为其平稳分布，并以速率mm指数混合。

###### 定理1 (Langevin扩散的指数混合).  
令U∈C2(Rd)U \in C^2(\mathbb{R}^d)满足∇2U(x)⪰mI\nabla^2 U(x) \succeq mI对于所有x∈Rdx \in \mathbb{R}^d成立，其中m>0m > 0，且∇U\nabla U是全局Lipschitz的。令(Xt)t≥0(X_t)_{t \ge 0}满足过阻尼Langevin随机微分方程(SDE)  
dXt=−∇U(Xt)dt+2dWt,dX_t = -\nabla U(X_t) dt + \sqrt{2} dW_t,  
其中WtW_t表示布朗运动。那么μ\mu是(Xt)(X_t)的一个不变分布，并且对于每个初始分布ν0∈P2(Rd)\nu_0 \in \mathcal{P}_2(\mathbb{R}^d)，  
W2(νt,μ)≤e−mtW2(ν0,μ), νt:=L(Xt),W_2(\nu_t, \mu) \le e^{-mt} W_2(\nu_0, \mu), \quad \nu_t := \mathcal{L}(X_t),  
其中L(Xt)\mathcal{L}(X_t)表示XtX_t的分布。不变分布μ\mu在P2(Rd)\mathcal{P}_2(\mathbb{R}^d)中是唯一的。

定理1给出了第一个判据。如果曲率参数mm很大，那么到目标分布的距离按e−mte^{-mt}缩小。因此，mm越大意味着收缩越快，这对于扩散是可取的，因为去噪应该迅速将噪声样本返回到数据分布。

仅有快速收缩是不够的。在实践中，分数是未知的，需要由神经网络估计。定理2给出了第二个判据。如果分数误差至多为ε\varepsilon，那么由此引起的分布误差以ε/m\varepsilon/m为界。因此，mm越大对应于在分数估计不完美时越稳定。

###### 定理2 (分数扰动下不变测度的稳定性).  
令μ∈P2(Rd)\mu \in \mathcal{P}_2(\mathbb{R}^d)具有密度pp，并定义U(x):=−logp(x)U(x) := -\log p(x)。假设U∈C2(Rd)U \in C^2(\mathbb{R}^d)是mm-强凸的，且∇U\nabla U是全局Lipschitz的。令s(x):=∇logp(x)=−∇U(x)s(x) := \nabla \log p(x) = -\nabla U(x)，并令s^:Rd→Rd\hat{s}: \mathbb{R}^d \to \mathbb{R}^d是全局Lipschitz的，且满足supx∈Rd‖s^(x)−s(x)‖≤ε\sup_{x \in \mathbb{R}^d} \|\hat{s}(x) - s(x)\| \le \varepsilon。考虑两个SDE：  
dXt=s(Xt)dt+2dWtdX_t = s(X_t) dt + \sqrt{2} dW_t 和 dX^t=s^(X^t)dt+2dWtd\hat{X}_t = \hat{s}(\hat{X}_t) dt + \sqrt{2} dW_t。  
假设第二个SDE存在一个不变分布μ^\hat{\mu}。那么μ^∈P2(Rd)\hat{\mu} \in \mathcal{P}_2(\mathbb{R}^d)且  
W2(μ^,μ)≤εm.W_2(\hat{\mu}, \mu) \le \frac{\varepsilon}{m}.

定理1和定理2共同表明，曲率是收敛速度和分数估计误差下稳定性的有用代理。然而，仅凭曲率并不能捕捉表示是否易于建模：方差可能仍然分布在各方向上。因此，我们使用有效秩作为维度的代理。如果激活集中在低维流形附近，扩散只需要建模少数有意义的方向。此处，tr(Σ)\mathrm{tr}(\Sigma)是总方差，‖Σ‖\|\Sigma\|是最大协方差特征值，因此reff(Σ)=tr(Σ)/‖Σ‖r_{\mathrm{eff}}(\Sigma) = \mathrm{tr}(\Sigma) / \|\Sigma\|衡量活跃方差方向的有效数量。

###### 引理1 (近似流形支持意味着低有效秩).  
令XX是一个ℝd\mathbb{R}^d值的随机变量，协方差Σ:=Cov(X)\Sigma := \mathrm{Cov}(X)。假设存在一个kk维C2C^2流形M⊂Rd\mathcal{M} \subset \mathbb{R}^d和一个可测映射Π:Rd→M\Pi: \mathbb{R}^d \to \mathcal{M}，使得  
tr(Cov(Π(X)))≤C1k\mathrm{tr}(\mathrm{Cov}(\Pi(X))) \le C_1 k, E‖X−Π(X)‖2≤C2(δ2+η)\mathbb{E}\|X - \Pi(X)\|^2 \le C_2(\delta^2 + \eta), 且‖Σ‖≥c>0\|\Sigma\| \ge c > 0。  
则  
reff(Σ):=tr(Σ)‖Σ‖≤2C1ck+2C2c(δ2+η).r_{\mathrm{eff}}(\Sigma) := \frac{\mathrm{tr}(\Sigma)}{\|\Sigma\|} \le \frac{2C_1}{c} k + \frac{2C_2}{c} (\delta^2 + \eta).  
特别地，如果δ,η\delta, \eta是受控常数，且‖Σ‖\|\Sigma\|有常数上下界，那么reff(Σ)=O(k)r_{\mathrm{eff}}(\Sigma) = O(k)。

引理1证实用reff(Σ)r_{\mathrm{eff}}(\Sigma)作为操作性的内在维度代理是合理的：靠近一个kk维流形且离流形误差可控时，有效秩受kk控制，而非环境维度dd。

定理3将这种维度控制与定理1和2中的曲率条件结合起来，将低有效维度与表示集中联系起来，同时曲率控制围绕均值的波动。集中部分遵循标准的Bakry-Émery和Herbst论证(Bakry et al., 2014; Ledoux, 2001)（同上）。

###### 定理3 (内在维度与有效表示复杂度).  
令μ∈P2(Rd)\mu \in \mathcal{P}_2(\mathbb{R}^d)具有密度p(x)=Z−1e−U(x)p(x) = Z^{-1} e^{-U(x)}，其中U∈C2(Rd)U \in C^2(\mathbb{R}^d)，并令Σ:=Cov(μ)\Sigma := \mathrm{Cov}(\mu)。假设对于所有x∈Rdx \in \mathbb{R}^d，∇2U(x)⪰mI\nabla^2 U(x) \succeq mI成立，其中m>0m > 0，且∇U\nabla U是全局Lipschitz的。此外，假设存在一个kk维C2C^2流形M⊂Rd\mathcal{M} \subset \mathbb{R}^d和一个可测映射Π:Rd→M\Pi: \mathbb{R}^d \to \mathcal{M}。对于X∼μX \sim \mu，假设  
tr(Cov(Π(X)))≤C1k\mathrm{tr}(\mathrm{Cov}(\Pi(X))) \le C_1 k, E‖X−Π(X)‖2≤C2(δ2+η)\mathbb{E}\|X - \Pi(X)\|^2 \le C_2(\delta^2 + \eta), 且‖Σ‖≥c0>0\|\Sigma\| \ge c_0 > 0。  
则  
reff(Σ)≤2C1c0k+2C2c0(δ2+η),r_{\mathrm{eff}}(\Sigma) \le \frac{2C_1}{c_0} k + \frac{2C_2}{c_0}(\delta^2 + \eta),  
从而  
(E‖X−EX‖2)1/2=tr(Σ)≤‖Σ‖1/2(2C1c0k+2C2c0(δ2+η))1/2.\bigl(\mathbb{E}\|X - \mathbb{E}X\|^2\bigr)^{1/2} = \sqrt{\mathrm{tr}(\Sigma)} \le \|\Sigma\|^{1/2} \left(\frac{2C_1}{c_0} k + \frac{2C_2}{c_0}(\delta^2 + \eta)\right)^{1/2}.  
此外，存在一个绝对常数c>0c > 0，使得对于所有t≥0t \ge 0，  
P(∣‖X−EX‖−E‖X−EX‖∣≥t)≤2exp(−cmt2).\mathbb{P}\!\left(\left|\|X - \mathbb{E}X\| - \mathbb{E}\|X - \mathbb{E}X\|\right| \ge t\right) \le 2\exp(-cm t^2).

定理3将引理1与强对数凹性下的集中性质结合起来。它表明低维集中通过有效秩控制有效表示复杂度，而曲率参数mm通过测度集中控制围绕均值的波动。综合来看，这些结果并非旨在作为Transformer激活的保证，而是作为扩散友好表示应具备何种特征的理论动机。由于反向扩散使用噪声边际分布的时间相关分数，而过阻尼Langevin动力学使用固定的目标分布分数，我们仅将这些结果用于定性启发：类似收缩的行为、对分数估计误差的鲁棒性以及低有效复杂度。因此，一个好的层应表现出类似曲率的收缩以实现去噪稳定，以及低有效维度以便于建模。由于真实的密度、Hessian和流形结构不可获取，我们使用经验谱代理来近似这些想法：局部协方差浓度、基于全局精度的刚度，以及有效秩。所有证明见附录A (https://arxiv.org/html/2605.14368#A1)。

### 3.2 定位：寻找扩散友好的层  

这些理论结果作为替代动机，而非假设隐藏状态是全局强对数凹的。与其提供保证