Elamur Research⌕
☰
home / papers / 2026 / You Only Need One Single Token to Refine Safety Alignment

You Only Need One Single Token to Refine Safety Alignment

Paper · 2026 · Wenqian Yu, Shuo Chen, Zhijiang Li et al.

Summary

Safety-aligned LLMs often over-refuse: they reject prompts that look dangerous but are actually benign. Existing training-free fixes include prompting, contrastive decoding and activation steering. These add inference cost or depend on the specific architecture, and they correct outputs after the fact without changing how the model separates harmful from benign inputs. The authors propose Single Token Alignment (STA). It keeps the base model frozen and trains only one continuous prefix embedding, about 4,096 parameters. The training objective combines an SFT loss with a SimPO preference loss. It also uses two weighting schemes. Hard weighting is strict data filtering: responses are labeled by the DeepSeek API as Direct Refusal, Full Compliance or Safe Partial Compliance, and only unambiguous preference pairs are kept. Soft weighting is a focal term that emphasizes difficult examples. Across 9 models and 10 datasets, covering LLMs, MLLMs and reasoning models, STA improves the balance between helpfulness and safety. It costs less than prefix tuning or LoRA and adds no inference overhead. This makes it a cheap, model-agnostic way to adjust the alignment of already-aligned models.

Key points

  • Efficiency: On a single A40 GPU, STA trains 4,096 parameters using 17.57 GB VRAM at 61.70 s per 100 steps. Prefix tuning needs 819K parameters, 19.10 GB and 73.99 s. LoRA needs about 4.2M parameters, 25.89 GB and 136.85 s.
  • Token length: Ablations on Llama-2 with prefix lengths N ∈ {2, 4, 16, 256} showed no consistent gain over N = 1. Removing any single component (SFT, SimPO, hard weighting or soft weighting) lowered USR.
  • Multimodal models (MOSSBench): On Llava-mistral, USR rose from 22.5% (base) to 100%, compared with 26.25% for Magic Image. On QwenVL, the gain was small: 78.75% to 79.58%.
  • Reasoning models: Qwen3-4B and Qwen3-8B were trained on FalseReject data with chain-of-thought responses. The largest safety gains were on DNA, with +6.93 and +5.00 points.
  • Jailbreaks and limitations: After brief adversarial training, USR on jailbreak tests reached 85–100%. The authors note that results depend on the quality of the training signals. STA has not yet been tested on unaligned base models or against newer jailbreak techniques.

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 OxfordOxford · United Kingdom122 works · 160,277 citations🎓Wuhan UniversityWuhan · China47 works · 1,246 citations🏢LinkedInSunnyvale · United States3 works · 91 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksSocial Sciences · 1,954 worksSafety Research · 626 works

🧭 Topics 1Research subjects and languages the work is about.Bias, fairness & safety · 587 works

Cite

Cite as Wenqian Yu, Shuo Chen, Zhijiang Li, Zhipeng Wang and Jindong Gu. 2026. You Only Need One Single Token to Refine Safety Alignment. 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.662

@inproceedings{yu2026you,
  title = {{You Only Need One Single Token to Refine Safety Alignment}},
  author = {Wenqian Yu, Shuo Chen, Zhijiang Li, Zhipeng Wang and Jindong Gu},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.662},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.662}
}

References

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

  1. To explore jailbreak robustness of STA, we conducted a preliminary investigation using the GCG jailbreak method. The initial STA model shows vulnerability to jailbreak(43% USR). However, our Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, and Neel Nanda. 2024. Ref… — Refusal in Language Models Is Mediated by a Single Direction
  2. Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, and 1 others. 2022. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862. find
  3. Jiaming Ji, Mickel Liu, Josef Dai, Xuehai Pan, Chi Zhang, Ce Bian, Boyuan Chen, Ruiyang Sun, Yizhou Wang, and Yaodong Yang. 2023a. Beavertails: Towards improved safety alignment of llm via a human-preference dataset. Advances in Neural Information Processing Systems, 36:24678–24704.
  4. Rishabh Bhardwaj and Soujanya Poria. 2023. Red-teaming large language models using chain of utterances for safety-alignment. arXiv preprint arXiv:2308.09662.
  5. Jiaming Ji, Tianyi Qiu, Boyuan Chen, Borong Zhang, Hantao Lou, Kaile Wang, Yawen Duan, Zhonghao He, Jiayi Zhou, Zhaowei Zhang, and 1 others. 2023b. Ai alignment: A comprehensive survey. arXiv preprint arXiv:2310.19852.
  6. Federico Bianchi, Mirac Suzgun, Giuseppe Attanasio, Paul Röttger, Dan Jurafsky, Tatsunori Hashimoto, and James Zou. 2023. Safety-tuned llamas: Lessons from improving the safety of large language models that follow instructions. arXiv preprint arXiv:2309.07875.
  7. Zouying Cao, Yifei Yang, and Hai Zhao. 2025. Scans: Mitigating the exaggerated safety for llms via safety-conscious activation steering. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 23523–23531. find
  8. Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. 2017. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30. find
  9. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, and 1 others. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  10. Justin Cui, Wei-Lin Chiang, Ion Stoica, and Cho-Jui Hsieh. 2024. Or-bench: An over-refusal benchmark for large language models. arXiv preprint arXiv:2405.20947.
  11. Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, and 1 others. 2023. Mistral 7b. arXiv preprint arXiv:2310.06825. find
  12. Xiang Lisa Li and Percy Liang. 2021. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190. find
  13. Xirui Li, Hengguang Zhou, Ruochen Wang, Tianyi Zhou, Minhao Cheng, and Cho-Jui Hsieh. 2024. Mossbench: Is your multimodal language model oversensitive to safe queries? arXiv preprint arXiv:2406.17806. — [arXiv:2406.17806](https://arxiv.org/abs/2406.17806) find
  14. Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, and 1 others. 2024a. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437. find
  15. Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee. 2024b. Improved baselines with visual instruction tuning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 26296–26306.
  16. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, and 1 others. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  17. Mantas Mazeika, Long Phan, Xuwang Yin, Andy Zou, Zifan Wang, Norman Mu, Elham Sakhaee, Nathaniel Li, Steven Basart, Bo Li, and 1 others. Harmbench: A standardized evaluation framework for automated red teaming and robust refusal. arXiv preprint arXiv:2402.04249.
  18. Javid Ebrahimi, Anyi Rao, Daniel Lowd, and Dejing Dou. 2018. Hotflip: White-box adversarial examples for text classification. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 31–36.
  19. Yu Meng, Mengzhou Xia, and Danqi Chen. 2024. Simpo: Simple preference optimization with a reference-free reward. Advances in Neural Information Processing Systems, 37:124198–124235. find
  20. Rima Hazra, Sayan Layek, Somnath Banerjee, and Soujanya Poria. 2024. Safety arithmetic: A framework for test-time safety alignment of language models by steering parameters and activations. arXiv preprint arXiv:2406.11801. find
  21. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.
  22. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2016. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843.
  23. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, and 1 others. 2022. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744.
  24. Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. 2023. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36:53728–53741.
  25. Ruchira Ray and Ruchi Bhalani. 2024. Mitigating exaggerated safety in large language models. arXiv preprint arXiv:2405.05418. — [arXiv:2405.05418](https://arxiv.org/abs/2405.05418) find
  26. Nina Rimsky, Nick Gabrieli, Julian Schulz, Meg Tong, Evan Hubinger, and Alexander Turner. 2024. Steering llama 2 via contrastive activation addition. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15504–15522. find
  27. Paul Röttger, Hannah Kirk, Bertie Vidgen, Giuseppe Attanasio, Federico Bianchi, and Dirk Hovy. 2024. Xstest: A test suite for identifying exaggerated safety behaviours in large language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational L…
  28. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. find
  29. Leheng Sheng, Changshuo Shen, Weixiang Zhao, Junfeng Fang, Xiaohao Liu, Zhenkai Liang, Xiang Wang, An Zhang, and Tat-Seng Chua. 2025. Alphasteer: Learning refusal steering with principled null-space constraint. arXiv preprint arXiv:2506.07022. — [arXiv:2506.07022](https://arxiv.org/abs/2506.07022) find
  30. Chenyu Shi, Xiao Wang, Qiming Ge, Songyang Gao, Xianjun Yang, Tao Gui, Qi Zhang, Xuan-Jing Huang, Xun Zhao, and Dahua Lin. 2024. Navigating the overkill in large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages… find
  31. Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. 2020. Learning to summarize with human feedback. Advances in neural information processing systems, 33:3008– 3021. find
  32. Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, and 1 others. 2024. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.08295. find
  33. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, and 1 others. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.
  34. Peng Wang, Shuai Bai, Sinan Tan, Shijie Wang, Zhihao Fan, Jinze Bai, Keqin Chen, Xuejing Liu, Jialin Wang, Wenbin Ge, and 1 others. 2024a. Qwen2-vl: Enhancing vision-language model’s perception of the world at any resolution. arXiv preprint arXiv:2409.12191.
  35. Xinpeng Wang, Chengzhi Hu, Paul Röttger, and Barbara Plank. 2024b. Surgical, cheap, and flexible: Mitigating false refusal in language models via single vector ablation. arXiv preprint arXiv:2410.03415. — [arXiv:2410.03415](https://arxiv.org/abs/2410.03415) find
  36. Yuxia Wang, Haonan Li, Xudong Han, Preslav Nakov, and Timothy Baldwin. 2023. Do-not-answer: A dataset for evaluating safeguards in llms. arXiv preprint arXiv:2308.13387.
  37. Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. 2021. Finetuned language models are zero-shot learners. arXiv preprint arXiv:2109.01652.
  38. Yotam Wolf, Noam Wies, Oshri Avnery, Yoav Levine, and Amnon Shashua. 2023. Fundamental limitations of alignment in large language models. arXiv preprint arXiv:2304.11082.
  39. Yifan Xia, Guorui Chen, Wenqian Yu, Zhijiang Li, Philip Torr, and Jindong Gu. 2025. Reimagining safety alignment with an image. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 9600–9614. find
  40. Zhangchen Xu, Fengqing Jiang, Luyao Niu, Jinyuan Jia, Bill Yuchen Lin, and Radha Poovendran. 2024. Safedecoding: Defending against jailbreak attacks via safety-aware decoding. arXiv preprint arXiv:2402.08983.
  41. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
  42. Zhehao Zhang, Weijie Xu, Fanyou Wu, and Chandan K Reddy. 2025. Falsereject: A resource for improving contextual safety and mitigating over-refusals in llms via structured reasoning. arXiv preprint arXiv:2505.08054. — [arXiv:2505.08054](https://arxiv.org/abs/2505.08054) find
  43. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, and 1 others. 2023. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in neural information processing systems, 36:46595–46623.
  44. Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J Zico Kolter, and Matt Fredrikson. 2023a. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.15043.
  45. Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr, J Zico Kolter, and Matt Fredrikson. 2023b. Universal and transferable adversarial attacks on aligned language models. arXiv preprint arXiv:2307.15043.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/yu-2026-you-only-need-one-single/