Tapered Language Models

Modern language models, including transformer, recurrent, and memory-based variants, share a common chassis: a stack of identical layers in which parameters are allocated uniformly across depth. This is a default inherited from the original transformer and largely unchanged since, yet a growing body of evidence suggests that layers contribute non-uniformly to the final output, with later layers refining the residual stream rather than transforming it. We ask whether parameter capacity should reflect this asymmetry. Our controlled experiment shows that, under a fixed budget, allocating more capacity to earlier layers and less to later layers improves perplexity over a uniform-width baseline, while the reverse allocation hurts. Building on this result, we introduce Tapered Language Models (TLMs), an architectural principle in which a parameter-bearing component is monotonically tapered across depth under a fixed total budget. MLPs are the natural site for this instantiation: they dominate parameter count across all modern LM families and expose width as a single, clean axis of variation. Across three model scales and four architectures (Transformer, Gated Attention, Hope-attention, and Titans), tapering MLP width via a smooth cosine schedule consistently improves perplexity and downstream benchmark performance over uniform baselines, at no additional parameter or compute cost. These findings establish depth-aware capacity allocation as a simple, architecture-agnostic axis of language model design, a free lever hidden in plain sight.

한 줄 요약

Tapered Language Models (TLMs)는 MLP 중간 차원을 코사인 스케줄로 점진적으로 감소시켜 언어 모델 성능을 향상시키는 새로운 아키텍처 원칙이다.

핵심 기여도

MLP 중간 차원 $ d_{\text{ff}} $를 코사인 스케줄로 감소시켜 성능 향상
440M, 760M, 1.3B 파라미터 규모에서 4가지 아키텍처(Transformer, Gated Attention, Hope-attention, Titans)에서 일관된 성능 개선
440M 모델에서 uniform baseline 대비 0.32 포인트의 perplexity 개선 (15.96 → 15.64)
파라미터 증가 없이 성능 향상 (parameter 및 compute cost 증가 없음)

핵심 아이디어

기존 언어 모델은 모든 레이어에 동일한 파라미터를 할당하는 uniform layer 구조를 따르지만, 실제 레이어의 기여도는 비균일하다. 특히, 후반 레이어는 잔차 스트림을 정제하는 역할을 하며, 새로운 정보를 생성하는 데 기여도가 낮다. 이에 따라, Tapered Language Models (TLMs)는 파라미터를 레이어의 깊이에 따라 단조적으로 할당하는 새로운 원칙을 제안한다. MLP의 중간 차원 $ d_{\text{ff}} $는 파라미터 할당의 주요 축이므로, 이를 코사인 스케줄로 감소시켜 앞 레이어에 더 많은 용량을 할당하는 방식을 채택한다. 이는 FFN이 키-벨류 메모리 역할을 하며, 깊이가 깊어질수록 기존 잔차와의 일치도가 높아진다는 해석과도 일치한다.

기술적 접근법

MLP 중간 차원 $ d_{\text{ff}} $를 코사인 스케줄로 감소
총 파라미터 예산은 유지 (parameter 및 FLOPs 증가 없음)
3가지 스케줄 (linear, cosine, sigmoid) 실험, 코사인 스케줄이 가장 효과적
440M, 760M, 1.3B 파라미터 규모에서 4가지 아키텍처 (Transformer, Gated Attention, Hope-attention, Titans)에 적용
FFN 파라미터 수는 $ 2d d_{\text{ff}} $ (SwiGLU는 $ 3d d_{\text{ff}} $)로 계산

주요 결과

440M 모델에서 uniform baseline 대비 0.32 포인트의 perplexity 개선 (15.96 → 15.64)
뒷 레이어에 용량을 집중할 경우 perplexity가 1.33 포인트 악화 (17.29)
중간 레이어에 용량을 집중할 경우 0.68 포인트 악화 (16.61)
3가지 스케줄 중 코사인 스케줄이 가장 효과적
모든 모델 규모와 아키텍처에서 일관된 성능 향상 (parameter 및 compute cost 증가 없이)

의의 및 한계

Tapered Language Models는 파라미터 할당을 깊이에 따라 조절하는 새로운 설계 원칙으로, 기존 uniform layer 구조의 한계를 극복한다. 특히, MLP가 전체 파라미터의 대부분을 차지하므로, 이를 단조적으로 감소시키는 것은 실질적인 성능 향상으로 이어진다. 그러나 이 연구는 MLP에만 적용되었으며, 다른 파라미터 할당 축 (예: attention head 수, memory slot 수 등)에도 동일한 효과가 있는지는 추가 실험 필요. 또한, 모든 아키텍처에 일반화 가능한지는 더 많은 모델 실험을 통해 검증해야 한다.

실용적 활용

Tapered Language Models는 기존 언어 모델을 개선하는 저비용 설계 원칙으로, 파라미터 예산이 제한된 상황에서 특히 유용하다. 특히, 대규모 언어 모델의 파라미터 효율성 향상, 모델 압축, 비용 효율적인 추론 환경에서 활용 가능하다. 비전 트랜스포머나 다모달 모델에도 적용 가능성이 있다.