Elamur Research⌕
☰
home / papers / 2026 / One Token Is Enough: Improving Diffusion Language Models with a Sink Token

One Token Is Enough: Improving Diffusion Language Models with a Sink Token

Paper · 2026 · Zihou Zhang, Zheyong Xie, Li Zhong et al.

Summary

Diffusion Language Models (DLMs) such as LLaDA and Dream generate text in parallel using bidirectional attention. Unlike autoregressive models, they have no fixed first token to absorb surplus attention. As a result, their attention sinks wander between positions, often landing on masked tokens, as the diffusion steps proceed. The authors study this "moving sink" effect through the value space. They find that sink tokens have unusually small value-vector norms and act as a safeguard against over-mixing of information. Because their location keeps changing, however, they make inference less stable.

As a fix, they add one extra [SINK] token to the sequence. Through a modified attention mask, this token sees only itself, while every other token can attend to it. The token is used during both training and inference. They test it on DLMs initialized from Qwen2.5-0.5B and Qwen2.5-1.5B and report consistent benchmark improvements. The gain does not depend on where the token is placed or on what it encodes. This makes it a cheap way to stabilize DLM attention without adding architectural complexity.

Key points

  • In LLaDA-Instruct and Dream-Instruct, the authors sampled 10K sink and 10K normal tokens during GSM8K and HumanEval inference. Sinks had much lower mean value L2 norms: 2.36 vs 7.60 for LLaDA and 3.15 vs 7.79 for Dream.
  • Placing the sink token at the start or the end of the sequence gives similar improvements. This suggests the benefit comes from its role as a fixed attractor, not from positional bias.
  • An ablation with 1, 2 and 4 sink tokens shows that one token delivers nearly all of the gain, with more tokens adding little.
  • An analysis of 100K GSM8K inference steps on the 0.5B and 1.5B models shows the pattern in the middle layers. There, the learned sink's norm stays far below that of the remaining tokens, even as it draws a substantial portion of the attention.
  • Forcing the sink's value states to zero vectors (Qwen2.5-0.5B) preserves the gains, e.g. ARC-e 55.68 vs 55.47 and PIQA 64.09 vs 63.44. The experiments cover only the 0.5B and 1.5B scales. Code is available at github.com/skywalker0523/OneTokenIsEnough.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓University of Chinese Academy of SciencesBeijing · China147 works · 8,480 citations🏢Xiaohongshu (Little Red Book)Shanghai · China11 works · 107 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Zihou Zhang, Zheyong Xie, Li Zhong, Haifeng Liu and Shaosheng Cao. 2026. One Token Is Enough: Improving Diffusion Language Models with a Sink Token. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.323

@inproceedings{zhang2026one,
  title = {{One Token Is Enough: Improving Diffusion Language Models with a Sink Token}},
  author = {Zihou Zhang, Zheyong Xie, Li Zhong, Haifeng Liu and Shaosheng Cao},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.323},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.323}
}

References

41 references, 15 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Federico Barbero, Alvaro Arroyo, Xiangming Gu, Christos Perivolaropoulos, Michael Bronstein, Petar Veličković, and Razvan Pascanu. 2025. Why do llms attend to the first token? arXiv preprint arXiv:2504.02732. — [arXiv:2504.02732](https://arxiv.org/abs/2504.02732) find
  2. Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. 2020. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439.
  3. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, and 39 others. 202… — Evaluating Large Language Models Trained on Code
  4. En Ci, Shanyan Guan, Yanhao Ge, Yilin Zhang, Wei Li, Zhenyu Zhang, Jian Yang, and Ying Tai. 2025. Describe, don’t dictate: Semantic image editing with natural language intent. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 19185–19194. find
  5. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457.
  6. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  7. Xin Dai, Pengcheng Huang, Zhenghao Liu, Shuo Wang, Yukun Yan, Chaojun Xiao, Yu Gu, Ge Yu, and Maosong Sun. 2026. Revealing the attention floating mechanism in masked diffusion models. arXiv preprint arXiv:2601.07894. find
  8. Yuntian Deng, Kiran Prasad, Roland Fernandez, Paul Smolensky, Vishrav Chaudhary, and Stuart Shieber. 2023. Implicit chain of thought reasoning via knowledge distillation. arXiv preprint arXiv:2311.01460. find
  9. Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. 2024. Model tells you what to discard: Adaptive kv cache compression for llms. In International Conference on Learning Representations (ICLR).
  10. Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, Peng Hao, and Kong Lingpeng. 2025. Scaling diffusion language models via adaptation from autoregressive models. In The Thirteenth International Conference on Learning Representat… find
  11. Xiangming Gu, Tianyu Pang, Chao Du, Qian Liu, Fengzhuo Zhang, Cunxiao Du, Ye Wang, and Min Lin. 2024. When attention sink emerges in language models: An empirical view. arXiv preprint arXiv:2410.10781. — [arXiv:2410.10781](https://arxiv.org/abs/2410.10781) find
  12. Ishaan Gulrajani and Tatsunori B Hashimoto. 2023. Likelihood-based diffusion language models. Advances in Neural Information Processing Systems, 36:16693–16715. find
  13. Samar Khanna, Siddhant Kharbanda, Shufan Li, Harshit Varma, Eric Wang, Sawyer Birnbaum, Ziyang Luo, Yanis Miraoui, Akash Palrecha, Stefano Ermon, Aditya Grover, and Volodymyr Kuleshov. 2025. Mercury: Ultra-fast language models based on diffusion. arXiv preprint arXiv:2506.17298. find
  14. Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. 2017. Race: Large-scale reading comprehension dataset from examinations. arXiv preprint arXiv:1704.04683.
  15. Nianyi Lin, Jiajie Zhang, Lei Hou, and Juanzi Li. 2025. Boundary-guided policy optimization for memory-efficient rl of diffusion large language models. arXiv preprint arXiv:2510.11683. find
  16. Ruikang Liu, Haoli Bai, Haokun Lin, Yuening Li, Han Gao, Zhengzhuo Xu, Lu Hou, Jun Yao, and Chun Yuan. 2024. IntactKV: Improving large language model quantization by keeping pivot tokens intact. In Findings of the Association for Computational Linguistics: ACL 2024, pages 7716–7741, Bangkok, Thailan… find
  17. Ilya Loshchilov and Frank Hutter. 2017. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. find
  18. Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. 2025a. Scaling up masked diffusion models on text. In The Thirteenth International Conference on Learning Representations. find
  19. Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. 2025b. Large language diffusion models. arXiv preprint arXiv:2502.09992.
  20. Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc-Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The lambada dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association fo…
  21. Guilherme Penedo, Hynek Kydlícek, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin A. Raffel, Leandro von Werra, and Thomas Wolf. 2024. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:30811–30849.
  22. Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang, Suozhi Huang, Liu Dayiheng, Zhou Jingren, and Lin Junyang. 2025. Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free. arXiv preprint arXiv:2505.06708. — [arXiv:2505.06708](https://arxiv.org/abs/2505.06708) find
  23. Qwen, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, Junyang Lin, and 24 others. 2025. Qwen2.5 Technical Report. 10.48550/arXiv.2412.151… find
  24. Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 2022. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695. find
  25. Maximo Eduardo Rulli, Simone Petruzzi, Edoardo Michielon, Fabrizio Silvestri, Simone Scardapane, and Alessio Devoto. 2025. Attention sinks in diffusion language models. arXiv preprint arXiv:2510.15731.
  26. Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. 2019. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint arXiv:1904.09728.
  27. Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. find
  28. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. find
  29. Yuerong Song, Xiaoran Liu, Ruixiao Li, Zhigeng Liu, Zengfeng Huang, Qipeng Guo, Ziwei He, and Xipeng Qiu. 2025. Sparse-dllm: Accelerating diffusion llms with dynamic cache eviction. arXiv preprint arXiv:2508.02558. — [arXiv:2508.02558](https://arxiv.org/abs/2508.02558) find
  30. Llama Team. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  31. Zeqing Wang, Gongfan Fang, Xinyin Ma, Xingyi Yang, and Xinchao Wang. 2025. Sparsed: Sparse attention for diffusion language models. arXiv preprint arXiv:2509.24014. — [arXiv:2509.24014](https://arxiv.org/abs/2509.24014) find
  32. Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. 2025. Fast-dllm: Training-free acceleration of diffusion llm by enabling kv cache and parallel decoding. arXiv preprint arXiv:2505.22618. — [arXiv:2505.22618](https://arxiv.org/abs/2505.22618) find
  33. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2023. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453. find
  34. Yi Xin, Qi Qin, Siqi Luo, Kaiwen Zhu, Juncheng Yan, Yan Tai, Jiayi Lei, Yuewen Cao, Keqi Wang, Yibin Wang, Jinbin Bai, Qian Yu, Dengyang Jiang, Yuandong Pu, Haoxing Chen, Le Zhuo, Junjun He, Gen Luo, Tianbin Li, and 13 others. 2025. Lumina-dimoo: An omni diffusion large language model for multimodal… — [arXiv:2510.06308](https://arxiv.org/abs/2510.06308) find
  35. Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang. 2025. Mmada: Multimodal large diffusion language models. arXiv preprint arXiv:2505.15809. — [arXiv:2505.15809](https://arxiv.org/abs/2505.15809) find
  36. Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. 2025. Dream 7b: Diffusion large language models. arXiv preprint arXiv:2508.15487. find
  37. Zebin You, Shen Nie, Xiaolu Zhang, Jun Hu, Jun Zhou, Zhiwu Lu, Ji-Rong Wen, and Chongxuan Li. 2025. Llada-v: Large language diffusion models with visual instruction tuning. arXiv preprint arXiv:2505.16933.
  38. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830. find
  39. Siyue Zhang, Yilun Zhao, Liyuan Geng, Arman Cohan, Anh Tuan Luu, and Chen Zhao. 2025. Diffusion vs. autoregressive language models: A text embedding perspective. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 4273–4303, Suzhou, China. Association for…
  40. Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. 2025. d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint arXiv:2504.12216. — [arXiv:2504.12216](https://arxiv.org/abs/2504.12216) find
  41. Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. 2025. Llada 1.5: Variance-reduced preference optimization for large language diffusion models. arXiv preprint arXiv:2505.19223.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhang-2026-one-token-is-enough-improving/