연구진은 멀티모달 Mixture-of-Experts (MoE) 모델에서 시각적 토큰과 텍스트 토큰의 균형을 유지하는 기하학 가이드형 토큰 라우팅 방식인 ReBA routing을 소개했습니다. 초기 실험 결과, 이 방식은 기존 라우터들이 흔히 어려움을 겪는 이미지 해상도 급증 상황에서도 학습을 안정적으로 유지하는 것으로 나타났습니다.
시각-언어 모델에 새로운 라우터가 필요한 이유
시각-언어 모델은 매우 다른 두 가지 스트림을 입력받습니다. 수백 개에 달할 수 있는 이미지 패치와 소수의 단어 토큰입니다. 표준 MoE 라우터는 모든 토큰을 동일한 종류의 데이터로 취급하기 때문에, 이미지 패치가 특정 전문가(expert)에게 몰리는 반면 텍스트 토큰은 유휴 상태로 남게 됩니다. 기존 연구들은 각 전문가에게 할당된 토큰의 총합만을 동일하게 맞추는 보조 손실(auxiliary loss)을 통해 부하를 해결하려 했습니다. 하지만 방대한 이미지 부하가 미미한 텍스트 부하를 상쇄해 버릴 수 있기 때문에, 이 손실 함수는 성능이 저하될 때까지 실제 불균형을 가려버리는 문제가 있습니다.
ReBA가 라우팅 방식을 바꾸는 방법
ReBA는 라우팅 프로세스에 **모달리티 경계(modality boundary)**를 추가합니다. 패치와 단어를 섞어버리는 단일 풀(pool) 대신, 이미지 토큰의 기하학적 배치를 존중하는 별도의 경로를 구축합니다. 각 이미지 인스턴스는 동일한 전체 가중치를 부여받아 특정 전문가가 병목 현상이 되는 것을 방지합니다. 패치의 공간적 배치를 가이드로 삼아, 입력 해상도가 변하더라도 시스템은 안정성을 유지합니다.
저자들이 보고한 주요 장점은 다음과 같습니다:
- 시각적 토큰과 언어적 토큰 사이의 명확한 분리를 강제합니다.
- 배치 내 각 이미지의 균등한 기여를 보장합니다.
- 특정 모달리티에 의해 전문가가 과부하되는 것을 방지합니다.
- 이미지 패치 수가 증가해도 균형을 유지합니다.
여러 벤치마크 설정에서 ReBA는 배치에 패치가 얼마나 추가되든 관계없이 전문가 풀을 고르게 활용하며, 기존의 보조 손실 방식보다 뛰어난 성능을 보였습니다.
한계 및 남은 과제
본 연구는 속도나 메모리 소비에 대한 수치를 제공하지 않으므로, 추가된 라우팅 로직이 숨겨진 비용을 발생시키는지 여부는 알 수 없습니다. 또한 저자들은 단일 이미지의 모든 패치가 하나의 단위로 작동한다고 가정했는데, 이 가정은 서로 다른 해상도의 이미지가 섞여 있거나 일부가 마스킹된 영역을 포함하는 배치에서는 무너질 수 있습니다.
향후 주목해야 할 점
- 성능 대 효율성의 트레이드오프: 향후 연구에서는 ReBA가 추가하는 오버헤드를 정량화해야 합니다.
- 혼합 배치 동작: 고해상도와 저해상도 이미지가 결합된 배치에서의 테스트를 통해 모달리티 경계가 잘 유지되는지 확인할 수 있을 것입니다.
- 대규모 파이프라인 도입: 라우팅 안정성이 이미지 캡셔닝(image captioning)이나 시각적 질의응답(visual question answering)과 같은 다운스트림 작업의 성능 향상으로 이어진다면, 개발자들은 표준 보조 손실 대신 ReBA를 채택할 수 있습니다.
시각-언어 MoE 파이프라인을 구축 중이라면, 이미지 해상도를 높일 때 기본 라우터를 ReBA로 교체함으로써 더욱 균일한 전문가 사용 패턴을 얻을 수 있습니다. 변경 후에는 토큰 분포 지표를 주의 깊게 살펴보십시오. 분포가 더 평탄해진다면(flatter distribution) 모달리티 경계가 제 역할을 하고 있다는 신호입니다.