💬 NLP

MARS: Enabling Autoregressive Models Multi-Token Generation

MARS is a lightweight fine-tuning method that enables autoregressive language models to generate multiple tokens per forward pass without architectural changes or extra parameters, achieving 1.5–1.7x throughput while maintaining baseline accuracy and offering real-time speed adjustment via confidence thresholding.

Original authors: Ziqi Jin, Lei Wang, Ziwei Luo, Aixin Sun

Published 2026-04-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ziqi Jin, Lei Wang, Ziwei Luo, Aixin Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a professional chef (the AI model) tasked with writing a recipe.

The Old Way (Standard Autoregressive Models):
Traditionally, this chef writes the recipe one word at a time. Even if the next word is obvious—like writing "salt" after "pinch of"—the chef stops, thinks, writes "salt," stops again, thinks, and writes "pepper." It's a very careful, step-by-step process. It's accurate, but it's slow because the chef never takes a shortcut, even when the path is clear.

The Problem with Current "Fast" Methods:
Other teams tried to speed this up by hiring a sous-chef (Speculative Decoding) or giving the chef extra arms (Medusa/EAGLE).

The sous-chef tries to guess the next few words, but you have to keep two chefs in the kitchen, doubling the space and equipment needed.
The extra arms require building new limbs onto the chef, which is expensive and complicated.

The MARS Solution:
The authors of this paper, MARS, came up with a clever trick. They didn't hire a new chef or grow new arms. Instead, they gave the same chef a new training exercise.

They taught the chef to look ahead and say, "I'm 95% sure the next three words are 'salt', 'pepper', and 'sugar'." If the chef is confident, they write all three at once. If they are unsure (like when writing a complex, creative story), they slow down and write just one word at a time, exactly like before.

How It Works (The Analogy)

Think of the chef's workspace as a sliding window.

The Training (The "Masked" Practice):
During training, the authors put a "blindfold" (a [MASK] token) over the next few words on the recipe card. They ask the chef: "Based on what you've written so far, what goes in these blank spots?"
- Crucial Twist: They also made the chef write the original recipe alongside the blindfolded one. This ensures the chef doesn't forget how to write normally. It's like practicing a new dance move while still keeping your balance on one foot.
The Four Rules (Closing the Gaps):
Previous attempts to do this failed because they broke the rules of how chefs work. MARS fixed three broken rules:
- Rule 1 (Causality): The chef must only look forward, never backward. (Don't peek at the finished dish to guess the ingredients).
- Rule 2 (Order): The chef must write left-to-right. (Don't write the last word before the first).
- Rule 3 (Logic): The chef must use the same "brain" (parameters) for everything.
- Rule 4 (The Blindfold): The only thing that changes is that the chef gets to guess a block of words at once, but only if they are confident.
The Result (The "Confidence Knob"):
When the chef is writing a boring, predictable sentence ("The sky is blue"), they might write 3 or 4 words in a single breath.
When they are writing a complex math problem or a tricky joke, they switch back to writing one word at a time.
- The Best Part: You can adjust a "confidence knob" (a threshold) in real-time. If the kitchen is super busy, you tell the chef, "Be a bit more confident, write 2 words at a time!" If the customer is picky, you say, "Slow down, write one word at a time." You don't need to swap chefs or restart the kitchen.

Why This is a Big Deal

No New Hardware: You don't need a bigger kitchen or a second chef. It's the exact same model, just smarter about when to speed up.
No Quality Loss: When the chef writes one word at a time, they are just as good as before. In fact, the practice of guessing blocks actually made them slightly better at the standard tasks.
Real Speed: In tests, this method made the AI 1.5 to 1.7 times faster without losing accuracy. It's like the chef suddenly found a way to chop vegetables twice as fast during the boring parts of the prep, but still carefully plated the fancy dish.

In a nutshell: MARS teaches an AI to "batch" its thinking. It learns to sprint when the path is clear and walk when the path is tricky, all without changing its brain or needing extra help.

1. Problem Statement

Autoregressive (AR) language models generate text one token at a time, incurring a uniform computational cost regardless of how predictable the next token is. This "one-token-per-forward-pass" paradigm is inefficient for predictable sequences (e.g., formulaic phrases or simple reasoning steps).

Existing solutions for multi-token generation face significant deployment hurdles:

Speculative Decoding: Requires maintaining a separate, smaller "draft" model alongside the target model, doubling memory footprint and adding orchestration complexity.
Multi-Head Approaches (e.g., Medusa, EAGLE): Require architectural modifications, adding extra prediction heads and parameters, which complicates production serving pipelines.
Block Diffusion: Attempts to convert AR models to predict blocks of tokens but often suffers from severe quality degradation, particularly in reasoning and coding tasks, due to mismatches in attention patterns and generation order.

The Goal: Develop a method that enables an AR model to generate multiple tokens per forward pass without architectural changes, extra parameters, or a separate draft model, while preserving the model's original single-token generation quality.

2. Methodology: MARS (Mask AutoRegreSsion)

MARS is a lightweight fine-tuning method that teaches an instruction-tuned AR model to predict multiple tokens per step by analyzing and closing the gaps between standard AR models and block-masked prediction.

A. Gap Analysis

The authors identify four gaps between AR models and block diffusion approaches. They argue that three are eliminable design choices, while one is inherent:

Token Masking (Inherent): Future tokens must be replaced with [MASK] placeholders to allow parallel prediction. This is unavoidable.
Attention Pattern (Eliminable): Many diffusion approaches use bidirectional attention within blocks. MARS enforces strictly causal (left-to-right) attention everywhere, maintaining the AR property.
Logits Alignment (Eliminable): AR models predict $x_{t+1}$ from position $t$ (right-shifted). MARS preserves this convention.
Generation Order (Eliminable): Diffusion often unmask tokens out of order. MARS enforces strict left-to-right acceptance of tokens.

By closing gaps 2–4, the model remains a fully functional AR model that only sees [MASK] placeholders within the current block.

B. Training Strategy

MARS starts from a standard AR SFT checkpoint and performs continued fine-tuning on the same instruction data using a dual-stream approach:

Clean Stream: Processes the original sequence with standard causal attention and next-token prediction loss ( $L_{AR}$ ).
Noisy Stream: Replaces blocks of $B$ tokens with [MASK] placeholders. The model predicts these masked tokens using the clean prefix from previous blocks as context, with a masked cross-entropy loss ( $L_{mask}$ ).
Combined Loss: The total loss is $L = L_{mask} + L_{AR}$ $L = L_{ma s k} + L_{A R}$ .
- Key Insight: The $L_{AR}$ term is crucial. Without it, the "AR signal" decays as $1/B$ (as block size increases), causing quality degradation in reasoning tasks. The combined loss ensures the AR-equivalent signal fraction remains above 50% regardless of block size, stabilizing performance.

C. Inference Mechanism

Sliding Window: At each step, $B$ [MASK] tokens are appended to the current prefix.
Confidence Thresholding: The model runs a single forward pass to generate logits for all $B$ positions. Tokens are accepted consecutively from left to right as long as the maximum probability $p(x_t) \geq \tau$ (confidence threshold).
Adaptive Throughput:
- If $\tau \to 1.0$ , the model accepts only 1 token (behaving exactly like standard AR).
- If $\tau < 1.0$ , the model accepts multiple tokens (e.g., 1.5–1.7 tokens/forward) for predictable content.
- This allows real-time speed adjustment per request without model swapping or retraining.

D. Block-Level KV Caching

To achieve actual wall-clock speedups in batch inference, MARS introduces a block-level KV caching strategy. Instead of recomputing the full sequence for every step, the cache is updated only when a full block of tokens is accepted across all samples in a batch. This mitigates the $O(T^2)$ recomputation cost that would otherwise occur.

3. Key Contributions

Gap Analysis & Recovery: Demonstrated that quality degradation in block-masked prediction is caused by unnecessary design choices (bidirectional attention, wrong logits, out-of-order generation), not the concept of multi-token prediction itself.
Architecture-Free Multi-Token Generation: Proposed MARS, which requires no extra parameters or architectural changes, reusing the standard language model head and SFT data.
SFT Loss Necessity: Identified that adding a standard AR loss on the clean stream is essential to prevent performance decay at larger block sizes, effectively decoupling AR competence from block size.
Block-Level KV Caching: Developed a caching strategy for batch inference that yields significant wall-clock speedups.

4. Experimental Results

Experiments were conducted on Qwen2.5-0.5B and Qwen2.5-7B models across six benchmarks (IFEval, BBH, MMLU-Pro, GPQA, GSM8K, HumanEval).

Quality Preservation (One-Token Mode):
- When configured to generate one token per step ( $\tau=1.0$ ), MARS matches or exceeds the baseline AR SFT performance.
- On Qwen2.5-7B, MARS improved the average score from 56.6 (AR SFT) to 58.1, with notable gains in GSM8K (+4.5) and HumanEval (+3.0).
- In contrast, Block Diffusion (bidirectional) collapsed on reasoning tasks (e.g., BBH dropped to 7.5).
Multi-Token Throughput & Tradeoff:
- With $\tau=0.95$ , MARS achieves 1.5–1.7× tokens per forward pass with minimal accuracy loss (e.g., 7B model drops only 1.3 points on average).
- The speed-quality frontier is smooth; there is no "cliff" where quality suddenly collapses.
Wall-Clock Speedup:
- Using block-level KV caching on Qwen2.5-7B with a batch size of 4, MARS achieved a 1.71× wall-clock speedup (161.2s vs. 276.2s) over standard AR with KV cache.
- Without caching, MARS was slower than AR due to full-sequence recomputation.

5. Significance

Deployment Simplicity: MARS eliminates the need for complex draft models or architectural modifications, making multi-token generation feasible for existing production pipelines.
Dynamic Latency Control: The confidence threshold $\tau$ acts as a "knob" for operators to trade off latency and quality on the fly based on system load, without restarting services.
Theoretical Insight: The paper clarifies that causality and multi-token generation are compatible, provided the model maintains strict causal attention and generation order, challenging the prevailing assumption that bidirectional attention is required for parallel text generation.

In summary, MARS proves that lightweight fine-tuning alone is sufficient to unlock multi-token generation capabilities in AR models, offering a practical, high-performance, and quality-preserving solution for accelerating LLM inference.