扩散模型应该在语言模型的哪个位置介入？几何引导的隐藏状态替换方法

DiHAL用几何特征定位扩散友好层，通过重建隐藏状态提升扩散语言模型性能

前置知识

扩散模型

扩散模型是一类生成模型，通过逐步向数据添加噪声，然后学习逆转这个加噪过程来生成新样本。在连续空间中，这个过程可以用随机微分方程描述，其中关键是要学习分数函数，即对数概率密度的梯度。扩散模型在图像生成领域取得了巨大成功，但在文本生成中面临独特挑战，因为文本本质上是离散的

本文的核心是将扩散模型应用于语言模型，理解扩散模型的工作原理（特别是连续扩散）是理解本文方法的基础

Langevin动力学

Langevin动力学是描述粒子在势能场中运动的随机微分方程，形式为 $dX_t = -\nabla U(X_t) dt + \sqrt{2} dW_t$，其中 $U(x) = -\log p(x)$ 是势能函数，$W_t$ 是布朗运动。这个过程的平稳分布就是目标分布 $p(x)$。当 $U$ 是强凸函数时，解会以指数速度收敛到目标分布，这为扩散模型提供了理论保证

论文利用Langevin动力学的理论结果（特别是收敛性和稳定性）来定义什么是扩散友好的表示空间，这是本文几何代理设计的理论基础

有效秩

有效秩 $r_{\text{eff}}(\Sigma) = \text{tr}(\Sigma) / \|\Sigma\|$ 是一个衡量数据内在维度的指标，其中 $\text{tr}(\Sigma)$ 是协方差矩阵的迹，$\|\Sigma\|$ 是最大特征值。它表示有多少个有效的主成分方向承载了大部分方差。如果数据集中在低维流形上，有效秩会远小于原始维度

论文使用有效秩作为衡量表示空间复杂度的代理，扩散友好的空间应该有较低的有效秩，意味着变异集中在少数几个方向上，更容易建模

强对数凹性

一个概率分布的密度函数 $p(x)$ 如果其对数的负函数 $U(x) = -\log p(x)$ 是强凸的，即存在 $m > 0$ 使得 $\nabla^2 U(x) \succeq mI$ 对所有 $x$ 成立，则称该分布是 $m$-强对数凹的。强对数凹性保证了分布有良好的集中性质，Langevin动力学会以 $e^{-mt}$ 的速度收敛到目标分布

论文将强对数凹性作为扩散友好空间的理想性质之一，虽然Transformer隐藏状态不一定全局满足这个性质，但可以用曲率相关的代理来近似

研究动机

连续扩散语言模型在性能上明显落后于自回归Transformer模型，特别是在连续扩散设置中（Jo and Hwang, 2026）。这个性能差距不能简单地归因于文本的离散性，因为Transformer语言模型也使用离散Token，但大部分计算发生在连续隐藏状态空间中。现有的扩散语言模型方法（如Diffusion-LM、SED、LD4LG、CoDAR）都需要将去噪后的连续向量映射回离散Token，这会引入恢复误差：表示空间中的小错误可能导致恢复出完全不同的Token（Li et al., 2022; Zhang et al., 2025）。例如，在Qwen3-8B上，Diffusion-LM的Gen.PPL高达683.43，SED达到778.82，而最新的CoDAR方法也有144.83的Gen.PPL，远高于原始自回归模型。

本文的目标是本文提出了一个根本性的问题：扩散应该在预训练语言模型的哪个位置介入？具体来说，是否存在某个内部隐藏状态空间是扩散友好的，即在这个空间中应用扩散模型能够更有效地进行去噪和重建？我们的目标是通过几何特征来识别这样的层，然后用一个条件扩散桥替换下层的Transformer，同时保留上层和原始LM head，从而避免直接的连续到离散Token恢复问题。

与已有工作不同的是，与现有方法不同，本文不将扩散直接应用于Token嵌入或学习到的潜变量，而是研究Transformer内部的隐藏状态空间。现有工作要么将扩散应用于Token嵌入（Diffusion-LM, SED），要么需要学习专门的潜变量（LD4LG），要么尝试从连续状态恢复离散Token（CoDAR）。本文的独特切入点是：Transformer已经包含多个连续隐藏空间，我们只需要找到其中几何性质适合扩散的那一个，然后让扩散在这个空间中工作，Token预测仍然由保留的上层和LM head完成。这本质上是将连续到离散的接口选择问题，而不是一个无法避免的瓶颈。

核心方法

DiHAL（Diffusion-Transformer Hybrid Architecture for Language Generation）采用Locate-and-Replace两步策略：首先，基于几何特征定位Transformer中扩散友好的层；其次，用一个条件扩散桥替换下层的Transformer，保留上层和原始LM head。直觉是：如果在某个隐藏层，表示空间具有良好的几何性质（如强曲率、低有效维度），那么在这个层应用扩散应该更容易去噪和重建。技术路线上，我们使用Langevin动力学和集中理论的结果定义理想的扩散友好性质，然后用可计算的几何代理来近似这些性质，最后训练一个UNet风格的扩散桥来重建选定层的隐藏状态。

DiHAL的核心创新是将连续扩散从Token级别的恢复问题转变为内部隐藏状态的重建问题。这意味着我们不需要将扩散输出直接解码为Token，而是重建一个Transformer已经知道如何解码的内部表示。这个思想带来了三个关键优势：第一，避免了连续到离散Token的直接映射，减少了恢复误差；第二，扩散只需要在局部优化的隐藏状态空间中工作，而不是从零开始学习整个语言模型；第三，可以充分利用预训练Transformer的上层和LM head，这些部分已经学习到了良好的语言表示和解码能力。

方法步骤详情

DiHAL方法的完整步骤如下：第一步，收集数据。对于每个8B规模的骨干模型，运行模型在300K个Dolma v1.7序列上，保存每个层的隐藏状态。第二步，计算几何代理。对于每个层的激活矩阵 $x \in \mathbb{R}^{M \times D}$（M个Token，D个隐藏维度），计算三个统计量：局部曲率代理 $\hat{m}_{\text{curv}}$ 是通过k-最近邻邻域的协方差计算，$\hat{m}^{(i)}_{\text{curv}} = 1/\lambda_{\max}(\Sigma^{(i)}_{\text{local}})$，取所有Token的中位数；单调性代理 $\hat{m}_{\text{mono}}$ 衡量全局方向刚性，使用规则的精度矩阵 $P = (\Sigma + \lambda I)^{-1}$，计算采样对的 $m_{ij} = (x_i - x_j)^\top P (x_i - x_j) / \|x_i - x_j\|^2$ 并取中位数；有效秩 $\hat{k} = r_{\text{eff}}(\Sigma) = \text{tr}(\Sigma) / \|\Sigma\|$ 衡量内在维度。第三步，计算选择分数并进行层选择。选择分数为 $\text{selection\_score}(\ell) = z(\log \hat{m}_{\text{curv}}(\ell)) + z(\log \hat{m}_{\text{mono}}(\ell)) - z(\log \hat{k}(\ell))$，其中 $z(\cdot)$ 表示层级别的z-score标准化。选择分数最高的层 $\ell^* = \arg\max_\ell \text{selection\_score}(\ell)$ 作为扩散插入点。第四步，训练扩散桥。对于选定的层 $\ell^*$，用UNet风格的扩散桥 $D_\theta$ 替换下层Transformer $F_{1:\ell^*}$，保留上层 $F_{\ell^*+1:L}$ 和LM head。扩散桥以嵌入输出作为条件 $c(x)$，目标是重建选定层的隐藏状态 $\hat{h}_{\ell^*} = D_\theta(c(x))$。第五步，训练和推理。在训练时，优化扩散损失 $L_{\text{diff}} = \mathbb{E}_{t,\epsilon}[\|\hat{\epsilon}_\theta(z_t, t, c) - \epsilon\|_2^2]$、重建损失 $L_{\text{rec}} = \|\hat{h}_{\ell^*} - h_{\ell^*}\|_2^2$、下一个Token损失 $L_{\text{LM}}$ 和知识蒸馏损失 $L_{\text{KD}}$ 的加权和。在推理时，跳过下层，直接用扩散桥从条件生成 $\hat{h}_{\ell^*}$，然后通过上层和LM head生成Token。

技术新颖性

DiHAL的技术新颖性体现在三个方面。第一，理论驱动的几何选择。虽然Transformer有很多层，但本文不是暴力搜索，而是基于Langevin动力学理论（Theorem 1-3）定义了扩散友好空间应该具备的性质：强曲率保证快速收敛和分数估计误差下的稳定性，低有效秩保证表示复杂度低。然后提出了三个可计算的几何代理来近似这些理想性质。第二，隐藏状态级别的扩散。与现有的Token嵌入扩散或潜变量扩散不同，DiHAL将扩散应用于Transformer内部的隐藏状态，这是首次系统性地研究扩散在预训练语言模型中的插入位置问题。第三，验证固定分数的有效性。实验表明，固定几何分数（不需要针对特定任务调整）就能与验证桥接损失有强相关性（Spearman $\rho = 0.9143 \pm 0.0069$ on Llama-3.1-8B, $\rho = 0.9267 \pm 0.0157$ on Qwen3-8B），这意味着可以用低成本的计算来预测哪些层适合扩散替换，而不需要为每个层训练完整的桥接。

**Figure 1**: Locate-and-Replace framework. Layer-wise geometric proxies score transformer layers, select an insertion point, and guide replacement with a diffusion block.

实验结果

实验在两个8B规模的解码器骨干模型上进行：Llama-3.1-8B-Instruct（32层，隐藏维度4096）和Qwen3-8B（36层，相同隐藏维度）。首先，层级别几何分析发现Transformer隐藏表示的几何性质随深度系统性变化，输入邻接层的局部曲率较大，全局单调性和有效秩则呈现不同的深度依赖趋势。这表明Transformer隐藏状态不是同等适合扩散的。对于Llama-3.1-8B，几何分数选择了第3层，分数为2.360，$\hat{m}_{\text{curv}} = 553.166$，$\hat{m}_{\text{mono}} = 227.324$，$\hat{k} = 1.332$，验证桥接损失为0.331；对于Qwen3-8B，选择了第2层，分数为2.044，$\hat{m}_{\text{curv}} = 168.963$，$\hat{m}_{\text{mono}} = 317.682$，$\hat{k} = 4.544$，验证桥接损失为0.060。固定预算层扫描显示几何分数与验证桥接损失有强相关性，在30次重复的500样本分数估计上，Llama-3.1-8B的Spearman $\rho = 0.9143 \pm 0.0069$，Kendall $\tau = 0.8011 \pm 0.0086$，Pearson $r = 0.8687 \pm 0.0027$；Qwen3-8B的Spearman $\rho = 0.9267 \pm 0.0157$，Kendall $\tau = 0.8208 \pm 0.0256$，Pearson $r = 0.8605 \pm 0.0098$。最佳预测层和最佳观测层的秩差距仅为2和1。在与连续扩散方法的诊断性匹配预算比较中，DiHAL在Qwen3-8B上达到了最低的Gen.PPL（136.02）和最高的多样性（0.5913），相比CoDAR的Gen.PPL 144.83和多样性0.4777，DiHAL将Gen.PPL降低了6.1%，多样性提高了23.8%。完全训练后的最终评估显示，DiHAL几何选择层在Llama-3.1-8B上的NLL为4.91，PPL为135.64，KL为0.73；在Qwen3-8B上的NLL为4.97，PPL为144.03，KL为0.53。相比CoDAR的NLL 5.18和PPL 177.87，DiHAL在PPL上降低了约19%。几何选择层与验证损失预测的oracle层相当，甚至在Llama-3.1-8B上超越了oracle（NLL 5.11 vs 4.91，PPL 165.67 vs 135.64），这证明了固定几何分数的有效性。

**Table 1**: Fixed-budget layer sweep results. We compare the geometry-selected layer against representative baselines under the same bridge-training budget. Validation bridge loss measures bridgeability.

**Table 2**: Agreement between fixed geometry score and bridgeability under the fixed-budget sweep. Correlations use negative validation bridge loss; parentheses report standard deviations over 30 repeated 500-example score runs. Geometry proxies use 100 repeated 3K-example subsamples.

**Table 3**: Diagnostic same-budget comparison of continuous diffusion target spaces and recovery interfaces. This is not a pretraining-matched comparison: DiHAL reuses a pretrained transformer suffix and LM head, while some baselines use smaller standalone recovery modules.

**Table 4**: Final evaluation of DiHAL insertion layers against CoDAR on the combined WikiText-103 and held-out Dolma evaluation sets. CoDAR is a recent strong continuous-diffusion baseline re-evaluated in the Qwen3-8B setting under the same evaluation pipeline.

**Figure 2**: Layer-wise geometry of hidden representations for Llama-3.1-8B-Instruct (left) and Qwen3-8B (right). Curvature-related and effective-rank statistics vary substantially across depth. The selected insertion layer is determined by balancing local compactness, global stiffness, and effective representation complexity, rather than by maximizing a single proxy.

查看结构化数据

任务	指标	本文	基线	提升
文本生成质量	Gen.PPL（GPT-2评估器困惑度）	136.02	CoDAR	从144.83降低到136.02，提升6.1%
文本生成多样性	多样性（Distinct-1到Distinct-4的乘积）	0.5913	CoDAR	从0.4777提升到0.5913，提升23.8%
语言建模困惑度	PPL（在WikiText-103和Dolma上）	135.64（Llama-3.1-8B），144.03（Qwen3-8B）	CoDAR	从177.87降低，分别降低约24%和19%
层选择预测能力	Spearman相关系数	0.9143±0.0069（Llama-3.1-8B），0.9267±0.0157（Qwen3-8B）	N/A	N/A（几何分数与验证桥接损失有强相关性）

局限与改进

作者承认DiHAL有几个重要局限性。首先，DiHAL不是一个独立的扩散语言模型，Token预测仍然依赖于保留的Transformer后缀和LM head，这意味着它没有完全脱离自回归框架。其次，由于计算限制，本文没有完全探索更大的桥接结构、更长的训练或更深的替换，这可能限制了性能的上限。第三，几何代理只是理论性质的近似，不保证Transformer隐藏状态全局满足强对数凹性等理想假设，本文只是在预训练模型上观察到了相关性。从我的观察来看，还有其他局限性：首先，DiHAL选择的插入层都很浅（第2-3层），这意味着替换的下层Transformer较少，对原始模型结构的影响有限，可能无法充分利用扩散的优势。其次，几何代理的计算需要保存所有层的隐藏状态，这在大规模模型上可能带来显著的存储开销。第三，论文只在8B规模的模型上进行了验证，在更大规模模型（如70B、100B）上的有效性还需要验证。第四，评估主要使用困惑度和多样性指标，缺乏对下游任务（如问答、摘要、翻译）的评估，这些任务更能反映实际生成质量。

独立分析的弱点

DiHAL的几个潜在弱点可以改进。第一，浅层插入限制了替换的范围。当前选择的第2-3层只替换了很少的Transformer层，这意味着大部分计算仍然保留。改进方向是研究如何使深层隐藏状态更加可桥接，例如将几何代理直接整合到桥接训练中，或者使用自适应的桥接架构来处理更复杂的表示。第二，几何代理的计算成本。当前方法需要保存所有层的隐藏状态并计算k-最近邻等统计量，在大规模模型上可能不可行。改进方向是开发在线计算的几何估计方法，或者使用更稀疏的采样策略。第三，缺乏下游任务评估。当前评估主要在语言建模指标上，但这不能完全反映实际应用的性能。改进方向是在多个下游任务（如摘要、问答、代码生成）上进行全面评估，以验证方法的实际价值。第四，与SOTA方法的比较不够充分。虽然论文比较了几个基线，但缺乏与最新的自回归方法或其他混合方法的深入比较。改进方向是包括更多基线，特别是在相同训练预算下的公平比较。

未来方向

作者提出的未来工作方向包括将几何代理直接整合到桥接训练中，使深层隐藏状态更加可桥接，从而可能替换更大的Transformer前缀。基于本文成果，可以延伸几个研究方向：第一，自适应几何学习。当前使用固定的几何代理，未来可以学习任务相关的几何特征，或者动态调整插入层的位置。第二，多头扩散桥。当前使用单个扩散桥，未来可以考虑在不同层使用多个桥接，或者使用层次化的扩散结构。第三，跨模型泛化。当前方法在两个8B模型上验证，未来可以研究几何分数是否可以跨模型迁移，即在一个模型上计算的分数是否可以指导另一个模型的选择。第四，与其他生成范式的结合。DiHAL将扩散与自回归结合，未来可以探索与生成对抗网络、变分自编码器等其他生成范式的混合。第五，理论扩展。当前理论主要基于强对数凹性等理想假设，未来可以发展更适合Transformer隐藏状态的理论框架，或者证明几何代理在更宽松条件下的有效性。

复现评估

论文的复现性评估：作者提供了详细的实现细节在附录中，包括超参数、网络架构和训练设置。但是，论文没有明确说明代码和数据的开源情况。从描述来看，实验使用了Dolma v1.7数据集，这是一个公开可用的数据集，但具体使用的300K序列子集可能没有公开发布。计算方面，论文提到了40 H100小时的预算用于扩散/恢复训练，这表明需要相当可观的GPU资源。复现难度主要在于：第一，需要完整的8B规模模型（Llama-3.1-8B-Instruct和Qwen3-8B）及其隐藏状态提取，这需要大量内存和存储；第二，几何代理的计算涉及k-最近邻等操作，在大规模数据上可能需要优化；第三，桥接训练需要仔细的调参和监控。总体而言，如果有足够的计算资源和明确的实现细节，中等研究团队应该能够复现核心结果，但完整复现所有实验可能需要相当大的投入。

扩散模型应该在语言模型的哪个位置介入？几何引导的隐藏状态替换方法 Where Should Diffusion Enter a Language Model? Geometry-Guided Hidden-State Replacement