PaperlayerView exact revision on arXiv

Source-linked AI summary

MARS: Enabling Autoregressive Models Multi-Token Generation

Ziqi Jin, Lei Wang, Ziwei Luo, Aixin Sun

arXiv:2604.07023v1cs.CL

TL;DR

AR models spend one forward pass per token even when continuations are predictable. MARS fine-tunes an instruction-tuned AR model for optional multi-token prediction without architectural additions, preserving AR behavior while improving throughput. It achieves baseline-level accuracy with 1.5–1.7× throughput and up to 1.71× wall-clock speedup with block-level KV caching, subject to training and serving limitations.

  • Problem

    AR decoding assigns one forward pass to every token, motivating a way to exploit predictable continuations without sacrificing standard AR behavior.

  • Method

    MARS is lightweight fine-tuning that preserves causal AR-compatible design choices while teaching an instruction-tuned AR model to optionally predict multiple tokens per forward pass.

  • Results

    MARS matches or exceeds AR quality in one-token mode and achieves 1.5–1.7× throughput in multi-token mode, with up to 1.71× wall-clock speedup using block-level KV caching.

  • Takeaways & Limitations

    A single MARS checkpoint can retain standard AR serving while enabling confidence-controlled acceleration without swapping models or restarting.

  • Takeaways & Limitations

    The 7B evaluation uses only B=4, aggressive thresholds can cause substantial quality loss, and block-level caching requires synchronization that limits gains at large batch sizes.

Abstract

from arXiv · show

Autoregressive (AR) language models generate text one token at a time, even when consecutive tokens are highly predictable given earlier context. We introduce MARS (Mask AutoRegreSsion), a lightweight fine-tuning method that teaches an instruction-tuned AR model to predict multiple tokens per forward pass. MARS adds no architectural modifications, no extra parameters, and produces a single model that can still be called exactly like the original AR model with no performance degradation. Unlike speculative decoding, which maintains a separate draft model alongside the target, or multi-head approaches such as Medusa, which attach additional prediction heads, MARS requires only continued training on existing instruction data. When generating one token per forward pass, MARS matches or exceeds the AR baseline on six standard benchmarks. When allowed to accept multiple tokens per step, it maintains baseline-level accuracy while achieving 1.5-1.7x throughput. We further develop a block-level KV caching strategy for batch inference, achieving up to 1.71x wall-clock speedup over AR with KV cache on Qwen2.5-7B. Finally, MARS supports real-time speed adjustment via confidence thresholding: under high request load, the serving system can increase throughput on the fly without swapping models or restarting, providing a practical latency-quality knob for deployment.

1 Introduction

MARS addresses the inefficiency of one-token-at-a-time AR decoding by using lightweight fine-tuning to enable confidence-based multi-token generation without auxiliary deployment components. It preserves AR quality while enabling faster inference and batch wall-clock speedups.

  • Results: MARS preserves baseline quality in one-token mode and achieves 1.5–1.7× throughput when multi-token generation is enabled.The model adaptively batches predictable continuations while retaining token-by-token generation for novel content.
  • Analysis: MARS analyzes four AR-to-block-diffusion gaps and identifies three as eliminable design choices rather than inherent limitations.Closing these gaps is reported as sufficient to recover baseline quality without architectural changes.
  • Method: MARS uses lightweight fine-tuning to let an instruction-tuned AR model predict multiple tokens per forward pass without architectural changes or additional parameters.The same SFT data is reused, and the resulting model remains callable like the original AR model.
  • Training: An auxiliary SFT loss on the clean input stream preserves performance at larger block sizes, where the AR-like training signal would otherwise decay.This loss is identified as the key ingredient for maintaining AR behavior during block prediction training.
  • Batch inference: A block-level KV caching strategy delivers wall-clock speedups of up to 1.71× over AR with KV cache on Qwen2.5-7B.The strategy targets batch inference, where algorithmic token speedup alone does not determine production throughput.

2 Background and Related Work

Block-masked prediction offers a route to parallel generation with the same backbone and language-model head, but it creates mismatches with standard AR training and inference. MARS minimizes these mismatches by retaining causal behavior, AR logit alignment, and left-to-right generation.

  • Background: AR decoding requires T serial forward passes to generate T tokens, while block-masked prediction trains one forward pass to recover a contiguous block of B masked future tokens.The masked sequence is conditioned on clean tokens from preceding blocks.
  • Background: Direct AR-to-block-diffusion conversion can degrade quality because partially masked inputs, altered attention patterns, and different generation orders mismatch AR training.The paper treats this mismatch as the motivation for its gap analysis.
  • Related work: Prior multi-token methods such as MTP, Medusa, and EAGLE require additional prediction heads, parameters, or architectural modifications.MARS instead pursues multi-token prediction through fine-tuning without extra heads.
  • Related work: Jacobi and Lookahead decoding use fixed-point iteration from initialized future tokens, whereas MARS uses [MASK] tokens and fine-tuning for multi-token generation.The approaches are presented as complementary strategies.
  • MARS positioning: MARS retains causal attention, right-shifted logits, and strictly left-to-right generation, leaving token masking as the sole inherent gap.The method aligns with AR on gaps (2)–(4).

3 Method

MARS preserves a fully functional autoregressive model while adding masked multi-token prediction through lightweight fine-tuning. Its causal attention, right-shifted logits, left-to-right acceptance, and combined clean/noisy training objective retain AR competence while enabling confidence-controlled acceleration.

  • Design principle: MARS closes three eliminable gaps—attention pattern, logits alignment, and generation order—while retaining token masking as the sole inherent difference from AR generation.It uses causal attention everywhere, right-shifted logits, and strictly left-to-right token acceptance.
  • Training setup: The model processes clean and masked copies of each sequence in one shared forward pass, training standard AR prediction alongside masked-block prediction.The noisy stream replaces blocks of B tokens with [MASK] placeholders and uses clean prefixes from earlier blocks as context.
  • Training setup: The structured attention mask gives clean tokens causal visibility, noisy tokens causal visibility within their block, and noisy blocks access to clean tokens from preceding blocks.This visibility pattern supplies prefix context while preserving causal intra-block behavior.
  • Preserving autoregressive competence: 25% for B=4, 12.5% for B=8, and 6.25% for B=16 are the fractions of masked-training positions with fully clean context before adding the clean-stream loss.This fraction decreases with block size, and the passage reports degradation on reasoning and coding tasks at larger block sizes.
  • Preserving autoregressive competence: 62.5% for B=4 and 53.1% for B=16 are the combined-loss AR-equivalent signal fractions, remaining above 50% as block size grows.The clean-stream SFT loss prevents the AR training signal from decaying with block size; the default model also has a variant trained without this loss.
  • Inference: Confidence thresholding lets MARS accept tokens consecutively from the left, append replacements to maintain a B-token window, and adjust throughput during serving without retraining or model swapping.As τ approaches 1.0, the procedure accepts at most one token per step and recovers exact AR behavior; lower τ accepts more tokens at some quality cost.

4 Experiments

Experiments show that MARS preserves one-token AR quality, stabilizes performance across larger blocks with an auxiliary SFT loss, and enables controllable multi-token and wall-clock speedups. Block-level KV caching is essential for translating token-level gains into batch-inference acceleration.

  • 4.2 MARS Preserves AR Quality in One-Token Mode: MARS matches or exceeds AR SFT in one-token mode across six benchmarks at both model scales.At 0.5B, MARS scores 30.4 average versus 28.7; at 7B, it scores 58.1 versus 56.6.
  • 4.2 MARS Preserves AR Quality in One-Token Mode: 10-epoch compute-matched AR SFT falls from 28.7 to 26.4 average, whereas MARS gains from masked prediction rather than extra training alone.The continued-AR baseline also drops on MMLU-Pro from 11.9 to 9.3 and GSM8K from 32.0 to 28.3.
  • 4.3 Why Larger Blocks Work: Validating the Signal Decay Hypothesis: 6.2 points: without SFT loss, increasing block size from 4 to 16 drops average performance from 28.4 to 22.2.The degradation is concentrated in reasoning and coding tasks, including GSM8K and HumanEval.
  • 4.3 Why Larger Blocks Work: Validating the Signal Decay Hypothesis: 0.7 points: with SFT loss, the same block-size increase reduces average performance only from 30.4 to 29.7.GSM8K improves from 32.8 to 33.8, while HumanEval decreases from 40.2 to 36.6; the full Pareto frontier also improves.
  • 4.4 A Smooth Speed–Quality Frontier via Confidence Thresholding: 1.68 tokens per forward pass: at 7B, multi-token mode loses 1.3 average points while retaining 56.8 versus the AR baseline’s 56.6.The threshold-controlled frontier provides a smooth speed–quality tradeoff, with no model swap or retraining required.
  • 4.5 Wall-Clock Speedup with Block-Level KV Cache: 1.71× wall-clock speedup: block-level KV caching reduces batch-size-4 completion time to 161.2s versus AR’s 276.2s.MARS also achieves 1.60× at batch size 8 and 1.34× at batch size 16, while accuracy remains within 2–3 points of AR.

5 Conclusion

MARS enables multi-token generation through lightweight fine-tuning while preserving AR capability, but its evaluation and serving trade-offs leave concrete boundaries for deployment and future work.

  • MARS remains a strict superset of AR, matching or exceeding baseline quality in one-token mode and reaching 1.5–1.7× throughput in multi-token mode.Block-level KV caching provides up to 1.71× wall-clock speedup over AR with KV cache on Qwen2.5-7B.
  • MARS closes three eliminable gaps between AR and block-masked prediction while retaining the inherent masking gap.The method preserves causal intra-block attention, right-shifted logits, and strict left-to-right generation.
  • The clean-stream SFT loss preserves AR competence during masked-prediction training, keeping the AR signal above 50% regardless of block size.
  • At 7B, only B=4 is evaluated, so whether the observed block-size speed–quality pattern holds at larger scales remains future work.The 0.5B experiments show similar Pareto frontiers across block sizes, with B=4 marginally ahead by less than 0.5pp in average accuracy.
  • MARS training doubles effective sequence length, while aggressive thresholds and block-boundary synchronization constrain efficiency and throughput.Aggressive thresholds τ < 0.7 cause substantial quality loss, and synchronization limits gains at large batch sizes.

A Training Details

MARS uses a clean-to-MARS training pipeline with matched hyperparameter settings, but its concatenated data construction substantially increases training cost over AR SFT.

  • Both AR SFT and MARS stages use identical settings per model size, with matched effective batch sizes across scales.
  • MARS doubles effective sequence length by concatenating a clean and noisy copy of each sequence, increasing training cost relative to standard SFT.Training costs are 33 versus 15 H200-hours at 0.5B and 202 versus 100 H200-hours at 7B; peak GPU memory is approximately 1.5× AR SFT.

B Threshold Sweep Details

The threshold sweep shows that MARS trades acceptance speed against accuracy: moderate thresholds increase tokens per forward pass with limited quality loss, while lower thresholds degrade quality.

  • At τ=0.95, MARS with B=4 reaches 1.51 tokens per forward on GSM8K with only 1.7pp accuracy loss relative to τ=1.0.
  • Lowering τ accepts more tokens but may reduce accuracy, especially for larger block sizes.Table 7 sweeps τ from 1.0, equivalent to one token per step, down to 0.5.

C Jacobi Decoding Baseline

The Jacobi baseline provides lower speedup than MARS and benefits from knowing the target length, but that advantage is task-dependent.

  • Jacobi initializes all future positions simultaneously and iteratively updates them with causal forward passes until convergence.
  • 1.07× average speedup for Jacobi decoding is lower than MARS’s 1.46× on the same AR SFT checkpoint.
  • Random-token prefixes limit Jacobi convergence, whereas MARS trains on [MASK] placeholders and consequently obtains higher acceptance rates.
  • Jacobi’s fixed generation length improves GSM8K and HumanEval scores over AR SFT, but the advantage does not apply uniformly across format-sensitive and reasoning tasks.The cited comparison reports 36.5 versus 32.0 on GSM8K and 42.1 versus 35.4 on HumanEval.

D Acceptance Metric Sensitivity

MARS’s sliding-window inference accepts tokens left-to-right while confidence exceeds a threshold, and this speed–quality trade-off is robust across confidence metrics. Entropy and top-2 margin degrade slightly more gracefully than raw probability, but differences are small.

  • Tokens are accepted left-to-right while a confidence score exceeds threshold τ; the main experiments use the top-token probability.The confidence score is maxv p(v | ·).
  • Acceptance metrics: Entropy accepts while H ≤τ, with lower entropy indicating higher confidence.
  • Acceptance metrics: Top-2 margin accepts while ptop1 −ptop2 ≥τ; a larger gap indicates higher confidence.
  • Speed–quality comparison: All three metrics trace similar Pareto curves, confirming robustness to the specific confidence measure.The comparison uses MARS with B=4 on GSM8K.
  • Speed–quality comparison: Entropy and top-2 margin show marginally smoother degradation at comparable speedups, while differences remain small.Probability is used in the main paper for simplicity.
Loading 2604.07023v1…