Transformer 언어 모델의 MLP 형태 재고: Hourglass FFN의 효율성

초록

본 논문은 기존 Transformer에서 널리 사용되는 좁‑넓‑좁(narrow‑wide‑narrow) 형태의 FFN을 대신해, 좁‑넓‑좁 구조의 서브‑MLP를 여러 개 쌓은 hourglass FFN를 제안한다. 동일 파라미터 예산 하에 113 M ~ 1 B 규모 모델을 실험한 결과, hourglass FFN가 특히 400 M 이하에서는 성능 우위를 보이며, 큰 규모에서도 기존 설계와 동등한 결과를 얻는다. 또한 FFN 파라미터를 절감해 attention 차원이나 레이어 깊이를 늘리는 재배분 전략이 전반적인 효율성을 높임을 확인한다.

상세 분석

이 연구는 Transformer의 핵심 구성 요소인 Feed‑Forward Network(FFN)의 설계 규칙을 근본적으로 재검토한다. 기존의 “좁‑넓‑좁” MLP는 입력 차원 d_model을 확장 비율 r (보통 2~4)만큼 넓힌 뒤 다시 축소하는 구조로, 전체 파라미터 중 60 % 이상을 차지한다. 저자들은 최근 이론적·실험적 결과에서 “wide‑narrow‑wide”(hourglass) 형태가 동일 파라미터 수에서 더 높은 함수 근사 능력을 가짐을 인용하고, 이를 Transformer FFN에 적용한다. 구체적으로, 하나의 레이어에 K개의 hourglass 서브‑MLP를 순차적으로 쌓고, 각 서브‑MLP는 d_h < d_model인 bottleneck 차원을 사용한다. 이렇게 하면 파라미터는 K × (3 · d_h · d_model) 정도로, d_h를 작게 잡아도 K를 늘려 깊이를 확보할 수 있다.

실험 설계는 두 축을 중심으로 진행된다. 첫 번째는 동일 d_model과 레이어 수(L) 하에 d_h와 K를 변형해 “FFN 폭 vs. 깊이” 트레이드오프를 탐색한 것이다. 결과는 d_h를 1/41/2 수준으로 축소하고 K를 510 정도로 늘릴 때, 검증 손실과 퍼플렉시티가 최소화됨을 보여준다. 두 번째는 전체 파라미터 예산을 고정하고, hourglass FFN가 절감한 파라미터를 attention 차원(d_attn) 혹은 레이어 수(L) 증가에 재투자하는 전략을 검증한다. 113 M 모델에서는 attention 차원을 28 M→30 M 정도 늘리거나 레이어를 12→14로 확대했을 때, 기존 narrow‑wide‑narrow 대비 0.3 %~0.5 % 퍼플렉시티 개선을 기록했다.

스케일링 실험에서는 113 M, 400 M, 800 M, 1 B 파라미터 모델을 각각 baseline과 hourglass 변형으로 학습했다. 400 M 이하에서는 hourglass가 평균 0.2~0.4 % 낮은 퍼플렉시티를 달성했으며, 1 B 규모에서는 차이가 거의 사라졌다. 이는 hourglass 구조가 작은‑중간 규모에서 파라미터 효율성을 크게 높이지만, 대규모에서는 기존 설계와 비슷한 표현력을 제공한다는 점을 시사한다.

또한, 저자들은 “U‑shaped” 폭‑깊이 트레이드오프를 발견했다. 즉, d_model을 크게 잡고 레이어를 적게 하면 계산량이 급증하지만, d_model을 중간 수준으로 유지하고 레이어를 늘리는 것이 가장 효율적이었다. 이 최적점은 hourglass FFN에서 특히 두드러졌으며, K = 6~8, d_h ≈ 0.4 · d_model 조합이 전반적인 성능을 최적화했다.

이 논문은 기존 Transformer 설계가 “FFN 폭 중심”이라는 암묵적 가정을 가지고 있었음을 지적하고, hourglass FFN을 통해 파라미터를 보다 균형 있게 배분할 수 있음을 실증한다. 특히, 파라미터가 제한된 상황(모바일, 엣지 디바이스)이나 연구 단계에서 빠른 프로토타이핑이 필요한 경우, 넓은 attention 모듈과 깊은 FFN를 결합한 hourglass 설계가 실용적인 대안이 될 수 있다.