기술 요약: 온디바이스 환경을 위한 소프트 프롬프트를 이용한 안전한 LLM 시스템의 증류(Distillation)

문제 정의

자원 제한적인 에지 디바이스(예: 스마트폰)에 대규모 언어 모델(LLM)을 배포하는 것은 중대한 안전성 과제를 안겨줍니다. 베이스 LLM과 출력의 안전성을 분류 및 필터링하는 별도의 "가드 모델(guard model)"을 결합한 이중 모델 시스템은 강력한 안전 보장을 제공하지만, 메모리와 연산 비용이 매우 높습니다. 이러한 방식의 순차적 구조(응답 생성 후 분류)는 두 번의 순전파(forward pass)를 유발하여 지연 시간(첫 토큰 생성 시간)과 메모리 점유율을 크게 증가시키며, 이는 에지 배포에 부적합하게 만듭니다. 양자화(quantization)나 직접적인 미세 조정(fine-tuning)과 같은 기존의 대안들은 종종 일반화 능력을 저해하거나, 안전성과 모델의 유용성 사이의 균형을 효과적으로 맞추지 못합니다.

방법론

본 논문은 학습 가능한 소프트 프롬프트(soft prompts)를 갖춘 단일 LLM에 이중 모델 시스템의 안전 동작을 증류하는 매개변수 효율적 미세 조정(PEFT) 접근 방식인 **TV-DiSP (Total Variation Distillation via Soft Prompts)**를 제안합니다.

핵심 프레 프레임워크

1. 시스템 공식화: 저자들은 베이스 LLM $p(y|x)$와 가드 모델 $p(s|x, y)$가 결합된 안전한 LLM 시스템 $p(r|x)$를 정의합니다. 가드 모델이 응답 $y$를 안전하지 않다고 판단하면($s=0$), 시스템은 사전 정의된 거절 메시지 $y_r$을 출력하고, 그렇지 않으면 $y$를 출력합니다.
2. 증류 목적 함수: 목표는 연속적인 프롬프트 임베딩 집합 $W$(소프트 프롬프트)를 학습하여, 증류된 모델 $q(r|x, W)$가 안전한 시스템 $p(r|x)$의 출력 분포를 모방하도록 하는 것입니다.
3. Total Variation(TV)을 통한 최적화:
   • 저자들은 TV 거리(Total Variation distance)를 최소화하는 것이 안전한 시스템과 증류된 모델 간의 다운스트림 태스크 성능 격차를 제한한다는 이론적 보장(정리 3.1)을 제시합니다.
   • 최적화 목적 함수(식 4)는 가드 모델의 안전 확률에 따라 가중치가 부여된 정답 출력(안전한 입력에 대한 베이스 LLM의 응답 또는 안전하지 않은 입력에 대한 거절 메시지)의 로그 확률을 최대화합니다.
   • 학습에는 프롬프트 데이터셋만 필요하며, 가드 모델이 학습 과정 중에 프롬프트를 안전 또는 안전하지 않음으로 동적으로 라벨링합니다.
4. 대안적 스킴: 연구에서는 다음 사항들도 평가합니다:
   • KL-DiSP: 더 느슨한 이론적 보장을 제공하는 Kullback-Leibler 발산 최소화.
   • 베이스라인: 퍼플렉서티(Perplexity) 최적화, REINFORCE(정책 경사), 그리고 LoRA 및 스티어링 벡터(Steering Vectors)에 적용된 증류 프레임워크 변형.

실험 설정

• 모델: 에지 제약을 시뮬레이션하기 위해 4비트로 양자화된 네 가지 인스트럭션 튜닝 LLM(Llama3-1B/3B, Qwen2-1.5B, Gemma2-2B)을 사용했습니다.
• 가드 모델: 증류 학습에는 LlamaGuard3-1B를 사용하였고, 평가는 LlamaGuard3-8B를 사용했습니다.
• 데이터셋: 학습에는 Beavertails 및 Toxigen 데이터셋을 활용했습니다. 평가는 분포 외(out-of-distribution) 벤치마크(HarmBench, Detect-JailBreak)와 유용성 벤치마크(IFEval, GSM8K, MMLU)를 포함했습니다.
• 하드웨어: Qualcomm Snapdragon 8 Elite Gen 5 칩셋이 탑재된 스마트폰에서 온디바이스 측정을 수행했습니다.

주요 결과

1. 안전성-연산량 트레이드오프 (Safety-Compute Trade-off)

• 성능: TV-DiSP는 이중 모델 시스템의 안전성을 단일 모델로 성공적으로 증류합니다. 예를 들어, Llama3-1B 모델의 경우, 안전성(Safety Guard Score, SGS로 측정)이 71%(베이스)에서 99%(이중 시스템)로 향상되었으며, TV-DiSP는 최소한의 오버헤드로 이에 상응하는 안전성 향상을 달성했습니다.
• 효율성: 메모리와 연산을 두 배로 요구하는 이중 모델 시스템과 비교했을 때, TV-DiSP는 1% 미만의 추가 메모리 오버헤드와 10% 미만의 추가 연산 오버헤드(토큰당 FLOPs)를 도입합니다.
• 일반화: 본 방법은 비적대적(non-adversarial) 데이터셋으로 학습되었음에도 불구하고, 분포 외 적대적 프롬프트(HarmBench, Detect-JailBreak)에 대해 강력한 견고성을 보여줍니다.

2. PEFT 방법론과의 비교

• 소프트 프롬프트 vs. LoRA/SV: TV-DiSP(소프트 프롬프트)는 다양한 학습률에 걸쳐 TV-DiLoRA 및 TV-DiSV(Steering Vectors)보다 일관되게 우수한 성능을 보였습니다. 스티어링 벡터는 혼합된 결과를 보였는데, 때때로 더 큰 모델에서 성능을 저하시켰으며, LoRA는 에지 환경에 너무 침습적(intrusive)인 것으로 간주되었습니다.
• 최적화 목적 함수:
  • TV vs. KL: TV와 KL 증류 모두 높은 안전 점수를 달 Achieve했습니다. 그러나 TV-DiSP는 안전성과 유용성 사이에서 더 우수한 균형을 제공했습니다.
  • Perplexity & REINFORCE: 퍼플렉서티 최적화는 과적합을 유발하여 유용성을 크게 떨어뜨렸습니다. REINFORCE는 유용성은 유지했으나 실질적인 안전성 향상을 제공하지 못했습니다.

3. 유용성 보존 (Usefulness Preservation)

• 온디바이스 평가(IFEval 및 GSM8K) 결과, TV-DiSP와 KL-DiSP 모두 높은 유용성을 유지함을 보여주었습니다.
• 과잉 거절 (Over-refusal): 중요한 발견은 KL-DiSP가 안전한 프롬프트에 대해서도 높은 과잉 거절률(84.6%)을 보여 무해한 질의를 잘못 거절했다는 점입니다. 반면, TV-DiSP는 이중 모델 시스템과 유사한 수준의 과잉 거절률(36.1%)을 유지하여 안전한 질문에 답하는 모델의 능력을 보존했습니다.

의의 및 주장

본 논문은 소프트 프롬프트 증류가 온디바이스 LLM 배포를 위한 안전 정렬(safety alignment)의 선호되는 접근 방식임을 입증합니다. 주요 기여는 다음과 같습니다:

1. 우월한 트레이드오프: 소프트 프롬프트가 안전성, 유용성 및 계산 효율성의 균형을 맞추는 데 있어 LoRA, 스티어링 벡터 및 대안적 학습 목적 함수보다 일관되게 우수함을 보여줍니다.
2. 이론적 보장: Total Variation 기반 증류를 도입하여 다운스트림 성능에 대한 확률적 보장을 제공함으로써, 증류된 모델이 안전한 시스템의 동작에서 크게 벗어나지 않음을 보장합니다.
3. 실용적 타당성: 다양한 아키텍처와 실제 에지 하드웨어에서 검증된 이 방법은 이중 모델 시스템에 필적하는 안전성을 달성하면서도 무시할 만한 자원 오버헤드(<1% 메모리, <10% 연산)만을 발생시켜, 자원이 제한된 기기에 안전한 LLM을 배포하기 위한 실질적인 솔루션을 제공합니다.

저자들은 어떤 방법도 적대적 공격에 대한 완전한 면역을 제공할 수는 없지만, TV-DiSP는 전통적인 가드 모델 아키텍처에 비해 공격 표면과 계산 부담을 크게 줄이는 강력하고 효율적이며 일반화 가능한 방어 메커니즘을 제공한다고 결론짓습니다.