Ares: Loss-Free Mixture-of-Experts Routing for Bidirectional Protein Encoders

Motivation

• Known: 단백질 언어 모델의 스케일링은 주로 dense transformer의 파라미터 수를 늘리는 방식에 의존해왔으며, ESM-2, ProGen2, ProGen3, ProtTrans, Ankh 등이 이러한 스케일링을 통해 구조/기능 예측 성능을 향상시켜왔다. MoE는 NLP와 vision 분야에서 성공적으로 적용되었고, ProGen3가 autoregressive 단백질 생성 모델에 MoE를 적용한 유일한 사례이다.
• Gap: Soft-MoE와 Expert-Choice와 같이 bidirectional(non-autoregressive) 구조에 자연스럽게 적합한 loss-free routing 전략들이 단백질 encoder에는 전혀 탐구되지 않았으며, 두 전략 간 및 MoE layer 배치(interleaved vs consecutive) 방식에 따른 학습 안정성 비교도 이루어진 바 없다.
• Why: 보조 balancing loss 없이 discrete top-k routing의 구조적 한계(load imbalance, token dropping, 하이퍼파라미터 튜닝 부담)를 회피하면서도, dense 모델 대비 컴퓨트 효율적으로 대규모 단백질 encoder를 스케일링할 수 있는 경로를 제시한다는 점에서 중요하다.
• Approach: UniRef50 데이터로 TPU pods 상에서 동일 compute budget 하에 4개의 4B-parameter MoE 변형(Expert-Choice consecutive/interleaved, Soft-MoE consecutive with/without router L2 normalization)을 학습시켜 routing 전략과 layer 배치에 따른 학습 안정성 및 downstream 성능을 체계적으로 비교하였다.

Achievement

1. 최초의 loss-free MoE 적용: Soft-MoE와 Expert-Choice를 bidirectional 단백질 encoder에 처음 적용하여 동일 데이터·컴퓨트·평가 조건 하에서 통제된 비교를 제공하였다.
2. routing 불안정성 규명: Expert-Choice가 MoE layer를 연속 배치(consecutive)할 경우 cascading routing instability로 조기 종료를 초래하고, dense layer와 interleave할 경우 반복적인 transient collapse를 유발함을 발견하였으며, 이를 straight-through estimator 기반 top-k routing의 gradient mismatch가 depth 차원에 적용된 결과로 설명하였다.
3. 컴퓨트 효율적 성능: 39B 토큰(ESM-2 대비 약 1/25 budget)만으로 학습한 Soft-MoE 변형이 Fluorescence와 GB1에서 ESM-2 3B를, Stability에서 ESM-2 650M을 능가하거나 대등한 성능을 보였다.
4. router L2 normalization ablation: Soft-MoE의 router-side L2 normalization을 제거해도 dispatch 단계가 사실상 무력화됨에도 downstream 성능이 유지됨을 보여, 이 스케일에서 dispatch가 representation quality에 기여하는 정도가 기존 Soft-MoE 이론이 예측하는 것보다 작다는 점을 시사하였다.

How

• 20-layer, d=1024 encoder-only transformer(bias-free linear, SwiGLU FFN, GQA, pre-norm RMSNorm, RoPE)를 기반 아키텍처로 사용, 약 4.2B parameter 규모로 구성
• MoE layer 배치를 interleaved(홀수 layer가 MoE)와 consecutive(10 dense + 10 MoE)로 나누어 비교, 각 MoE layer당 32 experts 사용
• Soft-MoE는 expert당 64 slots, noise level 0.05로 설정하고 원 논문의 router-side L2 normalization 유무에 따른 ablation을 수행
• Expert-Choice는 capacity factor 2.0을 사용하고, 원 논문과 달리 routing score를 expert 축이 아닌 token 축에 대해 softmax로 정규화하는 변형을 채택
• masked language modeling validation accuracy, per-layer token drop rate, expert co-selection(Jaccard overlap), per-layer log ratio 및 selectivity 지표 등을 통해 학습 안정성을 분석
• Fluorescence, GB1, Stability, Remote Homology, Secondary Structure, ProteinGym 등 구조/기능 벤치마크에서 ESM-2와 성능을 비교

Originality

• 단백질 언어 모델 분야에서 그동안 탐구되지 않았던 Soft-MoE와 Expert-Choice라는 두 가지 loss-free routing 전략을 bidirectional encoder에 최초로 도입
• MoE layer 배치(interleaved vs consecutive)가 routing 전략별로 서로 다른 안정성 결과를 낳는다는 placement-driven asymmetry를 실증적으로 규명하고, 이를 top-k routing의 gradient mismatch가 depth 차원으로 확장된 결과로 이론적으로 해석
• Expert-Choice의 routing score normalization을 원 논문의 expert-axis에서 token-axis로 변경한 변형을 명시적으로 구분하여 평가
• Soft-MoE의 핵심 요소로 알려진 router L2 normalization이 실제로는 dispatch 기능을 무력화시키면서도 성능에 영향이 없다는 반직관적 발견을 통해 기존 프레임워크의 이론적 가정에 도전

Limitation & Further Study

• 39B 토큰이라는 상대적으로 적은 학습 budget으로 인해 Remote Homology, Secondary Structure, ProteinGym 등 일부 벤치마크에서는 ESM-2 대비 성능 격차가 존재하며, 이 격차가 순수히 routing 전략의 한계인지 토큰 수 부족 때문인지 완전히 분리되지 않음
• Expert-Choice의 경우 원 논문의 expert-axis normalization이 아닌 token-axis normalization만 평가하여, canonical Expert-Choice 형태에 대한 결론으로 일반화하기는 어려움
• Soft-MoE interleaved 변형은 학습하지 않아 placement-invariance 가설이 실험적으로 완전히 검증되지 않음
• 후속 연구로 더 큰 스케일과 토큰 budget에서의 검증, DeepSeek-style bias 기반 balancing이나 Mixture-of-Tokens 등 다른 loss-free 대안과의 비교가 필요함

Evaluation

총평: 단백질 encoder에 대한 MoE routing 전략을 체계적이고 통제된 방식으로 최초 비교한 실용적이고 시의적절한 연구로, 특히 routing 안정성에 대한 흥미로운 실증적 발견과 router L2 normalization에 대한 반직관적 ablation 결과가 후속 연구에 유용한 통찰을 제공한다. 다만 workshop paper 성격상 스케일과 토큰 budget의 제약, 일부 미완성된 실험 설계(Soft-MoE interleaved 미학습, Expert-Choice 변형 한정 평가)는 향후 보완이 필요하다.