ML Security PapersToken-level Data Selection for Safe LLM Fine-tuning
Yanping Li 1, Zhening Liu 1, Zijian Li 1, Zehong Lin 2, Jun Zhang 1
Published on arXiv
2603.01185
Transfer Learning Attack
OWASP ML Top 10 — ML07
Prompt Injection
OWASP LLM Top 10 — LLM01
Key Finding
Token-level data selection significantly outperforms sample-level defense methods in both safety and utility across custom fine-tuning datasets with varying harmful data ratios
TOSS (Token-level data Selection for Safe LLM fine-tuning)
Novel technique introduced
Fine-tuning large language models (LLMs) on custom datasets has become a standard approach for adapting these models to specific domains and applications. However, recent studies have shown that such fine-tuning can lead to significant degradation in the model's safety. Existing defense methods operate at the sample level and often suffer from an unsatisfactory trade-off between safety and utility. To address this limitation, we perform a systematic token-level diagnosis of safety degradation during fine-tuning. Based on this, we propose token-level data selection for safe LLM fine-tuning (TOSS), a novel framework that quantifies the safety risk of each token by measuring the loss difference between a safety-degraded model and a utility-oriented model. This token-level granularity enables accurate identification and removal of unsafe tokens, thereby preserving valuable task-specific information. In addition, we introduce a progressive refinement strategy, TOSS-Pro, which iteratively enhances the safety-degraded model's ability to identify unsafe tokens. Extensive experiments demonstrate that our approach robustly safeguards LLMs during fine-tuning while achieving superior downstream task performance, significantly outperforming existing sample-level defense methods. Our code is available at https://github.com/Polly-LYP/TOSS.
Key Contributions
- First systematic token-level diagnosis showing that safety-degrading and utility-enhancing signals are intertwined at the token level during fine-tuning, revealing fundamental limitations of sample-level defenses
- TOSS framework that scores each token's safety risk using the loss difference between a safety-degraded reference model and a utility-oriented reference model, enabling precise removal of unsafe tokens
- TOSS-Pro progressive refinement strategy that iteratively improves the safety-degraded model's ability to identify unsafe tokens using increasingly higher-quality supervision
🛡️ Threat Analysis
The paper directly addresses safety degradation that occurs specifically during LLM fine-tuning (transfer learning). TOSS defends against attacks exploiting the fine-tuning process to erode safety alignment, including adversarial harmful data and benign-yet-harmful data — both operating through the transfer/fine-tuning pipeline.