Anchor-token masking

Never mask the retrieved prefix, always subject the target to the noise schedule. On identical data the fine-tune generalises 1.67× better out of distribution than the standard objective, and the paper shows where in the network the gain lives.

Retrieval-grounded generation hands a diffusion language model tokens it already knows and asks it to write around them. The standard training objective ignores that asymmetry and masks every position alike. This work trains rank-16 LoRA adapters on a 1.3B-parameter masked-diffusion model under three objectives that differ only in which positions may be masked, then measures how well each extends a locked anchor on held-out domains it never saw. The mechanism is tested four ways, ending in a direct ablation.

Three masking strategies are compared for fine-tuning masked-diffusion language models on retrieval-grounded data: anchor-token masking, where the prefix question is never masked and the answer statement is always subject to the noise schedule; random-position masking, uniform over both; and a reverse-asymmetric control that masks only the prefix. Rank-16 LoRA adapters are trained on Cassandra T1, a 1.3B-parameter model at epoch-5 cross-entropy 2.26, for 300 steps on 144 question and answer pairs, and evaluated by the corpus-vocabulary overlap of the text generated after a locked anchor, on 15 in-domain queries and 48 held-out queries across astronomy, medicine, law and cooking.

Anchor-token masking produces a 1.67× pooled held-out advantage over random-position masking, with a 95% bootstrap interval of [1.51, 1.85] and P(anchor > random) = 1.000. The in-domain advantage is 1.26× with interval [1.18, 1.34]; the two intervals do not overlap, so the effect favours generalisation rather than fit. The reverse control lands below random at P = 0.006, so the direction of the asymmetry carries the effect. A sweep over six training budgets shows the advantage peaking at 1.97× to 2.19× at 100 steps and closing to 1.07× at 1000 steps on 144 examples. Four mechanism probes, ending in a direct ablation, locate the effect in the value and output projections: a V/O-only LoRA outperforms the full LoRA at half the trainable parameters, and a Q/K-only LoRA falls below the random control.

A training example is a question Q and an answer statement A, padded to sequence length L. Each training step draws a mask ratio \gamma(t) uniformly from [0.5, 0.9] and then draws a mask m_i per position. The three strategies differ only in which class of position the Bernoulli draw may touch, and the loss is computed over the masked positions of that class.

At inference the retrieved locus statement s is tokenised and locked at positions [0, |\mathrm{tokenize}(s)|) of the answer slot. The unmasking loop excludes locked positions from the masked set at every denoising step, so locked tokens are preserved by construction; this mechanism is identical across arms and untouched by training. The coherence metric is \mathrm{corpus\_overlap}(\mathrm{text}, \mathrm{vocab}), the fraction of word tokens in the post-anchor text whose lowercase form is in the corpus vocabulary, the union of tokens over all locus statements in the relevant corpus. Confidence intervals come from 10,000 bootstrap resamples of the per-query values, percentile method, and P(a > b) is the fraction of resamples where the mean of a exceeds the mean of b.

The account the probes converge on: under anchor-token masking the prefix is fully visible at every step, so the adapter never has to model masked-prefix reconstruction and its rank-16 budget is free for the value and output projections. Those projections decide what content is read out of the positions the model already attends to. The ablation confirms that this reallocation is the effect, not a by-product of it. The mechanism is not that the model learns to look at anchors more; it is that the model learns to extract anchor-extending content from positions it was already attending to. Under out-of-distribution prompts that bias produces topical extension, while the random-mask adapter, having spent its capacity on query and key updates, falls back to training-corpus regurgitation.

• One base model, Cassandra T1 v1, undertrained at cross-entropy 2.26; one LoRA configuration, rank 16; 144 training pairs; 63 evaluation queries over five corpora.
• A fine-tuning regime only. From-scratch pretraining under the objective is proposed as the next experiment, not measured here.
• corpus_overlap is a vocabulary proxy for topical coherence. It does not score word order or grammar; the qualitative samples show it tracks an interpretable behaviour.
• The held-out corpora were authored for the study, with publicly verifiable facts, rather than drawn from an external benchmark.
• The 1.67× headline is a mid-budget measurement. The sweep's maximum was 2.19× at the lowest budget tested and the advantage closes at 1000 steps on 144 examples.
• The V/O ablation was run at one data, step and rank configuration.

The repository carries the three training runners, the four evaluation runners, the attention and ablation scripts, the sweep, the bootstrap analysis, the base checkpoint with its SHA-256, the three 24 MB adapters, the seven corpora in LTMi-XT format, the query sets, and the per-query JSONL outputs for all 504 generations. Seeds are 42 for training and 1234 for generation. Total reproduction time is about 60 minutes on a single consumer GPU.

@techreport{garren2026anchortokenmasking,
  author      = {Garren, Thomas},
  title       = {Anchor-Token Masking Produces an Out-of-Distribution
                 Generalization Advantage in Retrieval-Grounded
                 Masked-Diffusion Language Models},
  institution = {SOPHIA XT LLC},
  year        = {2026},
  month       = {May},
  doi         = {10.5281/zenodo.20496783},
  url         = {https://sophiaxt.com/research/anchor-token-masking}
}