MARS: Enabling Autoregressive Models Multi-Token Generation

Abstract: Autoregressive (AR) LLMs 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.

• The paper introduces a fine-tuning method that enables multi-token generation in AR models without modifying the underlying architecture.
• It uses a dual-stream training strategy, combining a clean stream with a masked block prediction to balance conventional AR loss and maintain generation order.
• Experimental results show up to 1.71× speedup in throughput with stable output quality, offering adjustable latency–quality tradeoffs for deployment.

MARS: A Fine-Tuning Approach for Multi-Token Generation in Autoregressive LLMs

Introduction and Motivation

Most transformer-based LMs use AR decoding, which generates output sequentially, one token per forward pass. This introduces inefficiencies, as even highly predictable tokens require a full forward pass. Existing methods that attempt to enable multi-token generation (e.g., speculative decoding, Medusa, EAGLE) incorporate architectural changes or additional models/parameters, increasing system complexity and deployment costs. In contrast, "MARS: Enabling Autoregressive Models Multi-Token Generation" (2604.07023) proposes a parameter-free, architecture-preserving fine-tuning routine that teaches instruction-tuned AR LMs to optionally generate multiple tokens per decoding step, achieving substantial throughput gains with no loss in one-token quality.

Methodology

Gap Analysis and Training Objectives

The authors identify four distinct gaps between classic AR generation and block-masked multi-token prediction schemes: token masking, attention pattern, logits alignment, and generation order. Most prior approaches sacrifice AR compatibility by deviating on attention and output conventions. MARS closes all but the inevitable masking gap, retaining strictly causal attention, right-shifted logits, and left-to-right decoding order, ensuring the derived model remains usable as a standard AR LM.

During fine-tuning, MARS leverages a dual-stream paradigm: each training sample is processed both in its original form (the clean stream) and with consecutive blocks replaced by [MASK] tokens (the noisy stream). The model is trained to predict all masked tokens using context only from the already-generated (unmasked) blocks and from [MASK] tokens within the current block, enforced by a carefully designed attention mask. The key innovation is the inclusion of a standard SFT loss on the clean stream, which preserves the model's core AR competence and counteracts the AR signal decay that plagues block-masked approaches as block size increases. The overall loss is a balanced sum of conventional AR NLL and masked block-wise NLL.

Figure 2: MARS attention mask and inference schematic for $L{=}8$, $B{=}4$; the training mask structure preserves causal attention and aligns block-wise prediction with AR conventions.

Inference Procedure

At test time, MARS operates via a left-to-right sliding window of size $B$: after generating a prefix, $B$ [MASK]s are appended, and the model predicts all $B$ masked tokens. Tokens are greedily accepted from left to right as long as their confidence (e.g., maximum softmax probability) exceeds a tunable threshold $\tau$. At minimum, one token per pass is always accepted, ensuring no regression over standard AR decoding. Speed–quality behavior is thus seamlessly adjustable per input/request by modifying $\tau$ alone, without model switches.

Experimental Results

Preservation of AR Quality

MARS achieves strict AR compatibility: in one-token mode, accuracy matches or exceeds the AR SFT baseline across six standard reasoning, instruction, and programming tasks (average +1.7 points at 0.5B, +1.5 points at 7B parameters; HumanEval increases by up to 4.8 points). Control experiments with compute-matched AR SFT indicate the gains are not merely an artifact of longer training.

Stability at Scale: Block Size and SFT Loss

The SFT loss is essential. Without it, increasing block size ($B$) causes sharp degradation (e.g., 6.2-point average drop from $B=4$ to $B=16$ at 0.5B), especially on reasoning/coding tasks. With the SFT loss, performance remains stable at all block sizes, with the AR signal maintained above 50%, and the speed-quality tradeoff curve consistently superior.

Figure 1: Speed–quality Pareto curves comparing MARS (with and without SFT loss) to baseline SFT for GSM8K and HumanEval; MARS with SFT loss dominates across the full range of operating points.

Speed–Quality Tradeoff

The main practical benefit is the throughput/latency control MARS provides. At $B{=}4$0, a 7B model accepts on average 1.68 tokens/forward, preserving accuracy at or above AR SFT levels (e.g., GSM8K +4.5pp, HumanEval +3.0pp rel. to AR SFT). The degradation with lower $B{=}4$1 is monotonic and small; the Pareto frontier is smooth, giving a robust latency–quality knob that can be dynamically tuned at inference time.

Efficient Deployment via Block-Level KV Caching

To obtain real wall-clock improvements, MARS employs a block-level KV cache scheme, amortizing residual computation and minimizing full-sequence recomputations. The results demonstrate up to 1.71$B{=}4$2 wall-clock speedup on Qwen2.5-7B relative to AR decoding with a vanilla cache, robust across moderate to large batch sizes.

Figure 3: Block-level KV cache operation for batch inference; visualization shows cache advancement and synchronization behavior as different samples fill their current block.

Implications and Future Directions

Practical implications: MARS enables deployers to serve a single model that dynamically balances speed and quality per request. It requires no extra heads, draft models, or architecture changes, and the fine-tuning overhead is limited to several further SFT epochs. Integration into production pipelines is thus straightforward.

Theoretical implications: The work reveals that causal attention is not fundamentally incompatible with parallel (multi-token) prediction—even when naively-masked block training degrades quality, this is largely due to removable design mismatches. The SFT loss ensures that AR competence is preserved even as the model learns to operate under block-masked contexts.

Comparison to prior art: MARS compares favorably to speculative decoding, Medusa/EAGLE, and Jacobi decoding. The block-masked approach, combined with SFT loss and strict causality/generation order, yields higher acceptance rates and does not incur secondary parameter overhead or inference-time model orchestration.

Limitations and future work: The main overhead of MARS is doubled sequence length during training and increased GPU peak memory, but this is minimal relative to pre-training. The batch KV cache incurs synchronization delays at block boundaries, and there are open questions regarding optimal block/adaptive window sizing, cursor-based cache management, and synergistic combinations with speculative decoding. The generalization of these results to very large (e.g., 30B+, 100B+) models, and the analysis of aggressive acceptance thresholds, are identified future directions.

Conclusion

MARS offers a lightweight, parameter-efficient method for augmenting AR LMs with robust multi-token generation capabilities. Empirically, it attains high speedup (1.5–1.7$B{=}4$3 throughput/wall-clock improvement) with minimal or negligible loss in output quality, and remains strictly compatible with legacy AR pipelines. Its design principles clarify the structural requirements for block-masked prediction without sacrificing AR utility, providing a formally grounded yet practical contribution to efficient LLM inference and deployment.