기술 요약: SDG-MoE (서명된 토론 그래프 혼합 전문가)

문제 제기

희소 혼합 전문가 (MoE) 아키텍처는 각 토큰에 대해 소수의 '전문가'만 활성화함으로써 대규모 언어 모델 (LLM) 의 확장 비용을 해결합니다. 표준 Top-K MoE 에서 선택된 전문가들은 토큰을 독립적으로 처리하며, 그 출력은 가중 합을 통해 결합됩니다. 이러한 집계 전략은 전문가들을 고립된 예측기로 취급하여, 최종 출력이 생성되기 전에 그들 간의 상호작용이나 숙고를 시뮬레이션하지 못합니다. 이전 연구들은 어텐션 기반 퓨전이나 보조 통신 메커니즘을 탐구해 왔지만, 활성화된 라우팅된 전문가들 간의 직접적이고 반복적인 상호작용은 여전히 충분히 탐구되지 않았습니다. 구체적으로, 기존 방법들은 종종 중앙 집중식 메모리 노드나 훈련 시 증류에 의존하며, 활성화된 전문가 집합이 추론 중에 상충되는 신호를 교환하고 합의를 도출할 수 있도록 하는 데는 미치지 못합니다.

방법론: SDG-MoE

저자들은 라우팅 단계와 최종 집계 사이에 경량의 반복적 숙고 단계를 삽입하는 새로운 희소 MoE 레이어인 SDG-MoE를 제안합니다. 이 아키텍처는 $K$개의 활성 전문가를 독립적인 단위가 아니라 서로를 지지하거나 비판할 수 있는 소규모 상호작용 그룹으로 취급합니다.

핵심 구성 요소

1. 개인 상태 및 공유 상태:
   라우팅 시, 각 활성 전문가 $i$는 은닉 표현 $r_i$를 계산합니다. 이는 두 가지 명확한 상태로 분할됩니다:

   • 개인 상태 ($h_{i,priv}$): 통신 그래프를 완전히 우회하여 전문가의 전문 지식을 보존하는 고차원 벡터입니다.
   • 공유 상태 ($h_{i,shr}$): 숙고 과정에 참여하는 저차원 벡터 ($d_s \ll d$) 입니다.

2. 서명된 토론 그래프:
   각 숙고 라운드 $t$에서, 활성 전문가들은 두 개의 학습된 방향 그래프를 구성합니다:

   • 지지 그래프 ($A^+$): 전문가들 간의 강화 영향을 포착합니다.
   • 비판 그래프 ($A^-$): 교정적 영향을 포착합니다. 표적화된 비판을 보장하기 위해, 이 그래프는 각 전문가당 상위 $m^-$개의 가장 강한 음의 간선만 유지하도록 희소화됩니다.
     이러한 그래프는 학습된 전문가 식별 임베딩으로 보강된 공유 상태의 쿼리/키 프로젝션을 통해 계산됩니다.

3. 게이트된 숙고 및 앵커링:
   상호작용 강도는 불일치 게이트에 의해 동적으로 제어됩니다. 불일치 점수 $D^{(t)}$는 활성 전문가들의 정규화된 공유 상태 간의 코사인 거리를 기반으로 계산됩니다.

   • 적응형 단계 크기: 사회적 업데이트 단계 $\lambda^{(t)}$는 $D^{(t)}$의 함수로, 전문가들이 현저히 불일치할 때만 증가합니다.
   • Friedkin–Johnsen 앵커링: 업데이트는 새로운 공유 상태가 업데이트된 상태와 초기 라우팅된 상태 ($h^{(0)}_{i,shr}$) 의 가중 조합이 되는 동역학 모델을 따릅니다. 매개변수 $\beta$는 '고집' 계수 역할을 하여 전문가들이 초기 전문성에서 너무 멀리 벗어나는 것을 방지합니다.
   • 서명된 대비: 업데이트 MLP 는 현재 상태, 지지 메시지, 그리고 서명된 대비 항 ($m^+ - \gamma m^-$) 을 수신하여 확인 증거와 반증 증거를 명시적으로 구분합니다.

4. 집계:
   $T$회의 숙고 라운드 후, 전문가 $i$의 최종 출력은 개인 상태와 최종 공유 상태 모두의 프로젝션입니다. 이들은 원래 라우터 가중치를 사용하여 집계됩니다.

이론적 분석

이 논문은 숙고 동역학의 안정성과 효율성을 확립하는 이론적 분석을 제공합니다:

• 안정성: 저자들은 특정 조건 (유계 업데이트 노름과 불일치 게이트의 리프시츠 연속성) 하에서 숙고 과정이 축소 사상임을 증명합니다. 이는 전문가 상태가 유계로 유지되고 고유한 고정점으로 수렴하여 통제되지 않은 이동을 방지함을 보장합니다.
• 복잡도: 통신 오버헤드는 활성 집합 크기 $K$에 대해 저차임이 입증되었습니다. 토큰당 비용은 $O(T(Kb^2 + K^2b))$이며, 여기서 $b$는 병목 차원을 나타냅니다. 작은 $K$와 $T$의 경우, 이는 전체 추가 전문가 레이어가 아닌 컴팩트한 보정을 추가합니다.

실험 결과

저자들은 Qwen 과 유사한 디코더 전용 트랜스포머 아키텍처를 사용한 통제된 사전 학습 설정에서 SDG-MoE 를 평가했습니다.

• 성능: 세 가지 시드 사전 학습 실행에서 SDG-MoE 는 검증 퍼플렉시티를 63.14(일반 MoE) 에서 48.03으로 감소시켜, 가장 강력한 베이스라인 대비 상당한 개선을 보였습니다.
• 외부 벤치마크: SDG-MoE 는 비교된 모든 시스템 중에서 WikiText-103, C4, Paloma 등 외부 데이터셋에서 가장 낮은 퍼플렉시티를 달성했습니다.
• 어블레이션:
  • 서명된 의미: 지지와 비판 간의 구분을 제거 (Dual-Unsigned) 하거나 비판 채널을 완전히 제거 (Zero-Neg) 하면 성능이 저하되어, 서명된 통신이 필수적임을 확인했습니다.
  • 적응형 게이팅: 적응형 불일치 게이트를 고정 계수로 대체하면 퍼플렉시티가 현저히 악화되어, 모델이 전문가 불일치에 기반하여 선택적으로 숙고해야 함을 나타냅니다.
  • 숙고 깊이: 두 번의 숙고 라운드 ($T=2$) 가 가장 좋은 결과를 낳았으며, $T=4$로 증가하면 성능이 저하되어 과도한 반복이 과교정으로 이어짐을 시사합니다.

중요성 및 주장

이 논문은 집계 전에 활성 라우팅된 전문가들 간의 직접적이고 서명된 통신을 가능하게 하는 최초의 희소 MoE 레이어를 도입했다고 주장합니다. 그 중요성은 다음과 같습니다:

1. 새로운 상호작용 메커니즘: 독립적인 전문가 투표나 중앙 집중식 조정을 넘어, 사회적 의견 동역학 (특히 Friedkin–Johnsen 모델) 에서 영감을 받은 피어 투 피어 숙고 모델로 이동합니다.
2. 구조화된 숙고: 전문가들이 서로를 지지하고 비판하게 허용하면서도 초기 전문 상태에 앵커링함으로써, 희소 라우팅의 효율성을 희생하지 않으면서 더 풍부한 표현을 산출함을 입증합니다.
3. 이론적 기반: 숙고 과정이 표준 MoE 의 제어된 교란임을 보여주는 안정성 보장을 제공하여, 유계 상태와 낮은 계산 오버헤드를 보장합니다.

저자들은 성능 향상이 단순히 증가된 용량이나 일반적인 통신 때문이 아니라, 지지, 비판, 앵커링, 적응형 게이팅의 균형 잡힌 상호작용에서 비롯된 것이라고 결론지었습니다. 그들은 이 접근법이 라우팅 후 전문가 상호작용을 위한 새로운 설계 공간을 열지만, 더 큰 규모의 평가와 완전히 사전 학습된 모델로의 통합은 향후 과제로 남아 있음을 인정합니다.