안전한 LLM 미세 조정(Fine-tuning)을 위한 토큰 수준 데이터 선택

Token-level Data Selection for Safe LLM Fine-tuning

대규모 언어 모델(LLM)을 특정 도메인 및 애플리케이션에 맞게 조정하는 표준적인 방법으로, 사용자 정의 데이터 세트를 사용한 미세 조정이 널리 사용되고 있습니다. 그러나 최근 연구에 따르면 이러한 미세 조정은 모델의 안전성에 상당한 저하를 초래할 수 있습니다. 기존의 안전 방어 방법은 샘플 수준에서 작동하며, 종종 안전성과 유용성 사이에서 만족스럽지 못한 균형을 보입니다. 이러한 제한 사항을 해결하기 위해, 우리는 미세 조정 과정에서 발생하는 안전성 저하를 체계적으로 토큰 수준에서 분석합니다. 이를 바탕으로, 우리는 각 토큰의 안전 위험을 안전성 저하 모델과 유용성 중심 모델 간의 손실 차이를 측정하여 정량화하는 새로운 프레임워크인 안전한 LLM 미세 조정을 위한 토큰 수준 데이터 선택(TOSS)을 제안합니다. 이러한 토큰 수준의 세분성은 안전하지 않은 토큰을 정확하게 식별하고 제거하여, 귀중한 작업 관련 정보를 보존할 수 있도록 합니다. 또한, 우리는 안전성 저하 모델이 안전하지 않은 토큰을 식별하는 능력을 반복적으로 향상시키는 점진적인 개선 전략인 TOSS-Pro를 소개합니다. 광범위한 실험 결과, 우리의 접근 방식은 미세 조정 과정에서 LLM을 안정적으로 보호하면서 기존의 샘플 수준 방어 방법보다 우수한 하위 작업 성능을 달성함을 보여줍니다. 저희의 코드는 https://github.com/Polly-LYP/TOSS 에서 확인할 수 있습니다.

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.