소수 토큰으로 안전을 지키다 파인튜닝 시 안전 토큰 제약 기법

초록

본 논문은 파인튜닝 과정에서 안전 정렬이 손상되는 문제를 해결하기 위해, 안전에 핵심적인 소수의 토큰에만 신뢰도를 유지하도록 제약하는 PACT 방법을 제안한다. 안전 토큰을 식별하고, 파인튜닝 중에 이 토큰들의 출력 확률을 원본 안전 정렬 모델과 일치시키는 정규화를 적용함으로써, 전반적인 성능 저하 없이 안전성을 보존한다.

상세 분석

본 연구는 대형 언어 모델(LLM)의 파인튜닝 시 발생하는 안전 정렬 드리프트 현상을 토큰 수준에서 분석한다. 먼저, 안전 정렬된 모델(M_safe)과 기본 사전학습 모델(M_base) 사이의 토큰별 확률 차이를 구해 Δ_t(v)=p_Msafe(v|x,y< t)−p_Mbase(v|x,y< t) 로 정의하고, 이를 전체 데이터와 위치에 평균하여 전역 차이 점수 d(v)를 산출한다. 상위 K=50개의 토큰을 ‘안전 토큰’으로 선정했으며, 여기에는 “I”, “cannot”, “sorry” 등 거부 표현에 자주 등장하는 어휘가 포함된다.

안전 토큰의 역할을 검증하기 위해 로그잇에 상수 α=5.0을 더해 신뢰도를 인위적으로 상승시키거나, −10⁹으로 억제하는 실험을 수행했다. 안전 토큰을 강화하면 공격 성공률이 33.5%에서 0.5%로 급감했으며, 반대로 억제하면 41.5%까지 상승했다. 이는 안전 정렬이 전체 어휘가 아니라 소수 토큰에 집중되어 있음을 실증한다.

다음으로, 유해 데이터가 섞인 파인튜닝(예: GSM8K + Adversarial 데이터) 동안 에포크별 안전 토큰 평균 신뢰도와 안전 점수(1−ASR)를 추적했다. 유해 데이터 비중이 증가할수록 두 지표가 동시에 감소하는 경향을 보였으며, 이는 안전 토큰 신뢰도의 저하가 정렬 손실의 주요 원인임을 시사한다.

이러한 관찰을 바탕으로 제안된 PACT는 두 단계로 구성된다. ① 안전 토큰 가중 정규화: KL 발산 손실 L_KL=∑_v w(v)·KL(p_Mft(v)||p_Msafe(v))에서 w(v)=d(v) 혹은 사전 정의된 안전 가중치를 부여해 안전 토큰에 강한 제약을 가한다. ② 안전 신호 보정: 학습 초기에 전체 컨텍스트 기반 분포와 응답 전용 분포를 혼합해 안전 토큰의 기준을 안정화한다. 비안전 토큰은 일반적인 작업 손실 L_task와 함께 최적화되므로, 모델은 안전성을 유지하면서도 다운스트림 성능을 크게 희생하지 않는다.

실험은 Qwen‑2.5‑7B‑Instruct, Llama‑3.1‑8B, Gemma‑2‑9B 등 네 가지 모델군에 대해 GSM8K(수학), SST‑2(감성), AGNEWS(뉴스 분류) 세 작업에 적용했다. 유해 데이터 비율을 0%~10%로 변동시키면서 PACT와 기존 방법(SafeLoRA, Lisa, 전체 KL 제약 등)을 비교했으며, PACT는 StrongReject에서 5.75‑9.27%, HarmBench에서 13.5‑29.5%의 공격 성공률을 달성하면서, 작업 정확도는 베이스 파인튜닝과 거의 동일하거나 약간 상회했다. 특히, 안전 토큰 수가 50개에 불과함에도 불구하고 전역 정규화 대비 2‑3배 높은 안전‑유틸리티 트레이드오프를 보여준다.

결론적으로, 안전 정렬은 어휘 전반이 아니라 핵심 토큰에 집중되어 있으며, 이 토큰에 대한 신뢰도만을 정밀하게 제어하면 파인튜닝 과정에서 안전성을 효과적으로 보존할 수 있다. PACT는 파라미터‑레벨 제약보다 더 미세하고 유연한 접근법으로, 다양한 모델과 작업에 적용 가능하며, 안전‑유틸리티 간의 균형을 최적화한다.