기술 요약: 트랜스포머 언어 모델의 형상 관습 재고

문제 정의

밀집(Dense) 트랜스포머 언어 모델은 역사적으로 경직된 아키텍처 관습을 준수해 왔습니다: 바로 피드포워드 네트워크(FFN)가 "좁음–넓음–좁음(narrow–wide–narrow)" 형태의 다층 퍼셉트론(MLP) 구조를 사용하는 것입니다. 이 설계에서는 모델 차원($d_{model}$)이 일반적으로 2에서 4 사이의 확장 비율을 가지는 더 넓은 중간 차원($d_h$)으로 확장되었다가 다시 투영됩니다. 이러한 구조는 파라미터 할당의 대부분을 차지하며, 모델의 파라미터와 부동 소수점 연산(FLOPs)의 대다수를 점유합니다.

저자들은 이 특정 형상이 필수적인 것인지, 아니면 단지 편리한 구현체일 뿐인지에 대해 의문을 제기합니다. 최근의 이론적 연구는 잔차(residual)가 있는 "넓음–좁음–넓음(hourglass)" 형태의 MLP가 최적의 함수 근사치 역할을 할 수 있으며, 잠재적으로 더 적은 파라미터로 더 우수한 표현력을 제공할 수 있음을 시사합니다. 또한, 생성형 비전 모델과 MoE(Mixture-of-Experts) 트랜스포머의 경험적 증거는 hourglass 구조가 표준 설계보다 뛰어난 성능을 보일 수 있음을 보여주었습니다. 핵심 문제는 기존의 "좁음–넓음–좁음" FFN이 필수적인 제약인지, 아니면 더 깊고 좁은 hourglass FFN이 어텐션(attention)과 FFN 모듈 사이의 용량 할당을 더 효율적으로 제공할 수 있는지 여부입니다.

방법론

본 논문은 기존의 FFN을 Hourglass FFN으로 대체하는 트랜스포머 변형 모델을 제안합니다. 방법론은 다음과 같습니다:

1. 아키텍처 재설계:

   • Hourglass FFN은 단일한 넓은 확장 레이어를 $K$개의 hourglass 형태 서브-MLP 스택으로 대체하며, 이들은 잔차 경로(residual pathways)로 연결됩니다.
   • 각 서브-MLP($MLP_i$)는 압축-확장 구조를 따릅니다: 입력을 $d_{model}$에서 $d_h$ (여기서 $d_h < d_{model}$)의 병목 차원으로 투영한 후, 비선형 활성화 함수(SwiGLU)를 적용하고 다시 $d_{model}$로 투영합니다.
   • 레이어의 출력은 입력을 $K$개의 이러한 잔차 서브 블록을 통과시킨 결과입니다.

2. 파라미터 재배분 전략:

   • 중간 차원 $d_h$를 줄임으로써(FFN을 더 "좁게" 만듦), 저자들은 일반적으로 FFN의 너비에 할당되었던 파라미터를 절약합니다.
   • 절약된 파라미터는 다음 중 하나로 재할당됩니다:
     • 모델 차원($d_{model}$) 및 어텐션 용량 증가.
     • FFN 스택의 깊이($K$) 증가.
     • 트랜스포머 레이어의 수($L$) 증가.

3. 실험 설정:

   • 베이스라인: 표준 LLaMA 스타일 트랜스포머 및 최첨단 OLMo-2 아키텍처와 비교합니다.
   • 규모(Scales): 113M에서 1B 파라미터 범위의 모델 크기에 걸쳐 실험을 수행합니다.
   • 제약 조건: FFN 형상의 효과를 격리하기 위해 총 파라미터 예산(0.001% 이내 차이)과 고정된 어텐션 크기를 기준으로 실험을 수행합니다.
   • 탐색 공간: $K$ (hourglass 스택의 깊이), $L$ (트랜스포머 레이어 수), $d_{model}$, 그리고 $d_h/d_{model}$ 비율을 변화시키며 아키텍처 탐색을 수행합니다.

주요 기여

본 논문은 네 가지 주요 기여를 합니다:

1. 표준에 대한 도전: 사실상의 "좁음–넓음–좁음" FFN 형상에 도전하며, 잔차로 연결된 스택형 서브-MLP를 통해 표현력을 유지하는 효율적인 "넓음–좁음–넓음" Hourglass FFN을 제안합니다.
2. 파라미터 재할당: FFN 너비에서 어텐션 너비나 내부 FFN 깊이($K > 1$)로 파라미터를 전환함으로써 성능이 향상됨을 입증하며, 특히 작은 규모(113M)에서 이를 확인했습니다.
3. 아키텍처 통찰: $d_{model}/L$ 및 $d_h/d_{model}$에 대한 최적의 비율을 찾아내는 U자형 너비-깊이 트레이드오프를 식별하고, 더 깊은 Hourglass FFN 구조가 견고함을 입증합니다.
4. 확장성 검증: 113M에서 1B 파라미터까지 Hourglass FFN을 검증하여, 특정 구성이 기존 모델보다 작은 규모에서 더 우수한 성능을 보이고 큰 규모에서는 그들과 대등한 성능을 보임을 보여줍니다.

결과

경험적 검증을 통해 다음과 같은 결과가 도출되었습니다:

• 고정된 어텐션에서의 성능: 113M 규모에서 고정된 어텐션 파라미터를 사용할 때, Hourglass FFN은 기존 FFN보다 낮은 검증 손실(validation loss)과 퍼플렉시티(PPL)를 달s성했습니다. 예를 들어, $K=8$일 때 PPL은 36.441(기존)에서 36.179로 감소했습니다.
• 최적의 구성:
  • 깊이 vs 너비: $d_h/d_{model} \approx 0.4$와 $K=4$ 레이어의 비율이 가장 낮은 검증 손실을 기록했습니다.
  • 파라미터 이동: 113M 규모에서 어텐션 모듈이 지배적이 되도록 파라미터를 이동시킨 경우(예: 어텐션 66M vs FFN 47M), 가장 좋은 성능(PPL 35.101)을 보였으며, 이는 작은 모델에서는 정적인 사실적 인출보다 문맥적 처리가 더 가치 있음을 시사합니다.
  • 너비-깊이 트레이드오프: $d_{model}/L$ 비율에 대해 뚜렷한 U자형 곡선이 관찰되었으며, 100과 250 사이에서 최적의 지점이 발견되었습니다.
• 확장성:
  • 소규모에서 중규모(113M – 906M): Hourglass 변형 모델은 검증 손실, 퍼플렉시티 및 다운스트림 추론 작업(예: Arc Easy, HellaSwag, PIQA)에서 기존 베이스라인보다 일관되게 우수한 성능을 보였습니다.
  • 대규모(1B): 1B 파라미터 규모에서 Hourglass 변형 모델은 강력한 기존 베이스라인과 대등한 성능(Val PPL 20.082 vs 20.002)을 달성했으며, OLMo-2 구성(20.466)을 능가했습니다.
• 어블레이션 연구(Ablation Studies): hourglass 깊이 $K$를 늘리는 것은 성능을 일관되게 향상시켰으나, $K=6$ 이후에는 한계 효용이 감소했습니다. $d_h/d_{model}$을 0.4 미만으로 줄이면 성능이 크게 저하되어, FFN 압축의 하한선이 존재함을 나타냈습니다.

의의 및 주장

본 논문은 트랜스포머의 오랜 "좁음–넓음–좁음" MLP 형상이 필수적이지 않으며 중복될 수 있다고 주장합니다. Hourglass FFN을 채택함으로써 저자들은 다음을 입증했습니다:

1. 효율성: 모델은 더 효율적인 파라미터 분포를 통해 경쟁력 있거나 더 우수한 성능을 달성할 수 있습니다.
2. 자원 재할당: Hourglass 구조는 어텐션 메커니즘이 파라미터 소비의 주체가 될 수 있도록 하는 근본적인 전환을 가능하게 하며, 이는 전통적인 FFN 중심의 예산 배분에 도전합니다.
3. 설계의 재평가: 본 연구의 결과는 어텐션과 FFN 사이의 균형에 대한 재고를 촉구하며, 더 깊고 좁은 FFN과 더 넓은 어텐션 메커니즘을 결합하는 것이 효율적이고 표현력 있는 현대적 언어 모델을 위한 더 우수한 설계 전략일 수 있음을 시사합니다.

저자들은 수십억 또는 수백억 개의 파라미터를 가진 모델로의 확장성에 대해서는 신중한 태도를 유지하며, 113M 규모에서 얻은 결과를 바탕으로 추론했음을 언급했습니다. 또한, 더 깊은 아키텍처에서 발생하는 "깊이의 저주(signal propagation issues)"에 대한 추가 조사가 필요하다고 밝혔습니다. 또한, Hourglass FFN과 고급 어텐션 메커니즘(예: Group-Query Attention) 사이의 상호작용은 아직 탐구되지 않은 영역임을 명시했습니다.