Accepted at AACL-IJCNLP 2026

AXON Supportive Token Revealing for Fast Diffusion Language Model Decoding

A training-free plug-in for diffusion language models. When parallel decoding stalls, AXON reveals anchors: confident masked tokens that uncertain positions attend to, giving the rest of the sequence the context it needs.

Giries Abu Ayoub1, Mario Barbara1, Lluís Pastor-Pérez2, Tanja Bien2, Aneesh Barthakur2, Alaa Maalouf1, Loay Mualem2,3

1Department of Computer Science, University of Haifa  ·  2Institute for AI, University of Stuttgart  ·  3IMPRS-IS

Paper Code

How AXON works

Gate → select anchors → reveal supportive context.

A lightweight gate checks every denoising step. When uncertain tokens lack support, AXON picks anchors with a submodular coverage objective and reveals them; otherwise the base decoder's proposal goes through unchanged.

Pipeline: a partially noised sequence 'The [MASK] Tower [MASK] [MASK] [MASK] .' goes to a parallel base decoder that proposes tokens. AXON's gate checks whether anchors are required: with enough support the base proposal 'is in' is kept; otherwise anchor selection picks 'Paris'. The next state reveals the selected tokens. Panels below detail the gate, anchor selection by a utility matrix and greedy coverage, and anchor injection lowering uncertainty over denoising steps.Click to enlarge
Fig. 1. Gate → select anchors → reveal supportive context.

Highlights

Fewer forward passes. No retraining.

  • Training-free plug-in; the backbone stays frozen.
  • Gated intervention using existing attention and confidence signals.
18.6%
fewer forward passes on HumanEval with LLaDA-1.5
DAWN + AXON (k = 1)
NFE 77.3 → 62.9
Accuracy 42.07 → 42.68%
14.8%
fewer forward passes on HumanEval with Dream-v0-Instruct-7B
DAWN + AXON (k = 1)
NFE 76.2 → 64.9
Accuracy 54.88 → 55.60%
29/36
settings where AXON needs fewer forward passes than its base decoder
3 backbones × 4 benchmarks × 3 base decoders
26/36
settings where AXON is more accurate than its base decoder
same 36 settings

Forward passes per sample: base decoder vs. + AXON

Benchmark
Base decoder
0120 forward passes (NFE)

NFE = forward passes of the model per sample (fewer is faster). AXON injects k = 1 anchor per gated step; AXONCVR adds anchors until their coverage matches the base proposal's. Full table in the paper.

AXON in action

Same problem, same decoder. Half the steps.

Both animations decode HumanEval/47 (median) with LLaDA-1.5 and the DAWN base decoder. DAWN alone needs 107 steps and its output fails the unit test; with AXON, the gate detects the stalls and reveals anchors (blue), un-blocking the surrounding tokens.

DAWN107 steps · 2.39 tok/step · fails unit test
DAWN decoding HumanEval/47 token by token: 107 steps, 2.39 tokens per step, gate fired 0 times.
DAWN + AXON54 steps · 4.74 tok/step · gate 7× · passes unit test
DAWN plus AXON decoding the same problem: 54 steps, 4.74 tokens per step, gate fired 7 times, with anchors highlighted.