Skip to content
Xalp & Peyto's
Go back

MARS: Several Tokens per Forward Pass, Without Giving Up Autoregression

The short version. Take an instruction-tuned AR model. Fine-tune it for 5 epochs on the same SFT data it already saw, with one extra loss term that costs no extra compute. You get back one checkpoint that does two things:

No new parameters. No second draft model. No architecture change. No separate deployment path: the same weights serve both modes, and the fallback to plain AR decoding is exact and free.

One checkpoint, two modes: at τ=1.0 tokens are emitted one per forward pass at full quality; at τ=0.95 the same weights commit runs of 2–4 tokens per pass for a 1.56× speedup.

Figure 1. One checkpoint, two modes. The dial τ is the only thing that changes between them; this is not two models.


The question behind the work

Diffusion language models generate several tokens at once, and there is a well-trodden recipe for converting a pretrained AR model into one. The recipe works in the sense that you get parallel generation. It also reliably damages the model: reasoning and coding degrade, sometimes catastrophically.

The usual reading is that this is the price of parallelism. We think that reading is wrong, and the paper’s first contribution is to say why.

When you convert an AR model into a block-diffusion model, you change four things at once:

  1. Token masking: some positions are [MASK] placeholders rather than real tokens.
  2. Attention pattern: intra-block attention becomes bidirectional.
  3. Logit alignment: position i predicts token i instead of token i+1.
  4. Generation order: tokens inside a block are filled in confidence order rather than left to right.

Only the first is inherent to predicting multiple tokens in one pass. The other three are conventions inherited from diffusion, and each one discards something the pretrained model spent its entire pretraining budget acquiring. Bidirectional attention throws away the causal structure of every attention head. Changing logit alignment invalidates the output head’s calibration. Out-of-order filling breaks the left-to-right factorization the model’s internal representations are organised around.

MARS keeps gaps 2, 3 and 4 aligned with the AR model and pays only for gap 1.

Four columns, one per gap. Token masking is unavoidable when emitting multiple tokens per pass; attention, logit alignment, and generation order are diffusion conventions MARS keeps on the AR side.

Figure 2. The thesis in one image: three of the four changes are conventions, not requirements. Keep the AR side of them.


The construction, and the part that turned out to matter

Training runs one forward pass over the concatenation [x ; x̃]: the clean response, followed by a copy in which each block of B tokens is replaced by mask placeholders. A structured attention mask enforces the right visibility: the clean half attends causally to itself, exactly as in AR training, and each masked block attends to the clean tokens of earlier blocks, which is the context it needs to predict its own.

Here is the part that is easy to miss and that shapes everything downstream: the clean tokens have to be in that forward pass anyway. You cannot ask a model to fill block k from placeholders unless the real tokens of blocks 1..k−1 are present as keys and values; otherwise the target is undefined and the task degenerates into unconditional generation. The clean stream is not a design choice. It is the conditioning context.

Which raises the actual design question. Those clean positions are computed regardless, and because they attend causally to one another, their logits are exactly standard AR next-token logits. Do we attach a loss to logits we already have?

Attaching it costs zero additional transformer compute. It also turns out to be the difference between a model that works and one that falls apart.

0.5B, six-benchmark averageB=4B=8B=16
with the AR loss30.129.730.1
without it28.025.521.7

With the loss, the block-size sweep is flat. Without it, quality decays monotonically, and at 7B the same ablation is far more violent: removing the AR loss at block-16 drops GSM8K from 75.3 to 31.5. The damage is not uniform. Instruction-following and knowledge benchmarks barely move, while long-chain arithmetic is destroyed. That signature is the tell: what is being lost is autoregressive competence specifically, in the tasks that depend on it most.

Six-benchmark average against block size B. With the free AR loss the curve is flat near 30; without it, it slides from 28.0 to 21.7, and at 7B GSM8K collapses from 75.3 to 31.5.

Figure 3. The money plot: the free loss is the load-bearing one.

A correction we made during review. Our initial framing was that the single-forward, two-stream design was necessary. A reviewer asked us to test the obvious alternatives, and they were right to ask: simply mixing AR-only and masked-only batches also works (30.3 at B=16, against 30.0 for MARS and 21.8 with no AR signal at all), and computing the two losses in separate forward passes is provably identical to MARS. In float64 the clean-stream logits agree bit-for-bit and the gradients to under 7e-7, at twice the compute.

So the honest claim is not about the architecture. The insight is that masked-block training erodes AR competence at a rate governed by how much AR-anchored signal remains, and any schedule that keeps that signal present prevents the collapse. The single forward is simply the cheapest way to keep it: it harvests both losses from one pass, while mixed batches spend roughly 1.5× the forward compute for the same training signal. The sharpened claim is more useful than the one we started with.


Inference: a knob, not a mode switch

At generation time the model proposes B tokens per forward pass and commits them left to right while its confidence holds, stopping at the first token below a threshold τ. Because the model never abandoned the AR conventions, τ=1.0 recovers exact single-token AR decoding: same weights, same code path, no conversion step.

Two findings are worth stating on their own.

The acceptance rule matters, and probability is not the best one. On a format-heavy benchmark such as IFEval, the natural gate (top-1 probability) degrades first, which is what you would expect if high-confidence formatting tokens are exactly the ones a probability gate happily commits. A margin gate, how far ahead the top token is rather than how confident it is, holds up better and reaches operating points probability cannot: between 1.25 and 1.35 tokens per forward it is consistently 0.6 to 0.8 points better, and it keeps producing usable points out to 2.57 tokens per forward, where probability’s range ends at 1.39.

Batching is where methods like this usually die, and mostly we fixed it. The original block-level KV cache advanced at the batch-minimum acceptance, so one slow sample stalled everyone. We rebuilt the sampler to defer KV commitment into the next proposal forward, with per-sample cursors and a cache-validity bitmask, which removes the extra forward, the cache copy and the synchronization barrier together:

batchpaper’s samplerdeferred-KV sampler
41.64×2.16×
81.55×2.06×
161.27×1.70×
321.01×1.17×
640.67×1.00× (falls back to AR)

The last row is the interesting one. With W proposals per step, W=1 appends no mask slots at all, so the step is computationally identical to AR. A load-aware controller that keeps W=1 in its candidate set is therefore bounded below by AR throughput by construction: under heavy load the system stops proposing and degrades into a plain AR server rather than into something slower than one. The optimal W anneals monotonically with load: 8, 8, 4, 2, 1.

Wall-clock speedup against batch size for the old and deferred-KV samplers. Both shrink with load, but the deferred-KV sampler lands on a 1.0× AR floor at batch 64 rather than falling through it.

Figure 4. The advantage shrinks with load and lands on a floor rather than falling through it.


Where it stands against the alternatives

We compared against a trained draft head (EAGLE-style) built on our own checkpoint at a matched budget of added parameters, sweeping head width and training length so that neither size nor undertraining could explain the outcome. The result splits cleanly across two axes, and both are worth reporting:

added paramstokens / target forwardwall-clockGSM8K
AR SFT (the target model)01.001.00×68.4
best EAGLE-style head19.3M2.001.32×68.4 (lossless)
MARS (τ=0.95)01.7851.56×70.7

A trained drafter proposes better than MARS. It converts that into less end-to-end speed, because every proposal costs a forward pass through the head, each ending in a 152k-vocabulary projection, outside the target forward, while MARS’s proposals ride inside a forward it was going to run anyway. Scaling the head to EAGLE’s own 7B width makes both axes worse: 1.48 tokens per forward and 0.96×, slower than plain AR.

There is a second asymmetry worth naming. A lossless method is pinned to its target: every draft-head row scores exactly 68.4, because that is what the target scores. MARS is not pinned, and at τ=0.95 it lands above the model it was trained from. The trade being weighed here is not speed for accuracy.


What we do not claim


Why we think this direction has room

Three observations point the same way.

The effect grows with scale. The speedup at 7B is larger than at 0.5B. If the number of tokens a single forward can safely commit grows with model capacity, then the interesting experiment is not 7B but 70B and beyond, and the open question is where that saturates. It is the measurement we most want to run next.

The cost model is unusually friendly. MARS adds one extra SFT run (100 → 202 GPU-hours at 7B) and nothing to the serving stack: no second module to load, batch, or keep in sync. Compare that to any draft-model approach, where the drafter is a permanent operational dependency. For a team that already runs SFT, this is close to free.

It composes rather than competes. Used as the draft inside standard speculative decoding, with verification untouched and the output distribution exactly matching the target, MARS is 1.40× faster than standard speculative decoding on the same model pair. It upgrades both endpoints of a speculative pipeline instead of replacing it.

The obvious next steps: a production engine (our numbers come from a research implementation on HuggingFace Transformers, and vLLM or SGLang would change the constant factors on both sides), the load-aware W scheduler as a first-class serving policy, and the cross-family replication we could not fit into this cycle.

Decision diagram: if you already run SFT and serve at small-to-moderate batch, MARS costs one extra SFT run and nothing at serving time; very large batches route to AR, and exact-distribution needs use MARS as the draft in speculative decoding.

Figure 5. When to reach for MARS, with its two honest off-ramps.


The one-line takeaway. Multi-token generation does not require abandoning autoregression. It requires not throwing away the three things about autoregression that were never the problem, and keeping enough AR-anchored training signal that the model stays the model you started with.


Share this post:

Next Post
How Much Do Diffusion LLMs Memorize? (vs Autoregressive)