토큰의 딜레마: 거대 비전 언어 모델의 지속 학습을 위한 드리프트 인식 토큰 할당 기반 동적 MoE

새로운 지식을 배울 때 기존 지식을 잊어버리는 '치명적 망각' 문제를 해결하기 위해 MoE 구조를 확장할 때 발생하는 토큰 배정 오류를 분석하고 해결책을 제시한다. 전문가를 추가할 때 기존 데이터가 엉뚱한 전문가에게 배정되어 성능이 떨어지는 현상을 막아 모델의 지속적인 성장을 가능하게 한다.

Mixture of Experts(MoE)는 입력 토큰을 가장 적합한 전문가에게 보내는 라우팅 메커니즘을 핵심으로 한다. 하지만 새로운 작업을 배울 때 라우터가 아직 덜 학습된 상태에서 기존 지식과 관련된 토큰을 신규 전문가에게 잘못 보내는 '라우팅 드리프트'가 발생하며, 이것이 이전 지식을 잊게 만드는 주범이다.

본 논문은 이를 '토큰의 딜레마'로 정의한다. 새로운 데이터셋 안에는 신규 지식을 담은 토큰뿐만 아니라, 기존 지식과 유사하거나 어느 쪽인지 모호한 토큰들이 섞여 있다. 이 모호한 토큰들이 최적화되지 않은 라우터에 의해 신규 전문가에게 배정되면, 라우터의 정책이 오염되어 결국 기존 지식의 인출 경로가 망가진다.

LLaVA-DyMoE는 학습 과정에서 각 토큰의 라우팅 점수를 분석해 성격을 파악한다. 모호하거나 기존 지식에 가까운 토큰은 신규 전문가 근처에 가지 못하도록 강제로 차단하고 기존 전문가에게 보내 지식을 보존한다. 반면 확실히 새로운 패턴을 가진 토큰만 신규 전문가가 학습하게 하여 전문성을 극대화한다.

LoRA 기반의 동적 MoE 구조를 채택하여 각 작업마다 새로운 LoRA 전문가를 추가한다. 기존 전문가와 라우터 파라미터는 동결하여 지식을 보존하고, 신규 전문가와 확장된 라우터 부분만 학습시키는 방식을 취한다.

TAG(Token Assignment Guidance) 메커니즘은 토큰별 라우팅 점수의 차이($D_{rel}$)를 계산한다. 두 전문가 그룹의 최대 점수 차이를 합산 점수로 나누어 상대적 차이를 구하고, 이 값이 임계값보다 작으면 모호한 토큰으로 분류한다. 모호한 토큰은 신규 전문가로의 경로를 $-\infty$로 마스킹하여 기존 전문가 그룹으로 유도한다.

RSR(Routing Score Regularization)은 두 가지 손실 함수를 사용한다. Exclusivity Loss는 기존 그룹과 신규 그룹의 출력 곱을 최소화하여 토큰이 한쪽 그룹에만 확실히 배정되게 한다. Specialization Loss는 신규 토큰이 신규 전문가를 적극적으로 활용하도록 유도하여 모델의 가소성을 확보한다.

CoIN 벤치마크의 8개 VQA 태스크에서 실험한 결과, LLaVA-DyMoE는 평균 최종 정확도(MFN) 57.03%를 기록하여 기존 MoELoRA(43.93%)나 IncMoELoRA(49.68%)를 크게 앞질렀다.

망각 정도를 나타내는 Backward Transfer(BWT) 지표에서 -4.67%를 기록했다. 이는 기존 방식들이 -16%에서 -23% 수준의 심각한 망각을 보인 것과 대조적으로 지식 보존 능력이 탁월함을 입증한 결과이다.

Ablation Study를 통해 TAG와 RSR이 각각 정확도 향상과 망각 방지에 기여함을 확인했다. 특히 TAG가 모호한 토큰의 잘못된 배정을 막아 망각 방지의 핵심적인 역할을 수행함을 수치로 증명했다.

LLaVA-v1.5-7B 백본에 CLIP ViT-L/14 인코더를 결합한 구조이다. 각 레이어의 MLP 블록을 16개의 rank-4 LoRA 전문가로 구성된 MoE 레이어로 대체한다. 라우팅은 Top-K 방식을 사용하며 학습 시 $K=16$으로 설정하여 전체 전문가 중 상위 점수를 받은 전문가들에게 토큰을 배분한다. 총 손실 함수는 표준 다음 토큰 예측 손실에 $\alpha=1e-3$ 가중치를 적용한 정규화 손실 항들을 더해 구성한다.

대규모 모델로의 확장성 및 더 복잡하고 현실적인 시나리오에서의 성능 검증이 과제로 남아 있다.