MASCing: Configurable Mixture-of-Experts Behavior via Activation Steering Masks

"Activation Steering Masks 를 통한 구성 가능한 전문가 혼합 (MoE) 행동"이라는 제목의 논문 "MASCing"에 대한 상세한 기술 요약입니다.

1. 문제 제기

대규모 언어 모델 (LLM) 은 희소 활성화 (sparse activation) 를 통해 추론 비용을 낮게 유지하면서 파라미터 수를 확장할 수 있는 능력 덕분에 전문가 혼합 (Mixture-of-Experts, MoE) 아키텍처를 활용하여 지배적인 위치를 차지하게 되었습니다. 그러나 이러한 희소 활성화 패러다임은 고유한 안전성 문제를 야기합니다:

• 결합된 행동 (Coupled Behavior): 모델의 행동은 동적 라우팅 결정과 긴밀하게 결합되어 있습니다. 토큰당 '전문가'의 부분 집합만 활성화되므로, 적대자가 라우터를 조작하여 안전 정렬 (safety-aligned) 전문가를 피하도록 유도하면 안전 메커니즘이 우회될 수 있습니다.
• 현재 방어 수단의 경직성: 전통적인 안전 정렬은 전체 파인튜닝 (fine-tuning) 또는 재학습에 의존합니다. 이러한 방법들은 대규모 MoE 모델에게는 계산적으로 불가능하며, 빠르게 변화하는 안전 정책 (예: 새로운 잭브레이크 전략이나 변화하는 콘텐츠 가이드라인) 에 민첩하게 적응할 수 있는 유연성이 부족합니다.
• 격차 (The Gap): 모델 가중치를 재학습하지 않고도 특정 안전 시나리오 (예: 잭브레이크에 대한 방어 강화 또는 특정 도메인에 대한 제한 완화) 를 위해 MoE 행동을 유연하게 재구성할 수 있는 경량화되고 학습이 필요 없는 메커니즘이 부족합니다.

2. 방법론: MASCing 프레임워크

**MASCing(MoE Activation Steering Configuration)**은 모델 행동을 유도하기 위해 라우팅 게이트 (routing gate) 수준에서 개입하는 3 단계 학습 불필요 프레임워크입니다. 이는 모델 가중치를 전혀 수정하지 않습니다.

1 단계: 행동의 순차적 모델링

• 대리 모델 (Surrogate Model): 이산적인 top-$k$ 전문가 선택 (비미분 가능) 을 최적화하는 대신, MASCing 은 **연속적인 비정규화 라우팅 로지트 (continuous, unnormalized routing logits)**에 대해 경량 LSTM 대리 모델을 학습시킵니다.
• 입력: LSTM 은 토큰 시퀀스에 걸쳐 모든 레이어의 모든 전문가에 대한 라우팅 로지트를 처리합니다.
• 목표: LSTM 은 이러한 라우팅 패턴을 목표 행동 (예: '거부' 대 '수용') 에 매핑하도록 학습합니다. 연속 로지트를 사용함으로써, 선택되지는 않았지만 여전히 라우팅 분포에 영향을 미치는 전문가에 대한 세밀한 정보를 보존합니다.
• 학습: LSTM 은 로지트 시퀀스가 특정 행동으로 이어지는지 여부를 분류하도록 (이진 분류) 학습됩니다.

2 단계: 조향 마스크 생성

• 회로 식별: 학습된 LSTM 을 미분 가능한 프록시로 사용하여, 목표 행동을 담당하는 전문가를 식별하기 위해 **조향 행렬 (steering matrix, $S$)**을 최적화합니다.
• 최적화: 행렬은 다음으로 구성된 복합 손실 함수를 최소화하도록 최적화됩니다:
  1. 이진 교차 엔트로피 (BCE): 대리 모델의 예측을 목표 클래스로 유도하기 위함.
  2. $L_1$ 정규화: 희소성을 강제하여 중요하지 않은 가중치를 0 으로 만들기 위함.
• 적응형 스케일링: 레이어 간 분산을 처리하기 위해, 조향 값은 각 레이어의 조향되지 않은 로지트 표준 편차로 스케일링됩니다.
• 가지치기: 최적화된 행렬에 대칭 크기 게이트 (임계값 $\tau$) 를 적용하여 최종 희소 조향 마스크를 생성합니다. 이 마스크는 특정 전문가를 증폭 (양수 값) 하거나 억제 (음수 값) 할 대상이 되는 전문가를 식별합니다.

3 단계: 조향 마스크 적용

• 추론 개입: 추론 중 정적 조향 마스크가 pre-forward 후크를 통해 대상 MoE 모델의 라우팅 로지트에 직접 적용됩니다.
• 메커니즘: 개입은 로지트에 스케일된 편향을 추가합니다: $g'_{l,t} = g_{l,t} + \alpha(\sigma_l \cdot \hat{S}_l)$.
• 결과: 이는 라우터의 확률 분포를 이동시켜, 실제로 top-$k$ 선택을 대체하여 특정 전문가 회로를 활성화하거나 비활성화합니다. 개입이 희소하고 표적화되어 있으므로 모델은 일반적인 언어 유틸리티를 유지합니다.

3. 주요 기여

1. 새로운 프레임워크: 재학습 없이 MoE 안전 행동을 유연하게 양방향으로 재구성 (강화 및 완화) 할 수 있는 최초의 프레임워크인 MASCing 의 도입.
2. 연속 로지트 분석: 이산적인 top-$k$ 선택 분석에서 연속 라우팅 로지트 분석으로의 방법론적 전환. 이는 비미분 가능한 병목 현상을 우회하고 선택되지 않은 전문가에 잠재된 안전 정보를 포착합니다.
3. 대리 모델 기반 회로 발견: 복잡한 교차 레이어 라우팅 의존성을 특정 행동 회로에 매핑하기 위해 LSTM 대리 모델을 사용하여, 최소한의 고타격 전문가 개입을 식별할 수 있게 함.
4. 양방향 유연성: 동일한 모델 아키텍처에서 안전성 강화(잭브레이크 방어) 와 제약 완화(도메인별 정책 준수) 모두를 수행할 수 있는 능력 입증.

4. 실험 결과

이 프레임워크는 DeepSeek-MoE, Mixtral, Qwen, GPT-OSS 를 포함한 7 개의 오픈소스 MoE 모델에 대해 두 가지 대조적인 시나리오에서 평가되었습니다:

A. 멀티턴 잭브레이크 방어

• 목표: 여러 턴에 걸쳐 모델을 해로운 콘텐츠로 점진적으로 유도하는 적대적 프롬프트에 대한 방어를 강화.
• 성능:
  • 베이스라인: 평균 방어 성공률 52.5%.
  • MASCing: 평균 성공률을 **83.9%**로 향상.
  • 최고 성능: 최대 89.2% (Qwen3-30B-A3B-Instruct-2507).
• 비교: 이산적 전문가 선택에 의존하고 대화 역사 컨텍스트를 무시하는 최첨단 SteerMoE(평균 58.4%) 보다 훨씬 우수한 성능을 보임.

B. 성인 콘텐츠 생성 (정책 완화)

• 목표: 적절한 컨텍스트에서 성인 콘텐츠 생성을 가능하게 하여 (정책 전환 시뮬레이션) '과도한 거부 (over-refusal)'를 줄임.
• 성능:
  • 베이스라인: 평균 생성 성공률 52.6%.
  • MASCing: 평균 성공률을 **82.0%**로 향상.
  • 최고 성능: 최대 93.0% (Phi-3.5-MoE-Instruct).

C. 유틸리티 보존

• 벤치마크: MMLU(사실적 지식) 및 GSM8K(수학적 추론) 에서 평가.
• 결과: MASCing 은 유틸리티에서 미미한 평균 감소 (약 4.1%) 만 초래했습니다. 모델은 복잡한 추론 및 일반 작업을 수행할 수 있는 능력을 유지하여, 조향이 파국적 망각 (catastrophic forgetting) 을 유발하지 않음을 입증했습니다.

D. 효율성

• 학습: LSTM 대리 모델은 단일 NVIDIA H100 GPU 에서 약 5 분 내에 학습 가능.
• 추론: 로지트에 대한 요소별 덧셈으로 인해 지연 시간이 미미하게 추가되며, 모델 가중치를 위한 추가 메모리가 필요 없음.

5. 중요성 및 영향

• 배포를 위한 실용성: MASCing 은 안전 구성을 위한 '플러그 앤 플레이' 솔루션을 제공합니다. 개발자는 전체 재학습의 비용과 시간 없이도 새로운 정책이나 emerging 위협에 모델을 신속하게 적응시킬 수 있습니다.
• 메커니즘적 통찰: 이 논문은 MoE 모델의 안전성이 단일체가 아니라 특정 전문가 회로에 분산되어 있으며, 라우팅 로지트를 통해 격리되고 조작될 수 있음을 보여줍니다.
• 양방향 제어: 안전 제약이 보안 목적으로 강화되거나 특정 사용 사례를 위해 완화될 수 있음을 보여주면서, 안전 (안전 대 불안전) 에 대한 이분법적 관점에 도전하여 모델 무결성을 유지합니다.
• 이산적 방법 대비 우위: 결과는 희소 아키텍처에서 복잡하고 멀티턴인 행동을 이해하고 제어하기 위해 이산적 top-$k$ 분석보다 **연속 라우팅 정보 (로지트)**를 활용하는 것이 우월함을 확인시켜 줍니다.

결론적으로, MASCing 은 MoE 모델에서 구성 가능한 안전성을 위한 새로운 패러다임을 정립하여, 동적 배포 환경에서 모델 유틸리티와 안전 정렬 간의 균형을 관리하기 위한 경량화되고 효율적이며 매우 효과적인 메커니즘을 제공합니다.