Coupling Models for One-Step Discrete Generation¶

会议: NeurIPS 2026
arXiv: 2605.07193
代码: https://github.com/pengzhangzhi/Coupling-Models
领域: LLM 其他
关键词: 离散生成模型, 单步生成, 归一化流, 掩码扩散模型, 可控引导

一句话总结¶

提出两阶段离散生成范式 Coupling Models，先通过似然归一化流将离散序列与高斯潜变量对齐建立可采样耦合，再训练并行解码器实现单步离散序列生成与高效引导。

研究背景与动机¶

离散结构生成是语言建模、生物序列设计（如 DNA 增强子设计）以及离散图像建模的核心基础。当前在这些任务中占据统治地位的生成模型高度依赖多步迭代计算：自回归模型（如 LLaMA、Qwen）通过逐步条件分解生成，虽然表达能力极强且具备优异的可扩展性，但其固有的串行推理机制成为了生成速度的主要瓶颈；离散扩散模型、掩码精炼模型以及离散流匹配模型虽然能在各位置实现并行更新，但为了建立全局 token 间的复杂依赖关系，通常需要数十乃至数百次模型前向去噪与精炼步骤。

由此催生了对真正单步（one-step）并行离散生成模型的迫切需求：从简单先验分布采样一次，经模型单次前向传播即直接输出完整离散序列。然而，近期的加速研究主要集中在“轨迹压缩”路径上，例如通过连续流蒸馏、轨迹自蒸馏、端点对齐（Flow Maps）或修正流（Rectified Flow）将多步采样轨迹强行压缩到极少数步乃至单步。这类方法往往在连续单纯形（simplex）或独热嵌入空间构建复杂的几何路径与端点映射，训练不仅极易受到多步累积误差与失真影响，而且直接面对离散序列在单步条件下的表示瓶颈。更根本的理论障碍在于：一个直接输出离散类别的单步解码器若缺乏共享全局隐变量，其在各位置独立采样的输出分布必然退化为完全相互独立的乘积分布，自由度仅为 \(L(V-1)\)，从数学上根本无法表示完整序列单纯形中多达 \(V^L-1\) 自由度的复杂交叉依赖关系。

因此，单步离散生成的核心难题不仅是优化问题，更是表示问题。本文的核心切入角度在于：不应强行压缩多步扩散或单纯形流的动力学轨迹，而应在数据空间与采样先验之间显式构建一个易于采样的连续隐空间耦合。核心 idea：通过两阶段耦合范式，先利用似然归一化流将离散序列嵌入对齐到标准高斯先验诱导成对监督，再训练并行解码器单步反演该耦合，以共享高斯隐变量的混合分布彻底突破独立发射的表达瓶颈。

方法详解¶

整体框架¶

Coupling Models 摒弃了基于轨迹模拟与流图蒸馏的思路，构建了两阶段生成管线：Stage A 负责构建离散序列与标准高斯潜变量之间的双向可采样耦合，为解码器提供高质量配对监督数据；Stage B 则在该配对数据上训练一个单步并行分类解码器。在推理时，模型完全无需运行 Stage A，只需抽取高斯噪声 \(z \sim \mathcal{N}(0, I)\)，通过并行解码器单次前向传播即可同时输出全部位置的类别概率分布并采样完成。

%%{init: {'theme': 'base', 'themeVariables': {'primaryColor': '#EEF2FF', 'primaryBorderColor': '#6366F1', 'lineColor': '#4F46E5', 'fontSize': '13px'}, 'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    subgraph S1["Stage A: 离散序列到高斯潜变量的耦合构建"]
        direction TB
        X["离散序列输入 x"] --> Enc["可重参数化编码器 E_psi<br/>生成连续表征 u"]
        Enc --> Flow["连续归一化流 NF_phi<br/>对齐到高斯潜变量 z"]
        Enc --> DecA["轻量辅助解码器 D<br/>重构离散序列 x_hat"]
    end

    subgraph S2["Stage B: 潜变量到序列的逆映射学习"]
        direction TB
        DataZ["诱导配对监督数据集 D_Z = (z, x)"] --> Gen["并行一阶解码器 G_theta<br/>以全局 z 为条件输出全部 logits"]
    end

    subgraph S3["推断阶段: 单步生成与少步扩展"]
        direction TB
        Prior["先验采样 z ~ N(0, I)"] --> DecOne["单步并行发射<br/>Categorical(softmax(l/tau))"]
        Prior --> MDM["少步掩码扩散扩展<br/>固定全局 z 进行局部 P2-self 精炼"]
        Prior --> DirectGuidance["单步可控引导<br/>CFG / 分类器梯度 / 奖励微调"]
    end

    S1 -->|冻结编码器与流参数| DataZ
    Gen -->|推理时单次评估| DecOne

关键设计¶

1. 离散到高斯的两阶段连续耦合构建

针对直接单步离散解码器独立分类发射导致的 \(L(V-1) \ll V^L-1\) 表达力瓶颈，模型不能直接在离散序列上拟合连续概率流。Stage A 将离散序列 \(x = (x_1, \ldots, x_T)\) 首先通过可重参数化编码器 \(E_\psi(x, \epsilon)\) 映射到连续表征 \(u\)（其中 \(\epsilon \sim \mathcal{N}(0, I)\)），再通过可逆归一化流 \(z = \mathrm{NF}_\phi(u)\) 将 \(u\) 变换为采样潜变量 \(z\)。同时引入轻量级辅助解码器 \(D\) 从 \(u\) 中重构原始序列，避免连续隐表征丢失离散序列的关键上下文细节。Stage A 的联合优化目标由序列重构损失、潜空间 KL 散度约束和流的最大似然对数雅可比项组成：

\[\mathcal{L}_{A} = \lambda_{\mathrm{rec}} \left[-\sum_{t=1}^{T} \log p_{D}(x_t \mid u)\right] + \lambda_{\mathrm{KL}} \, \mathrm{KL}(q_\psi(u \mid x) \,\|\, \mathcal{N}(0, I)) + \lambda_{\mathrm{flow}} \left[-\log p_Z(\mathrm{NF}_\phi(u)) - \log \left|\det \frac{\partial \mathrm{NF}_\phi(u)}{\partial u}\right|\right]\]

通过这种显式分离，编码器负责捕获充分的离散重构信息，归一化流则精确对齐潜在边际分布与标准高斯先验 \(p_Z = \mathcal{N}(0, I)\)。由此生成的成对样本 \((z, x)\)，使得连续隐变量 \(z\) 既包含了生成目标 \(x\) 的完整全局依赖，又在边缘上与推理时的采样先验严格吻合。

2. 潜变量对齐误差控制与单步并行解码

传统 VAE 结构在单步生成中常因先验“空洞”或边际分布失配而在推理时产生严惩的分布外漂移。论文在理论上给出了全变差（Total Variation）误差上界：若解码器的条件 KL 散度不超过 \(\varepsilon_{\mathrm{dec}}\)，流模型诱导潜变量边际分布与高斯先验的 KL 散度不超过 \(\varepsilon_{\mathrm{flow}}\)，则真实数据分布与生成边缘分布的距离受控：

\[\mathrm{TV}(p_{\mathrm{data}}, p_{\theta}^{\mathrm{gen}}) \leq \sqrt{\frac{\varepsilon_{\mathrm{dec}}}{2}} + \sqrt{\frac{\varepsilon_{\mathrm{flow}}}{2}}\]

在冻结 Stage A 参数后，Stage B 收集配对数据集 \(\mathcal{D}_Z = \{(z^{(i)}, x^{(i)})\}_{i=1}^N\)，训练一个参数化的并行 Transformer 解码器 \(G_\theta(z) \in \mathbb{R}^{T \times V}\)。虽然条件发射在给定 \(z\) 下具有独立乘积形式 \(G_\theta(x \mid z) = \prod_{t=1}^T \mathrm{softmax}(\ell_t)[x_t]\)，但生成边缘分布为潜在混合模型 \(p_\theta^{\mathrm{gen}}(x) = \int G_\theta(x \mid z) p_Z(z) dz\)，完全承载了跨位置的全局长程关联。

3. 少步掩码扩散扩展（P2-self 精炼）

耦合隐空间不仅支持纯单步直接生成，还能无缝拓展至保留全局连续先验的少步掩码去噪。在 Stage B 中，可以用以全局潜变量 \(z\) 为条件的掩码去噪器 \(H_\theta(x_t, z)\) 替代纯并行解码器。训练时对序列以随机时间 \(t \sim \mathrm{Unif}(0, 1)\) 进行掩码加噪，仅在被掩码的位置计算交叉熵：

\[\mathcal{L}_{\mathrm{MDM}}(\theta) = -\mathbb{E}_{(z, x) \sim \mathcal{D}_Z, t \sim \mathrm{Unif}(0, 1), m} \left[\frac{1}{\sum_{i=1}^T m_i} \sum_{i: m_i=1} \log p_\theta(x_i \mid x_t, z)\right]\]

推理时，全局潜变量 \(z \sim \mathcal{N}(0, I)\) 在采样全过程中保持恒定，为序列生成提供全局语义锚点；解码器结合置信度评估执行少步 P2-self 采样，利用重掩码强度 \(\eta\) 自适应撤销低置信度已生成 token 并重新填补，用极少步数（2~8 步）显著提升局部语法一致性。

4. 终端预测层直接可控引导

在多步离散扩散模型中，引导信号（如分类器引导或奖励信号）必须穿透长达数十步乃至数百步的离散非可微转移链，导致梯度估计方差极大、计算开销昂贵。Coupling Models 由于具备清晰的潜变量与端点结构，直接在最终预测层构建引导： - 无分类器引导（CFG）：在训练阶段以一定概率丢弃类别条件 \(y\)，推理时直接在最终 logits 上进行单步线性外推：\(\ell^{\mathrm{cfg}} = \ell^u + s(\ell^c - \ell^u)\)； - 潜变量分类器引导：固定解码器 \(G_\theta\)，利用 Softmax 或 Gumbel-Softmax 连续松弛 \(\rho(\ell)\)，在潜空间直接对初始先验 \(z\) 进行 \(K\) 步轻量梯度上升：\(z^{k+1} = z^k + \eta \nabla_{z^k} R(\rho(G_\theta(z^k, y)), y)\)； - 奖励微调（Reward Fine-Tuning）：通过优化带锚点惩罚的奖励目标 \(\mathcal{L}_{\mathrm{FT}}(\theta) = -\lambda_{\mathrm{reward}} \mathbb{E}_{z, y} R(\rho(G_\theta(z, y)), y) + \lambda_{\mathrm{anchor}} \mathcal{R}(\theta, \theta_0)\) 直接微调解码器。

损失函数 / 训练策略¶

Stage A 采用 AdamW 优化器，在图像和 DNA 任务上联合优化式 (1)；在 LM1B 文本生成中，使用冻结的 Qwen2.5-0.5B（隐藏层 896 维）作为编码器，降维至 256 维后接含 12 个耦合层的 RealNVP 架构，流权重项在训练前 5% 从 0.5 退火至 1.0。Stage B 在文本任务中采用 12 层 Transformer 解码器，并额外引入来自冻结 Qwen2.5-0.5B 教师模型的 KL 蒸馏损失（权重 1.0，温度 2.0）；为了缓解训练与先验采样分布微小差异，Stage B 训练期间有 30% 的概率用缩放高斯纯噪声 \(0.90 \times \mathcal{N}(0, I)\) 代替 Stage A 的输出潜变量。

实验关键数据¶

主实验¶

论文在受控二值图像生成（MNIST-Binary）、生物序列设计（DeepFlyBrain 增强子）以及开放域文本生成（LM1B）三个代表性领域全面评估了 Coupling Models 的单步生成能力。

数据集	任务/指标	Coupling Model (本文)	前序最强单步基线	相对提升 / 优势
MNIST-Binary	FID \(\downarrow\) (单步 1 step)	5.50	10.10 (CFM) / 12.90 (PAIRFLOW+ReDi)	FID 下降 45.5%
Fly Brain (DNA)	FB FBD \(\downarrow\) (单步 1 step)	12.9 ± 0.3	15.8 (Distilled Dirichlet FM)	FBD 降低 18.4%
LM1B (Text)	Gen-PPL \(\downarrow\) (低困惑度点)	45.82 (Entropy 3.61)	68.11 (DFM-ESD, Entropy 3.79)	困惑度降低 32.7%
LM1B (Text)	Gen-PPL / 真实参考熵 (0.65T)	98.96 (Entropy 4.32)	119.34 (FMLM, Entropy 4.16)	在更高熵 (4.32 vs 4.16) 下 PPL 降低 17.1%

消融实验：引导机制与推理效率对比 (MNIST-Binary)¶

论文在受控 MNIST-Binary 条件生成任务上对比了 Coupling Models 的单步引导与 512 步掩码扩散模型（MDM）在生成质量、分类精度与推理速度上的权衡表现。

模型	引导方式	参数量 (M)	FID \(\downarrow\)	准确率 (%) \(\uparrow\)	函数评估次数 (NFE) \(\downarrow\)	每千样本耗时 (秒) \(\downarrow\)
MDM 基线	CFG	73.74	8.72	99.8	512	356.82
MDM 基线	分类器引导	73.74	7.95	97.2	256	244.68
MDM 基线	奖励微调	73.74	8.82	99.7	256	203.04
Coupling Model	CFG	24.07	5.64	99.1	2	1.72 (加速 207×)
Coupling Model	分类器引导	24.07	5.26	98.5	6	13.20 (加速 18.5×)
Coupling Model	奖励微调	24.07	5.09	99.5	1	0.63 (加速 322×)

关键发现¶

单步质量超越迭代蒸馏：在 MNIST-Binary 上，仅用单步前向传播的 Coupling Model 取得了 5.50 的 FID，显著优于经过复杂自蒸馏/修正流加速的基线（CFM 的 10.10、PAIRFLOW+ReDi 的 12.90），甚至逼近多步连续流模型（64 步 PAIRFLOW 5.17，64 步 UDLM 5.01）。
语言生成前沿帕累托位移：在 LM1B 评测中，前序单步方法往往通过降低生成多样性（熵骤降至 3.10 甚至更低）来换取虚假的低困惑度；而 Coupling Model 在宽泛的温度扫描中（0.50~0.70）展现出平滑的质量-多样性权衡，将单步生成前沿全面推向更低 PPL、更高熵的理想区域。
少步精炼迅速收敛：在 2~8 步掩码精炼下，Coupling Model 在 LM1B 上的 Gen-PPL 分别达到 66.2、69.3 和 64.3，全程优于所有基于蒸馏与离散流的少步基准模型（DFM-ESD、DFM-PSD 及 FMLM）。
极短控制路径大幅降低计算开销：在引导实验中，Coupling Model 将原本离散扩散需要 256~512 次网络调用的极长控制链路简化至 1~6 次 NFE，生成耗时缩短了 18~320 倍，且 FID 质量指标全面优于强扩散基线。

亮点与洞察¶

范式转变（Aha Moment）：离散单步生成的瓶颈不是求解器步长或动力学轨迹拟合不足，而是底层缺乏能承载全局多词依赖的连续空间表示。通过将多步动力学模拟替换为“两阶段可采样隐空间耦合构建”，彻底解耦了序列信息保真与先验高斯化两个正交目标。
表达能力界限破解：严格论证了单步 token 独立预测模型无法建模一般离散序列联合分布的理论缺陷，通过先验高斯隐变量建立潜在混合模型，以最低的参数增长和单次前向实现了全局上下文感知。
可迁移的离散控制范式：在离散生成中利用最终输出层的连续松弛直接对潜变量或 logits 做优化，绕过了长期困扰离散扩散的多步随机梯度传递难关，为蛋白质序列定向突变、代码条件修饰等离散生成控制提供了通用轻量解法。

局限与展望¶

模型尺度与任务规模：当前实验集中在中等规模模型（如 Qwen2.5-0.5B 表征）与常规基准数据集上，尚未在大规模前沿大语言模型（数十亿参数以上）、超长上下文生成或大型复杂蛋白质结构生成上进行 Scaling 验证。
文本评估指标的固有偏差：生成困惑度（Gen-PPL）由外部 GPT-2 Large 评定，难以全面反映开放文本的真实逻辑自洽性与人类可读性，存在类似图像生成中似然度失真的问题。
未来改进方向：探索端到端联合微调 Stage A 与 Stage B 的稳定性策略；进一步研究在多模态离散 token（如图像离散 codebook、语音离散 token）中的通用适配与大模型预训练整合。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 明确跳出主流轨迹压缩思维定式，从离散序列联合分布表达力的第一性原理出发提出两阶段耦合生成范式。
实验充分度: ⭐⭐⭐⭐☆ 跨越二值图像、生物 DNA 增强子序列和开放文本三大异质领域，且涵盖单步、少步和引导实验。
写作质量: ⭐⭐⭐⭐⭐ 动机清晰，数学论证严格，结构严谨且实验数据翔实。
价值: ⭐⭐⭐⭐⭐ 为离散序列的极速并行生成开辟了极具潜力的新方向，显著降低了可控生成与推理部署的算力成本。