Spherical Flows for Sampling Categorical Data¶

会议: NeurIPS 2026
arXiv: 2605.05629
论文: https://arxiv.org/abs/2605.05629 https://openreview.net/forum?id=uVMZotQcJ7
代码: https://github.com/jannis-chemseddine/spherical-flows
领域: 其他
关键词: 连续扩散模型、流匹配、von Mises-Fisher 分布、离散数据生成、预测-校正采样

一句话总结¶

针对离散序列生成中欧氏连续流存在的模长漂移与维度缩放问题，本文在单位球面 $\mathbb{S}^{d-1}$ 上构建以 von Mises-Fisher（vMF）分布为条件路径的流模型，利用径向对称性将连续性方程化简为一维标量 ODE 并推导解析分数，单模型统一了 ODE、SDE 与预测-校正采样。

研究背景与动机¶

自回归解码长期主导着离散序列（如自然语言与结构化代码）的生成建模，但其固有的串行解码机制限制了推理吞吐，并容易累积误差。为了实现全序列并行非自回归生成，将连续空间的扩散模型与流匹配（Flow Matching）拓展至离散数据成为重要探索方向。现有离散生成流模型主要采用两类范式：一类是直接在有限离散状态空间中通过掩码或马尔可夫跳跃过程建模；另一类则是将离散 Token 嵌入到连续空间（如欧氏空间 $\mathbb{R}^d$ 或概率单纯形 $\Delta^N$）中构建连续动力学流，代表方法如 CDCD、LangFlow 等。

在连续嵌入范式中，模型通常将 Token 嵌入约束在单位范数上，以方向区分语义、以模长控制置信度。然而，欧氏空间中主流的线性插值（VP 路径）或高斯扰动（VE 路径）天然允许采样轨迹偏离球面，使隐变量在扩散初期大幅向原点收缩，破坏了几何约束；更关键的是，高维高斯噪声会导致不同 Token 间的几何相对关系随着维度 $d$ 的增加发生剧烈的信噪比扭曲。此外，球面上朴素的测地线大圆插值（Slerp 路径）虽受球面流匹配启发，却缺乏闭式解析黎曼分数，无法自然支持随机化校正器（Predictor-Corrector）与 SDE 采样，且在固定时刻下期望余弦相似度与空间维度 $d$ 解耦，无法随维度增加提供渐进的信息感知范围。

针对上述几何失配与可计算性瓶颈，本文的核心切入点是直接在乘积球面 $(\mathbb{S}^{d-1})^L$ 上定义扩散过程，引入内禀方向统计中的 von Mises-Fisher（vMF）分布族作为条件测度路径。核心 idea：利用 vMF 分布的径向对称性将球面连续性方程约化为余弦相似度的一维标量常微分方程，求解出解析可算的条件速度场并闭式导出黎曼分数，仅凭交叉熵损失学习逐位后验概率，即可无缝统一流 ODE 与朗之万预测-校正（Predictor-Corrector）采样。

方法详解¶

整体框架¶

本文针对长度为 $L$、词表大小为 $N$ 的离散序列生成任务。首先为每个词表 Token 分配一个可学习且严格约束在单位球面上的嵌入向量 $w_k \in \mathbb{S}^{d-1}$。在正向过程中，真实序列的嵌入点受到 vMF 噪声扰动，其浓度参数从 $t=0$ 时的 $\kappa_0=0$（球面均匀分布）单调递增至 $t \to 1$ 时的无穷大（狄拉克点测度）。在反向生成过程中，系统从全序列球面均匀先验出发，通过求解连续速度场驱动的流 ODE，并在时间步间交替执行基于黎曼分数的朗之万校正步，最终通过球面内积分类器解码为离散 Token。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["先验采样<br/>x₀ ~ 𝒰(𝕊ᵈ⁻¹)ᴸ"] --> B["神经网络后验预测<br/>p_θ,t(w_k | x) = Softmax(⟨w_k, x̂⟩ + b_k)"]
    B --> C["1. 条件流与速度场标量约化<br/>v_t = κ̇_t ∑ p_k ψ̃_t(⟨w_k,x⟩) P_x(w_k)"]
    C --> D["流 ODE 预测步 (Euler 推进)<br/>x_{i+1} = Retract(x_i + Δt · v)"]
    D --> E["2. 闭式黎曼分数与预测-校正<br/>g = κ ∑ p_k P_x(w_k) 朗之万校正"]
    E --> F["3. 自条件反馈调节<br/>前序均值嵌入叠加修正输入特征"]
    F -->|反向时间积分推进| G["4. 渐进信息感知维度缩放<br/>浓度调度 κ_t 随 d 自适应调谐"]
    G --> H["球面内积解码<br/>ŷ = argmax ⟨w_k, x̂₁⟩ + b_k"]

关键设计¶

1. 条件流与速度场标量约化：破解流匹配在黎曼流形上的不可积瓶颈

在流匹配框架中，由连续性方程（Continuity Equation, CE）求解非欧流形上的条件速度场通常面临偏微分方程解析不可积的数学障碍。针对球面 $\mathbb{S}^{d-1}$ 上的条件概率路径 $p_t(x|w)$，若其相对于均值方向 $w$ 具有径向对称性（即密度只依赖于内积标量 $s = \langle w, x \rangle$），本文证明其切空间速度场必定呈现 $v_t(x|w) = \psi_t(\langle w, x \rangle) \mathrm{P}_x(w)$ 形式，其中 $\mathrm{P}_x(w) = w - \langle w, x \rangle x$ 为到切空间的投影算子。以此为基础，高维球面连续性方程严格等价于一维标量通量方程（Flux Equation）： $$ \partial_t f_t(s) + \partial_s \big(f_t(s) \psi_t(s) (1 - s^2)\big) = 0, \quad s \in (-1, 1) $$ 其中 $f_t(s) = C_d(\kappa_t)(1 - s^2)^{(d-3)/2}\exp(\kappa_t s)$ 为投影到内积标量轴上的边际密度。通过变量分离，关于速度权重 $\psi_t$ 的偏微分系统被彻底约化为常微分方程（ODE），且在端点 $s \in [-1, 1]$ 具有唯一有界光滑解。进一步解耦得到 $\psi_t(s) = \dot{\kappa}_t \tilde{\psi}_t(s)$，其中无量纲因子 $\tilde{\psi}_t(s)$ 仅由当前浓度 $\kappa_t$ 决定而完全独立于时间调度，可通过数值稳定的积分表示高效预计算。边缘速度场最终化简为词表上的后验加权切向和： $$ v_{\theta, t}^l(\mathbf{x}) = \dot{\kappa}t \sum}^N p_{\theta, t}^l(w_k|\mathbf{x}) \tilde{\psit(\langle w_k, x^l \rangle) \mathrm{P}(w_k) $$

2. 闭式黎曼分数与预测-校正：统一确定性流与随机性扩散

传统球面测地线流由于缺乏解析边际密度导数，无法构造评分匹配（Score Matching），推理只能依赖确定性 ODE。然而在离散 Token 生成中，确定性数值积分误差会随着序列步数累积导致模式崩溃。vMF 分布密度的对数关于点 $x$ 呈现纯线性的内积形式 $\log p_t(x|w) = \log C_d(\kappa_t) + \kappa_t \langle w, x \rangle$。将其投影至球面切空间后，条件黎曼分数具有极简闭式：$\nabla_{\mathbb{S}^{d-1}} \log p_t(x|w) = \kappa_t \mathrm{P}_x(w)$。该项与速度场享有完全相同的后验权重与投影方向，边际黎曼分数为： $$ \nabla_{\mathbb{S}^{d-1}, x^l} \log p_{\theta, t}(\mathbf{x}) = \kappa_t \sum_{k=1}^N p_{\theta, t}^l(w_k|\mathbf{x}) \mathrm{P}_{x^l}(w_k) $$ 单一网络只需输出后验分布，便能以零额外参数开销同时获得漂移速度与切向分数。这使得模型在推理时能够执行朗之万预测-校正（Predictor-Corrector, PC）以及流 SDE 采样：在流 ODE 沿切向推进一步后，利用黎曼分数执行 $k$ 步朗之万随机游走，驱动扰动样本在当前等概率超曲面上自我纠错，极大平抑积分误差。

3. 自条件反馈调节：利用历史状态估计强化非自回归长程一致性

非自回归模型在全序列一次性生成时，各个 Token 之间缺乏显式的串行依赖，容易在全局语篇一致性或数独逻辑约束上面临协调困难。本文将自条件机制（Self-Conditioning, SC）迁移至球面流系统中：在推理与训练过程中，将当前预测的词表后验期望嵌入 $\sum_k p_t^l(w_k|\mathbf{x}) w_k$ 通过可学习线性映射投影为特征修正偏置，叠加回 Transformer Backbone 的输入端。在训练时以 $p=0.25$ 的概率激活无梯度预前向计算；在推理采样时则直接复用上一时间步的后验分布，不增加额外的模型评估次数（NFE）。

4. 渐进信息感知维度缩放：解决高维空间信噪比坍塌

在测地线球面插值（Slerp）中，当空间维度 $d \to \infty$ 时，任意固定时间 $t > 0$ 的期望余弦相似度均收敛至与维度无关的固定常数 $\sin(\frac{\pi}{2} t)$，这意味着高维几何下信号呈现突变式揭示，模型缺乏渐进感知的动态范围。而对于 vMF 分布，期望余弦相似度由修正贝塞尔函数之比严格确定： $$ \mathbb{E}_{x \sim \mathrm{vMF}(w, \kappa_t)}[\langle w, x \rangle] = A_d(\kappa_t) = \frac{\kappa_t}{d} + \mathcal{O}(\kappa_t^3) $$ 达到相同的信噪比水平所需的浓度参数 $\kappa_t$ 严格正比于空间维度 $d$。这一理论特性使 vMF 路径能够随着隐藏层嵌入维度的扩展，自然维持对 Token 语义信息平滑、可控的渐进揭示，为高保真度文本与复杂逻辑推理提供了几何上的尺度不变性。

损失函数 / 训练策略¶

网络的优化目标仅为标准的逐位置多分类交叉熵损失： $$ \mathcal{L}(\theta) = - \mathbb{E}{t \sim \mathcal{U}(0,1),\, \mathbf{w} \sim p}},\, \mathbf{xt \sim p_t(\cdot|\mathbf{w})} \left[ \sum_t) \right] $$ 理论分析证明，该交叉熵损失的全局唯一极小值点严格等价于真实后验概率分布 }^L \log p_{\theta, t}^l(w^l|\mathbf{x$p_t^l(w^l|\mathbf{x}_t)$，避免了高方差的连续切向量回归损失。词表嵌入 $\{w_k\}_{k=1}^N$ 在训练中保持单位范数约束并随网络参数联合优化。实际训练中，均匀时间采样 $t \sim \mathcal{U}(0,1)$ 被自适应交叉熵累积分布采样替换（遵循 CDCD 方案），使训练算力集中在难区分的时间区域。

实验关键数据¶

主实验¶

各方法在数独求解、大规模语言建模以及数学逻辑推理任务上的核心指标对比如下表所示。

任务 / 数据集	评估指标	本文 vMF (PC / SC)	基线 VP (欧氏)	离散掩码流 / 参考 SOTA	相对提升 / 核心表现
Sudoku-Extreme (NFE=128)	解约束正确率 Valid. (↑)	74.5% (PC)	69.1%	56.7% (MDM $p{=}0.5$)	较欧氏提升 +5.4%，大幅超越离散扩散
LM1B 语言建模 (NFE=128)	生成困惑度 gen. PPL (↓)	27.0 ($H{\ge}4.17$)	83.1 ($H{\ge}4.33$)	75.82 (DFM-Diag) / 66.7 (AR)	在匹配熵下达到连续流生成最前沿 PPL
TinyGSM 推理 (NFE=1024)	GSM8K 测试精度 Acc. (↑)	43.5% (PC+SC)	28.5% (PC+SC)	18.0% (MDLM / $\mathbb{S}$-FLM)	相比欧氏连续流绝对提升 +15.0%，大幅打破离散基准
OpenWebText (NFE=1024)	生成困惑度 gen. PPL (↓)	36.0 ($H{=}5.20$)	36.3 ($H{=}5.22$)	36.5 (LangFlow) / 47.1 (DFM)	逼近自回归模型基线 (AR 35.9)

消融实验¶

针对采样器设计（ODE vs. PC）、几何空间选择以及自条件机制（Self-Conditioning, SC）的消融对比（以 Sudoku-Extreme 与 TinyGSM 任务为主）：

实验模块 / 路径配置	数独有效率 Valid. (%)	TinyGSM 精度 Acc. (%)	说明
vMF 完整模型 (ODE 采样)	65.8%	6.4% (无SC) / 23.7% (含SC)	基础流动流积分基线
vMF + PC 朗之万校正	74.5%	33.0% (无SC) / 43.5% (含SC)	引入分数校正显著改善离散分布支撑
测地线插值 (Geodesic, 仅ODE)	49.4%	5.0% (无SC) / 14.9% (含SC)	缺乏闭式分数，且维度尺度不可调
欧氏保方差路径 (VP + PC)	69.1%	12.7% (无SC) / 28.5% (含SC)	缺乏球面约束，在推理任务中落后 vMF 15%
欧氏方差爆炸路径 (VE + PC)	46.5%	1.9% (无SC) / 4.1% (含SC)	离散流下校正器完全失效
vMF 去除自条件 (w/o SC, PC)	74.5%	33.0%	去掉 SC 后长程逻辑与多步推理精度明显衰退

关键发现¶

预测-校正器是 vMF 性能爆发的决定性引擎：在纯 ODE 采样下，vMF 与欧氏 VP 性能互有胜负；然而一旦开启基于黎曼分数的朗之万校正（PC 采样），vMF 能够无缝沿几何切空间抑制累积偏离，其在数独、语言和数学推理上的收益远高于欧氏 VP，在 TinyGSM 上带来接近 20 个百分点的绝对增益。
高维下欧氏空间的模长漂移与信噪比失调：VE 与 VP 路径允许状态在欧氏全空间移动，未考虑 Token 范数恒定的内禀结构，使得校正步经常将状态推离高概率流形；vMF 保持模长严格恒定，避免了由模长波动引起的置信度畸变。
自条件与延迟接入控制熵崩塌：长序列（如 OpenWebText $L=1024$）自条件生成随着 NFE 增加会发生生成熵衰减（低于真实语料熵）。通过引入延迟激活策略（$\tau \in \{1/3, 1/2\}$），在扩散初期保持纯净噪声驱动，既锁定了低困惑度，又完全保住了文本的多样性。

亮点与洞察¶

切向投影与标量化 ODE 的数学化简极为优雅：研究没有停留在高维黎曼几何的抽象符号上，而是精准抓住 vMF 分布径向对称的物理特性，通过一维标量通量方程彻底解耦时间调度与速度场计算，消除了数值积分中复杂的几何联络运算。
“一个后验，两路复用”的极简学习范式：模型仅凭标准的监督交叉熵目标即可完成训练，输出的词表概率既是终点反投影的语义概率，又是流 ODE 速度场的加权系数，更是 PC 采样黎曼分数的加权系数，无需额外训练打分网络或方差估计器。
从连续流跨越到离散复杂推理的可行性证明：在极其依赖长程精准依赖的 TinyGSM 推理任务上，连续球面流配合校正机制达到了 43.5% 的高精度，颠覆了以往“连续扩散模型无法做严格符号/数学推理”的固有认知。

局限与展望¶

高维高浓度极限下的数值下溢风险：在 $t \to 1$ 阶段，浓度 $\kappa_t \to \infty$，径向边际密度 $f_t(s)$ 极度尖锐，若未做对数偏移（Log-shift）与细致的截断，容易导致浮点数欠载与除零异常。
步数开销仍高于自回归与离散掩码模型：在数学推理中达到最佳性能依赖 NFE=1024，相比自回归模型（512 步解码）在单次采样吞吐上尚未展现出纯粹的速度优势。
未来方向：探索与引导采样（Classifier-Free Guidance）的闭式球面结合，以及将球面流匹配推广到多模态跨表征对齐与离散程序代码合成。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出球面 vMF 连续流匹配理论并给出闭式解与解析分数，数学推导极为扎实。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖离散数独、小/大语言模型及高难度数学符号推理，并覆盖 ODE、SDE 与 PC 丰富消融。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰，几何图景自然，逻辑闭环极佳。
价值: ⭐⭐⭐⭐⭐ 为离散数据的连续扩散/流生成奠定了内禀几何新范式，具有很强的学术参考与拓展价值。