【AI论文】语言建模的连续扩散模型-CSDN博客

本文链接：https://blog.csdn.net/m0_66899341/article/details/145744310

摘要：扩散模型已成为建模离散分类数据时，替代自回归模型的一种有前景的选择。然而，直接在离散数据空间上工作的扩散模型并未充分利用迭代细化的能力，因为信号在离散状态之间的转换过程中会丢失。与离散方法相比，现有针对离散数据的连续扩散模型性能有限，且两者之间关系不明确，这限制了离散数据扩散模型的发展。在本文中，我们提出了一种用于语言建模的连续扩散模型，该模型融入了底层分类分布的几何特性。我们在离散扩散和统计流形上的连续流之间建立了联系，并基于这种类比，引入了一种简单的扩散过程设计，该设计推广了以往的离散扩散模型。我们进一步提出了一种基于径向对称的无模拟训练框架，以及一种解决流形高维性的简单技术。在语言建模基准测试和其他模态上的综合实验表明，我们的方法优于现有的离散扩散模型，且性能接近自回归模型。代码可在https://github.com/harryjo97/RDLM获取。Huggingface链接：Paper page，论文链接：2502.11564

1. 引言

扩散模型概述：扩散模型作为自回归模型的替代方案，在建模离散分类数据方面展现出巨大潜力。然而，直接在离散数据空间上工作的扩散模型未能充分利用迭代细化的优势，因为在离散状态转换过程中信号会丢失。
现有问题：与离散方法相比，现有的针对离散数据的连续扩散模型性能有限。此外，离散扩散与连续扩散模型之间的联系尚不明确，这限制了离散数据扩散模型的发展。
研究动机：本文旨在提出一种用于语言建模的连续扩散模型，该模型能够融入底层分类分布的几何特性，以解决上述问题。

2. 背景

2.1 离散扩散模型

定义与过程：离散扩散模型使用马尔可夫链直接在离散结构上定义扩散过程。前向过程描述从当前状态到其他状态的转移，通过乘以转移矩阵Q_t来实现。
具体形式：例如，Austin等人（2021）提出了几种转移矩阵的设计，包括掩码（吸收状态）扩散和均匀扩散。连续时间马尔可夫链（CTMC）进一步将框架扩展到连续时间。

2.2 统计流形与分类分布

统计流形：离散数据空间X={1,...,d}和(d-1)维概率单纯形∆_(d-1)之间的对应关系，构成了分类分布的统计流形P(X)。
几何特性：统计流形配备了Fisher-Rao度量，并且存在一个从P(X)到(d-1)维球面的正象限S_(d-1)^+的微分同胚π。

2.3 黎曼扩散混合框架

定义：黎曼扩散混合框架提供了一种在一般流形上进行生成建模的简单方法。它通过桥接过程Q_z在流形M上定义扩散过程，该过程将初始分布传输到数据分布。
过程建模：扩散过程被建模为桥接过程的混合，其漂移由神经网络η_θ通过桥接匹配目标进行回归。

3. 黎曼扩散语言模型

3.1 离散扩散的连续重参数化

连续重参数化：为了融入底层分类分布的几何特性，我们利用统计流形对离散数据进行参数化。每个流形上的点对应于离散样本空间X上的分类分布的参数。
重参数化问题：Fisher-Rao度量在流形的边界上定义不明确，导致数值问题。因此，我们使用微分同胚π将流形映射到球面的正象限，从而保留了分类分布的几何特性。

3.2 超球面上的生成过程

生成过程重构：通过重参数化，建模离散数据的分布可以重构为在超球面上建模分布。我们基于黎曼扩散混合框架构建了一个在超球面上的生成过程。
桥接过程：从任意点u∈S_(d-1)^+到第k个令牌e_k的桥接过程可以推导出来，并用于构建超球面上的扩散过程。

3.3 收敛性分析

径向过程：我们分析了桥接过程的收敛性，通过其径向过程r_k^t=d_g(X_t,e_k)来描述。对于σ_0>σ_T，径向过程在早期时间步长内迅速收敛，使得神经网络难以准确近似。
几何调度：我们经验性地发现，使用几何调度σ_t=σ_0(1-t/T)^σ可以实现逐渐收敛。

3.4 掩码扩散与均匀扩散

掩码扩散：将初始分布固定为掩码令牌e_m，得到一个混合过程，该过程概括了掩码离散扩散过程。
均匀扩散：将起点设置为单纯形的质心，并扩展到允许向令牌子集S的过渡具有不同概率ζ的均匀扩散。

3.5 混合路径

混合路径定义：我们通过混合具有相同噪声调度σ_t的生成过程的概率路径，推导出一个新的生成过程家族。
具体实例：例如，创建一个由掩码桥和均匀桥混合的路径，其初始分布为λ_0δ(e_m)+(1-λ_0)δ(dd e_i)。

3.6 推广流匹配

流匹配推广：我们的框架推广了在统计流形上的流匹配方法。通过设计噪声调度σ_t:=σ_0→0，我们获得流匹配模型的条件向量场。

4. 基于径向对称的无模拟训练

4.1 似然界

ELBO推导：我们使用紧流形上的Girsanov定理推导出我们的模型的似然界。对于流形S_d上的点z，我们可以将模型负对数似然的上界表示为KL散度。

4.2 参数化与目标

漂移参数化：由于直接近似漂移函数具有挑战性，我们提出使用神经网络s_θ对概率p_T|t(X_T|X_t)进行建模，然后通过该概率表示漂移。
训练目标：基于ELBO，我们推导出最大似然训练目标，该目标可以解释为最小化漂移近似的均方误差。此外，我们还引入了基于交叉熵的训练目标，该目标在训练中收敛更快，性能更好。

4.3 重要性采样

时间分布调整：通过调整时间分布以专注于特定区间，我们改进了训练目标。我们使用一个简单的密度q(t)=1-ϵ（如果t∈[a,b]，否则为ϵ）来实现这一点。

4.4 过渡分布近似

Riemannian Normal：由于超球面上的扩散过程不产生可处理的过渡分布，我们需要模拟该过程，这在训练中成为瓶颈。因此，我们引入了一种近似采样方法，该方法基于切空间上的高斯分布的指数映射的推进测度（即Riemannian Normal），实现了无模拟训练。

4.5 连接到投影过程

参数推导：虽然Riemannian Normal的参数α_t和ρ_t通常难以处理，但我们通过一维投影过程c_w^t:=⟨X_t,w⟩（对于w=X_0和X_1）来推导它们。

5. 令牌序列的生成

5.1 序列建模

产品流形：由于序列中的每个令牌都被重参数化为d维球面，因此长度为n的序列被建模在产品流形(S_d)^n上。
依赖关系：每个超球面上的扩散过程相互依赖，由一组SDE描述。

5.2 统计流形的维度拆分

维度拆分定义：对于大型词汇集，相应的统计流形具有高维度，导致两个挑战：突然收敛和大型输入维度。我们引入维度拆分技术来解决这些问题。
技术细节：将第k个令牌映射到产品流形(S_˜_b)^m，其中˜_b=b（对于掩码扩散）或b-1（否则），m=⌈log_b d⌉。

6. 相关工作

6.1 离散扩散模型

模型概述：离散扩散模型直接在离散数据空间上定义马尔可夫链。D3PM（Austin等人，2021）引入了离散时间马尔可夫前向过程，SEDD（Lou等人，2024）提出了学习离散状态的得分熵。

6.2 针对离散数据的连续扩散模型

早期尝试：早期工作通过完全将离散数据放松到连续空间（Han等人，2022）或将令牌嵌入潜在空间（Li等人，2022；Dieleman等人，2022）来尝试。然而，这些方法未能准确建模分类分布的离散性。
流匹配方法：最近的工作利用统计流形上的流匹配方法（Cheng等人，2024；Davis等人，2024），但这些方法的性能落后于离散扩散模型。

7. 实验

7.1 文本生成

数据集：我们在Text8和One Billion Words数据集上评估了我们的黎曼扩散语言模型（RDLM）。
基线模型：我们与最先进的自回归和扩散模型进行了比较，包括Multinomial Diffusion、D3PM、SEDD、MDLM、MD4、Plaid、BFN等。
结果：RDLM在Text8数据集上取得了优于所有先前扩散模型的结果，并在One Billion Words数据集上显著优于现有连续扩散模型。

7.2 像素级图像建模

数据集：我们将模型应用于顺序无关的图像数据，其中每个图像被表示为具有256个词汇的离散令牌集。
结果：RDLM在CIFAR-10数据集上实现了最低的每维度比特数（BPD），优于离散扩散模型和自回归模型。

7.3 DNA序列设计

任务：我们在启动子DNA序列设计任务上评估了我们的方法，该任务旨在根据转录谱生成有效的启动子DNA序列。
结果：RDLM实现了最低的均方误差（MSE），优于流匹配方法和离散扩散模型。

7.4 分析

过渡分布近似：我们测量了模拟过渡分布和近似分布之间的最大平均差异（MMD）距离，结果表明近似是可靠的。
训练目标：我们验证了基于交叉熵的损失函数的有效性，与均方误差损失相比，它提供了更快的收敛速度和更好的负对数似然。
维度拆分：对于具有大型词汇集的数据集，拆分维度的流形显著提高了模型性能。

8. 结论

主要贡献：我们提出了黎曼扩散语言模型（RDLM），一种用于语言和离散数据的连续扩散模型。我们建立了离散扩散和统计流形上连续流之间的联系，并引入了一种简单的扩散过程设计。
实验验证：在语言建模基准测试和其他模态上的综合实验表明，RDLM优于现有的离散扩散模型，并且性能接近自回归模型。
未来方向：我们计划扩展RDLM以进行可控文本生成，并探索其在图生成和分子合成等领域的应用。