MARS: Enabling Autoregressive Models Multi-Token Generation

会议: NeurIPS 2026
arXiv: 2604.07023
代码: https://github.com/Xalp/MARS
领域: LLM效率
关键词: 自回归模型, 多Token生成, 块掩码预测, 投机解码, 效率优化

一句话总结

MARS 提出一种无需改动架构与额外参数的轻量微调方法，通过在因果注意力与右移 Logits 下联合优化干净流与块掩码流，使标准自回归模型在保持单 Token 生成能力的同时获得可调节的多 Token 并行生成能力，实现高达 1.71× 的端到端推理加速。

研究背景与动机

标准自回归（AR）语言模型在推理时对每一个输出 Token 消耗均等的计算量：无论预测的是确定性极高的固定转折词，还是真正需要开放探索的逻辑推理关键词，模型都必须执行一次完整的单步前向传播。这种严格串行的逐 Token 解码机制不仅造成大量冗余算力浪费，而且使服务吞吐量直接受到输出长度的严密制约。

为了打破这一串行瓶颈，现有的多 Token 生成范式普遍需要重构推理部署栈。以投机解码为代表的方法需要部署独立的轻量草稿模型，造成显存占用成倍增加与编排调度复杂度骤升；而 Medusa 和 EAGLE 等多头预测方案则要求为骨干网络额外附加多组专用预测头并单独训练，破坏了基础架构的通用性。另一类基于离散块扩散（Block Diffusion）的方法试图直接训练模型预测块内多个 Token，但在数学推理与代码生成任务上往往遭遇灾难性的性能崩塌。

深入剖析 AR 模型与块扩散之间的本质差异可以发现，现有方案退化的根源并非多 Token 生成本身，而是人为引入了破坏因果性与自回归对齐的冗余改动——即块内双向注意力破坏因果掩码、修改 Logits 对齐约定、以及无序填空破坏自左向右生成顺序。本文的核心 idea 是：完全保留因果注意力、右移 Logits 与自左向右解码顺序，仅保留必要的块内 [MASK] 占位符，并通过干净流辅助 SFT 损失在整个序列中解耦 AR 监督信号衰减，使单个模型在完全不退化原 AR 能力的前提下自由切换多 Token 加速。

方法详解

整体框架

MARS 的设计宗旨是让多 Token 生成成为现有指令微调自回归模型的附加增益能力，而非替代模型原有的自回归特性。在训练阶段，输入序列被复制为两路拼接序列并行输入：前半部分为保留完整原始 Token 的干净流，执行标准单步因果因果预测；后半部分为将固定大小块内 Token 替换为 [MASK] 占位符的掩码流。两者通过专门构造的二维分块注意力掩码矩阵实现单次前向传播，保证块内与跨流的严格因果可见性。

在推理阶段，MARS 采用自左向右的滑动窗口机制。生成光标后追加一组 [MASK] 占位符并通过单次前向计算获得所有位置预测分布，随后从左至右连续采纳概率高于置信度阈值的 Token 并滑入历史前缀。配合针对批处理推理设计的块级 KV 缓存管理，MARS 实现了算法步数加速向实际墙上时钟加速的高效转化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 x<br/>(指令与回复上下文)"] --> B["双流并行拼接与因果掩码构建<br/>clean: x 与 noisy: [MASK] 块拼接"]
    B --> C["双流联合损失训练<br/>L = L_mask + L_AR 稳定因果信号"]
    C --> D["自左向右滑动窗口解码<br/>置信度阈值 tau 连续接纳"]
    D --> E["块级 KV 缓存批推理加速<br/>同批次块边界同步更新"]

关键设计

1. 消除三大设计裂痕与纯因果双流注意力掩码：让块预测完全兼容自回归结构

针对以往块扩散方法在微调中破坏 AR 模型基础因果先验导致推理崩塌的痛点，MARS 将 AR 与块扩散之间的四个差异进行解耦： Token 遮蔽（Token Masking）作为并行预测的物理要求是唯一不可避免的固有差异，而注意力模式（因果 vs 双向）、Logits 对齐（右移 vs 各异）以及生成顺序（从左到右 vs 乱序）皆为应当消除的人为偏差。MARS 在结构上全面向 AR 看齐，设计了维度为 \(2L \times 2L\) 的复合注意力掩码 \(\mathbf{M} \in \{0, -\infty\}^{2L \times 2L}\)，作用于由干净流 \(\mathbf{x}\) 与掩码流 \(\tilde{\mathbf{x}}\) 拼接而成的向量 \(\mathbf{z} = [\mathbf{x}; \tilde{\mathbf{x}}]\)：

\[ M_{ij} = \begin{cases} 0 & \text{if } i, j \leq L, \; j \leq i \quad \text{(干净流标准因果自注意力)} \\ 0 & \text{if } i, j > L, \; \beta(i-L) = \beta(j-L), \; j \leq i \quad \text{(掩码流块内单向因果注意力)} \\ 0 & \text{if } i > L, \; j \leq L, \; \beta(j) < \beta(i-L) \quad \text{(掩码流向前序干净块跨流可见)} \\ -\infty & \text{otherwise} \end{cases} \]

其中 \(\beta(t) = \lceil t/B \rceil\) 为块索引。该掩码保证了干净流在前半区完全复现标准 AR 训练，掩码流在后半区仅允许各位置单向查看本块内先前的 [MASK] 占位符，且仅向前跨流读取先于当前块的历史干净块内容。这不仅保持了因果注意力的严格左向流动与右移 Logits 约定，更彻底杜绝了位置泄漏与未来依赖。

2. 干净流辅助 SFT 损失：打破大块尺寸下自回归监督信号的衰减困境

单纯使用掩码损失 \(\mathcal{L}_{\text{mask}}\) 会随块大小 \(B\) 的增大而面临 AR 监督信号被快速稀释的问题。在单块长度为 \(B\) 的序列中，仅有块内第 1 个位置 \(t=1\) 能够观察到全干净的历史上下文，其余位置 \(t=2,\dots,B\) 均被迫依赖含有部分未知 [MASK] 的残缺上下文进行预测。若仅依赖掩码流，拥有完全干净上下文的 AR 信号比例仅为 \(r_{\text{AR}}^{\text{(mask only)}} = 1/B\)。当 \(B=16\) 时，有效 AR 监督比例暴跌至 6.25%，导致模型在复杂代码和数学推理任务中因丧失基础因果建模能力而显著退化。

为了从根本上解耦块大小与 AR 能力的负相关性，MARS 在同一次前向传播中无缝复用干净流前半段的 Logits，引入等权重的辅助自回归损失：

\[ \mathcal{L} = \mathcal{L}_{\text{mask}} + \mathcal{L}_{\text{AR}} = -\sum_{t \in \mathcal{M}} \log p_\theta(x_t \mid \tilde{\mathbf{x}}, \mathbf{x}_{<\beta(t)}) - \sum_{t=1}^L \log p_\theta(x_t \mid \mathbf{x}_{<t}) \]

引入该损失后，整体序列获得的等效 AR 信号占比提升至 \(r_{\text{AR}}^{\text{(combined)}} = \frac{L + L/B}{2L} = \frac{1 + 1/B}{2}\)。即使在极大的块尺寸 \(B=16\) 下，有效 AR 信号占比依然稳定保持在 53.1%（始终大于 50%）。这使得模型在学会依据不完整掩码上下文并行推断未来 Token 的同时，永久保留强悍的下一 Token 基础预测能力。

3. 自左向右自适应滑动窗口推理：提供平滑连续的延迟-质量服务旋钮

传统的块级无序填空生成在处理复杂语言逻辑时极易产生前后语义矛盾，而投机解码则需要频繁回退草稿并重新纠错。MARS 摒弃复杂的非自回归验证机制，在推理阶段每步向已有前缀后追加 \(B\) 个 [MASK] 占位符并执行单步前向传播。算法严格自左向右顺序核验最高置信度概率 \(\max_v p(x_t = v)\)，只要其大于等于预设阈值 \(\tau\) 便予以连续接纳，直至遇到首个低于阈值的位置；同时设立底线约束，强制每步至少采纳第 1 个 Token。随后已接纳的 \(N\) 个真实 Token 移入历史前缀，窗口右移并补齐新的 [MASK] 占位符。

这种自适应判定机制赋予了部署系统极大的灵活性：阈值 \(\tau\) 成为在线服务中无需重载模型或停机重启即可动态调节的「油门」。当系统负载平稳且注重精确度时，设 \(\tau=1.0\) 严格退化为逐字生成的纯 AR 模式，质量与原生模型完全持平甚至更优；当瞬时并发请求飙升时，系统可动态调小 \(\tau\)（如 0.95 或 0.90）立即换取 1.5–1.7× 的加速吞吐，且质量无断崖式跌落。

4. 块级 KV 缓存管理机制：消除批处理全序列重计算与多样本步调冲突

针对标准自回归依靠逐 Token 增量 KV 缓存而块生成缺乏固定单步边界的问题，直接无缓存推理会引发 \(O(T^2)\) 的灾难性序列重计算，导致并发批处理吞吐急剧恶化。MARS 设计了细粒度的块级 KV 缓存机制：以预设的块缓存大小 \(B_{\text{cache}}\)（如 16 或 32）为同步基准，前缀 KV 缓存在跨块时仅通过单次完整前向计算一次性冻结；在块内进行内部滑动预测时，前向计算仅需输入当前待预测的 \(B\) 个 [MASK] 向量并挂载前缀缓存。

针对多样本批推理中不同样本接纳长度不一致的问题，批处理状态由当前步所有样本中接纳 Token 数的最小值进行保守同步扩展，进度较快的样本在块边界短暂休眠等待，直至全批次均填满该块后再统一将整个块写入持久前缀缓存。这种设计有效均摊了前缀重计算代价，彻底释放了 GPU 显存带宽与并行计算潜能。

实验关键数据

主实验

论文在 Qwen2.5-0.5B-Instruct 与 Qwen2.5-7B-Instruct 两个模型尺度上，评测了包括 IFEval（指令遵循）、BBH（逻辑推理）、MMLU-Pro（综合知识）、GPQA（专业问答）、GSM8K（数学推理）和 HumanEval（代码生成）在内的 6 大基准。

在单 Token 评估模式下（\(\tau=1.0\)），MARS 在零质量牺牲的前提下甚至小幅超越原始 AR SFT 基线；而在多 Token 模式下（\(\tau=0.95\)），模型实现了极低损耗的高吞吐生成。

模型尺度与方法	运行模式 (阈值 / 块大小)	IFEval	BBH	MMLU-Pro	GPQA	GSM8K	HumanEval	平均分	单步平均Token数
Qwen2.5-0.5B
Base (无微调)	AR 单步	27.8	6.1	2.4	13.6	27.4	26.2	17.2	1.00
AR SFT (5 ep)	AR 单步	48.4	26.3	11.9	17.9	32.0	35.4	28.7	1.00
AR SFT (10 ep 等算力)	AR 单步	47.8	26.3	9.3	14.1	28.3	32.3	26.4	1.00
Block Diffusion (\(B=4\))	块生成单步	47.1	7.5	2.0	17.9	30.6	31.7	22.8	1.00
MARS w/o SFT loss (\(B=4\))	单步 (\(\tau=1.0\))	48.5	27.4	12.3	19.0	29.5	33.5	28.4	1.00
MARS (本文完整版, \(B=4\))	单步 (\(\tau=1.0\))	51.3	26.6	12.4	19.4	32.8	40.2	30.4	1.00
MARS (本文完整版, \(B=4\))	多Token (\(\tau=0.95\))	45.5	26.6	11.5	19.9	31.1	39.6	29.0	1.46
Qwen2.5-7B
Base (无微调)	AR 单步	63.4	24.0	10.5	11.6	57.3	82.9	41.6	1.00
AR SFT	AR 单步	67.0	54.0	43.9	27.5	68.7	78.7	56.6	1.00
MARS (本文完整版, \(B=4\))	单步 (\(\tau=1.0\))	68.2	54.6	44.4	26.6	73.2	81.7	58.1	1.00
MARS (本文完整版, \(B=4\))	多Token (\(\tau=0.95\))	63.0	54.3	44.2	27.7	71.0	80.5	56.8	1.68

在实际批推理墙上时间评测中（GSM8K 256题，Qwen2.5-7B，\(\tau=0.95\)），块级 KV 缓存展现出决定性的加速效能：

并发 Batch Size	评估方法	缓存粒度 \(B_{\text{cache}}\)	生成速度 (tok/s)	运行总耗时 (s)	相比 AR KV-Cache 加速比
Batch=4	AR (标准 KV 缓存)	–	143.4	276.2	1.00×
Batch=4	MARS (无缓存重计算)	–	127.1	306.5	0.90×
Batch=4	MARS + 块级缓存	32	241.9	161.2	1.71×
Batch=8	AR (标准 KV 缓存)	–	236.7	169.1	1.00×
Batch=8	MARS + 块级缓存	32	368.4	105.6	1.60×
Batch=16	AR (标准 KV 缓存)	–	434.3	91.8	1.00×
Batch=16	MARS + 块级缓存	16	566.0	68.7	1.34×

消融实验

消融实验深入验证了辅助 SFT 损失在不同块尺寸扩展下的决定性支撑作用（0.5B 模型，\(\tau=1.0\)）：

模型配置	块尺寸 \(B\)	IFEval	BBH	MMLU-Pro	GPQA	GSM8K	HumanEval	全局平均分	说明
AR SFT 基线	–	48.4	26.3	11.9	17.9	32.0	35.4	28.7	标准 5 轮次 SFT
MARS 无 SFT 损失	\(B=4\)	48.5	27.4	12.3	19.0	29.5	33.5	28.4	仅掩码损失，信号占比 25%
MARS 无 SFT 损失	\(B=8\)	48.5	24.3	11.1	20.8	24.9	28.7	26.4	信号衰减至 12.5%，数学显著滑坡
MARS 无 SFT 损失	\(B=16\)	42.6	21.7	10.6	16.3	21.0	20.7	22.2	信号仅剩 6.25%，推理代码全面崩溃
MARS 含完整 SFT 损失	\(B=4\)	51.3	26.6	12.4	19.4	32.8	40.2	30.4	信号比例提升至 62.5%
MARS 含完整 SFT 损失	\(B=8\)	49.6	26.9	12.0	19.6	32.8	37.2	29.7	性能极稳，数学不掉点
MARS 含完整 SFT 损失	\(B=16\)	50.7	27.0	12.1	17.9	33.8	36.6	29.7	信号稳定在 53.1%，大块能力完全保全

关键发现

• SFT 损失是抑制大块退化的生命线：若缺乏干净流辅助 SFT 损失，随着块尺寸从 4 扩展至 16，模型在 GSM8K 与 HumanEval 上的表现分别断崖式下跌 8.5 与 12.8 点，均分下挫 6.2 点；而加入 SFT 损失后，大块 \(B=16\) 均分依然保持在 29.7，与 \(B=4\) 的 30.4 几乎持平，甚至在 GSM8K 上反超至 33.8。
• 架构一致性远胜于盲目打破因果：采用双向块内注意力的 Block Diffusion 在 BBH 和 MMLU-Pro 上分别崩塌至 7.5 和 2.0，甚至劣化至未经指令微调的 Base 模型水准；而 MARS 坚守因果掩码和单向生成，不仅推理能力毫发无损，甚至由于掩码预测提供了天然的数据正则化，在单步模式下比原 SFT 在 0.5B 和 7B 上分别提升 1.7 与 1.5 点。
• 格式严格型任务对阈值更为敏感：在多 Token 模式下，绝大部分推理任务在 \(\tau=0.95\) 时掉点极微（通常 \(<1\) 点），而在 IFEval 这种需要精确计算段落数量和严格排版的格式约束任务上掉点约 5 点，说明模型在并行跳进时容易跨过排版关键标点。

亮点与洞察

• 极简主义与纯粹的超集扩展：不引入额外的辅助小模型、不增加任何新的前向预测头、不修改任何模型参数量，仅通过训练时双流拼接与因果掩码微调，就让现成的单一自回归检查点拥有了任意切换的「单步保真」与「多步加速」双重形态。
• 训练端因果流信号占比理论的提出：论文精确揭示了块扩散方法在自回归迁移中遭遇性能滑坡的数学机理（有效 AR 上下文比例按 \(1/B\) 快速衰减），并通过构造等长干净流将因果信号占比永久锁定在 50% 以上，为后续非自回归与自回归混合训练提供了通用的设计准则。
• 生产环境零代价动态负载自适应：在实际服务编排中无需准备多套针对不同负载水平的模型容器，通过在线动态调节每个请求的置信度阈值 \(\tau\)，即可瞬间完成吞吐量与精度的连续平滑权衡。

局限与展望

• 训练序列长度翻倍开销：由于在训练阶段必须沿长度维度拼接干净流与掩码流（\([ \mathbf{x} ; \tilde{\mathbf{x}} ]\)），有效输入长度扩大至 \(2L\)，导致单步前向反向计算耗时约为普通 SFT 的 2.0–2.2×，显存峰值约上升 1.5×。
• 激进加速阈值下的质量劣化：当置信度阈值设置得过于激进（如 \(\tau < 0.7\)）时，生成的文本质量会出现较大幅度衰退，说明基于单 Token 贪心概率的左向连续截断机制在低置信场景下仍存在改进空间。
• 批处理块边界同步的吞吐上限：在极大规模 Batch Size 场景下，由于不同样本生成速度存在差异，较快样本在块边界处的等待空转会削弱加速比（Batch=16 时加速比从 1.71× 回落至 1.34×），未来有待结合动态游标异步无锁缓存机制进行深度解耦。

相关工作与启发

• vs Block Diffusion：Block Diffusion 引入了块内双向自注意力与基于置信度的乱序填空，彻底破坏了预训练 AR 模型的因果先验与 Logits 映射，导致复杂推理能力大幅暴跌；MARS 坚持纯因果注意力掩码与严格自左向右接纳，使推理精度完全不退化。
• vs Medusa / EAGLE：Medusa 与 EAGLE 通过挂载多套额外的轻量预测头来实现多 Token 预测，引入了额外的网络参数、特定的前向图改造与多头同步训练负担；MARS 完全复用原本单一的语言模型预测头，零参数新增，且兼容任何现有部署引擎。
• vs 传统投机解码 (Speculative Decoding)：投机解码需要常驻一套独立的较小 Draft 模型并维持复杂的交互通信流水线，带来了翻倍的显存占用与调度痛点；MARS 依托单模型自提速，无需维护第二套模型权重。
• vs Jacobi / Lookahead 解码：Jacobi 解码采用全序列随机初始化并依赖不动点迭代收敛，由于模型未在错误前缀上微调过，平均每前向步仅能接纳 1.07 个 Token；MARS 引入明确的 [MASK] 标记并经专门微调，在相同基线上达到 1.46 的单步接纳率。

评分

• 新颖性: ⭐⭐⭐⭐☆ [敏锐发现并消除了块扩散与 AR 之间的三大冗余人为裂痕，以极简因果双流方案实现免加头多 Token 生成]
• 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 0.5B/7B 尺度、6 大权威基准、算力对齐基线、消融实验以及细致的块级 KV 缓存实际吞吐实测]
• 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密收敛，动机清晰坦诚，公式节制且直击机理，图文对照高度自洽]
• 价值: ⭐⭐⭐⭐⭐ [为当前大模型线上服务部署提供了一种无需改动架构与参数、低成本微调即可获得弹性加速的极高实用价值方案]

相关论文

• [NeurIPS 2026] Multi-Token Residual Prediction
• [NeurIPS 2026] Orthrus: Memory-Efficient Parallel Token Generation via Dual-View Diffusion
• [ICLR 2026] Self-Speculative Decoding Accelerates Lossless Inference in Any-Order and Any-Subset Autoregressive Models
• [NeurIPS 2026] Dynamic Delayed Tree Expansion for Improved Multi-Path Speculative Decoding
• [NeurIPS 2026] N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation