Token-level Data Selection for Safe LLM Fine-tuning

이 논문은 미세 조정 중 발생하는 안전성 저하를 해결하기 위해 샘플 단위가 아닌 토큰 단위로 위험을 정량화하고 제거하는 새로운 프레임워크 TOSS 와 점진적 정제 전략 TOSS-Pro 를 제안하여 안전성과 유용성을 동시에 향상시키는 방법을 제시합니다.

핵심

🍳 비유: "요리사 (AI) 의 레시피 수정"

생각해 보세요. 이미 훌륭한 요리사 (안전하게 훈련된 AI) 가 있습니다. 이 요리사는 유해한 재료를 절대 쓰지 않고, 손님에게 해로운 음식을 주지 않도록 훈련받았습니다.

이제 우리는 이 요리사에게 특정 지역 (예: 한국) 의 맛을 내게 하고 싶어서, 새로운 레시피 (데이터) 를 가르치려 합니다. 하지만 문제는 이 새로운 레시피 속에 실수로 독이 섞인 재료가 있거나, 유해한 조리법이 숨어있을 수 있다는 점입니다.

기존의 방법들은 이렇게 했습니다:

• 과거의 방법 (샘플 단위 삭제): "이 레시피 카드 한 장에 독이 섞인 글자가 하나라도 있으면, 그 카드 전체를 쓰레기통에 버려라!"
  • 문제점: 독이 섞인 글자가 하나뿐인데, 그 카드를 통째로 버리면 맛있는 소스나 중요한 조리법도 함께 사라져서 요리사가 한국 음식을 제대로 못 만들게 됩니다. (안전은 되지만, 유용성이 떨어짐)

💡 이 논문의 해결책: "TOSS (알맹이만 골라내는 정교한 체)"

이 연구팀은 **"레시피 카드 전체를 버릴 필요 없이, 독이 섞인 글자 (토큰) 하나하나만 골라내서 제거하자"**고 제안합니다.

1. 핵심 아이디어: "글자 단위 (Token-level) 의 정밀 검사"

레시피 한 장을 볼 때, "이 카드가 위험하니까 다 버려"가 아니라, **"이 카드의 10 번째 글자는 '독'이지만, 11 번째 글자는 '맛있는 소스'야. 10 번째 글자만 잘라내고 11 번째는 살리자"**는 방식입니다.

2. 어떻게 작동할까요? (두 명의 심사위원)

이 시스템은 두 명의 가상의 심사위원을 고용합니다.

• 심사위원 A (위험 감수자): "유해한 음식"을 만드는 데 특화된 요리사입니다. (위험한 패턴을 잘 아님)
• 심사위원 B (맛 탐구자): "맛있는 한국 음식"을 만드는 데 특화된 요리사입니다. (유용한 정보를 잘 아님)

이제 새로운 레시피의 글자 하나하나를 두 심사위원에게 보여줍니다.

• 만약 어떤 글자가 **심사위원 A(위험)**에게는 "아, 이거 우리 레시피랑 딱 맞아!"라고 생각되는데, **심사위원 B(맛)**에게는 "이건 우리 레시피랑 상관없어"라고 생각된다면?
• 결론: 그 글자는 위험하니까 잘라내자!
• 반대로 두 심사위원 모두 "이건 좋은 글자야"라고 하면 그 글자는 살려두자!

이렇게 글자 하나하나의 위험도를 계산해서, 해로운 부분만 '가위'로 잘라내고 유용한 부분은 그대로 남깁니다.

3. TOSS-Pro: "점점 더 똑똑해지는 검사관"

처음엔 검사관이 완벽하지 않을 수 있습니다. 그래서 TOSS-Pro라는 업그레이드 버전이 있습니다.

• 1 차 검사에서 위험한 글자를 찾아내서 제거하고, 그 정보를 바탕으로 검사관 A 를 다시 훈련시킵니다.
• 이렇게 반복적으로 훈련을 시키면, 검사관이 점점 더 정교하게 "위험한 글자"를 찾아내게 되어, 더 완벽한 요리를 만들 수 있게 됩니다.

---

🌟 왜 이 방법이 좋은가요?

1. 안전하면서도 맛있습니다: 유해한 글자만 잘라내서, 요리사는 여전히 한국 음식의 맛 (유용성) 을 잘 내면서도 독 (위험) 은 내지 않습니다.
2. 데이터 낭비가 없습니다: 레시피 카드 전체를 버리지 않으므로, 소중한 정보 (맛있는 소스) 를 잃지 않습니다.
3. 정밀합니다: "무조건 다 버리는" 기존 방법보다 훨씬 정교하게 문제를 해결합니다.

📝 요약

이 논문은 **"AI 를 수정할 때, 나쁜 부분 때문에 좋은 부분까지 다 버리지 말고, 나쁜 글자 (토큰) 하나하나만 골라내서 제거하는 정교한 기술 (TOSS)"**을 개발했다고 말합니다. 마치 나쁜 과일은 도려내고 좋은 과일은 그대로 남겨서 맛있는 주스를 만드는 것과 같습니다.

이 방법을 쓰면 AI 는 여전히 똑똑하고 유용하면서도, 해로운 말은 하지 않는 안전한 친구가 될 수 있습니다.

기술 요약

1. 문제 정의 (Problem Definition)

대형 언어 모델 (LLM) 은 특정 도메인이나 애플리케이션에 맞게 **파인튜닝 (Fine-tuning)**을 거치면 성능을 크게 향상시킬 수 있습니다. 그러나 최근 연구에 따르면, 커스텀 데이터셋으로 파인튜닝하는 과정에서 모델의 **안전성 (Safety)**이 심각하게 저하될 수 있습니다.

• 기존 방어 방법의 한계: 기존 데이터 중심 방어 전략 (데이터 혼합, 샘플 단위 필터링 등) 은 샘플 (Sample) 단위로 작동합니다.
  • 데이터 혼합: 안전 관련 데이터를 추가하면 안전성은 높아지지만, 모델이 과도하게 거절하거나 (Over-refusal) 타스크 성능이 떨어지는 안전성 - 유용성 (Safety-Utility) 트레이드오프 문제가 발생합니다.
  • 샘플 단위 필터링: 유해한 샘플 전체를 제거하는 방식은 안전성을 일부 개선하지만, 유해한 샘플 내부에도 유용한 작업 정보가 포함되어 있는 경우 이를 함께 버리게 되어 유용성 (Utility) 손실이 발생합니다.
• 핵심 가설: 안전성 저하의 원인은 전체 샘플이 아니라, 토큰 (Token) 단위로 존재합니다. 겉보기에 안전해 보이는 (Benign) 데이터라도 특정 토큰이 모델의 안전 정렬 (Safety Alignment) 을 무너뜨릴 수 있으며, 해로운 신호와 유용한 신호가 동일한 샘플 내에 혼재되어 있습니다.

2. 방법론 (Methodology)

저자들은 토큰 단위 데이터 선택 (TOSS) 프레임워크를 제안하여, 안전성을 해치는 토큰만 정밀하게 제거하고 유용한 토큰은 보존하는 방식을 도입했습니다.

A. 핵심 메커니즘: 손실 차이 (Loss Difference) 기반 점수화

TOSS 는 두 개의 참조 모델 (Reference Models) 을 학습시켜 각 토큰의 안전성 위험을 정량화합니다.

1. 안전성 저하 모델 (Safety-degraded model, $f_{\theta h}$): 유해 데이터 ($D_h$) 로 파인튜닝된 모델로, 유해 패턴을 인식하는 전문가 역할을 합니다.
2. 유용성 지향 모델 (Utility-oriented model, $f_{\theta u}$): 고품질 유틸리티 데이터 ($D_u$) 로 학습된 모델로, 하류 작업 (Downstream task) 에 필요한 정보를 반영합니다.

토큰 점수 계산:
커스텀 데이터셋의 각 토큰 $y_{i,j}$에 대해 두 모델의 손실 (Loss) 차이를 계산합니다.
$$S(y_{i,j}) = -\log P(y_{i,j} | \dots; \theta_u) + \log P(y_{i,j} | \dots; \theta_h)$$

• 높은 점수: 안전성 저하 모델에서는 낮은 손실 (유해 패턴에 적합), 유용성 모델에서는 높은 손실 (유용한 정보와 거리가 먼) 을 의미함 $\rightarrow$ 유해 토큰으로 간주하여 제거.
• 낮은 점수: 유용성 모델과 잘 맞고 안전성 모델과 거리가 먼 경우 $\rightarrow$ 보존.

이 점수를 기준으로 전체 데이터셋의 토큰을 전역적으로 순위 매겨 (Global Ranking), 상위 $d\%$의 토큰을 마스킹 (제거) 하고 나머지 토큰으로 선택적 파인튜닝을 수행합니다.

B. TOSS-Pro: 점진적 정제 전략 (Progressive Refinement)

초기 안전성 저하 모델의 품질이 선택 성능에 영향을 미칠 수 있다는 점을 보완하기 위해 TOSS-Pro를 제안했습니다.

• 반복적 개선: 초기 유해 데이터셋으로 학습된 안전성 저하 모델을 기반으로 토큰 점수를 계산합니다.
• 고품질 유해 샘플 추출: 점수가 높은 (가장 유해한) 토큰이 포함된 샘플들을 선별하여 유해 데이터셋에 추가합니다.
• 모델 업데이트: 확장된 데이터셋으로 안전성 저하 모델을 재학습시켜, 더 정확한 유해 토큰 식별 능력을 갖도록 점진적으로 정제합니다. 이 과정을 $T$회 반복한 후 최종 마스킹을 적용합니다.

3. 주요 기여 (Key Contributions)

1. 토큰 단위 안전성 진단: 파인튜닝 중 안전성 저하가 샘플 단위가 아닌 토큰 단위에서 발생함을 체계적으로 분석하고 실증했습니다. 특히 초기 응답 토큰뿐만 아니라 중간/후반 토큰에서도 안전성 저하 신호가 발견됨을 보였습니다.
2. TOSS 프레임워크 제안: 안전성 저하 모델과 유용성 모델을 결합한 손실 차이 메트릭을 도입하여, 해로운 토큰만 '외과 수술'처럼 제거하는 정밀한 데이터 선택 방식을 구현했습니다.
3. TOSS-Pro 전략: 점진적 정제를 통해 안전성 저하 모델의 식별 능력을 향상시켜, 초기 모델의 품질에 의존하지 않고 더 강력한 방어 성능을 달성했습니다.
4. 안전성 - 유용성 최적 균형: 기존 샘플 단위 방어 방법 (SEAL 등) 보다 월등히 우수한 안전성 유지와 동시에 하류 작업 성능을 보존하는 결과를 입증했습니다.

4. 실험 결과 (Results)

Llama-3-8B-Instruct 및 Llama-2-7B-Chat-hf 모델을 대상으로 한 실험 결과, TOSS 는 기존 방법들을 압도했습니다.

• 성능 비교 (Win Rate 기준):
  • 안전성 (HEx-PHI, ANTHROPIC HH): TOSS-Pro 는 Llama-3-8B 에서 평균 83.83% 의 승률을 기록하여, 기존 최상위 방법인 SEAL(61.48%) 보다 약 22% 이상 높은 안전성을 보였습니다.
  • 유용성 (SLIMORCA): 안전성을 유지하면서도 유용성 점수 (68.85%) 를 SEAL(57.41%) 보다 크게 상회하며, 표준 SFT 대비도 우수한 성능을 유지했습니다.
• 전송 가능성 (Transferability): 한 모델 (Llama-3-8B) 에 대해 선택된 토큰 마스킹 패턴을 동일한 토크나이저를 사용하는 다른 모델 (Llama-3.2-1B/3B) 에 직접 적용했을 때도 높은 성능을 유지하여, 토크나이저 단위만 선택하면 된다는 실용적 장점을 입증했습니다.
• Ablation Study:
  • 전역 순위 (Global Ranking): 샘플 내 고정 비율 제거 (Local Ranking) 보다 전역적 순위 매기기가 훨씬 효과적이었습니다.
  • 두 모델의 상호보완성: 안전성 모델만 또는 유용성 모델만 사용할 경우 각각 유용성 저하나 안전성 실패가 발생했으나, 두 모델을 결합했을 때 최적의 균형을 이뤘습니다.
  • 점진적 정제: TOSS-Pro 는 TOSS 대비 안전성 성능을 추가적으로 6% 까지 향상시켰습니다.

5. 의의 및 중요성 (Significance)

이 연구는 LLM 파인튜닝의 안전성 문제를 해결하기 위한 **패러다임의 전환 (Sample-level $\rightarrow$ Token-level)**을 제시합니다.

• 정밀한 제어: 불필요한 데이터 전체를 버리는 것이 아니라, 해로운 신호가 있는 토큰만 선별적으로 제거함으로써 데이터의 가치 (Utility) 를 최대한 보존하면서도 안전성 (Safety) 을 강력하게 방어합니다.
• 실용성: 복잡한 알고리즘 수정이나 추가적인 추론 비용 없이, 데이터 전처리 단계에서 효율적으로 적용 가능하여 실제 산업 환경에서의 맞춤형 LLM 배포에 큰 기여를 할 것으로 기대됩니다.
• 안전성 저하의 근본 원인 해결: 파인튜닝 후 모델을 '수리'하는 방식이 아니라, 학습 데이터 자체에서 유해 요인을 제거하여 안전성 저하를 원천 차단하는 접근법을 제시했습니다.

결론적으로, TOSS 는 맞춤형 LLM 의 안전성과 유용성 사이의 긴장 관계를 해결하는 새로운 표준을 제시하며, 신뢰할 수 있는 AI 시스템 구축에 중요한 기여를 하고 있습니다.