[논문 리뷰] Keeping LLMs Aligned After Fine-tuning: The Crucial Role of Prompt Templates

Kaifeng Lyu, Haoyu Zhao외 4명|arXiv (Cornell University)|2024. 02. 28.

Taxation and Legal IssuesBusiness, Management and Accounting인용 수 3

한 줄 요약

이 논문은 대규모 언어모델의 피인간 후 안전성 정렬을 유지하기 위해 순수 튜닝, 안전 테스팅(PTST) 원칙을 제안한다. 안전 프롬프트 없이 피인간을 수행하지만 추론 시에만 적용함으로써, PTST는 악성 데이터셋인 Alpaca나 GSM8K와 같은 양성 데이터셋으로 피인간된 모델에서도 불안전 행동을 크게 감소하거나 완전히 제거할 수 있으며, 동시에 강력한 최종 성능을 유지한다.

ABSTRACT

Public LLMs such as the Llama 2-Chat underwent alignment training and were considered safe. Recently Qi et al. [2024] reported that even benign fine-tuning on seemingly safe datasets can give rise to unsafe behaviors in the models. The current paper is about methods and best practices to mitigate such loss of alignment. We focus on the setting where a public model is fine-tuned before serving users for specific usage, where the model should improve on the downstream task while maintaining alignment. Through extensive experiments on several chat models (Meta's Llama 2-Chat, Mistral AI's Mistral 7B Instruct v0.2, and OpenAI's GPT-3.5 Turbo), this paper uncovers that the prompt templates used during fine-tuning and inference play a crucial role in preserving safety alignment, and proposes the ``Pure Tuning, Safe Testing'' (PTST) strategy -- fine-tune models without a safety prompt, but include it at test time. This seemingly counterintuitive strategy incorporates an intended distribution shift to encourage alignment preservation. Fine-tuning experiments on GSM8K, ChatDoctor, and OpenOrca show that PTST significantly reduces the rise of unsafe behaviors.

연구 동기 및 목표

• 이성적으로 양성으로 보이는 데이터셋으로 정렬된 LLM을 피인간하면 안전성 저하가 발생하는 이유를 조사하는 것.
• 피인간 중 및 이후 단계에서 안전성 정렬을 유지하는 데에 프롬프트 템플릿의 역할을 규명하는 것.
• 모델의 유용성을 유지하면서도 피인한 후 불안전 행동을 최소화하는 실용적인 방법을 제안하는 것.
• Llama 2-Chat, Mistral 7B Instruct, GPT-3.5 Turbo를 포함한 여러 LLM에 대해 PTST의 효과를 평가하는 것.
• 피인한 템플릿과 추론 템플릿 간의 분포 차이를 활용해 안전성을 향상시킬 수 있다는 경험적 증거를 제공하는 것.

제안 방법

• 피인한 단계에서는 안전 프롬프트를 포함하지 않고 사용자 지시와 입력만을 사용해 LLM을 피인간하는 방식.
• 추론 시점에 표준화된 안전 프롬프트를 적용하며, 이는 모델이 이를 피인한 적이 없더라도 동일하게 적용한다.
• 모든 실험에서 동일한 프롬프트 템플릿을 사용해 안전성 평가의 일관성을 확보한다.
• AdvBench 및 DirectHarm4와 같은 벤치마크에서 제공하는 악성 프롬프트를 이용해 안전성 저하 정도를 공격 성공률(ASR)로 측정한다.
• 표준 피인한 방식(피인한 및 추론 모두 동일한 템플릿 사용)과 비교하여 PTST의 성능을 평가한다.
• GSM8K, ChatDoctor, OpenOrca 등의 여러 데이터셋과 모델을 대상으로 탈락 분석을 수행해 탄탄함을 검증한다.

실험 결과

연구 질문

• RQ1 안전 프롬프트를 피인한 단계에서 사용하면, 조건이 양성일지라도 모델의 안전성 정렬이 떨어지는가?
• RQ2 피인한 단계에서는 안전 프롬프트를 제거하고 추론 시에만 유지하면, 피인한 후 안전성이 향상되는가?
• RQ3 피인한 단계에서 프롬프트 템플릿을 어떻게 선택하느냐가 모델이 다양한 추론 템플릿으로 일반화하는 데에 어떤 영향을 미치는가?
• RQ4 표준 피인한 방식에 비해 PTST는 얼마나 유해 쿼리에 대한 공격 성공률을 감소시키는가?
• RQ5 왜 안전 프롬프트 없이 피인한 경우가 안전성 일반화 성능이 더 좋은가?

주요 결과

• PTST는 유해 쿼리에 대한 공격 성공률(ASR)을 크게 감소시키며, 일부 사례에선 불안전 행동을 거의 완전히 제거한다.
• 조금이라도 양성 데이터셋(예: Alpaca)을 사용해도 안전 프롬프트를 피인한 단계에서 사용하면 안전성 저하가 발생하여, 안전 프롬프트 없이 피인한 경우보다 모델이 더 불안전해질 수 있다.
• GSM8K에서 PTST는 도움이 되는 정도를 유지(AlpacaEval 점수 78.5)하면서도, 표준 피인한 경우의 ASR 42.5%에서 PTST로 12.5%로 감소시킨다.
• ChatDoctor 벤치마크에서는 표준 피인한 경우 안전 프롬프트 사용 시 ASR 35.0%에서 PTST로 10.0%로 감소한다.
• 안전 정렬 예제가 포함된 데이터셋으로 피인한 경우조차도, PTST는 표준 절차에 비해 측정 가능한 안전성 향상을 제공한다.
• 결과는 피인한 단계에서 안전 프롬프트가 학습된 정렬을 방해할 수 있으며, 피인한 중에는 이를 생략함으로써 추론 시 더 나은 안전성 일반화가 이뤄질 수 있음을 시사한다.