Late Context Masking: Accuracy and Fidelity to Context-Aware Decoding
Year: 2026
Type: preprint
Abstract: Language models are often more likely to repeat what they memorized than to follow the text in their prompt. Context-aware decoding (CAD) is a common fix: it runs the model a second time without the context and pushes the prediction away from that run. We make the second run cheaper. The two runs share the model's lower layers, and only its last 16 layers, out of 32 to 48, are run again with the context hidden from them, a design we call the late fork. On MemoTrap, a benchmark that asks the model to end familiar quotes with new words, the late fork is 3.5 and 8.6 percentage points more accurate than CAD on Qwen3-4B-Base and Qwen3-8B-Base. The contrast weight sets how hard a decoder pushes away from the second run, and in this comparison both use 0.5, half the weight CAD's authors used on MemoTrap. When CAD's weight is instead tuned on a held-out development split, the 4B comparison is inconclusive. The two decoders choose different tokens: where CAD departs from greedy decoding, the late fork picks CAD's token only 43% and 62% of the time. On Mistral-7B the late fork is about as accurate as CAD, and on Yi-1.5-9B, where CAD adds only 2.5 points over greedy decoding, the late fork adds 4.5 more. On redefine-math it gives no gain over greedy decoding. We conclude that the late fork is a different decoder from CAD, and that its benefit depends on the task. Code and data The code, decision records and per-item outputs are in https://github.com/SeverinVisionary/late-context-masking at commit 35bbd0d. Disclosure The research questions, preregistrations and conclusions are the author's, and the author checked every result. Claude Opus 5.5 (Anthropic) wrote the experiment code and analysis. The author drafted this prose with AI assistance. Codex gpt-6-astra (OpenAI) reviewed the research plan and three of the later preregistrations, and ChatGPT reviewed the research plan and this draft. Experiments ran on one Colab A100 GPU. The author is responsible for the content.
Author: Yang, Hanyu
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W7225611059
Cites:
Cited by:
Related to: