Flow Map Language Models: One-step Language Modeling via Continuous Denoising¶
会议: NeurIPS 2026
arXiv: 2602.16813
OpenReview: https://openreview.net/forum?id=VqiKJfJ98m
代码: https://github.com/david3684/flm
领域: LLM 其他
关键词: 流匹配, 连续扩散语言模型, 流映射蒸馏, 单步生成, 概率单纯形
一句话总结¶
针对离散扩散语言模型在极少步采样下因因式分解误差而质量崩塌的瓶颈，提出在连续单热嵌入空间构建连续流语言模型（FLM），并通过流映射二元去噪重参数化将其蒸馏为流映射语言模型（FMLM），首次在单个前向评估下超越 8 步离散扩散生成质量。
研究背景与动机¶
现代前沿语言模型主要建立在自回归（Autoregressive, AR）机制之上，通过教师强制（teacher forcing）在训练阶段实现高度并行，但在推理生成阶段必须逐 token 串行生成。这种自回归生成带来的严重串行延迟，驱使学界广泛探索非自回归架构，其中基于离散扩散（Discrete Diffusion）的语言模型备受瞩目。离散扩散通过对完整序列施加加噪与去噪过程，理论上支持每一步并行输出或精炼多个 token。
然而，离散扩散模型在少步生成（few-step regime）乃至单步生成下普遍遭遇灾难性质变崩溃。这一困境的根本原因在于：文本离散序列的状态空间规模呈组合爆炸，逆向去噪概率分布无法在全离散空间直接建模，现有方法不得不依赖条件独立性假设对序列各位置进行因式分解近似。这一近似仅在连续时间极限或极多采样步数下勉强成立；一旦步数压缩至数步，模型便会完全忽略 token 间的联合相关性，导致生成文本退化为无序乱码或陷入低熵重复的死循环。
与之相对，在连续空间中建模的连续流模型具备确定性概率流 ODE，从数学上天然享有跨越任意时间区间的确定性流映射（Flow Map）结构，无需因式分解假设。长期以来学界普遍持有“连续扩散难以建模离散语言模态”的成见，忽视了流映射在离散数据上的巨大潜力。本文的核心 idea 是：将离散文本嵌入至无约束欧氏空间的标准单热编码，利用二元去噪器将流映射严格投影回概率单纯形，从而在连续空间建立具备显式流映射的语言模型，实现完全无因式分解误差的高保真单步生成。
方法详解¶
整体框架¶
本文构建了一套从连续流匹配预训练到高效流映射蒸馏的完整双阶管线。整体架构包含两大核心阶段： 1. 流语言模型（Flow Language Model, FLM）预训练：将长度为 \(L\) 的离散文本序列 \(\mathbf{y} \in V^L\) 无损映射到单热张量 \(\mathbf{x}_1 \in \mathbb{R}^{L \times |V|}\)，并在高斯噪声 \(\mathbf{x}_0 \sim \mathcal{N}(0, \mathbf{I})\) 与数据之间建立随机插值路径 \(I_t = (1-t)\mathbf{x}_0 + t\mathbf{x}_1\)。为了避免直接回归高维满秩速度场导致网络内部瓶颈欠拟合，FLM 学习预测后验干净数据的去噪器 \(\hat{D}_t(I_t)\)，并通过解码错误率重参数化时间调度。 2. 流映射语言模型（Flow Map Language Model, FMLM）蒸馏：流映射 \(X_{s,t}\) 负责将时间点 \(s\) 的状态直接输运至时间点 \(t\)。论文创新性地定义了二元去噪器 \(\delta_{s,t}\)，证明其在数学上恒落在概率单纯形内部，并利用半群性质构建交叉熵/KL 散度蒸馏目标，使得模型仅需单步 \(X_{0,1}\) 即可将纯高斯噪声一步映射为高质量单热词元表示。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["离散文本序列 y ∈ V^L<br/>单热编码映射为 x1 ∈ R^(L×|V|)"] --> B["高斯先验与随机插值<br/>It = (1-t)x0 + t x1"]
B --> C["解码错误率时间重参数化<br/>τ(t) 对齐 token 判决集中区"]
C --> D["单步去噪器与单纯形交叉熵<br/>Dt(x) 预测干净数据后验概率"]
D --> E["二元去噪器与单纯形映射<br/>δs,t(x) = x + (1-s) vs,t(x) ∈ Δ^(|V|-1)"]
E --> F["半群一致性流映射蒸馏<br/>多尺度三元组 (s,u,t) 凸组合约束"]
F --> G["单步/极少步确定性推理<br/>x^1 = X0,1(x0) 结合 argmax 解码"]关键设计¶
1. 单步去噪器与单纯形交叉熵：规避高维速度回归欠拟合
在常规连续流匹配中，模型通常直接拟合速度场目标 \(\dot{I}_t = \mathbf{x}_1 - \mathbf{x}_0\)。然而在词表维度 \(|V| \approx 30,000 \sim 50,000\) 的大词表场景下，高斯初始噪声 \(\mathbf{x}_0 \in \mathbb{R}^{L \times |V|}\) 属于高维满秩随机张量，而 Transformer 内部隐层特征维度 \(d \ll |V|\)，直接用均方误差（MSE）回归速度场会遭遇极严重的特征瓶颈与欠拟合。
针对该痛点，本文采用线性变量代换将速度场转为干净数据后验均值（即去噪器）： $\(D_t(\mathbf{x}) \coloneqq \mathbb{E}[\mathbf{x}_1 \mid I_t = \mathbf{x}], \quad b_t(\mathbf{x}) = \frac{D_t(\mathbf{x}) - \mathbf{x}}{1-t}\)$ 由于 \(\mathbf{x}_1^l\) 是离散词元的单热向量，其条件期望展开恰好为词表上的离散后验概率分布：\(D_t(\mathbf{x})^l = [p_{1|t}^l(\mathbf{x}_1^l = \mathbf{e}_1 \mid I_t = \mathbf{x}), \dots, p_{1|t}^l(\mathbf{x}_1^l = \mathbf{e}_{|V|} \mid I_t = \mathbf{x})]^\top \in \Delta^{|V|-1}\)。因此，模型输出层可以直接接入 token 级别的 softmax 归一化，将输出严格约束在概率单纯形上。此时可抛弃连续空间的 MSE 损失，完全采用适配单热几何的交叉熵损失进行训练： $\(\mathcal{L}_{\mathrm{CE}}(\hat{D}) = -\sum_{l=1}^L (\mathbf{x}_1^l)^\top \log \hat{D}_t(I_t)^l\)$ 论文从理论上严格证明，最小化该交叉熵等价于最小化预测后验与真实后验的 KL 散度，并推导出了在截断区间上由交叉熵超额风险（excess risk）严格控制生成分布与真实分布之间 Wasserstein-2 距离的泛化上界。
2. 解码错误率时间重参数化：解决单热几何下的训练信号极化
在连续图像生成中，时间变量通常采用均匀采样 \(t \sim \mathcal{U}[0, 1]\)。但在离散词元的连续单热插值过程中，词元决策呈现极端的非均匀性。定义在时刻 \(t\) 截断并执行 \(\mathrm{argmax}\) 解码的解码错误率为： $\(P_e(t) \coloneqq \frac{1}{L}\sum_{l=1}^L P(g^l(\mathbf{x}_t) \neq g^l(\mathbf{x}_1))\)$ 随着词表规模 \(|V|\) 扩大，错误率在 \(t \in [0, 0.8]\) 的宽广区间内几乎维持在 \(1 - 1/|V|\) 的高位平台期，而在 \(t \to 1\) 极窄的时间窗口内发生陡峭雪崩式下降。这意味着传统均匀采样把绝大部分算力浪费在了没有实质去噪判决的信息稀疏区，而在决定性别的窄窗口内采样极度欠缺，导致训练发散与崩溃。
为此，论文提出基于解码错误率下降进度的解析重参数化函数 \(\tau(t)\)，要求时间采样步长与去噪进度严格成正比： $\(\tau(t) = \frac{P_e(0) - P_e(t)}{P_e(0)} = 1 - \frac{|V|}{|V|-1} P_e(t)\)$ 在训练阶段从重参数化时间空间均匀采样 \(\tau(t) \sim \mathcal{U}[0, 1]\) 反解出物理时间 \(t = \tau^{-1}\)；在推理采样时构建均匀 \(\tau\) 网格 \(t_n = t(n/N)\)。该机制将训练梯度与推理计算动态聚集在词元判决发生的关键时段，成功将欧氏连续流稳定扩展到 \(|V| = 50,257\) 的大词表规模。
3. 二元去噪器与单纯形映射：构建确定性少步流映射
流映射 \(X_{s,t}(\mathbf{x})\) 是概率流 ODE 的演化算子，满足 \(X_{s,t}(\mathbf{x}_s) = \mathbf{x}_t\)。通常参数化为 \(X_{s,t}(\mathbf{x}) = \mathbf{x} + (t-s)v_{s,t}(\mathbf{x})\)，其中 \(v_{s,t}\) 为平均速度场。但由于高斯噪声的扩散效应，流映射输出张量 \(X_{s,t}\) 在行进过程中必然脱离概率单纯形。如果直接采用连续空间均方误差（如 MeanFlow 或一致性模型 MSE）学习 \(v_{s,t}\)，将失去离散几何的先验优势。
本文设计了革命性的“二元去噪器”（two-time denoiser）重参数化： $\(\delta_{s,t}(\mathbf{x}) \coloneqq \mathbf{x} + (1-s)v_{s,t}(\mathbf{x})\)$ 论文在数学上严格证明，\(\delta_{s,t}(\mathbf{x})\) 等价于沿流轨迹上单步去噪器 \(D_\tau\) 的积分加权均值： $\(\delta_{s,t}(\mathbf{x}) = \frac{(1-s)(1-t)}{t-s} \int_s^t \frac{D_\tau(X_{s,\tau}(\mathbf{x}))}{(1-\tau)^2} d\tau\)$ 由于单步去噪器恒在概率单纯形上，其非负权重的加权平均亦严格保证 \(\delta_{s,t}(\mathbf{x}) \in \Delta^{|V|-1}\)。基于此，不仅流映射可通过凸组合精确反解 \(X_{s,t}(\mathbf{x}) = \frac{1-t}{1-s}\mathbf{x} + \frac{t-s}{1-s}\delta_{s,t}(\mathbf{x})\)，二元去噪器还能天然配置 softmax 输出层，完全在单纯形流形上运作。
4. 半群一致性流映射蒸馏与边界注入：稳定单步生成收敛
基于二元去噪器，流映射的半群性质 \(X_{s,t} = X_{u,t} \circ X_{s,u}\) 转化为单纯形上的凸组合代数恒等式： $\(\delta_{s,t}(\mathbf{x}) = \gamma \delta_{s,u}(\mathbf{x}) + (1-\gamma) \delta_{u,t}(X_{s,u}(\mathbf{x})), \quad \gamma = \frac{(1-t)(u-s)}{(1-u)(t-s)}\)$ 当 \(s \le u \le t\) 时，权重 \(\gamma \in [0, 1]\)，使得右侧目标组合为一个合法的单纯形离散概率分布。蒸馏目标由此分解为对角锚定损失与非对角半群损失之和： $\(\mathcal{L}_{\mathrm{diag}} = -\sum_{l=1}^L \hat{D}(I_s)^l \cdot \log \hat{\delta}_{s,s}^l(I_s)\)$ $\(\mathcal{L}_{\mathrm{off}} = -\sum_{l=1}^L \bar{\delta}^l \cdot \log \hat{\delta}_{s,t}^l(I_s), \quad \bar{\delta} = \mathrm{sg}\left(\gamma \hat{\delta}_{s,u}(I_s) + (1-\gamma) \hat{\delta}_{u,t}(\hat{X}_{s,u})\right)\)$ 为了消除连续时间重采样在 \(s=0\) 处的极低采样密度，主算法引入以 \(p=1/32\) 的概率显式采样全局跳跃边界对 \((s,t)=(0,1)\) 的机制，强力确保了直接从初始高斯噪声单步映射到完整干净文本 \(X_{0,1}(\mathbf{x}_0)\) 的收敛精度。
损失函数 / 训练策略¶
- FLM 预训练阶段：采用 Adam 优化器，学习率 \(3 \times 10^{-4}\)，全局批量 512，在 LM1B 和 OWT 上充分预训练 100 万步（1M steps）。网络骨干统一采用 1.79 亿参数（179M）的 DiT-B（12 层 Transformer block，配备 RoPE 位置编码与 AdaLN 时间条件调制）。
- FMLM 蒸馏阶段：冻结预训练好的 FLM 教师模型，学生模型结构完全同款，以相同的优化器参数快速蒸馏 10 万步（100k steps）。蒸馏过程收敛速度比基础预训练快 10 倍。
实验关键数据¶
主实验¶
实验在 One Billion Word（LM1B，序列长 128，词表 30,522）与 OpenWebText（OWT，序列长 1024，词表 50,257）上展开全面对比。评估指标为 GPT-2 Large 测得的生成困惑度（Gen. PPL \(\downarrow\)）及单 token 一元熵（Entropy，真实基准 LM1B 为 4.31，OWT 为 5.44）。
表 1：全步长（1024 步）生成性能对比
| 模型 | 类型 | LM1B Gen. PPL (↓) | LM1B Entropy (4.31) | OWT Gen. PPL (↓) | OWT Entropy (5.44) |
|---|---|---|---|---|---|
| RDLM | 连续黎曼扩散 | 268.21 | 4.33 | - | - |
| CANDI | 离散-连续混合 | 120.99 | 4.35 | 143.13 | 5.71 |
| MDLM | 掩码离散扩散 | 109.21 | 4.32 | 105.15 | 5.63 |
| Duo | 离散对偶扩散 | 98.14 | 4.31 | 77.69 | 5.55 |
| FLM (本文) | 欧氏连续流 | 96.91 | 4.29 | 62.23 | 5.33 |
表 2：极少步（1~4 步）蒸馏模型生成性能对比
| 数据集 | 步数 | MDLM + SDTT | MDLM + Di4C | Duo + DCD | Duo + Di4C | FMLM (本文) |
|---|---|---|---|---|---|---|
| LM1B | 1 步 | 1429.48 (4.31) | 1217.10 (4.38) | 1224.52 (4.33) | 292.94 (3.79*) | 119.34 (4.16) |
| 2 步 | 602.14 (4.28) | 621.59 (4.37) | 520.08 (4.20) | 247.69 (3.87*) | 110.19 (4.21) | |
| 4 步 | 241.01 (4.28) | 247.32 (4.00) | 210.88 (4.23) | 150.67 (4.00) | 98.76 (4.21) | |
| OWT | 1 步 | 1260.86 (5.26) | 1298.80 (5.29) | 5743.29 (6.02) | 370.51 (3.92*) | 168.30 (5.17) |
| 2 步 | 877.22 (5.34) | 758.23 (5.35) | 891.16 (5.41) | 210.22 (4.63*) | 133.29 (5.25) | |
| 4 步 | 339.73 (5.38) | 239.27 (5.40) | 250.86 (5.37) | 154.67 (4.85*) | 111.31 (5.26) |
(注：星号标注表示该基准通过严重的 token 重复退化导致熵崩塌低于 4.0，虚假降低了 PPL)
消融实验¶
表 3：LM1B 上的核心设计消融结果
| 模块类别 | 消融配置 | 采样步数 | Gen. PPL (↓) | Entropy | 核心结论分析 |
|---|---|---|---|---|---|
| FLM 参数化与损失 | 速度场回归 \(b_t\), MSE | 1024 | 3801.36 | 4.85 | 噪声高维满秩导致特征瓶颈，模型完全发散不收敛 |
| 去噪器 \(D_t\), MSE | 1024 | 129.04 | 3.97 | 预测干净数据有效，但缺乏单纯形约束导致退化 | |
| 去噪器 + Softmax, MSE | 1024 | 120.16 | 4.28 | 单纯形约束有效防止了熵崩溃 | |
| 去噪器 + Softmax, CE (完整) | 1024 | 96.91 | 4.29 | 交叉熵紧密匹配单热离散几何，达成最优生成质量 | |
| 时间重参数化 | 无重参数化 (均匀 \(t\)) | 1024 | 149.18 | 4.29 | 采样点未匹配判决区间，梯度信号被严重稀释 |
| 学习型熵时间调度 | 1024 | 130.42 | 4.27 | 相比均匀分布有改善，但依赖在线学习不够稳定 | |
| 解码错误率重参数化 \(\tau(t)\) | 1024 | 106.98 | 4.30 | 解析显式对齐判决区间，训练效率与稳定性最强 | |
| 连续几何选择 | 黎曼流形扩散 | 1024 | 268.21 | 4.33 | 几何约束计算复杂度高且表现受限 |
| 单纯形扩散 (Simplex) | 1024 | 85.07 | 3.76 | 初始高维均匀分布缺乏多样性，发生严重模式坍塌 | |
| 无约束欧氏空间 (完整) | 1024 | 96.91 | 4.29 | 高斯球壳初始化覆盖所有方向，质量与多样性俱佳 | |
| FMLM 蒸馏设计 | 平均速度 \(v_{s,t}\), MSE | 1 | 199.01 | 3.93 | 缺乏离散先验，单步生成质量显著退化 |
| \(\delta\)-去噪器 + Softmax, MSE | 1 | 181.72 | 4.18 | 单纯形约束改善熵，但 MSE 目标次优 | |
| \(\delta\)-去噪器 + Softmax, CE | 1 | 119.34 | 4.16 | 二元去噪器配合交叉熵带来高达 80 点 PPL 巨幅提升 | |
| 无边界采样 (\(p=0\)) | 1 | 142.61 | 3.81 | 原点跨度缺乏训练信号，极少步生成显著劣化 | |
| 边界采样 (\(p=1/32\), 完整) | 1 | 119.34 | 4.16 | 显式强化全局输运信号，单步 PPL 优化至极致 |
关键发现¶
- 离散扩散蒸馏的致命困境：在 1 步或 2 步采样下，所有离散扩散蒸馏方法（如 MDLM+SDTT、Duo+DCD）均发生灾难性失败，其 PPL 飙升至 1200 以上（输出语法全无的随机组合乱码），或者如 Duo+Di4C 那样通过退化为极低熵（\(\sim 3.79\)）的无休止 token 重复来换取低 PPL。这从实验上坐实了因式分解近似在少步生成下的不可克服性。
- 单步匹敌 8~16 步：FMLM 单步生成的质量（LM1B PPL 119.34，OWT PPL 168.30）直接击败了现有离散扩散蒸馏模型在 4 步甚至 8 步的生成水平，在 OWT 文本续写条件生成上单步 PPL 达到 141.66，比 MDLM+SDTT 优越 2.6 倍。
- 组合逻辑推理的飞跃：在 81 步数独（Sudoku）端到端求解任务中，FMLM 在单次前向传播（1 步）下实现高达 39.84% 的谜题求解有效率，而所有离散扩散基线在 1 步下完全为 0.00%；结合拒绝采样后相比 81 步自回归模型带来 7.6 倍的确定性加速。
亮点与洞察¶
- 颠覆“连续空间不适合建模文本”的既有认知：论文不仅证明无约束欧氏空间的高斯插值流能够胜任语言建模，更指明了离散扩散存在“无法定义样本级确定性流映射”的数学盲区（命题 3.5），确立了连续流在单步非自回归生成中的独特理论霸权。
- 二元去噪器的数学巧思：巧妙利用积分因子将流映射 ODE 解耦为单步去噪器的加权积分，将脱离单纯形的 \(X_{s,t}\) 成功重构为严格守恒在单纯形内部的 \(\delta_{s,t}\)，从而将分类交叉熵与半群一致性完美融合。
- 免训练扩散引导的前瞻赋能：由于流映射提供了通往终点的可微前瞻（look-ahead）路径 \(X_{t,1}\)，引导分类器或奖励模型仅需在真实干净数据上预训练，无需像离散扩散那样在全部加噪轨迹上昂贵训练，大幅拓展了推理期导向（Autoguidance）和强化学习端点微调的前景。
局限与展望¶
- 单热张量的高显存开销：由于在训练中直接对全尺寸单热矩阵 \(\mathbb{R}^{L \times |V|}\) 执行反向传播，每步训练需要更新全部 \(|V| \times d\) 嵌入权重，相比只更新离散索引对应向量的嵌入扩散，显存和算力开销高出约 30%，亟待稀疏梯度或分块词表优化。
- 与自回归前沿模型的生成差距：尽管在多步和少步上均超越现有离散扩散，但在绝对生成复杂度与超长文本连贯性上，单步连续流与前沿大参数量自回归模型相比仍处于探索阶段。
相关工作与启发¶
- vs MDLM / Duo (离散扩散模型)：离散扩散模型通过离散马尔可夫链或掩码转移建模序列，必须对全序列联合概率使用条件独立性因式分解近似，导致少步生成严重崩溃；本文通过连续欧氏流与样本级流映射彻底摆脱因式分解假设，单步即实现高保真连贯生成。
- vs CDCD / Diffusion-LM (连续嵌入扩散)：以往连续语言扩散多采用可学习的稠密词嵌入或预训练嵌入，常遭遇表示崩塌、范数失控及复杂的重参数化调整；本文采用无损且无参数的规范单热编码与去噪器概率重参数化，架构更纯粹稳定。
- vs Consistency Models / MeanFlow (图像流映射蒸馏)：图像领域的流映射蒸馏完全建立在连续欧氏 MSE 损失上，无法适配离散分布；本文提出的二元去噪器将半群流映射收敛到概率单纯形，成功开辟了连续流蒸馏在离散序列生成上的通用范式。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立离散语言数据的单纯形二元去噪器流映射理论与算法。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 LM1B、OWT、数独逻辑求解、前瞻奖励引导及全面参数消融。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨自洽，对连续与离散扩散的深层矛盾剖析深刻透彻。
- 价值: ⭐⭐⭐⭐⭐ 为超低延迟大语言模型并行推理和单步生成提供了极具前景的突破性技术路径。