MARS: Enabling Autoregressive Models Multi-Token Generation

Ziqi Jin; Lei Wang; Ziwei Luo; Aixin Sun

cs.CLApr 8, 2026

MARS: Enabling Autoregressive Models Multi-Token Generation

Authors: Ziqi Jin, Lei Wang, Ziwei Luo, Aixin Sun

Organizations: Nanyang Technological University · Singapore Management University · Uppsala University

Abstract

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.

Figures & tables

Question: Toulouse has twice as many sheep as Charleston. Charleston has 4 times as many sheep as Seattle. How many sheep do Toulouse, Charleston, and Seattle have together if Seattle has 20 sheep?
MARS-7B (confidence threshold $\tau{=}0.95$ ), 66 forwards, 2.55 tokens/forward:
Let ’s break it down step by step: 1 . Seattle has 20 sheep. 2. Charleston has 4 times as many sheep as Seattle , so Charleston has 4 x 20 = 80 sheep. 3. Toulouse has twice as many sheep as Charleston, so T oulouse has 2 x 80 = 16 0 sheep. Now , let ’s add up … Total number of sheep : 20 + 80 + 16 0 = 1 00 + 160 = 260 sheep
Shaded regions = multiple tokens accepted in one forward pass. Unshaded tokens are generated one at a time. Darker = more tokens per step. Formulaic phrases and calculations are chunked; novel reasoning proceeds token-by-token.

Figure 1: Example MARS generation on GSM8K. The model adaptively generates 1–4 tokens per forward pass based on confidence: predictable continuations are batched together, while novel content proceeds token-by-token. This achieves 2.55

\times

token per forward over standard AR decoding.

	AR	MARS	Block Diffusion
(1) Token Masking	No masking	Masked within block	Masked within block
(2) Attention Pattern	Causal	Causal	Bidirectional (intra-block)
(3) Logits Alignment	Right-shifted	Right-shifted	Varies
(4) Generation Order	Left-to-right	Left-to-right	Confidence-based

Table 1: Gaps between AR and block diffusion. MARS aligns with AR on gaps (2)–(4), leaving only the inherent masking gap.

Figure 2: MARS attention mask and inference for $L{=}8$ , $B{=}4$ . Left: training mask with $[\mathbf{x}\mid\tilde{\mathbf{x}}]$ concatenation. The orange cells show that noisy positions attend to each other causally within each block, in contrast to Block Diffusion ( Arriola et al., 2025 ) which uses bidirectional attention within blocks. Right: sliding-window inference. The dashed line marks the generation cursor; $B$ [MASK] tokens are appended and filled via one forward pass. Accepted tokens (blue) slide into the prefix for the next step.

Qwen2.5-0.5B-Instruct
Model	IFEval	BBH	MMLU-Pro	GPQA	GSM8K	HumanEval	Avg
Base (no fine-tune)	27.8	6.1	2.4	13.6	27.4	26.2	17.2
AR SFT (5 ep)	48.4	26.3	11.9	17.9	32.0	35.4	28.7
AR SFT (10 ep)	47.8	26.3	9.3	14.1	28.3	32.3	26.4
Block Diffusion Arriola et al. (2025) ( $B{=}4$ )	47.1	7.5	2.0	17.9	30.6	31.7	22.8
MARS-0.5B w/o SFT loss ( $B{=}4$ )	48.5	27.4	12.3	19.0	29.5	33.5	28.4

Table 2: One-token mode (

\tau{=}1.0

): MARS vs. AR SFT, compute-matched AR SFT (10 epochs), and Block Diffusion. All models generate one token per forward pass. Bold : best per column within each scale.

MARS w/o SFT loss
Model	IFEval	BBH	MMLU-Pro	GPQA	GSM8K	HumanEval	Avg
AR SFT	48.4	26.3	11.9	17.9	32.0	35.4	28.7
$B{=}4$	48.5	27.4	12.3	19.0	29.5	33.5	28.4
$B{=}8$	48.5	24.3	11.1	20.8	24.9	28.7	26.4
$B{=}16$	42.6	21.7	10.6	16.3	21.0	20.7	22.2
MARS with SFT loss

Table 3: Effect of SFT loss across block sizes (0.5B,

\tau{=}1.0

). Without the SFT loss, larger blocks rapidly erode quality; with it, performance is stable across

B

Figure 3: Speed–quality Pareto curves on GSM8K (left) and HumanEval (right). Solid lines: MARS (with SFT loss). Dashed lines: w/o SFT loss. Dotted: AR SFT baseline. With SFT loss, MARS dominates at every operating point on both tasks.

Qwen2.5-0.5B-Instruct
Model	IFEval	BBH	MMLU-Pro	GPQA	GSM8K	HumanEval	Avg
MARS $B{=}4$	45.5 ( $-$ 5.8)	26.6 (0.0)	11.5 ( $-$ 0.9)	19.9 (+0.5)	31.1 ( $-$ 1.7)	39.6 ( $-$ 0.6)	29.0 ( $-$ 1.4)
tok/fwd	1.20	1.90	1.50	1.36	1.51	1.26	1.46
MARS $B{=}8$	44.0 ( $-$ 5.6)	28.3 (+1.4)	11.9 ( $-$ 0.1)	19.0 ( $-$ 0.6)	31.7 ( $-$ 1.1)	36.6 ( $-$ 0.6)	28.6 ( $-$ 1.1)
tok/fwd	1.23	1.96	1.49	1.46	1.54	1.27	1.49
MARS $B{=}16$	45.3 ( $-$ 5.4)	25.9 ( $-$ 1.1)	11.9 ( $-$ 0.2)	17.9 (0.0)	31.2 ( $-$ 2.6)	36.6 (0.0)	28.1 ( $-$ 1.6)

Table 4: Multi-token mode. Each cell shows accuracy with the absolute change from

\tau{=}1.0

\tau{=}0.95

in parentheses. Gray italic rows show tokens accepted per forward pass for each task. Full threshold sweep in Appendix Table 7 .

Figure 4: Block-level KV cache for batch inference (

B_{\text{cache}}{=}4

, batch size 3). Each step forwards

B

[MASK] tokens against the cached prefix. The cache advances by the minimum number of tokens accepted across all samples: after Step 1 (S1 accepts 4, S2 accepts 2, S3 accepts 1), one token is cached (green). S1 idles while S2 and S3 continue. Once all samples fill the block, the entire block is cached (yellow) and new [MASK] tokens are appended for the next block.

		Batch size=4			Batch size=8			Batch size=16
Method	$B_{\text{cache}}$	tok/s	time	$\times$	tok/s	time	$\times$	tok/s	time	$\times$
AR (KV cache)	–	143.4	276.2	1.00	236.7	169.1	1.00	434.3	91.8	1.00
MARS (no cache)	–	127.1	306.5	0.90	112.4	346.9	0.49	97.7	399.1	0.23
MARS + block cache	4	157.2	248.8	1.11	236.4	165.4	1.02	399.3	97.2	0.94
MARS + block cache	8	196.0	199.4	1.39	289.9	134.4	1.26	484.1	80.1	1.15
MARS + block cache	16	228.5	170.5	1.62	338.0	115.1	1.47	566.0	68.7	1.34

Table 5: Batch inference on GSM8K (256 questions) with Qwen2.5-7B at

\tau{=}0.95

B_{\text{cache}}

: KV cache synchronization granularity. For each batch size we report tokens/sec, total wall-clock time (seconds), and speedup relative to AR with KV cache (

\times

). Bold : best MARS configuration per batch size.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	0.5B	7B
Base model	Qwen2.5-0.5B-Instruct	Qwen2.5-7B-Instruct
Training data	Dolci-Instruct-SFT ( $\sim$ 2M examples)
Optimizer	AdamW
Learning rate	$5\times 10^{-6}$
LR schedule	Cosine decay
Warmup ratio	0.03

Appendix

Table 6: Training hyperparameters. Both stages (AR SFT

\to

MARS) use identical settings per model size. Effective batch sizes are matched across scales.

MARS
	GSM8K						HumanEval
	$B{=}4$		$B{=}8$		$B{=}16$		$B{=}4$		$B{=}8$		$B{=}16$
$\tau$	Acc	T/F	Acc	T/F	Acc	T/F	P@1	T/F	P@1	T/F	P@1	T/F
AR SFT	32.0	1.00	32.0	1.00	32.0	1.00	35.4	1.00	35.4	1.00	35.4	1.00
1.0	32.8	1.00	32.8	1.00	33.8	1.00	40.2	1.00	37.2	1.00	36.6	1.00
0.95	31.1	1.51	31.7	1.54	31.2	1.57	39.6	1.26	36.6	1.27	36.6	1.25

Appendix

Table 7: Speed–quality tradeoff on GSM8K (0-shot) and HumanEval (0-shot). Tok/Fwd: average tokens accepted per forward pass.

\tau{=}1.0

: one token per step (from Table 2 ). Lower

\tau

accepts more tokens but may reduce accuracy.

Method	IFEval	BBH	MMLU-Pro	GPQA	GSM8K	HumanEval	Avg	Tok/fwd
AR SFT	48.4	26.3	11.9	17.9	32.0	35.4	28.7	1.00
Jacobi (on AR SFT)	41.0	19.2	11.7	17.6	36.5	42.1	28.0	1.07
MARS ( $\tau{=}0.95$ , $B{=}4$ )	45.5	26.6	11.5	19.9	31.1	39.6	29.0	1.46

Appendix

Table 8: Jacobi decoding [ Teng et al., 2024 ] on the AR SFT checkpoint (0.5B, training-free). Jacobi uses fixed-point iteration: all future positions are initialized with random tokens and iteratively updated via causal forward passes until convergence. Tok/fwd: average tokens accepted per forward pass.

Figure 5: Speed–quality trade-off under three acceptance metrics (MARS

B{=}4

, GSM8K). All three metrics trace similar Pareto frontiers, indicating that the speed–quality trade-off is robust to the choice of acceptance criterion. Entropy and top-2 margin degrade slightly more gracefully than raw probability at comparable tokens per forward pass.

Multi-Token Prediction

Multi-token prediction trains models to forecast several future tokens at once, supporting more efficient generation and richer predictive representations.

See all 39 papers

K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling
Jun 9, 2026Zhiwei Tang, Yuanyu He, Yizheng Han +4LLM Serving Language Model Decoding
1DAMO Academy, Alibaba Group · 2Hupan Lab · 3Zhejiang University +1
BitLM: Unlocking Multi-Token Language Generation with Bitwise Continuous Diffusion
May 12, 2026Shaobin Zhuang, Yuang Ai, Jiaming Han +7Multi-Token Prediction Autoregressive Language Modeling
Shanghai Jiao Tong University · MMLab, The Chinese University of Hong Kong · Institute of Automation, Chinese Academy of Sciences +1
Efficient Training with Foresight: Multi-Token Auxiliary Supervision for Autoregressive Image Generation
Aug 26, 2026Guo Niu, Xiongfei Yao, Teng Wang +1Autoregressive Image Generation Multi-Token Prediction
Foshan University Foshan, China · The University of Hong Kong Hong Kong, China · Sun Yat-sen University Guangzhou, China

Jun 9, 2026J/K move · Enter open · S save

K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling

Zhiwei Tang, Yuanyu He, Yizheng Han, Wangbo Zhao+3

1DAMO Academy, Alibaba Group · 2Hupan Lab · 3Zhejiang University+1

LLM Serving Language Model Decoding Autoregressive Language Modeling LLM Inference Acceleration Language Model Distillation

Open paper AskPDF arXiv

Autoregressive (AR) language modeling is the dominant paradigm for text generation, yet its sequential token-by-token decoding makes inference memory-bound and inefficient. Existing acceleration approaches, such as speculative decoding and diffusion language models, can yield speedups under certain conditions but do not directly address high-load batch serving--the scenario most critical for industrial-scale deployment. We introduce K-Forcing, a push-forward language modeling paradigm for joint next-k-token decoding. K-Forcing distills an existing AR model into a conditional push-forward mapping--one that transforms independent uniform noise variables into a joint sample of multiple future tokens in a single forward pass. This design preserves fixed-length outputs, reuses the AR teacher backbone, and remains compatible with standard AR serving infrastructure. We train this mapping via progressive self-forcing distillation, which gradually expands the prediction window while enabling the student to closely match the sequence distribution of the AR teacher. We evaluate K-Forcing on LM1B and OpenWebText using a standard causal Transformer backbone. When aggressively configured to generate k = 4 tokens per forward pass, K-Forcing delivers approximately 2.4-3.5x speedup across different batch sizes, while incurring modest quality degradation relative to its AR teacher. As inference increasingly dominates the lifetime compute cost of modern LLMs, K-Forcing offers a promising route toward accelerating AR generation under real-world high-load deployment.

Abstract

Multi-Token Prediction

Explore similar work

K-Forcing: Joint Next-K-Token Decoding via Push-Forward Language Modeling

BitLM: Unlocking Multi-Token Language Generation with Bitwise Continuous Diffusion

Efficient Training with Foresight: Multi-Token Auxiliary Supervision for Autoregressive Image Generation