라우터 네트워크(router network)는 게이팅 함수(gating function) 또는 가중치 함수(weighting function)라고도 불리며, 혼합 전문가(mixture-of-experts, MoE) 머신러닝 모델의 구성 요소이다. 이는 입력을 받아 각 전문가 네트워크가 출력에 기여하는 정도를 결정하는 가중치 집합을 생성한다. 현대의 대규모 언어 모델에서 라우터 네트워크는 각 입력 토큰에 대해 소수의 전문가만 선택하는 데 자주 사용되는데, 이 기법을 희소 활성화(sparse activation)라고 하며, 모델 용량을 유지하면서 계산 비용을 줄인다.
MoE에서 라우팅의 개념은 위원회 기계(committee machines)와 적응형 로컬 전문가 혼합(adaptive mixtures of local experts)에 대한 초기 연구로 거슬러 올라간다. 전형적인 MoE 아키텍처에는 여러 전문가 네트워크가 있으며, 각각은 입력 공간의 서로 다른 영역에 특화될 수 있고, 라우터 네트워크는 입력을 적절한 전문가에 할당하는 방법을 학습한다. 라우터의 출력은 전문가에 대한 확률 분포 또는 전문가 출력과 가중 합을 통해 결합되는 비음수 가중치 집합일 수 있다.
기본 아키텍처
표준 MoE 설정에서 라우터 네트워크는 입력 \( x \)를 가중치 벡터 \( (w(x)_1, ..., w(x)_n) \)로 매핑하는 함수 \( w(x) \)이며, 여기서 \( n \)은 전문가 수이다. 최종 출력은 \( f(x) = \sum_{i=1}^n w(x)_i f_i(x) \)로 계산되며, 여기서 \( f_i \)는 \( i \)-번째 전문가이다. 라우터와 전문가는 일반적으로 평균 제곱 오차 또는 교차 엔트로피와 같은 손실 함수에 대한 경사 하강법을 통해 공동으로 훈련된다.
라우터 네트워크 자체는 종종 선형 계층 뒤에 소프트맥스 활성화가 이어지는 작은 신경망이다. 희소 MoE 모델에서 라우터는 top-k 선택 메커니즘을 사용할 수 있으며, 여기서 가장 높은 가중치를 가진 전문가만 활성화되고 나머지는 무시된다. 이러한 희소성은 계산 비용을 줄이는 핵심으로, 모델이 각 입력에 대해 총 파라미터의 일부만 사용할 수 있게 한다.
역사적 발전
게이팅 네트워크를 사용하여 여러 전문가를 결합하는 아이디어는 1990년대에 탐구되었다. 초기 사례 중 하나는 Hampshire와 Waibel이 보고한 메타-파이 네트워크(meta-pi network)로, 전문가 출력의 가중 합을 사용하고 음소 분류 작업에서 훈련되었다. 그들의 실험에서 그들은 여섯 명의 일본어 화자의 음성을 분류하기 위해 각각 시간 지연 신경망인 여섯 개의 전문가를 훈련했다. 그들은 라우터 네트워크가 다섯 개의 전문가를 다섯 명의 개별 화자에게 전담시키고, 여섯 번째 화자의 목소리는 다른 남성 화자에 대한 전문가 조합으로 처리하는 것을 학습했음을 관찰했다.
또 다른 영향력 있는 초기 모델은 가우시안 혼합 모델을 사용한 적응형 로컬 전문가 혼합이었다. 이 접근 방식에서 각 전문가는 고정 공분산을 가진 가우시안 분포를 예측했고, 라우터는 입력에 따라 가중치를 할당하는 선형-소프트맥스 함수였다. 이 모델은 라우팅이 입력 공간을 동질적인 영역으로 분할하는 데 사용될 수 있으며, 각 전문가가 로컬 영역에 특화됨을 입증했다.
현대 대규모 언어 모델에서의 역할
2020년대에 라우터 네트워크는 MoE 계층을 사용하는 대규모 언어 모델(LLM)의 핵심 구성 요소가 되었다. Google DeepMind 및 OpenAI와 같은 회사에서 개발한 모델은 추론 비용을 비례적으로 증가시키지 않고 파라미터 수를 확장하기 위해 희소 MoE 아키텍처를 채택했다. 이러한 모델에서 각 트랜스포머 계층은 여러 전문가를 포함할 수 있으며, 라우터 네트워크는 토큰당 몇 개의 전문가를 선택한다.
예를 들어, 트랜스포머 기반 LLM에서 입력 토큰 임베딩은 전문가에 대한 확률 분포를 계산하는 라우터 네트워크를 통과한다. 그런 다음 라우터는 top-k 전문가(예: k=2 또는 k=4)를 선택하고 해당 전문가 출력의 가중 합으로 출력을 계산한다. 이를 통해 모델은 수십억 개의 파라미터를 가지면서도 각 토큰에 대해 소수의 하위 집합만 활성화하여 훈련과 추론을 더 효율적으로 만든다.
희소 활성화와 부하 균형
희소 활성화는 MoE 모델에서 라우터 네트워크의 주요 이점이다. 입력당 몇 개의 전문가만 활성화함으로써 모델은 큰 총 파라미터 수를 가질 수 있지만 유효 계산 비용은 훨씬 작다. 그러나 이는 부하 균형 문제를 도입한다: 라우터가 일관되게 동일한 몇 개의 전문가를 선택하면 해당 전문가는 과부하되고 다른 전문가는 저활용된다. 이를 해결하기 위해 현대 MoE 모델은 라우터가 토큰을 전문가 간에 더 고르게 분배하도록 장려하는 보조 손실 함수를 포함하는 경우가 많다.
라우터 성능을 개선하기 위해 훈련 중 라우터의 로짓에 노이즈를 추가하여 탐색을 촉진하거나 미분 가능한 top-k 선택을 사용하는 등 다양한 기법이 제안되었다. 일부 모델은 계층적 라우팅 체계를 사용하는데, 여기서 1차 라우터가 전문가 그룹을 선택하고 2차 라우터가 그 그룹 내에서 선택한다.
다른 기법과의 관계
라우터 네트워크는 멀티 헤드 어텐션 및 크로스 어텐션과 같은 머신러닝의 다른 개념과 밀접하게 관련되어 있으며, 이들도 입력에 따라 다른 구성 요소에 가중치를 부여한다. 그러나 어텐션 메커니즘은 입력 시퀀스의 다른 부분에 가중치를 부여하는 반면, 라우터 네트워크는 일반적으로 독립적인 함수 근사기인 다른 전문가 네트워크에 가중치를 부여한다.
라우터 네트워크는 모델 가지치기와도 유사점을 공유하는데, 둘 다 계산 비용을 줄이는 것을 목표로 하지만, 가지치기는 파라미터를 영구적으로 제거하는 반면 라우팅은 각 입력에 대해 사용할 파라미터를 동적으로 선택한다. 또한 라우터 네트워크는 여러 모델의 출력을 결합하므로 앙상블 학습의 한 형태로 볼 수 있지만, 전통적인 앙상블과 달리 전문가는 라우터와 공동으로 훈련된다.
구현 및 훈련
실제로 라우터 네트워크는 입력 표현을 받아 각 전문가에 대한 로짓을 출력하는 선형 계층으로 구현된다. 로짓은 소프트맥스 함수를 통과하여 가중치를 생성한다. 훈련 중에 라우터와 전문가는 역전파를 사용하여 공동으로 업데이트된다. 손실 함수는 일반적으로 작업 손실(예: 언어 모델링의 교차 엔트로피)과 보조 부하 균형 손실을 모두 포함한다.
라우터 네트워크 훈련의 한 가지 과제는 전문가의 이산적 선택(예: top-k)이 미분 불가능하다는 점이다. 이를 극복하기 위해 많은 구현은 훈련 중에 소프트맥스 가중치를 직접 사용하거나 top-k 선택에 대해 직통 추정기(straight-through estimator)를 사용한다. 일부 모델은 각 계층에 대해 별도의 라우터 네트워크를 사용하여 서로 다른 계층이 서로 다른 유형의 라우팅 결정에 특화될 수 있게 한다.
언어 모델 이외의 응용
라우터 네트워크는 LLM에서 가장 두드러지지만 다른 영역에서도 사용된다. 예를 들어, 컴퓨터 비전에서 라우터가 있는 MoE 계층은 이미지 분류 및 생성 작업에 적용되었다. 음성 인식에서 초기 MoE 모델은 음향 공간을 분할하기 위해 라우터를 사용했다. 라우터 네트워크는 강화 학습에서도 사용되며, 여기서 서로 다른 전문가가 서로 다른 정책을 나타낼 수 있고 라우터는 상태에 따라 적절한 정책을 선택한다.
인공지능 연구의 맥락에서 라우터 네트워크는 라우팅 효율성, 해석 가능성 및 확장성 개선에 대한 지속적인 작업이 있는 활발한 연구 영역이다. Stanford AI Lab 및 Berkeley AI Research와 같은 기관의 연구자들은 대규모 모델에서 라우터의 동작을 이해하는 데 기여했다.
과제 및 향후 방향
성공에도 불구하고 라우터 네트워크는 여러 과제에 직면한다. 한 가지 문제는 라우터가 모든 입력을 처리하고 라우팅 결정을 신속하게 내려야 하므로 병목 현상이 될 수 있다는 점이다. 또 다른 과제는 라우터가 전체 작업에 최적이 아닌 방식으로 입력을 라우팅하는 방법을 학습하여 성능 저하를 초래할 수 있다는 것이다. 또한 보조 부하 균형 손실은 주요 작업 손실과 간섭할 수 있어 신중한 조정이 필요하다.
향후 연구는 입력 내용을 더 미묘한 방식으로 고려하는 학습된 라우팅 정책이나 각 수준에서 고려되는 전문가 수를 줄이는 계층적 라우팅과 같은 더 정교한 라우팅 메커니즘을 탐구할 수 있다. 또한 특정 입력이 특정 전문가로 라우팅되는 이유를 이해할 수 있도록 라우터를 더 해석 가능하게 만드는 데에도 관심이 있다.
요약하면, 라우터 네트워크는 MoE 모델의 기본 구성 요소로, 효율적인 확장과 특화를 가능하게 한다. 초기 게이팅 함수에서 LLM의 현대 희소 라우터로의 발전은 수십 년에 걸친 머신러닝 기법의 진화를 보여준다.