MASCing: 활성화 스티어링 마스크를 통한 Mixture-of-Experts 동작의 가변적 구성

Mixture-of-Experts(MoE) 모델은 효율적이지만 특정 전문가(Expert)가 활성화되는 경로에 따라 안전성 위험이 발생할 수 있다. MASCing은 모델의 가중치를 수정하는 비싼 재학습 과정 없이, 추론 시점에 라우팅 신호에 마스크를 씌우는 것만으로 모델의 행동을 정밀하게 제어하고 안전성을 강화한다.

MoE 아키텍처는 입력 토큰마다 소수의 전문가(Expert)만 선택하여 연산 효율을 높이지만, 이 라우팅 결정이 모델의 최종 답변 성격(안전성 등)을 결정짓는 핵심 병목이 된다. 기존의 안전성 정렬은 전체 모델을 미세 조정해야 하므로 비용이 크고 변화하는 정책에 빠르게 대응하기 어렵다는 한계가 있다.

MASCing은 라우팅 레이어에서 발생하는 logit 신호가 시간에 따른 패턴을 가진다는 점에 착안하여, 이를 LSTM으로 분석해 특정 행동을 유발하는 '전문가 회로'를 찾아낸다. Softmax를 거치기 전의 연속적인 logit 값을 활용함으로써, 단순한 전문가 선택 유무보다 훨씬 풍부한 정보를 보존하며 미분 가능한 최적화가 가능해진다.

결과적으로 특정 전문가의 활성화를 정밀하게 유도하거나 억제하는 '스티어링 마스크'를 생성한다. 이 마스크는 추론 시점에 라우팅 logit에 더해져 모델이 스스로 안전한 전문가를 선택하도록 유도하며, 이는 모델의 언어 능력 손실을 최소화하면서도 타겟 행동만 효과적으로 변화시킨다.

MASCing은 세 단계 파이프라인으로 구성된다. 첫째, 행동의 순차적 모델링 단계에서는 MoE 게이트 레이어의 비정규화된 라우팅 logit을 추출하여 LSTM 모델을 학습시킨다. 입력 토큰 시퀀스 t에 대해 각 레이어 l의 logit 벡터 x를 LayerNorm과 아핀 변환을 거쳐 z로 변환하고, 이를 LSTM에 입력하여 최종 hidden state h를 통해 해당 시퀀스가 타겟 행동(거부 등)으로 이어질 확률을 예측한다.

둘째, 스티어링 마스크 생성 단계에서는 학습된 LSTM을 미분 가능한 대리 모델로 사용하여 스티어링 행렬 S를 최적화한다. 손실 함수 L = BCE(y, y_hat) + lambda * ||S||_1을 사용하여 타겟 클래스로 유도하면서도 L1 정규화를 통해 마스크의 희소성(Sparsity)을 확보한다. 최적화된 S에 대칭적 크기 임계값(Symmetric magnitude gate) tau를 적용하여 최종 마스크 S_hat을 생성한다.

셋째, 마스크 적용 단계에서는 추론 시점에 MoE 모델의 라우팅 logit g에 마스크를 더한다. g' = g + alpha * (sigma * S_hat) 공식을 사용하며, 여기서 sigma는 레이어별 logit의 표준편차이고 alpha는 스티어링 강도를 조절하는 하이퍼파라미터이다. 수정된 logit g'을 바탕으로 표준 top-k 라우팅이 수행되어 최종 전문가가 선택된다.

다중 턴 탈옥 방어 실험에서 MASCing은 7개 모델 평균 방어 성공률을 52.5%에서 83.9%로 향상시켰다. 특히 Qwen3-30B 모델의 경우 47.3%에서 89.2%로 비약적인 개선을 보였다. 성인 콘텐츠 생성 제어 실험에서도 평균 순응률이 52.6%에서 82.0%로 증가하여 정책 변화에 따른 유연한 대응 능력을 입증했다.

기존 기법인 SteerMoE와의 비교에서 MASCing은 평균 83.9%의 성공률을 기록하여 SteerMoE(58.4%)를 크게 앞질렀다. 이는 LSTM을 통해 대화 문맥 전체의 라우팅 패턴을 파악하고 연속적인 logit 값을 활용한 결과로 분석된다.

모델의 일반적인 언어 능력 유지 측면에서 MMLU와 GSM8K 벤치마크를 측정한 결과, 평균적으로 약 4.1%의 성능 저하만이 관찰되었다. 이는 모델의 가중치를 직접 수정하지 않고 라우팅 경로만 미세하게 조정함으로써 치명적인 망각(Catastrophic forgetting) 없이 특정 행동만 제어할 수 있음을 보여준다.

MASCing은 MoE의 라우팅 메커니즘이 비미분적인 top-k 선택으로 인해 직접 최적화가 어렵다는 점을 LSTM 대리 모델로 해결한다. LSTM은 모든 전문가의 logit 분포를 입력으로 받아 레이어 간 종속성과 토큰 간의 시간적 관계를 모두 학습한다. 구현 시 NVIDIA H100 GPU 한 장에서 5분 이내에 LSTM 학습이 가능할 정도로 가볍다. 스티어링 마스크는 정적인 행렬로 생성되어 추론 시 오버헤드가 거의 없으며, 레이어별 logit의 표준편차를 활용한 적응형 스케일링(Adaptive scaling)을 통해 모델 내부의 값 변동에 강건하게 대응한다. 하이퍼파라미터 분석 결과, 스티어링 강도 alpha와 희소성 임계값 tau의 적절한 균형이 모델의 유틸리티 유지와 행동 제어 성능 사이의 Trade-off를 결정하는 핵심 요소임이 밝혀졌다.

LSTM 대리 모델이 매우 깊은 모델의 복잡하고 비선형적인 라우팅 역학을 완벽하게 근사하지 못할 수 있다. 또한 전문가 자체의 가중치를 수정하지 않으므로, 모델이 해당 안전 행동을 생성할 수 있는 내재적 능력을 이미 갖추고 있어야만 효과가 있다. 마지막으로 현재의 마스크는 정적(Static)이어서 분포를 벗어난 새로운 유형의 공격에는 취약할 수 있다.