Fixed-Point Masked Generative Modeling¶
会议: NeurIPS 2026
arXiv: 2605.31215
OpenReview: https://openreview.net/forum?id=7Ox69QdIwT
代码: https://github.com/andreamiele/fp-mgm/
领域: 图像生成
关键词: 掩码生成模型, 不动点求解器, 跨步一致性正则, 三状态重用, 离散扩散
一句话总结¶
本文提出面向掩码生成模型的不动点去噪框架 CoFRe，利用参数共享的不动点层替代中间固定深度的 Transformer 堆叠以支持自适应深度，并结合跨步一致性损失与 token 级三状态重用规则，在显著降低参数量和显存的同时大幅提升低前向预算下的生成质量。
研究背景与动机¶
掩码生成模型（Masked Generative Models, MGMs）如面向文本的离散扩散模型 MDLM 和面向图像的离散掩码生成模型 MaskGIT，通过迭代去噪被替换为掩码标记的离散序列实现并行解码，在跨模态生成中展现出强大的表现力。然而，传统 MGM 的去噪主干在每个迭代细化步都需要对整条序列执行完整的双向 Transformer 前向计算，导致训练过程极耗显存且耗时漫长；更为严峻的是，当在低前向预算（即限制 Transformer 块的总前向传播次数）下进行采样时，固定的网络深度使得分配给每步的去噪计算极其僵硬，生成质量出现断崖式下跌。
现有提升 MGM 效率的工作主要集中于改进采样调度策略或设计固定深度的轻量化骨干网络，但它们仍未脱离在每个去噪步为去噪器分配固定计算容量的范式。尽管深层平衡模型（DEQ）与连续不动点扩散模型（FPDM）证实了利用参数共享的不动点层可以解耦网络参数量与有效计算深度，但连续扩散方法无法直接移植到离散掩码生成领域。离散掩码去噪在时间步进过程中存在剧烈且非均匀的 token 级状态跃迁——部分 token 保持可见，部分维持被掩状态，而新揭示的 token 则面临局部条件的突变。
因此，掩码生成模型的不动点化面临两大核心矛盾：一是跨步输入离散突变导致隐表征空间非平滑漂移，低步数下去噪器预测剧烈震荡；二是前序不动点解无法在所有 token 位置无差别重用。本文的核心 idea 是构建完整的端到端训练与推理框架 CoFRe，将掩码去噪器中间层重构为参数共享的不动点求解器，在训练阶段引入跨步一致性正则化拉齐时间轨迹上的隐表征，并在采样阶段根据 token 离散转移类型实施三状态差异化热启动重用。
方法详解¶
整体框架¶
CoFRe 将标准的掩码去噪 Transformer 解耦为四个串联组件：显式预处理堆叠 \(P_{\theta_P}\)、输入条件投影 \(G_{\theta_G}\)、隐式不动点块 \(F_{\theta_F}\) 以及显式后处理堆叠 \(H_{\theta_H}\)。在每个去噪步 \(t\)，模型首先将当前带掩码的离散序列映射为初始隐状态，再通过参数共享的不动点块迭代逼近隐平衡态，最后由后处理层解码为 clean token 的对数几率（logits）。
在整个流程中，训练期采用随机无雅可比反向传播（SJFB）优化不动点参数，并利用嵌套相关掩码构建跨步一致性损失；推理期则依据历史平衡态与当前预处理特征，按 token 转移类型执行三状态热启动重用（3SR）。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["带掩码离散输入序列 z_t 与时间步 t"] --> B["显式预处理堆叠 P 与条件投影 G<br/>提取初始隐表征与注入信号"]
B --> C["不动点求解块 F<br/>权重共享自适应迭代逼近隐平衡态"]
C --> D["显式后处理堆叠 H<br/>映射平衡态隐表征输出去噪 Logits"]
C -.->|"前序平衡态 h*_{t+1}"| E["三状态重用机制 3SR<br/>依可见/仍掩/新揭示划分重用权重"]
E -.->|"热启动初值 h0_t"| C
D --> F["损失监督：基础掩码交叉熵 + 跨步一致性正则"]关键设计¶
1. 不动点掩码去噪网络结构：解耦参数量与去噪计算深度 传统 MGM 依赖固定层数的 Transformer 堆叠，参数量随层数线性增加，无法在不同采样预算下灵活调整各步计算量。本文用参数共享的不动点块替代去噪器中间的冗余层，将去噪器形式化为： $$ \mathbf{h}{\mathrm{pre},t}=P}}(\mathbf{z{t},t),\;\tilde{\mathbf{h}}}=G_{\theta_{G}}(\mathbf{h{\mathrm{pre},t}),\;\mathbf{h}}^{\star}=\operatorname{Fix}!\left(F_{\theta_{F}}(\cdot;\tilde{\mathbf{h}{t},t)\right),\;\boldsymbol{\ell}}(\mathbf{z{t},t)=H,t) $$ 若显式预处理、不动点块和后处理块的 Transformer 层数分别为 }}(\mathbf{h}_{t}^{\star\(K_{\mathrm{pre}}\)、\(K_{\mathrm{fp}}\) 和 \(K_{\mathrm{post}}\)，则在推理时仅需参数化 \(K_{\mathrm{pre}}+K_{\mathrm{fp}}+K_{\mathrm{post}}\) 个独立层，而在迭代求解 \(N\) 步时即可达到等效于 \(K_{\mathrm{pre}}+NK_{\mathrm{fp}}+K_{\mathrm{post}}\) 层的前向深度。在训练阶段，模型使用随机无雅可比反向传播（S-JFB），随机抽取 \(N_{\mathrm{ng}} \sim \mathcal{U}\{0,\dots,4\}\) 步无梯度迭代推进隐状态，紧接着抽取 \(N_{\mathrm{g}} \sim \mathcal{U}\{3,\dots,6\}\) 步带梯度跟踪的迭代计算损失梯度，既避免存储全求解轨迹激活值的巨额显存消耗，又使模型适应不同的迭代逼近深度。
2. 跨步一致性正则化：拉齐时序表征以实现跨时间自蒸馏 在有限前向预算下，掩码去噪步数大幅压缩，导致步间输入跳跃剧烈并引发严重的表征震荡与 logits 漂移。针对中间隐状态缺乏直接平滑时序监督的问题，本文引入跨步一致性损失 \(\mathcal{L}_{\mathrm{CONS}}\)。对于同一 clean 样本 \(\mathbf{x}\)，算法采样两个噪声水平 \(t_c < t_s\) 并使用同一组随机数生成嵌套的相关掩码，确保较干净的教师分支上下文严格包含较噪的学生分支上下文。随后对学生与教师分支在不动点块和后处理块之后的输出隐状态施加均方误差约束： $$ \mathcal{L}{\mathrm{CONS}}=\left|\mathbf{h}}-\mathrm{sg}(\mathbf{h{c})\right| $$ 其中 }^{2\(\mathrm{sg}(\cdot)\) 表示截断梯度。这一损失在表征空间充当跨时间自蒸馏信号，促使去噪器在面对噪声更重、掩码比例更高的学生输入时，能够提前预测出未来更干净阶段的稳定特征，有效抑制滞后 logits 漂移并锐化掩码位置的预测分布。
3. 三状态重用热启动机制：针对离散跃迁差异调制初值 数值求解器在推理时需要一个初始隐状态 \(\mathbf{h}_t^0\)。若直接从前序去噪步的平衡态完全重用（\(\mathbf{h}_t^0 = \mathbf{h}_{t+1}^\star\)），则默认了整条序列的不动点环境平稳；但在掩码生成中，不同 token 位置经历着异构的物理变化。实证分析表明，新揭示 token 的平衡态漂移极大，仍被掩码的 token 随着可见比例增加漂移逐渐收敛，而已保持可见的 token 漂移最小。为此，本文设计基于三状态调制的重用规则： $$ \mathbf{h}{t}^{0}=\boldsymbol{\gamma}}\odot\mathbf{h{t+1}^{\star}+(1-\boldsymbol{\gamma}})\odot\mathbf{h{\mathrm{pre},t} $$ 调制系数 \(\gamma_t^i\) 在隐藏维度广播，按 token 转移类型定义为： $$ \gamma $$ 其中 }^{i}=\begin{cases}1.0,&\text{若位置 } i \text{ 是未变化的可见 token} \ \gamma_{\mathrm{mask,min}}+(\gamma_{\mathrm{mask,max}}-\gamma_{\mathrm{mask,min}})v_t,&\text{若位置 } i \text{ 仍处于掩码状态} \ 0.2,&\text{若位置 } i \text{ 是当前步新揭示的 token}\end{cases\(v_t = \frac{1}{d}\sum_{j=1}^d \mathbb{I}[z_t^j \neq [\textsc{mask}]]\) 表示当前步可见 token 所占比例。该机制让稳定的可见 token 完整继承先验特征以加速收敛，让新揭示 token 更多依赖当前预处理输入注入，显著减少因粗暴完全重用带来的累积初始化误差。
4. 预训练掩码模型短程自适应：层映射与无痛架构迁移 针对无需从头预训练的应用需求，本文提出通过居中层相似性映射将标准 MGM 转换为 FP-MGM。中心核对齐（CKA）特征相似度分析显示，标准 12 层 MDLM 在第 5 层与第 6 层之间存在显著的表征几何跃迁，而第 6 至 12 层之间的表征高度自相似（CKA 达到 0.998）。因此，转换流程将第 1 层映射为显式预处理，抽取第 6 层初始化参数共享不动点块，将第 12 层映射为后处理。随后以原预训练模型为冻结教师，在新架构上施加温度缩放的 KL 蒸馏损失，仅需 40k 步自适应（约原预训练步数的 4%）即可完成高效架构迁移，性能甚至优于 1M 步从头训练的基线模型。
损失函数 / 训练策略¶
在从头训练阶段，总优化目标为基础掩码生成交叉熵损失与跨步一致性正则项的加权和： $$ \mathcal{L} = \mathcal{L}{\mathrm{MGM}} + \lambda \mathcal{L} $$ 其中 }\(\lambda\) 在后训练阶段引入，在前 5k 步从 0 线性预热至 0.1。为防止过度正则化导致模型过锐化和样本生成多样性（熵）崩溃，训练采用验证集困惑度上升 15% 的早停法则。对于检查点转换自适应阶段，优化目标结合基础交叉熵与学生掩码位置上的温度缩放 KL 散度： $$ \mathcal{L}=\mathcal{L}{\mathrm{base}}+\lambda\tau^{2}\frac{1}{|\mathcal{M}}|}\sum_{i\in\mathcal{M{s}}\mathrm{KL}!\left(p(i)\right) $$ 实验设定蒸馏温度 }^{\tau}(i)\,|\,p_{s}^{\tau\(\tau=1.5\)，权重 \(\lambda=0.1\)，噪声步长差距 \(\Delta \sim \mathcal{U}[0.05, 0.30]\)。
实验关键数据¶
主实验¶
论文在文本建模数据集 OpenWebText（上下文长度 1024）和图像生成基准 ImageNette（256×256 分辨率）上系统评估了 CoFRe 的质量与成本权衡。计算预算以 Transformer 块的前向计算次数严格对齐。
| 数据集 / 模态 | 前向预算 | 本文 CoFRe | 强基线 (MDLM+SDTT / MaskGIT) | 相对增益 / 说明 |
|---|---|---|---|---|
| OpenWebText (文本生成) | 96 步 | 101.791 (Gen PPL) | 193.050 (MDLM+SDTT) | 生成困惑度显著下降 47.3% |
| OpenWebText (文本生成) | 192 步 | 65.182 (Gen PPL) | 89.170 (MDLM+SDTT) | 生成困惑度下降 26.9% |
| OpenWebText (文本生成) | 384 步 | 48.755 (Gen PPL) | 62.290 (MDLM+SDTT) | 生成困惑度下降 21.7% |
| OpenWebText (文本生成) | 768 步 | 37.846 (Gen PPL) | 47.040 (MDLM+SDTT) | 生成困惑度下降 19.5% |
| ImageNette (图像生成) | 48 步 | 96.733 (FID) | 174.086 (MaskGIT-Large) | FID 降低 77.35 点，IS 由 9.29 增至 14.41 |
| ImageNette (图像生成) | 96 步 | 51.008 (FID) | 117.644 (MaskGIT-Large) | FID 降低 66.64 点，IS 由 13.37 增至 15.96 |
| ImageNette (图像生成) | 192 步 | 27.624 (FID) | 54.617 (MaskGIT-Large) | FID 减半，生成视觉质量明显占优 |
| ImageNette (图像生成) | 384 步 | 22.838 (FID) | 30.020 (MaskGIT-Large) | 达到全量更深模型的收敛水平 |
主实验资源开销对比：在 OpenWebText 上，CoFRe 较 MDLM 参数量减少 38.8%，训练时间从 ~139h 降至 ~123h，单卡 VRAM 从 112.4 GiB 降至 93.44 GiB；在 ImageNette 上，CoFRe 相比 MaskGIT-Large 训练时间缩短 48.6%（从 17h46m 降至 9h08m），峰值显存降低 50.7%（从 72.45 GiB 降至 35.74 GiB）。
消融实验¶
论文在 OpenWebText 上对 CoFRe 各核心组件进行了系统剥离分析（以下为采样预算 96 与 768 下的 GPT-2 Large 生成困惑度与熵表现）：
| 模型配置 / 消融项 | 预算 96 (PPL / Entropy) | 预算 768 (PPL / Entropy) | 说明 |
|---|---|---|---|
| MDLM 原始基线 | 830.82 / 5.91 | 143.88 / 5.70 | 极低预算下生成近乎崩溃 |
| FP-MDLM (仅不动点骨干) | 375.63 / 5.81 | 179.65 / 5.70 | 低预算暴降 54.8%，但高预算缺乏正则变差 |
| FP-MDLM + 完全重用 (Full Reuse) | 516.68 / 5.82 | 269.01 / 2.21 | 缺乏一致性时高预算出现严重熵坍缩崩溃 |
| FP-MDLM + 三状态重用 (3SR) | 454.10 / 5.80 | 254.26 / 5.38 | 缓解突变冲击但未解决根本 logits 漂移 |
| FP-MDLM + \(\mathcal{L}_{\mathrm{CONS}}\) (仅一致性) | 104.15 / 5.45 | 41.67 / 5.16 | 一致性损失是低预算生成质量跃升的最大推手 |
| 完整 CoFRe (FP-MDLM + \(\mathcal{L}_{\mathrm{CONS}}\) + 3SR) | 101.79 / 5.43 | 37.85 / 5.14 | 一致性与三状态重用高度互补，达到全局最优 |
| 40k 步自适应模型 (Adapted + 3SR) | 286.40 / 5.70 | 126.87 / 5.57 | 仅用 4% 训练量即大幅超越 1M 步基础模型 |
关键发现¶
- 一致性正则化是质量提升的核心引擎：消融数据显示，单纯引入不动点求解器虽然降低了参数量并使低预算性能翻倍，但在未加正则约束时，不动点状态存在累积漂移；加入 \(\mathcal{L}_{\mathrm{CONS}}\) 后，低预算 PPL 从 375.63 锐降至 104.15，滞后 logits KL 散度平均下降 15.2%（在 lag=1 处下降 19.0%），起到了跨时序自蒸馏的关键作用。
- 三状态重用与一致性具有强协同效应：在未经正则化的模型上，盲目全量重用会导致严重的模式坍塌（如在预算 768 步时熵骤降至 2.209）；而在经过 \(\mathcal{L}_{\mathrm{CONS}}\) 约束后，状态重用转变为正向赋能，3SR 凭借对未变化可见 token（\(\gamma=1.0\)）、仍掩码 token（随可见比例动态插值）与新揭示 token（\(\gamma=0.2\)）的精准区分，全面优于全量重用与无重用。
- 预训练权重自适应的高效性：CKA 相似度证实了标准 Transformer 深度方向存在明显的两阶段结构（前 5 层为特征提取，后 7 层高度自相似）。以此结构先验进行层映射再微调，仅需 40k 步就能达到 126.87 的高预算 PPL，比 1M 步从头训练的模型高出近 50 点。
亮点与洞察¶
- 离散掩码状态跃迁与数值求解热启动的深度解耦：连续扩散模型中的不动点技术因无法应对离散 token 剧烈突变而碰壁。本文通过精细统计不同 token 转移类别的平衡态位移，巧妙设计了三状态加权调制，将连续动态系统的隐式求解平滑地桥接到了离散序列生成中。
- 表征空间一致性替代复杂的输出端蒸馏：以往加速采样往往依赖复杂的多步蒸馏策略（如 SDTT）。本文证明仅需在同一序列的嵌套相关掩码间约束隐平衡态的 MSE 均方误差，便能在无外部复杂教师的前提下实现跨步对齐，大幅抑制步间 logits 漂移。
- 深层冗余性的分析指导架构重构：利用 CKA 深入剖析预训练双向 Transformer 各层表征的内在几何分布，以科学的数据支撑而非盲目调参来确定显式层与不动点层的切割边界，为大型基础模型的参数共享与瘦身提供了通用范式。
局限与展望¶
- 实验尺度的局限性：作者承认当前验证主要集中在 OpenWebText 级别的文本建模和 ImageNette 图像生成基准上，尚未扩展至百亿参数级的大型语言模型或超大规模文生图/视频等超大计算场景。
- 超参数敏感与早停依赖：\(\mathcal{L}_{\mathrm{CONS}}\) 后训练阶段容易引发过锐化导致生成多样性受损，当前依赖于验证集困惑度上升 15% 的经验性早停策略，且三状态系数仍需网格搜索，缺乏自适应动态学习机制。
- 非单调去噪轨迹的泛化疑问：三状态重用规则是专门为单调渐进揭示的掩码策略量身定制的，对于包含重新加掩（remasking）或双向修正机制的非单调采样器，状态转移更为复杂，需要设计更具弹性的重用准则。
相关工作与启发¶
- vs MDLM / MaskGIT: 经典掩码生成模型每个细化步必须完整前向穿透所有固定层，计算开销大且低预算下不可调。CoFRe 通过不动点块实现自适应深度分配，大幅节约显存与训练时间，在相同计算预算下生成质量显著胜出。
- vs SDTT (Self-Distillation Through Time): SDTT 针对离散模型训练多步至单步的压缩蒸馏，流程繁琐且依赖多阶段学生更新；CoFRe 采用紧凑的隐表征均方误差与嵌套掩码设计，充当轻量级跨步自蒸馏，训练流程更简洁。
- vs FPDM (Fixed-Point Diffusion Models): FPDM 面向连续高斯扩散过程，状态变化连续平滑；CoFRe 针对离散掩码模型的突变痛点，首创三状态重用策略与掩码自适应系数调制，解决了离散条件切换导致的不动点发散难题。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将不动点求解器深度适配到离散掩码生成模型，三状态重用与跨步一致性设计精巧。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖文本与图像双模态、从头预训练与短程自适应、细致的消融实验、CKA 特征几何分析以及滞后 logits 统计。
- 写作质量: ⭐⭐⭐⭐⭐ 结构清晰，逻辑严密，图表信息量极其丰富，问题剖析一针见血。
- 价值: ⭐⭐⭐⭐⭐ 为轻量化、自适应深度的离散扩散与掩码并行生成提供了全新且极具实用价值的范式。