다중폭·깊이 재활용 피드포워드 네트워크 VersatileFFN

초록

VersatileFFN은 하나의 공유 FFN 파라미터를 폭‑폭과 깊이‑깊이 두 축에서 재활용하는 구조로, 폭‑전문가 경로는 가벼운 가상 MoE를, 깊이‑재귀 경로는 토큰별 반복 적용을 제공한다. 난이도 인식 게이팅이 두 경로를 동적으로 결합해 “쉬운” 토큰은 빠르게, “어려운” 토큰은 깊게 처리한다. 파라미터는 그대로 유지하면서 연산량만 늘려 효율성을 높이며, 다양한 벤치마크와 모델 규모에서 기존 파라미터‑동등 혹은 FLOPs‑동등 방법들을 능가한다.

상세 분석

VersatileFFN은 기존 Transformer의 Feed‑Forward Network(FFN)를 두 개의 적응형 서브‑패스로 재구성한다. 첫 번째인 폭‑다양성(Width‑Versatile) 경로는 하나의 대형 FFN 가중치를 스트라이드 방식으로 비중첩 서브스페이스로 분할해 N개의 가상 전문가(virtual expert)를 만든다. 라우터 W_g가 토큰별로 Top‑K 전문가를 선택하고, 선택된 전문가들의 출력은 Softmax‑정규화된 게이팅 가중치로 가중합된다. 이 설계는 물리적인 파라미터 증가 없이 MoE와 유사한 전문화 효과를 제공한다. 두 번째인 깊이‑다양성(Depth‑Versatile) 경로는 동일한 FFN를 재귀적으로 L_max 번 적용한다. 토큰마다 필요한 반복 횟수를 예측하기 위해 W_loop 헤드가 Gumbel‑Softmax 기반의 차분 가능한 루프 카운트를 출력한다. 훈련 시에는 확률적 샘플링과 Straight‑Through Estimator를 사용해 미분 가능하게 만들고, 추론 시에는 argmax 로 결정된 반복 횟수만큼 실제 연산을 수행한다. 난이도‑인식 융합 단계에서는 깊이‑경로에서 얻은 기대 반복 횟수를 토큰 난이도 지표로 활용해 λ∈