Token-level Data Selection for Safe LLM Fine-tuning

TL;DR

This paper introduces TOSS, a token-level data selection framework for fine-tuning large language models that effectively balances safety and utility by identifying and removing unsafe tokens during training.

Contribution

The paper proposes a novel token-level safety risk measurement and a progressive refinement strategy for safer LLM fine-tuning, outperforming existing sample-level methods.

Findings

TOSS effectively identifies unsafe tokens during fine-tuning.

TOSS improves downstream task performance while enhancing safety.

The progressive refinement strategy further enhances safety detection accuracy.

Abstract

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.