MoE 모델의 라우팅과 어텐션을 결합한 새로운 방식 제안
Attention-Aware Routing: Coupling Routing and Attention in MoEs
arXiv기존 MoE 라우터가 제한적인 컨텍스트 정보만 사용하던 한계를 극복하기 위해, 어텐션 가중치에서 추출한 시간적/스펙트럼 특징을 활용하는 Attention-Aware Routing(AAR)이 제안되었습니다.
- AAR은 성능 향상을 입증했을 뿐 아니라, 라우팅과 어텐션 메커니즘이 서로 영향을 주고받는 '결합 회로'임을 밝혀내어 모델의 작동 원리를 심층적으로 분석했습니다.
- 이는 추론 과정에서 오답으로 인해 내용이 길게 이탈하는 현상(diverging generation)을 줄여주고, 정확한 답변 유지를 돕는 등 실질적인 안정성 개선 효과를 가져옵니다.
- 다만, AAR은 적용되는 레이어 깊이에 매우 민감하여 모든 층에 무분별하게 사용할 경우 사실 검색 능력이 저하될 수 있으므로 신중하고 선택적인 접근이 필요합니다.
기존 (MoE) 언어 모델에서 사용되는 라우터는 의 숨겨진 상태만을 기반으로 전문가를 선택하고 를 부여하는 한계가 있었습니다. 이에 연구진은 Attention-Aware Routing (AAR)을 제안했습니다. AAR은 슬라이딩 윈도우로 추출된 어텐션 가중치에서 시간적 및 스펙트럼 특징을 라우터에 추가하여, 숨겨진 상태와 분리된 모델의 컨텍스트 상태 요약을 활용합니다.
AAR은 기본 를 고정하고 라우팅 만 학습시켜 구현되었으며, OLMoE 기반에서 라우팅 전용 SFT 기준선 대비 GSM8K 성능을 +3.37 pp 향상시키는 결과를 보였습니다. 또한 연구진은 라우팅과 어텐션이 서로 영향을 주고받는 '결합 회로'를 형성함을 입증했습니다. 즉, l 레이어에서의 라우팅 변화가 잔차 스트림을 통해 다음 l+1 레이어의 어텐션 싱크를 증폭시키며 어텐션을 재구성합니다.
실질적인 측면에서 AAR은 오답으로 인해 내용이 길게 이탈하는 현상(long diverging generation)을 줄여주어 정확한 답변 유지를 돕습니다. 다만, AAR의 적용은 레이어 깊이에 매우 민감하여 모든 층에 무분별하게 사용할 경우 사실 검색 능력이 저하될 수 있습니다. 따라서 연구진은 각 레이어의 특성을 고려한 선택적 접근이 필요하다고 강조했습니다.
용어 풀이
- Mixture-of-Experts
- 여러 하위 모델 중 일부만 골라 계산하는 구조. 모델이 커도 실행 비용을 줄일 수 있어요.
- 토큰
- AI 모델이 글을 처리하는 단위. 단어보다 작은 조각이며 사용량과 요금을 셀 때 기준이 돼요.
- 가중치
- 학습으로 정해진 모델 내부의 숫자 값. 내려받는 모델 파일의 본체예요.
- 트랜스포머
- 오늘날 대부분의 언어 모델이 쓰는 신경망 구조. 문장 속 단어들의 관계를 한꺼번에 계산해요.
- 매개변수
- 모델이 학습하면서 조정한 내부 숫자. 개수가 많을수록 대체로 모델이 크고 무거워요.