기술 요약: 언어 모델의 선형화에서 중요한 것은 무엇인가?

문제 정의

선형화(Linearization)는 트랜스포머의 $O(L^2)$ 자기 주의(self-attention) 메커니즘을 서브쿼드라틱(subquadratic) 토큰 믹서(예: 선형 주의, 게이트형 재귀)로 대체함으로써 효율적인 언어 모델(LM)을 개발하기 위한 주요 전략으로 부상했습니다. 이 접근 방식은 대규모 트랜스포머의 성능을 유지하면서도 선형 시간 추론과 상수 메모리 오버헤드를 제공하는 것을 목표로 하지만, 이 분야를 이해하는 데에는 여전히 상당한 격차가 존재합니다.

현재의 문헌들은 모델 크기, 학습 예산, 증류(distillation) 목적 등이 상이한 이질적인 평가 조건으로 인해 어려움을 겪고 있습니다. 결과적으로, 어떤 특정 귀납적 편향(예: 게이팅, 감쇠, 또는 델타 규칙 업데이트)이 트랜스포머의 능력을 계승하는 데 가장 적합한지 불분명합니다. 또한, 파라미터 수와 토큰 예산에 따른 선형화된 모델의 스케일링 동작, 그리고 지시어 튜닝(instruction tuning)이나 긴 문맥 검색(long-context retrieval)과 같은 다운스트림 태스크에 대한 생존 가능성 역시 미해결 과제로 남아 있습니다.

방법론

이러한 격차를 해결하기 위해 저자들은 7가지 대표적인 서브쿼드라틱 아키텍처인 xLSTM, GLA, Gated DeltaNet, DeltaNet, RetNet, Kimi (KDA), 그리고 Linear Attention을 비교하는 통합적 경험적 프레임워크를 제안합니다.

실험 설정

본 연구는 사전 학습된 트랜스포머 교사 모델(SmolLM2 제품군)과 학생 모델을 정렬하되, 교사 모델, 데이터 분포, 학습 레시피를 일정하게 유지하는 3단계 증류 파이프라인을 사용하여 토큰 믹서를 주요 변수로 고립시킵니다.

1. 토큰 믹서(Attention-Output) 정렬: 학생의 토큰 믹서를 MSE 손실을 사용하여 교사의 어텐션 출력과 정렬합니다. 이는 전체 믹싱 경로에 걸쳐 조밀한 감독 신호를 제공합니다.
2. 은닉 상태(Hidden-State) 정렬: 교사와 학생의 중간 은닉 표현을 MSE 손실을 통해 정렬합니다.
3. 엔드 투 엔드 지식 증류: 표준 지식 증류(다음 토큰 예측 + KL 발산)를 사용하여 학생을 엔드 투 엔드로 학습시킵니다.

평가 차원

본 연구는 세 가지 별개의 실험적 축을 통해 모델을 평가합니다:

1. 스케일링 법칙(Scaling Laws): 성능 계층을 관찰하기 위해 세 가지 파라미터 규모(140M, 360M, 1.7B)에서 고정된 3B 토큰 예산으로 모델을 학습시킵니다.
2. 점근적 성숙도(Asymptotic Maturity): 모델을 약 1.7B 파라미터로 고정하고 10B 토로 확장된 예산으로 학습하여, 성능 격차가 학습 부족 때문인지 아니면 근본적인 아키텍처 한계 때문인지 결정합니다.
3. 다운스트림 적응(Downstream Adaptation): 증류된 학생 모델에 지시어 튜닝(사후 학습 SFT) 및 지시어 인지 증류를 수행하여, 지시어를 따르는 능력과 긴 문맥 검색(Needle-In-A-Haystack 태스크) 능력을 평가합니다.

주요 기여 및 결과

1. 성능 계층의 안정성 및 스케일링 지수

연구 결과, 성능 계층은 규모(140M에서 1.7B까지)에 관계없이 놀라울 정도로 안정적으로 유지됩니다.

• 게이트형 아키텍처: 게이팅 메커니즘이 있는 모델(xLSTM, Gated DeltaNet)은 게이트가 없는 베이스라인(Linear Attention, RetNet)보다 일관되게 우수한 성능을 보입니다.
• 교차점(Crossover Point): 중요한 발견은 1.7B 규모에서의 "교차"입니다. xLSTM이 작은 규모에서는 앞서지만, Gated DeltaNet이 1.7B에서 최고의 성능(교사 성능의 94.2% 회복)을 보이며 등장합니다. 이는 오차 수정 업데이트 규칙(Delta-rule)이 모델 용량이 증가함에 따라 순수 가산적(additive) 대안보다 더 우수한 스케일링 지수를 가짐을 시사합니다.
• 내재적 결함: Linear Attention과 게이트형 모델 간의 성능 격차는 규모가 커짐에 따라(140M에서 ~14%에서 1.7B에서 ~20%로) 넓어지며, 이는 가산적 선형 주의의 결함이 아키텍처 자체에 내재되어 있으며 단순히 파라미터 수를 늘리는 것만으로는 극복할 수 없음을 나타냅니다.

2. 상태 해상도(State Resolution)라는 근본적 병목 현상

학습을 10B 토큰까지 확장한 결과, 상대적인 성능 격차는 초기에 확립되며 모델이 점근적 한계에 도달하더라도 지속됨을 확인했습니다.

• 성숙도: 게이트형 모델의 성능은 6~7B 토큰 이후 정체되는 반면, 게이트가 없는 모델은 더 낮은 수준에서 더 빨리 포화됩니다. 이는 격차가 학습 부족 때문이 아니라, 서로 다른 믹서가 정보를 압축하는 방식의 근본적인 한계를 나타낸다는 것을 확인시켜 줍니다.
• 검색 능력: Needle-In-A-Haystack(NIAH) 태스크에서 대부분의 아키텍처는 시퀀스 길이가 길어짐에 따라 상태 해상도(state resolution)의 붕괴를 겪습니다. Linear Attention과 RetNet은 512 토큰 이상에서 거의 완전히 실패합니다.
• Gated Delta-Rule의 우월성: Gated DeltaNet은 확장된 문맥(예: 2k 토큰에서 81.8% 정확도)에서도 유의미한 검색 정밀도를 유지하는 유일한 아키텍처로, 게이트형 델타 규칙이 표적화된 상태 업데이트 및 상태 포화를 방지하는 데 더 효과적인 메커니즘을 제공함을 입증합니다.

3. 지시어 적응 및 증류

연구는 지시어 튜닝이 아키텍처의 약점을 보완할 수 있는지 평가합니다.

• 차이의 증폭: 사후 학습 지시어 튜닝은 모든 아키텍처의 성능을 향상시키지만, 기존의 차이를 증폭시킵니다. 강력한 모델(Gated DeltaNet, xLSTM)은 상당한 이득을 얻는 반면, 약한 모델은 제한적인 이득만을 보입니다.
• 증류의 한계: 지시어 튜닝된 교사를 사용하는 지시어 인지 증류는 표준 사후 학습 튜닝에 비해 일관된 이점을 제공하지 못하며, 어떤 경우에는 긴 문맥 검색 성능을 저하시키기도 합니다.
• 적응에 대한 결론: 지시어 적응은 선형화 과정에서 부과된 근본적인 아키텍처 제약을 변경하지 못합니다. 정보를 선택적으로 유지하고 덮어쓰는 능력이 다운스트림 성능을 결정하는 지배적인 요인으로 남습니다.

의의 및 주장

본 논문은 성공적인 선형화를 위해서는 아키텍처의 귀납적 편향이 주요 제약 사항이며, 이는 단순히 학습 컴퓨팅을 스케일링하거나 증류 토큰 예산을 늘리는 것만으로는 극복할 수 없다고 주장합니다.

• 설계 계층 구조: 본 연구는 서브쿼드라틱 설계의 명확한 계층 구조를 확립하며, 오차 수정 업데이트 규칙(Delta-rule)이 가산적 누적이나 채널별 게이팅보다 성공에 더 결정적임을 식별했습니다.
• 상태 해상도: 연구는 상태 해상도가 증류 예산보다 더 근본적인 병목 현상이라고 상정합니다. 내부 상태 정밀도를 관리하지 못하는(포화를 피하지 못하는) 모델은 학습 기간에 관계없이 효과적으로 스케일링하거나 긴 문맥을 처리하는 데 실패합니다.
• 재현성: 추가 연구를 촉진하기 위해 저자들은 모든 베이스, 지시어 튜닝 및 지시어 증류 모델을 공개하여 이러한 발견의 재현성을 보장합니다.

저자들은 서브쿼드라틱 모델을 안정적으로 학습할 수는 있지만, 토큰 믹서의 선택이 성능의 천장을 결정하며, 현재로서는 게이트형 델타 규칙(gated delta-rule) 형식이 효율적이고 긴 문맥 처리가 가능한 언어 모델을 향한 가장 강력한 경로를 제공한다고 결론짓습니다.