🪐 MARS: Enabling Autoregressive Models Multi-Token Generation

MARS converts any instruction-tuned autoregressive (AR) model into a multi-token predictor with zero architectural changes, zero additional parameters, and a single checkpoint. The AR capability is fully preserved; multi-token generation is available as an opt-in acceleration mode. A fine-tuned model checkpoint can be served at a quality-latency tradeoff chosen per-request. We release all our checkpoints and training/eval scripts to support further research.

• Lossless at τ=1.0. MARS matches or exceeds AR SFT on IFEval, BBH, MMLU-Pro, GPQA, GSM8K, and HumanEval when decoding one token at a time.
• 1.5–1.7× tokens-per-forward at τ=0.95 with minimal accuracy loss.
• Up to 1.69× wall-clock over AR-with-KV-cache on Qwen2.5-7B via block-level KV caching.
• One knob at inference time. The confidence threshold τ trades quality for speed per request. No separate draft model, no new heads.
• Zero architectural changes. MARS is a fine-tuning method, not a new architecture. Inherits all AR serving infrastructure.

• 2026-04-20 — Uploaded 7B blk8/16 models to HuggingFace.

MARS is trained with a [noisy | clean] concatenation and a structured attention mask that keeps the clean stream pure-causal (AR) while letting the noisy stream predict a full block of masked tokens in parallel, conditioned only on the clean tokens from previous blocks. At inference the sampler walks left-to-right, appending B mask tokens at the cursor, running one forward, and accepting the prefix of greedy predictions whose confidence passes threshold τ. When the model is confident, it commits multiple tokens per forward; when it is uncertain, it falls back to one-token AR.

L = CE(noisy_logits, labels)        # multi-token prediction
  + α · CE(clean_logits, labels)    # AR loss on the clean stream

The auxiliary AR term (ar_weight=1.0) prevents the AR-like training signal from decaying as 1/B as block size grows, and is the key ingredient behind MARS's stability at B=8 and B=16.

Solid lines: MARS (with SFT loss). Dashed: without SFT loss (0.5B only). Dotted: AR SFT baseline. Each point corresponds to a different confidence threshold τ ∈ {1.0, 0.95, 0.9, 0.8, 0.7, 0.6, 0.5}.

At τ=1.0 (leftmost point on each curve), MARS sits at or above the AR SFT line — a strong precondition: no quality is sacrificed for the option to accelerate. From there, lowering τ trades accuracy for tokens/forward along a smooth Pareto curve.

git clone https://github.com/Xalp/MARS.git
cd MARS
bash setup_env.sh
conda activate mars

All models available on HuggingFace.

MARS first trains the model with AR SFT (control step), then MARS training with the same dataset. This is a control step to isolate MARS's effect from the effect of additional data exposure.

• Stage 1: AR SFT — standard next-token prediction fine-tuning.
• Stage 2: MARS — masked block prediction with auxiliary AR loss, starting from the Stage 1 checkpoint.

export BASE_DIR=/path/to/experiment/root
bash scripts/train/train_0.5b.sh

# Stage 1: AR SFT
accelerate launch --config_file scripts/accelerate_configs/zero2.yaml \
    --num_processes 8 train/train_ar_sft.py \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --output_dir ${BASE_DIR}/models/ar_sft \
    --dataset_args allenai/Dolci-Instruct-SFT \
    --num_train_epochs 5 --learning_rate 5e-6 \
    --per_device_train_batch_size 48 --max_length 512 --bf16

# Stage 2: MARS (with SFT loss, block_size=4)
accelerate launch --config_file scripts/accelerate_configs/zero2.yaml \
    --num_processes 8 train/train_mars.py \
    --model_name_or_path ${BASE_DIR}/models/ar_sft/checkpoint-final \
    --output_dir ${BASE_DIR}/models/mars_blk4 \
    --dataset_args allenai/Dolci-Instruct-SFT \
    --num_train_epochs 5 --learning_rate 5e-6 \
    --per_device_train_batch_size 48 --max_length 512 \
    --block_size 4 --right_shift_logits --ar_weight 1.0 --bf16

No acceleration, isolates quality.

accelerate launch --num_processes 8 mars/eval_harness.py \
    --tasks gsm8k_cot --num_fewshot 0 \
    --model mars --apply_chat_template \
    --model_args "pretrained=${BASE_DIR}/models/mars_blk4/checkpoint-final,max_new_tokens=256,steps=256,block_size=4,cfg=0.0,right_shift_logits=True"

Commits the prefix of greedy predictions whose probabilities exceed τ.

accelerate launch --num_processes 8 mars/eval_harness.py \
    --tasks gsm8k_cot --num_fewshot 0 \
    --model mars --apply_chat_template \
    --model_args "pretrained=${BASE_DIR}/models/mars_blk4/checkpoint-final,max_new_tokens=256,steps=256,block_size=4,cfg=0.0,right_shift_logits=True,confidence_threshold=0.95"

python scripts/benchmark/bench_block_cached.py \
    --mars_model ${BASE_DIR}/models/mars_blk4/checkpoint-final \
    --ar_model ${BASE_DIR}/models/ar_sft/checkpoint-final \
    --batch_size 16 --block_sizes "4,8,16,32" --threshold 0.95

@misc{jin2026marsenablingautoregressivemodels,
      title={MARS: Enabling Autoregressive Models Multi-Token Generation}, 
      author={Ziqi Jin and Lei Wang and Ziwei Luo and Aixin Sun},
      year={2026},
      eprint={2604.07023},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2604.07023}, 
}