Tapered Language Models

기술 요약: 테이퍼드 언어 모델 (Tapered Language Models)

문제 정의
트랜스포머(Transformer), 순환 변형 모델(recurrent variants), 메모리 기반 아키텍처를 포함한 현대의 언어 모델(LM)들은 공통적인 구조적 기본값을 공유한다: 즉, 매개변수가 깊이에 따라 균등하게 할당된 동일한 레이어의 스택이다. 레이어가 출력에 기여하는 방식이 비균등하다는 증거가 늘어나고 있음에도 불구하고(초기 레이어는 주로 얕은 구문 패턴을 포착하고, 후기 레이어는 잔차 스트림(residual stream)을 변형하기보다 이를 정제하는 역할을 함), 이러한 균등한 할당 방식은 여전히 표준으로 남아 있다. 본 논문은 이러한 "일률적인(one-size-fits-all)" 접근 방식이 레이어 중요도의 비대칭성을 반영하지 못하며, 매개변수 활용도가 낮은 후기 레이어에서 잠재적으로 매개변수 용량을 낭비할 수 있다고 주장한다.

방법론
저자들은 고정된 총 예산 하에서 매개변수를 보유한 구성 요소가 깊이에 따라 단조롭게 줄어드는(monotonically tapered) 아키텍처 원칙인 **테이퍼드 언어 모델(TLM)**을 도입한다.

• 대상 구성 요소: 본 연구는 이 원칙을 **MLP 중간 너비($d_{ff}$)**에 적용한다. MLP를 선택한 이유는 현대적 언어 모델에서 매개변수 수를 지배하며, 토큰 혼합 모듈(token-mixing module)과 독립적으로 너비를 단일하고 깔끔한 변화 축으로 노출하기 때문이다.
• 테이퍼링 스케줄(Tapering Schedules): 불연속적인 블록 단위 할당 대신, 저자들은 $l \in \{0, \dots, L-1\}$ 레이어에 대해 $d_{ff}(l)$의 부드럽고 연속적인 감쇄 스케줄을 제안한다. 세 가지 스케줄이 평가되었다:
  1. 선형(Linear): 일정한 감쇄율을 가짐.
  2. 코사인(Cosine): 반 코사인 곡선을 따르며, 양 끝단에서는 느리게 감소하고 중간에서는 급격히 감소함.
  3. 시그모이드(Sigmoid): 중간 지점 주변의 좁은 대역에 전이를 집중시킴.
• 제약 조건: 테이퍼링은 엄격하게 예산을 보존한다. 레이어당 평균 너비는 균등한 베이스라인 너비($\frac{1}{L}\sum d_{ff}(l) = d_{ff}^{baseline}$)와 동일하도록 제한되어, 총 매개변수 수와 FLOPs(학습 및 추론 모두)가 균등한 베이스라인과 동일하게 유지된다.
• 실험 설정: 이 원칙은 네 가지 서로 다른 아키텍처(표준 트랜스포머, 게이트 어텐션, 호프-어텐션(Hope-attention), 타이탄스(Titans))와 세 가지 모델 규모(440M, 760M, 1.3B 파라미터)에 걸쳐 평가된다. 연구는 고정된 학습 설정(Llama 3 토크나이저, 4K 시퀀스 길이)을 사용하며, 테이퍼드 모델을 균등 베이스라인과 인-디스트리뷰션(in-distribution) 검증 퍼플렉시티(perplexity), 아웃-오브-디스트리뷰션(out-of-distribution) 퍼플렉시티(WikiText, LAMBADA), 그리고 8개의 상식 추론 벤치마크를 통해 비교한다.

주요 결과

1. 할당의 방향이 중요하다: 440M 트랜스포머를 대상으로 한 동기 부여 실험에서 레이어를 초기, 중기, 후기 그룹으로 나누었다. 용량을 초기 레이어에 더 많이 할당하는 방식("Wider-early")은 균등 베이스라인 대비 퍼플렉시티를 0.32 포인트 개선했다. 반대로, 용량을 후기 레이어에 더 많이 할당하는 방식("Wider-late")은 성능을 1.0 포인트 이상 저하시켜, 용량 할당의 방향이 매우 중요하다는 것을 입증했다.
2. 최적의 스케줄 및 비율: 440M 모델에 대한 너비 비율 및 스케줄 탐색 결과, 코사인 스케줄이 선형 및 시그모이드 스케줄보다 일관되게 우수한 성능을 보였다. 최적의 구성은 시작-종료 너비 비율이 1.50 $\to$ 0.50인 경우로 나타났으며, 균등 베이스라인의 16.28 대비 1.84 포인트 개선된 14.44의 검증 퍼플렉시티를 달로 달성했다.
3. 일반화: 고정된 최적 구성(코사인, 1.5/0.5)을 네 가지 아키텍처 전체의 760M 및 1.3B 모델에 적용했을 때, 퍼플렉시티와 다운스트림 상식 추론 정확도가 일관되게 향상되었다. 이러한 이득은 토큰 혼합 모듈이 소프트맥스 어텐션, 게이트 어텐션, 또는 순환/메모리 기반 메커니즘을 사용하는지 여부와 관계없이 관찰되었다.
4. 기계론적 통찰: 사전 학습된 트랜스포머 분석 결과, MLP 출력은 깊이가 깊어질수록 기존 잔차 스트림과 점진적으로 더 정렬되는 것(코사인 유사도로 측정)으로 나타났다. 이는 후기 MLP가 새로운 특징을 계산하기보다 기존 콘텐츠를 강화함을 시사한다. 테이퍼링은 용량이 가장 적게 사용되는 곳(후기 레이어)에서는 용량을 줄이고, 더 직교적이며 유용한 곳(초기 레이어)에는 용량을 재할당함으로써 이러한 패턴에 아키텍처를 정렬시킨다.

의의 및 주장
본 논문은 **깊이 인식적 용량 할당(depth-aware capacity allocation)**을 언어 모델을 위한 단순하고 아키텍처에 무관한 설계 레버로 확립한다. 저자들은 테이퍼드 MLP 너비가 매개변수 수나 연산 비용을 늘리지 않고도 퍼플렉시티와 다운스트림 작업 정확도를 높이는 "공짜(free)" 성능 향상을 제공한다고 주장한다.

그 의의는 트랜스포머 섀시(chassis)의 계승된 균일성에 도전한다는 데 있다. 연구 결과는 테이퍼링 원칙이 어텐션 메커니즘의 특수한 속성이 아니라, 깊이에 따른 일반적인 매개변수 할당에 관한 것임을 시사한다. 저자들은 특정 최적 스케줄(코사인, 1.5/0.5)이 440M 모델에서 도출되었으며 모든 규모에 대한 글로벌 최적값은 아닐 수 있다고 언급했으나, 다양한 아키텍처와 크기에서 관찰된 일관된 이득은 이 원칙이 견고함을 보여준다. 결론적으로 본 논문은 깊이 인식적 할당이 비전 및 멀티모달 트랜스포머를 포함하여 언어 모델링을 넘어 파운데이션 모델에 적용될 수 있는, 눈앞에 있는 "숨겨진 레버"라고 결론짓는다.