基于超球面流的语言建模

# 基于超球面流的语言建模
来源: https://arxiv.org/html/2605.11125 Justin Deschenaux EPFL 瑞士洛桑 [email protected] &Caglar Gulcehre EPFL, 瑞士洛桑 微软 AI

###### 摘要
离散扩散语言模型作为自回归（AR）模型的替代品取得了 rapid 进展，这得益于其并行生成能力。然而，为了计算可行性，离散扩散模型从分解分布中采样，其表达能力不如 AR 模型。最近的流语言模型（FLMs）将连续流应用于语言，通过确定性常微分方程（ODE）将噪声传输到数据，从而避免了分解采样。FLMs 操作于维度随词汇表大小缩放的一热向量（one-hot vectors），使得 FLMs 的训练成本高昂。此外，由于所有不同的一热嵌入在 $L_2$ 空间中距离相等，添加高斯噪声没有明确的语义解释（这与图像不同，在图像中高斯噪声会逐渐破坏结构）。我们引入 $\mathbb{S}$-FLM，这是一种位于超球面上的潜在流语言模型。$\mathbb{S}$-FLM 通过在学习到的速度场（使用交叉熵学习）上旋转 $\mathbb{S}^{d-1}$ 中的向量来生成序列，避免了实例化一热向量的开销。之前的 FLMs 在生成困惑度（Gen. PPL）上匹敌 AR 模型，但在数学和代码等可验证领域中，高似然样本并不一定正确。$\mathbb{S}$-FLM 大幅改进了大词汇表推理任务上的连续流语言模型表现，并在标准温度采样（$T=1$）下缩小了与掩码扩散模型的差距，而在优化的低温（$T=0.1$）解码下仍存在差距。

参见图注
参见图注
图 1: $T=1$ 时 GSM8K 上的准确率。左图：带有 $\mathbb{S}$-arch（第3.3节 https://arxiv.org/html/2605.11125#S3.SS3）的 $\mathbb{S}$-FLM 解码策略。精确速度（公式 15 https://arxiv.org/html/2605.11125#S3.E15）和随机解码（算法 3 https://arxiv.org/html/2605.11125#alg3, *Stoch.*）在 ~12% 附近趋于平稳。将速度限制为 $p_{1|t}^\theta$ 的 top-$k$ 条目可以提高准确率，其中 top-1 达到 ~18%。右图：$\mathbb{S}$-FLM（带 $\mathbb{S}$-arch）与 MDLM 和 Duo 的比较。使用精确速度时，$\mathbb{S}$-FLM 在 NFE $\leq 16$ 时优于这两个基线。

参见图注
图 2: $\mathbb{S}$-FLM 概览。训练（上图）：我们将每个标记嵌入为 $\mathbb{S}^{d-1}$ 上的单位范数向量。我们通过对干净嵌入和 $\mathbb{S}^{d-1}$ 上的随机向量进行 SLERP 获得噪声潜在变量 $z_t^\ell$。我们使用交叉熵训练去噪器 $p_{1|t}^\theta$。采样（下图）：$p_{1|t}^\theta$ 通过对指向每个干净嵌入 $e_v, v \in \mathcal{V}$ 的切向量进行边缘化来定义速度场。从 $\mathbb{S}^{d-1}$ 上的均匀噪声开始，我们沿速度场积分，并通过 $\arg\max_{v \in \mathcal{V}} p_{1|1}^\theta(v | z_1)$ 解码最终潜在变量。

## 1 引言
自回归（AR）模型目前主导着语言建模。得益于链式规则分解和 Transformer 架构 (vaswani2017attention, https://arxiv.org/html/2605.11125#bib.bib89)，AR 似然可以快速评估，且 AR 语言模型可以扩展到大规模 (kaplan2020scalinglawsneurallanguage, https://arxiv.org/html/2605.11125#bib.bib39; openai2024gpt4technicalreport, https://arxiv.org/html/2605.11125#bib.bib63; openai2024gptoss, https://arxiv.org/html/2605.11125#bib.bib64; grattafiori2024llama3herdmodels, https://arxiv.org/html/2605.11125#bib.bib57; geminiteam2025gemini, https://arxiv.org/html/2605.11125#bib.bib29; gemmateam2025gemma3technicalreport, https://arxiv.org/html/2605.11125#bib.bib31)。然而，在采样期间，AR 模型每个标记需要一个前向传递，且因果注意力可能会损害需要双向上下文的推理任务 (papadopoulos2024arrowstimelargelanguage, https://arxiv.org/html/2605.11125#bib.bib65; kitouni2024factorizationcursetokenspredict, https://arxiv.org/html/2605.11125#bib.bib44; zhangli2024reversenumberdecodingorder, https://arxiv.org/html/2605.11125#bib.bib98; nagarajan2025roll, https://arxiv.org/html/2605.11125#bib.bib60)。

离散扩散模型 (austin2023structureddenoisingdiffusionmodels, https://arxiv.org/html/2605.11125#bib.bib5; campbell2022continuoustimeframeworkdiscrete, https://arxiv.org/html/2605.11125#bib.bib9; saho2024simpleeffectivemaskeddiffusion, https://arxiv.org/html/2605.11125#bib.bib76; gat2024discreteflowmatching, https://arxiv.org/html/2605.11125#bib.bib28; saho2025diffusionduality, https://arxiv.org/html/2605.11125#bib.bib77; shi2025simplifiedgeneralizedmaskeddiffusion, https://arxiv.org/html/2605.11125#bib.bib82; nie2025scalingmaskeddiffusionmodels, https://arxiv.org/html/2605.11125#bib.bib61; vonruette2026scalingbehaviordiscretediffusion, https://arxiv.org/html/2605.11125#bib.bib93; saho2026scalingmaskeddiffusionlanguage, https://arxiv.org/html/2605.11125#bib.bib78; wu2025fastdllmtrainingfreeaccelerationdiffusion, https://arxiv.org/html/2605.11125#bib.bib95) 在生成困惑度（Gen. PPL）上接近 AR 模型，并具有并行生成和双向上下文的能力。然而，在每一步去噪过程中，标记是从分解的边缘分布中*采样*的，而不是联合采样。这种分解使得离散扩散在并行生成标记时比 AR 模型的表达能力弱。

使用流匹配（Flow Matching）训练的连续流 (lipman2023flowmatchinggenerativemodeling, https://arxiv.org/html/2605.11125#bib.bib47; liu2022flowstraightfastlearning, https://arxiv.org/html/2605.11125#bib.bib49; albergo2023buildingnormalizingflowsstochastic, https://arxiv.org/html/2605.11125#bib.bib1) 学习一个定义*常微分方程*（ODE）的速度场，将噪声样本传输到数据分布。因此，推理步骤联合更新所有位置，并避免了离散扩散的分解采样问题。最近的工作 (roos2026categoricalflowmaps, https://arxiv.org/html/2605.11125#bib.bib75; lee2026onesteplanguagemodelingcontinuous, https://arxiv.org/html/2605.11125#bib.bib45; potaptchik2026discreteflowmaps, https://arxiv.org/html/2605.11125#bib.bib68) 重新激发对基于流的*流语言模型*（FLMs）(li2022diffusionlmimprovescontrollabletext, https://arxiv.org/html/2605.11125#bib.bib46; dieleman2022continuousdiffusioncategoricaldata, https://arxiv.org/html/2605.11125#bib.bib22; gulrajani2023likelihoodbaseddiffusionlanguagemodels, https://arxiv.org/html/2605.11125#bib.bib33) 的兴趣。最近的 FLMs 将标记表示为一热向量，添加高斯噪声，并使用交叉熵（CE）训练去噪器。尽管它们匹敌了 AR 和离散扩散模型的 Gen. PPL，但这些 FLMs 有两个主要缺点。

(1) 首先，将标记表示为一热向量成本高昂。*大型语言模型*（LLMs）通常使用包含 100k-200k 标记的词汇表 (openai2024gptoss, https://arxiv.org/html/2605.11125#bib.bib64; qwen2025qwen25technicalreport, https://arxiv.org/html/2605.11125#bib.bib70)，因此大型 FLMs 需要为每个标记存储一个 >100k 维的向量。在添加高斯噪声后，去噪器将这些向量与嵌入矩阵相乘，而不是查找单个向量。因此，FLMs 的训练速度慢于离散扩散和 AR 模型。

(2) 其次，在图像中，高斯扩散首先平滑地退化高频成分。对于一热向量，添加高斯噪声的解释并不清晰。

#### 贡献
我们提出了*超球面流语言模型*（$\mathbb{S}$-FLM），这是一种在嵌入上的流，无需实例化一热向量。

(1) 回想一下，余弦距离比欧几里得距离更好地捕捉标记嵌入之间的相似度 (mikolov2013efficientestimationwordrepresentations, https://arxiv.org/html/2605.11125#bib.bib58; pennington2014glove, https://arxiv.org/html/2605.11125#bib.bib67; wang2020understandingcontrastiverepresentationlearning, https://arxiv.org/html/2605.11125#bib.bib94)。由于余弦距离由单位超球面上的弧长决定，我们将 $\mathbb{S}$-FLM 实现为 $\mathbb{S}^{d-1}$ 上的黎曼流。我们的前向过程将单位范数嵌入传输到均匀先验。

(2) $\mathbb{S}$-FLM 操作于 $d$ 维嵌入，而不是 $|\mathcal{V}|$ 维的一热向量。因此，假设相同的骨干网络，$\mathbb{S}$-FLM 具有与离散扩散模型相似的训练成本（不像基于一热向量的 FLMs，后者成本更高）。我们进一步引入了 $\mathbb{S}$-arch，其激活位于 $\mathbb{S}^{d-1}$ 上。将激活与输入对齐提高了 GSM8K 和 OpenWebText (OWT) (Gokaslan2019OpenWeb, https://arxiv.org/html/2605.11125#bib.bib30) 上的样本质量。

(3) 在 GSM8K 上，基于 TinyGSM (liu2023tinygsm, https://arxiv.org/html/2605.11125#bib.bib48) 训练的先前 FLMs 的准确率低于 1%。相比之下，$\mathbb{S}$-FLM 使用精确速度达到 ~12%，使用 top-1 速度达到 ~18%（第3.1节 https://arxiv.org/html/2605.11125#S3.SS1.SSS0.Px3）。在标准温度采样（$T=1$）下，$\mathbb{S}$-FLM 使用 top-1 速度在所有*函数评估次数*（NFE）预算下缩小了与 MDLM 和 Duo 的差距，并在 NFE $\leq 16$ 时使用精确速度优于它们（图1 https://arxiv.org/html/2605.11125#S0.F1）。在低温（$T=0.1$）解码下仍存在差距，此时 MDLM 和 Duo 达到 33-36%（图3 https://arxiv.org/html/2605.11125#S4.F3）。

## 2 背景

#### 符号
我们用 $\mathcal{V}$ 表示大小为 $|\mathcal{V}|$ 的有限词汇表。粗体字母表示长度为 $L$ 的标记序列 $\mathbf{x} \in \mathcal{V}^L$，其中 $x^\ell$ 是第 $\ell$ 个元素，或者表示 $\mathbb{R}^d$ 中的向量，含义由上下文清楚界定。我们将 $\mathbb{S}^{d-1} := \{\mathbf{x} \in \mathbb{R}^d : \|\mathbf{x}\| = 1\}$ 写为 $\mathbb{R}^d$ 中的单位超球面，并将 $\mathcal{U}(\mathbb{S}^{d-1})$ 写为 $\mathbb{S}^{d-1}$ 上的均匀分布。标记嵌入存储在查找表 $\mathbf{E} \in \mathbb{R}^{|\mathcal{V}| \times d}$ 中，我们将与 $v \in \mathcal{V}$ 关联的行写为 $\mathbf{e}_v \in \mathbb{R}^d$。

#### 语言建模和自回归模型
语言模型使用密度 $p_\theta(\mathbf{x})$ 近似序列上的数据分布 $p_{\text{data}}: \mathcal{V}^L \rightarrow [0,1]$。AR 模型使用概率链式规则分解 $p_\theta(\mathbf{x}) = \prod_{\ell=1}^L p_\theta(x^\ell | x^{<\ell})$。这种分解使得精确的似然训练成为可能，但意味着推理速度慢，因为我们逐个生成标记，且无法以未来标记为条件。

### 2.1 流生成建模
连续归一化流（CNFs）(Chen2018NeuralOD, https://arxiv.org/html/2605.11125#bib.bib12; grathwohl2018ffjordfreeformcontinuousdynamics, https://arxiv.org/html/2605.11125#bib.bib32) 是 $\mathbb{R}^d$ 上的生成模型。CNFs 学习从噪声分布 $p_0 = p_{\text{noise}}$ 到数据分布 $p_1 = p_{\text{data}}$ 的连续时间传输。依赖时间的速度场 $u_t^\theta: \mathbb{R}^d \rightarrow \mathbb{R}^d$ 通过积分诱导*流* $\phi_t: \mathbb{R}^d \rightarrow \mathbb{R}^d$：
$$
\frac{d}{dt}\phi_t(z_0) = u_t^\theta(\phi_t(z_0)), \quad \phi_0(z_0) = z_0, \tag{1}
$$
速度场 $u_t^\theta$ 由具有连续可导且有界导数的神经网络参数化，因此 ODE (1 https://arxiv.org/html/2605.11125#S2.E1) 有唯一解 (coddington1955theory, https://arxiv.org/html/2605.11125#bib.bib15)。对于 $t \in [0,1]$，推前 $p_t = [\phi_t]_\# p_0$ 定义中间密度 $p_t$，族 $\{p_t\}_{t \in [0,1]}$ 称为从 $p_0$ 到 $p_1$ 的*概率路径*。积分 (1 https://arxiv.org/html/2605.11125#S2.E1) 到 $t$ 产生样本 $z_t = \phi_t(z_0) \sim p_t$。

#### 流匹配
流匹配（FM）(lipman2023flowmatchinggenerativemodeling, https://arxiv.org/html/2605.11125#bib.bib47; liu2022flowstraightfastlearning, https://arxiv.org/html/2605.11125#bib.bib49; albergo2023buildingnormalizingflowsstochastic, https://arxiv.org/html/2605.11125#bib.bib1) 是一种学习将 $p_0$ 传输到 $p_1$ 的速度场 $u_t^\theta$ 的方法（附录A https://arxiv.org/html/2605.11125#A1）。*真实*速度 $u_t$ 通常表示为期望：
$$
u_t(z_t) = \int u_{t|1}(z_t | x) p_{1|t}(x | z_t) dx, \tag{2}
$$
其中 $u_{t|1}$ 是*条件*速度，以 $x \sim p_{\text{data}}$ 为条件，$p_{1|t}$ 是给定噪声示例 $z_t$ 的后验。由于 $p_{1|t}$ 通常难以处理，FM 训练 $u_t^\theta$ 对抗 $u_{t|1}$。令 $\psi_{t|1}$ 表示与 $u_{t|1}$ 关联的*条件*流。一个常见的选择是线性插值：
$$
\psi_{t|1}(z_0 | x) = z_t = \alpha_t x + (1 - \alpha_t) z_0, \tag{3}
$$
其中 $z_0 \sim p_0$ 且 $x \sim p_1$。$\alpha_t: [0,1] \rightarrow [0,1]$ 是单调递增的*噪声计划*，满足 $\alpha_0 = 0$ 和 $\alpha_1 = 1$。条件速度由 $\psi_{t|1}$ 的时间导数给出：
$$
u_{t|1}(z_t | x) = \dot{\alpha}_t (x - z_0). \tag{4}
$$
FM 的一个关键结果 (lipman2023flowmatchinggenerativemodeling, https://arxiv.org/html/2605.11125#bib.bib47) 是*条件流匹配*（CFM）损失的最小化器是边缘速度 $u_t$ (2 https://arxiv.org/html/2605.11125#S2.E2)：
$$
\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}[0,1], z_0 \sim p_0, x \sim p_1} \| u_t^\theta(z_t) - u_{t|1}(z_t | x) \|^2. \tag{5}
$$

### 2.2 超球面的几何
我们总结了在 $\mathbb{S}^{d-1}$ 上使用的原语。关于黎曼几何的深入处理，参见 do Carmo (docarmo1992riemannian, https://arxiv.org/html/2605.11125#bib.bib23)。$\mathbf{p}, \mathbf{q} \in \mathbb{S}^{d-1}$ 之间的*测地距离*是