기술 요약: Leviathan – 언어 모델에서 입력 및 출력 표현의 분리

1. 문제 제기

현대 대규모 언어 모델 (LLM) 은 일반적으로 결합 가중치 (tied-weight) 아키텍처를 사용합니다. 여기서 단일 임베딩 행렬 $E \in \mathbb{R}^{V \times D}$가 이중적인 역할을 수행합니다. 즉, 입력 토큰을 연속 표현으로 매핑하고, 숨겨진 상태를 다시 어휘 공간으로 투영하여 출력 예측을 수행합니다. 이러한 결합은 단일 행렬이 두 가지 서로 다른 목표를 동시에 충족하도록 강요합니다.

1. 입력 표현: 의미적으로 관련된 토큰이 표현 공간에서 인접한 영역을 차지하도록 하는 **부드러움 (smoothness)**이 필요합니다.
2. 출력 분류: 어휘 전체의 모든 토큰을 효과적으로 분리하는 **구별력 (discrimination)**이 필요합니다.

저자들은 이 결합이 아키텍처적 한계를 부과한다고 가정합니다. 완전히 분리된 (untied) 모델 (입력과 출력 행렬을 별도로 학습하는 것) 은 이론적으로 이러한 목표를 분리하지만, 실제로는 효과가 입증되지 않았습니다. 실험 결과, 완전히 분리된 모델은 매개변수 수가 약 50% 증가했음에도 불구하고 결합 베이스라인보다 성능이 낮았습니다. 이는 희귀 토큰을 포함한 이산 임베딩에 대한 최적화 어려움과 불충분한 학습 신호로 인해 추가 매개변수가 효과적으로 활용되지 못했음을 시사합니다.

2. 방법론: Leviathan 아키텍처

Leviathan 은 완전히 분리된 임베딩의 prohibitive 한 매개변수 비용을 치르지 않고도 입력 표현을 출력 예측으로부터 분리하는 Transformer 아키텍처를 도입합니다.

2.1 학습된 임베딩 벡터화 (LEV)

핵심 혁신은 LEV 레이어로, 밀집 입력 임베딩 행렬을 학습된 연속 함수 $G: \{0, \dots, V-1\} \to \mathbb{R}^D$로 대체합니다. 이 함수는 컴팩트한 매개변수 공유 메커니즘을 통해 토큰 인덱스를 임베딩으로 매핑합니다.

• 구성적 인덱싱 (Compositional Indexing): 어휘를 $b$진수 분해를 사용하여 $k$개의 구성 요소로 분해하여 토큰 인덱스 $i$를 튜플 $(i_1, \dots, i_k)$로 변환합니다.
• 공유 코드북 (Shared Codebooks): 이러한 인덱스는 공유 코드북 $C_1, \dots, C_k$에 매핑되어 시드 표현 $z(i)$를 생성합니다. 이를 통해 구조적으로 유사한 토큰이 사용하는 공유 항목을 통해 전달된 그래디언트 신호를 희귀 토큰도 활용할 수 있습니다.
• 연속 투영 (Continuous Projection): 시드는 선형 투영 및 정규화를 거쳐 유계 잠재 좌표 $\tilde{z}_\ell \in [0, 1]^{d_{seed}}$로 투영됩니다.
• 스플라인 확장 (Spline Expansion): 각 차원은 $\kappa$개의 매듭을 가진 단변수 2 차 B-스플라인 기저 확장으로 모델링됩니다.
• 텐서 곱 (Tensor Product): 차원들은 랭크-$r$ 텐서 곱을 통해 결합됩니다.
• 헤드 합산 (Head Summation): 최종 임베딩은 선형 투영 후 $h$개의 병렬 헤드로부터의 출력 합계입니다.

2.2 아키텍처 및 분리

• 분리된 출력 헤드: LEV 레이어는 출력 투영 헤드가 완전히 분리된 (fully untied) 상태 (입력과 독립적인 밀집 행렬) 로 유지되도록 합니다.
• 매개변수 효율성: LEV 레이어는 밀집 임베딩 테이블보다 훨씬 작기 때문에 분리된 출력 헤드의 오버헤드는 무시할 수 있습니다. 12 억 (1.2B) 매개변수 규모에서 총 매개변수 증가는 불과 **0.2%**입니다.
• 유도 편향 (Inductive Bias): 이 아키텍처는 토큰 공간에 부드러운 유도 편향을 부과하며, 이는 표준 룩업 테이블의 이산적이고 잠재적으로 이방적인 기하학과 대조됩니다.

3. 실험 설정

• 모델: 실험은 2 억 (200M), 4 억 (400M), 12 억 (1.2B) 매개변수의 세 가지 규모에서 디코더 전용 Transformer 로 수행되었습니다.
• 베이스라인:
  • TIED: 표준 결합 가중치 베이스라인.
  • UNTIED: 완전히 분리된 모델 (4 억 규모에서 6 억 1 천 1 백만 매개변수).
  • ALBERT: 분리된 출력 헤드를 가진 입력 매개변수 일치 분해 (선형 압축).
  • 절제 실험 (Ablations): 비선형성과 구성적 구조의 기여도를 분리하기 위해 매개변수 일치 평탄 MLP 및 분해된 MLP.
• 학습: 모든 모델은 동일한 하이퍼파라미터, 시드, 토큰 예산 (12 억 모델의 경우 최대 241 억 토큰) 으로 The Pile(저작권이 없는) 에서 학습되었습니다.
• 평가: The Pile 의 홀드아웃 검증 집합에 대한 언어 모델링 성능, Wikitext-103에 대한 제로샷 일반화, 그리고 하위 벤치마크 (HellaSwag, ARC, PIQA, WinoGrande, LAMBADA) 평가.

4. 주요 결과

4.1 언어 모델링 성능

Leviathan 은 모든 규모에서 결합 가중치 베이스라인을 일관되게 능가하며, 학습이 진행됨에 따라 모델 규모가 커질수록 성능 격차가 벌어집니다.

• 12 억 규모: Leviathan 은 결합 베이스라인 대비 검증 퍼플렉시티를 9.2%(0.087 nats) 감소시켰습니다.
• 샘플 효율성: 12 억 규모에서 Leviathan 은 결합 베이스라인의 최종 손실에 도달하는 데 2.1 배 적은 학습 토큰이 필요했습니다.
• 학습 역학: 작은 규모 (2 억) 에서 Leviathan 은 초기에는 베이스라인보다 성능이 낮았으나 이후 이를 추월했습니다. 12 억 규모에서는 학습 내내 베이스라인을 능가했습니다.

4.2 대안과의 비교

• 완전히 분리된 모델 대비: 완전히 분리된 모델 (UNTIED) 은 학습 내내 결합 베이스라인보다 성능이 낮아, 효율적인 매개변수화 없이 단순하게 분리하는 것은 최적이 아님을 확인시켜 주었습니다.
• 선형 분해 대비: ALBERT 스타일의 분해된 베이스라인 또한 결합 베이스라인보다 성능이 낮았으며, 이는 선형 압축만으로는 결합 기하학의 표현적 이점을 회복할 수 없음을 나타냅니다.
• Leviathan 대비: Leviathan 은 UNTIED 와 ALBERT 모두를 능가하여, Chinchilla 최적 연산 예산 내에서 비선형 연속 생성이 단순한 분리 및 선형 분해보다 우수함을 입증했습니다.

4.3 하위 작업 및 일반화

• 하위 작업: 12 억 규모에서 Leviathan 은 평가된 모든 6 개 하위 벤치마크에서 성능을 향상시켰습니다. 특히 LAMBADA에서 퍼플렉시티가 30% 감소했습니다.
• 제로샷 전이: 다른 분포인 Wikitext-103 에서 Leviathan 은 규모가 커짐에 따라 퍼플렉시티 개선 폭이 증가했습니다 (2 억 규모에서 2.4%, 12 억 규모에서 12.3%). 이는 연속 표현이 데이터 분포 간에 더 효과적으로 전이됨을 시사합니다.

4.4 빈도 계층화 분석

성능 향상은 토큰 분포의 **긴 꼬리 (long tail)**에 집중되어 있습니다.

• 희귀 토큰: 1 천만 분의 1 미만으로 나타나는 토큰의 경우, Leviathan 은 12 억 규모에서 퍼플렉시티를 81% 감소시켰습니다.
• 빈번한 토큰: 1 만 분의 1 이상으로 나타나는 토큰의 경우, 향상 폭은 거의 0 에 수렴했습니다 (가장 빈번한 토큰의 경우 약간 음수). 이는 학습 신호가 강할 때 이산 임베딩으로도 충분하기 때문입니다.
• 메커니즘: 연속 매개변수화는 구조적으로 유사한 토큰 간에 그래디언트를 공유하여, 데이터 희소성으로 인해 이산 룩업이 실패하는 영역에서 강력한 신호를 제공합니다.

4.5 견고성

• 어휘 크기: 5 만, 10 만, 20 만 어휘 크기에서 결과가 일관되었으며, 전반적으로 약 2% 의 상대적 향상을 보였습니다.
• 절제 실험: 향상 폭은 다음과 같이 분해되었습니다: 연속 표현 (+1.13%), 구성적 구조 (+0.37%), 분리 가능한 스플라인 부드러움 (+0.62%).

5. 의의 및 주장

이 논문은 가중치 결합의 표준 관행이 단순히 가중치를 분리함으로써 해결되지 않는 아키텍처적 한계를 부과한다고 주장합니다. 대신, 최적의 해결책은 입력 및 출력 표현을 분리하면서 연속 매개변수화를 통해 부드러움을 강제하는 것입니다.

• 아키텍처적 통찰: 입력 임베딩의 표현 기하학이 중요합니다. 연속 매개변수화는 특히 이산 임베딩이 어려움을 겪는 희귀 토큰의 경우, 이 기하학을 개선하는 데이터 및 연산 효율적인 수단입니다.
• 효율성: Leviathan 은 무시할 수 있는 매개변수 오버헤드 (12 억 규모에서 0.2%) 로 이러한 이점을 달성하여 현대 언어 모델에 실용적인 업그레이드가 됩니다.
• 확장성: LEV 레이어의 이점은 규모와 데이터에 따라 누적되므로, 모델이 커질수록 이산 룩업 대비 부드럽고 연속적인 입력 표현의 이점이 더 두드러집니다.

저자들은 빈도 의존적 개선 패턴이 상관관계일 수 있지만, 규모, 어휘, 아키텍처 변형 전반에 걸친 일관성이 희귀 토큰에 대한 이산 임베딩 학습과 관련된 최적화 어려움을 연속 매개변수화가 효과적으로 완화한다는 강력한 증거를 제공한다고 결론지었습니다.