MOSAIC: Composable Safety Alignment with Modular Control Tokens

이 논문은 다양한 맥락에 따라 안전 규칙을 유연하게 조합하고 적용할 수 있도록, 고정된 백본 모델 위에 학습 가능한 제어 토큰을 도입한 모듈형 안전 정렬 프레임워크 'MOSAIC'을 제안합니다.

핵심

🛡️ MOSAIC: AI 의 안전장치를 '레고 블록'처럼 자유롭게 조립하다

이 논문은 인공지능 (AI) 이 얼마나 안전한지, 그리고 그 안전 기준을 어떻게 유연하게 바꿀 수 있는지에 대한 새로운 아이디어를 제시합니다. 제목인 MOSAIC(모자이크) 는 마치 다양한 색상의 타일을 모아 하나의 그림을 만들 듯, AI 의 안전 규칙을 작은 조각들로 나누어 필요할 때만 조립한다는 의미를 담고 있습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 지금의 문제점: "단단하게 굳어버린 안전장비"

지금까지 AI 를 안전하게 만들기 위해 사용했던 방법들은 두 가지 큰 단점이 있었습니다.

• **방법 A **(모델 자체를 고치는 것) AI 의 두뇌 (파라미터) 자체를 다시 학습시켜 안전 규칙을 심어놓는 방식입니다.
  • 비유: 마치 아이에게 "술은 절대 마시지 마"라고 가르친 뒤, 그 기억을 지우지 않고 평생 가지고 다니게 만드는 것입니다. 만약 나중에 "성인이 된 아이에게는 술에 대한 설명도 필요해"라고 상황이 바뀐다면, 다시 아이를 학교에 보내고 (재학습) 모든 것을 다시 가르쳐야 합니다. 비용도 많이 들고, 다른 능력 (예: 수학 실력) 이 떨어질 위험도 있습니다.
• **방법 B **(메모를 붙이는 것) AI 에게 "술은 마시지 마"라고 긴 메모 (프롬프트) 를 붙여서 알려주는 방식입니다.
  • 비유: 매번 대화할 때마다 AI 에게 긴 규칙서를 읽어주는 것입니다. 규칙이 너무 길어지면 AI 가 지쳐서 규칙을 잊어버리거나, 엉뚱한 부분까지 막아버릴 수 있습니다.

2. MOSAIC 의 해결책: "안전 레고 블록"

저자들은 이 문제를 해결하기 위해 MOSAIC이라는 새로운 방식을 제안합니다.

• 핵심 아이디어: AI 의 두뇌는 건드리지 않고, **작은 '안전 레고 블록' **(학습 가능한 제어 토큰) 만 따로 만들어서 사용합니다.
• 어떻게 작동할까요?
  • 레고 블록: 각 안전 규칙 (예: '도박 금지', '술 금지', '폭력 금지') 마다 작은 레고 블록 하나씩이 있습니다.
  • 상황에 따른 조립:
    • 어린이가 질문할 때: '도박'과 '술' 레고 블록을 AI 앞에 붙입니다. AI 는 "이건 안 돼요!"라고 거절합니다.
    • 성인 전문가가 질문할 때: '폭력' 레고 블록만 붙입니다. '도박'이나 '술' 레고는 떼어내므로, 성인에게는 관련 정보를 알려줄 수 있습니다.
  • 장점: AI 의 두뇌 (기반 모델) 는 그대로 두기 때문에, 원래의 똑똑함은 유지된 채로 안전 규칙만 순간적으로 바꿀 수 있습니다.

3. 기술적인 마법: "과도한 거절"을 막는 비법

이 방식이 잘 작동하려면 두 가지 중요한 기술이 필요합니다.

1. **레고 조합 훈련 **(Combinational Task Sampling)
   • 레고 블록 하나하나가 혼자 있을 때도, 다른 블록들과 섞였을 때도 제 역할을 하도록 훈련시킵니다. 마치 레고 조각들이 서로 충돌하지 않고 잘 어울리도록 연습시키는 것과 같습니다.
2. **상상력 교정 **(Counterfactual Knowledge Distillation)
   • 문제: 안전 레고 블록을 붙였을 때, AI 가 "아, 안전 블록이 붙었네? 그럼 모든 질문을 거절해야겠다!"라고 오해해서, 안전한 질문까지 거절하는 경우가 생깁니다 (과도한 거절).
   • 해결: AI 에게 "안전 블록이 붙었을 때와 붙지 않았을 때의 대답을 비교해봐. 안전한 질문은 원래대로 대답하고, 위험한 질문만 거절해"라고 가르칩니다.
   • 비유: **스승이 "이 학생은 위험한 질문만 거절하고, 나머지는 평소처럼 친절하게 대답해"**라고 가르치는 것과 같습니다.

4. 실험 결과: "정확하고 유연한 안전"

저자들은 실제 상황에 가까운 새로운 테스트 데이터를 만들어 실험했습니다.

• 결과: 기존 방법들보다 **위험한 질문을 정확히 막아내는 능력 **(방어 성공률) 이 매우 높으면서도, **안전한 질문까지 막아버리는 실수 **(과도한 거절) 가 훨씬 적었습니다.
• 확장성: 새로운 안전 규칙 (예: '성인물 금지') 이 생겼을 때, AI 를 다시 학습시킬 필요 없이 해당 레고 블록만 추가하면 바로 적용됩니다.

5. 한 줄 요약

MOSAIC 은 AI 의 두뇌를 건드리지 않고, 상황에 따라 필요한 '안전 레고 블록'만 끼워 넣음으로써, 어린이에게는 엄격하고 성인에게는 유연한 AI 를 만들어주는 똑똑한 시스템입니다.

이 기술은 앞으로 AI 가 다양한 국가, 연령대, 상황에 맞춰 안전 기준을 유연하게 조절할 수 있는 길을 열어줄 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 의 안전성 정렬 (Safety Alignment) 은 현재 다음과 같은 한계점을 겪고 있습니다:

• 정적 정책의 비유연성: 기존 안전 정책은 모델 파라미터에 고정되어 있어, 사용자, 지역, 응용 분야에 따라 다른 안전 규칙을 적용하기 어렵습니다.
• 파라미터 수준 정렬의 문제: SFT(지도 미세조정) 나 RLHF(인간 피드백 강화학습) 와 같은 방법은 안전 행동을 일반 모델 능력과 엮어버려 (entangle), 특정 규칙만 업데이트하려면 비용이 많이 드는 재학습이 필요하며, 기존 성능에 치명적인 간섭을 일으킬 수 있습니다.
• 프롬프트 기반 방법의 한계: 자연어 지시를 통해 안전 규칙을 제어하는 방식은 불일치한 추론을 유발하고, 여러 제약 조건을 적용할 때 토큰 오버헤드가 커지며, 명시적인 제어 메커니즘이 부족합니다.
• 과도한 거부 (Over-refusal): 안전 규칙이 활성화되었을 때, 관련 없는 benign(안전한) 요청까지 거부하는 문제가 발생합니다.

2. 방법론 (Methodology)

저자들은 안전성 정렬을 표현 학습 (Representation Learning) 문제로 재정의하고, MOSAIC이라는 모듈형 프레임워크를 제안합니다.

핵심 아이디어: 모듈형 제어 토큰 (Modular Control Tokens)

• 학습 가능한 토큰: 각 안전 제약 조건 (예: 도박, 알코올, 폭력 등) 을 모델의 임베딩 공간에 학습 가능한 소수의 제어 토큰 (Control Tokens) 집합으로 인코딩합니다.
• 동적 활성화: 추론 시, 특정 안전 규칙을 적용하려면 해당 토큰을 입력 앞에 추가하고, 적용하지 않으면 제거합니다. 이는 모델의 기본 파라미터를 변경하지 않고도 안전 정책을 유연하게 조합 (Compose) 할 수 있게 합니다.

주요 기술적 구성 요소

1. 순서 기반 태스크 샘플링 (Order-based Task Sampling):
   • 여러 카테고리 토큰을 조합할 때 발생하는 지수적인 데이터 비용을 해결하기 위해, 활성화된 제약 조건의 수 (Order) 에 따라 태스크 풀을 구성합니다.
   • 각 Order(예: 1 개 카테고리, 2 개 카테고리 조합 등) 마다 고정된 학습 예산을 할당하여, 다양한 토큰 조합을 효율적으로 학습시키면서도 데이터 크기를 통제합니다.
2. 반사실적 지식 증류 (Counterfactual Knowledge Distillation):
   • 과도한 거부 완화: 안전한 요청 (Benign queries) 에 대해 제어 토큰이 있을 때와 없을 때의 모델 출력을 비교합니다.
   • 목표: 제어 토큰이 활성화되었을 때에도, 해당 요청이 안전 규칙 위반이 아닌 경우 기본 모델 (Frozen Backbone) 의 원래 응답 분포를 유지하도록 KL 발산 (KL Divergence) 을 최소화합니다. 이를 통해 제어 토큰이 불필요할 때는 개입하지 않도록 학습시킵니다.

3. 주요 기여 (Key Contributions)

1. 개념적 재정의: 안전성 정렬을 단일 파라미터 수정이 아닌, 조합적 표현 학습 (Compositional Representation Learning) 문제로 재정의하여 조건부 안전 제어를 모듈형 제약 활성화로 접근했습니다.
2. MOSAIC 프레임워크 제안:
   • 고정된 백본 모델 위에 학습 가능한 제어 토큰을 사용하여 안전성을 구현합니다.
   • 점진적 확장성: 새로운 안전 카테고리를 기존 모델을 재학습하지 않고 토큰만 추가하여 확장할 수 있습니다.
   • 과도한 거부 감소: 구조화된 조합 학습과 반사실적 지식 증류를 통해 안전하지 않은 요청은 차단하되, 안전한 요청은 정상적으로 처리하도록 최적화했습니다.
3. 새로운 평가 벤치마크 구축:
   • 기존 안전 벤치마크는 이미 정렬된 모델의 데이터와 겹치는 문제가 있어, 실제 배포 시나리오를 반영하기 어렵습니다.
   • 저자들은 1,500 개의 사용자 요청으로 구성된 새로운 벤치마크를 구축했습니다. 이 데이터는 기본 설정의 LLM 이 거부하지 않는 요청들을 포함하여, 선택적이고 조건부로 안전 규칙을 적용하는 능력을 정밀하게 평가할 수 있게 합니다.

4. 실험 결과 (Results)

Llama-3.1-8B 및 Llama-3.2-3B 모델을 기반으로 한 실험 결과는 다음과 같습니다:

• 방어 성공률 (DSR) 과 과도한 거부율 (OR):
  • MOSAIC 은 모든 작업 순서 (1 차, 2 차, 3 차 카테고리 조합) 에서 99% 이상의 높은 방어 성공률을 기록했습니다.
  • 기존 SFT 기반 방법은 DSR 은 높았으나 과도한 거부 (OR) 가 약 6% 수준이었으나, MOSAIC 은 이를 1.8%~2.9% 수준으로 획기적으로 낮췄습니다.
  • 특히 카테고리 토큰 수를 5 개로 늘리면 OR 가 더욱 감소하는 경향을 보였습니다.
• 모델 유틸리티 보존:
  • MMLU 와 같은 일반 언어 모델링 벤치마크에서 MOSAIC 은 기본 모델의 성능을 거의 저하시키지 않았습니다 (기존 모델과 유사한 점수 유지).
• 점진적 확장성:
  • 새로운 안전 카테고리를 추가할 때 기존 카테고리의 성능이 크게 저하되지 않았으며, DSR 은 99% 이상 유지되었습니다.
• 비교 분석:
  • In-context prompting 과 ORPO(LoRA 기반) 는 복잡한 조합 상황에서 성능이 떨어지거나 과도한 거부 문제가 발생했습니다.
  • Ablation Study: 멀티태스크 학습과 반사실적 KD 가 모두 결합되었을 때 가장 좋은 성능을 발휘함을 확인했습니다. 특히 반사실적 KD 는 과도한 거부를 줄이는 데 결정적인 역할을 했습니다.

5. 의의 및 결론 (Significance)

• 실용적 유연성: MOSAIC 은 연령대, 지역, 법적 요구사항 등에 따라 안전 정책을 동적으로 조정할 수 있는 실용적인 솔루션을 제공합니다.
• 비용 효율성: 전체 모델을 재학습할 필요 없이 소수의 토큰만 학습하여 새로운 안전 규칙을 추가할 수 있어, 배포 및 유지보수 비용을 크게 절감합니다.
• 정밀한 제어: "무조건적인 거부"가 아닌 "조건부 차단"을 가능하게 하여, 모델의 유용성 (Utility) 을 해치지 않으면서도 필요한 안전 장치를 구현할 수 있음을 입증했습니다.

이 연구는 LLM 의 안전성 정렬이 고정된 규칙이 아닌, 상황에 따라 유연하게 구성 가능한 모듈형 시스템으로 전환되어야 함을 시사하며, 향후 다중 가치와 다양한 안전 요구사항을 가진 환경에서의 LLM 배포에 중요한 기여를 할 것으로 기대됩니다.