Pruning Unsafe Tickets: A Resource-Efficient Framework for Safer and More Robust LLMs

1. 연구 배경 및 문제 정의 (Problem)

대규모 언어 모델 (LLM) 과 비전 - 언어 모델 (VLM) 이 실제 응용 분야에 빠르게 도입되고 있지만, 사전 학습 (Pre-training) 과정에서 학습된 위험한 행동 (Unsafe Behaviors) 이 여전히 존재합니다.

• 기존 방법의 한계: 현재 주류인 정렬 (Alignment) 기법인 SFT(지도 미세조정) 와 RLHF(인간 피드백을 통한 강화학습) 는 선호하는 응답을 생성하도록 모델을 유도하지만, 위험한 출력을 유발하는 하위 네트워크 (Unsafe Subnetworks) 를 명시적으로 제거하지는 않습니다.
• 자원 소모: 기존 정렬 방법은 대규모 데이터셋, 막대한 GPU 메모리, 긴 학습 시간이 필요합니다.
• 취약성: 이러한 모델들은 여전히 적대적 프롬프트 (Jailbreak) 나 분포 변화에 취약하며, 내부 모델 개입이나 프롬프트 엔지니어링과 같은 경량 대안들은 확장성이나 효율성 측면에서 한계가 있습니다.

2. 핵심 방법론 (Methodology)

저자들은 로또 티켓 가설 (Lottery Ticket Hypothesis, LTH) 에 기반하여, 정렬된 모델 내부에도 위험한 행동을 담당하는 희소 하위 네트워크인 "위험 티켓 (Unsafe Tickets)" 이 남아있다고 가정합니다. 이를 제거하여 안전성을 유지하면서도 성능을 보존하는 "안전 티켓 (Safety Tickets)" 을 찾아내는 자원 효율적 가지치기 (Pruning) 프레임워크 를 제안합니다.

프레임워크는 크게 4 단계로 구성됩니다:

1. 행동 프로파일링 (Behavior Profiling):

   • 모델의 실제 실패 사례를 반영하기 위해 모델에 다양한 프롬프트를 입력하여 안전/위험 응답 데이터를 생성합니다.
   • 외부 안전 분류기 (Llama-Guard) 를 사용하여 레이블을 자동화하고, 응답 길이를 고정하여 일관성을 유지합니다.

2. 귀속 분석 (Attribution Analysis):

   • 기존 경량화 기법인 Wanda를 확장하여 사용합니다.
   • Gradient-Free(경사 불필요) 방식: 역전파 없이 파라미터가 위험/안전 생성에 기여하는 정도를 추정합니다.
   • 응답 중심 마스크: 프롬프트 토큰이 아닌 응답 생성 단계 (Response tokens) 의 활성화 값만 고려하여, 위험한 행동이 발생하는 시점의 파라미터 중요도를 정확히 파악합니다.

3. 컴포넌트 스코어링 (Component Scoring):

   • 위험 데이터셋 ($S'_u$) 과 안전 데이터셋 ($S'_s$) 에 대한 중요도 점수의 비율을 계산하여 대비 중요도 (Contrastive Importance Score) 를 도출합니다.
   • $I = S'_u / (S'_s + \epsilon)$ 공식으로, 안전 생성에는 기여하지만 위험 생성에는 과도하게 기여하는 파라미터를 식별합니다.
   • 각 컴포넌트 내에서 Z-score 정규화를 수행하여 층 간 비교를 가능하게 합니다.

4. 반복적 가지치기 (Iterative Pruning):

   • 한 번에 모든 파라미터를 제거하는 것이 아니라, 반복적으로 가장 중요한 (위험한) 파라미터를 제거합니다.
   • Greedy Search: 매 Iteration 마다 가장 높은 점수를 가진 컴포넌트를 선택하여 빠르게 가지치기.
   • Beam Search: 여러 가지 가지치기 경로를 탐색하여 안전성과 유용성 (Utility) 의 최적 균형을 찾는 더 정교한 방법.

3. 주요 기여 (Key Contributions)

1. 자원 효율적 정렬 프레임워크: 미세조정 (Fine-tuning) 이나 추가 추론 비용 없이, 제한된 GPU 리소스 (경사 불필요) 로 안전성과 견고성을 동시에 향상시키는 모델-중립적 (Model-agnostic) 가지치기 프레임워크를 제안했습니다.
2. 새로운 관점 (LTH 기반): 모델 내에 존재하는 "위험 티켓"을 제거함으로써 "안전 티켓"을 드러낸다는 관점을 정립했습니다. 이는 모델이 본질적으로 안전성을 유지할 수 있는 하위 네트워크를 포함하고 있음을 시사합니다.
3. 광범위한 검증: 언어 모델 (Mistral 계열) 과 비전 - 언어 모델 (LLaVA) 에서 광범위한 평가를 수행하여, 위험한 생성을 대폭 줄이면서도 일반 성능을 유지함을 입증했습니다.

4. 실험 결과 (Results)

실험은 Mistral-7B, Mistral-Nemo, Mistral-Small, LLaVA-v1.6-Mistral-7B 등 다양한 모델에서 수행되었습니다.

• 안전성 향상:
  • Mistral-7B: 위험 응답 비율이 22.8% 에서 1.17% (Beam Search 기준) 로 감소했습니다.
  • LLaVA: 위험 응답 비율이 17.1% 에서 1.96% 로 감소했습니다.
  • 기존 방법 (DPO, Circuit Breaker, Goal Prioritization) 대비 위험 응답을 효과적으로 차단하면서도 과도한 거부 (Over-Refusal) 를 현저히 줄였습니다. (예: Circuit Breaker 는 위험 응답을 0% 로 만들지만 정상 질문도 100% 거부함).
• 견고성 (Robustness):
  • GCG, AutoDAN, PAIR, TAP 등 다양한 Jailbreak 공격에 대한 공격 성공률 (ASR) 을 크게 낮췄습니다. (예: Mistral-7B 에서 AutoDAN 공격 성공률을 45.7% → 0.67% 로 감소).
• 유용성 (Utility) 보존:
  • MT-Bench 등을 통한 유용성 점수 저하가 미미했습니다. (Mistral-7B 기준 8.07 → 7.13).
• 자원 효율성:
  • 학습 시간: Mistral-7B 기준 Greedy 방식은 455 초 만에 완료.
  • 메모리: 최대 18GB GPU 메모리만 요구 (기존 Circuit Breaker 는 36GB 이상 필요).
  • 추론 비용: 추가 토큰이나 메모리 오버헤드 없이 기존 모델과 동일한 속도로 추론 가능.
• 가치 있는 발견:
  • 가지치기는 모델의 위험 입력 인식 능력을 해치지 않으며, 출력 분포를 재조정하여 위험한 완성도를 줄이고 적절한 거절 응답을 늘립니다.
  • 가지치기가 집중된 영역은 Self-Attention Output Projection과 MLP 의 두 번째 레이어로, 이는 기존 연구와 일치하며 위험 행동의 핵심 회로임을 시사합니다.

5. 의의 및 결론 (Significance)

이 논문은 사후 (Post-hoc) 정렬 전략으로서 가지치기의 잠재력을 입증했습니다.

• 실용성: 고비용의 재학습 없이, 제한된 하드웨어 환경에서도 배포 가능한 경량 안전화 솔루션을 제공합니다.
• 확장성: 양자화 (Quantization) 된 모델이나 다양한 아키텍처 (Qwen, Gemma, LLaMA 등) 에서도 효과적으로 작동합니다.
• 이론적 통찰: LLM 의 안전성 문제가 단순히 "학습되지 않은" 것이 아니라, "제거해야 할 특정 하위 네트워크"의 문제일 수 있음을 보여주며, 모델 내부 구조를 조작하여 안전성을 확보하는 새로운 패러다임을 제시합니다.

결론적으로, 이 프레임워크는 리소스가 제한된 환경에서도 LLM 과 VLM 을 더 안전하고 견고하게 만드는 가볍고 효율적인 표준이 될 수 있음을 보여줍니다.