A Single-Layer Model Can Do Language Modeling

기술 요약: 단일 레이어 모델이 언어 모델링을 수행할 수 있다

문제 제기
현대 언어 모델은 일반적으로 깊이를 증가시켜 용량을 확장하며, 각 레이어가 자체 상태 (예: 트랜스포머의 레이어별 KV 캐시 또는 Mamba, GDN, RWKV, xLSTM 의 레이어별 행렬 상태) 를 유지하는 여러 레이어를 쌓아 올립니다. 이는 계산과 메모리를 스택 전체에 분산시킵니다. 반면, 생물학적 신경망은 깊게 쌓는 것보다 시간 경과에 따라 동일한 회로를 반복 방문하는 강한 재귀 (recurrence) 에 의존합니다. 본 논문은 언어 모델링에 있어 이러한 생물학적 형태가 얼마나 유효한지 조사합니다. 구체적으로, 깊은 스택에 상태를 분산시키는 대신 매 단계에서 단일 상태 벡터를 반복적으로 재방문할 때, 단일하고 얕은 재귀 레이어가 얼마나 멀리 나아갈 수 있는지 탐구합니다.

방법론: 그라운딩 예측 네트워크 (Grounded Prediction Networks, GPN)
저자들은 매 시간 단계에서 재방문되는 단일 상태 벡터로 정의된 아키텍처인 그라운딩 예측 네트워크 (GPN) 를 제안합니다. 핵심 루프는 세 가지 연산으로 구성됩니다:

1. 그라운딩 (Grounding): 이전 예측 상태 ($s^p_{t-1}$) 가 게이트된 잔차 (gated residual) 를 통해 현재 관측된 토큰 ($x_t$) 과 융합되어 그라운딩된 상태 ($s^g_t$) 를 생성합니다.
2. 예측 (Prediction): 그라운딩된 상태를 처리하여 다음 상태 ($s^p_t$) 를 예측합니다. 이 예측은 다음 단계의 토큰 분포로 디코딩됩니다.
3. 메모리 (GPN+M): 저자들은 중앙 집중식 행렬 메모리 ($M$) 를 추가하여 기본 GPN 을 GPN+M으로 확장합니다. 이 메모리는 게이트된 델타 규칙을 통해 업데이트되며, 피드포워드 네트워크 (FFN) 와 병렬로 읽힙니다. FFN(오프라인 파라미터) 과 메모리 읽기 (온라인 파라미터) 모두 상태 업데이트에 대한 가산 예측자 (additive predictors) 로 작용합니다.

레이어별 상태를 유지하는 표준 재귀 모델과 달리, GPN 은 매 단계에서 재방문되는 단일 공유 행렬 메모리를 가진 하나의 넓고 얕은 레이어를 사용합니다. 깊이는 공간 (레이어 쌓기) 이 아니라 시간 (루프 재방문) 을 통해 달성됩니다. 상태 벡터는 이중 역할을 수행합니다: 다음 단계에 대한 예측이자 지금까지 관측된 역사의 인코딩 역할을 하며, '그라운딩'은 다음 사이클을 위한 인코딩을 완성합니다.

주요 기여

1. 아키텍처: 단일 상태 벡터, 하나의 FFN, 그리고 매 단계에서 재방문되는 하나의 중앙 집중식 행렬 메모리를 활용하는 GPN 과 GPN+M 의 도입. 이는 트랜스포머와 선형 RNN 의 레이어별 쌓인 상태와 대조됩니다.
2. 개념적 관점: 다음 단계에 대한 상태의 예측이 해당 단계의 인코딩의 일부로 간주되는 프레임워크로, 단일 벡터가 두 가지 역할을 모두 수행합니다.
3. 실증적 증거: 1 억 3 천만 파라미터 기준 성능 벤치마크에서 1 레이어 GPN+M 이 깊은 베이스라인과 경쟁력 있는 퍼플렉시티를 달성하며, 2 레이어 변형은 격차를 더 좁혔습니다.
4. 관측 가능성: 전체 작업 컨텍스트가 단일 벡터이므로, 저자들은 잔류 스트림을 파헤칠 필요 없이 기하학적 구조를 직접 검사할 수 있습니다. 이를 통해 아키텍처에 명시적으로 설계되지 않은 등장적 구조를 식별했습니다.

결과
FineWeb-Edu 데이터셋 (1 억 3 천만 파라미터, 2048 컨텍스트 윈도우) 에서 평가:

• 1 레이어 GPN+M: 퍼플렉시티 (PPL) 18.06을 달성했습니다. 이는 12 레이어 Transformer++(PPL 16.05) 대비 13% 이내이며, 10 레이어 게이트드 델타넷 (GDN)(PPL 15.34) 대비 18% 이내입니다.
• 2 레이어 GPN+M: PPL 16.95를 달성하여 Transformer++ 대비 격차를 6% 로, GDN 대비 11% 로 좁혔습니다.
• LAMBADA 작업: 얕은 모델은 깊은 베이스라인에 비해 컨텍스트 의존적 작업에서 낮은 정확도를 보여, 토큰당 계산 깊이에 따른 비용을 강조합니다.
• 메모리 없음 (GPN): 행렬 메모리가 없는 버전은 훨씬 더 낮은 성능 (PPL 23.51) 을 보이며, 연관성 메모리 구성 요소의 필요성을 확인시켜 줍니다.

등장적 구조 (분석)
단일 벡터 상태는 직접적인 기하학적 분석을 가능하게 하여, 명시적 정규화 없이 자발적으로 등장한 세 가지 구조를 드러냈습니다:

1. 고차원 기하학 (High-Rank Geometry): 상태 벡터는 넓으며 ($d=2496$) 전체 너비를 활용합니다. 거의 상관관계가 없는 차원과 '죽은' 차원이 없으며, 안정된 랭크는 약 40 입니다.
2. 지속적인 기본 토큰 방향: 상태 벡터의 평균 ($\bar{s}$) 은 빈번한 기능 토큰 (예: "the", "and", ",") 으로 지배되는 분포로 디코딩되는 크고 고정된 오프셋입니다. 디코더에 명시적 편향을 추가하더라도 이 방향은 지속되므로, 단순한 디코딩을 넘어 구조적 역할 (예: RMSNorm 고정) 을 수행함을 시사합니다.
3. 빠른/느린 메모리 분할: 아키텍처적으로 동일한 15 개의 메모리 헤드가 훈련 중에 자발적으로 두 개의 풀로 나뉩니다: 13 토큰의 반감기를 가진 '빠른' 풀 (6 개 헤더) 과 5322 토큰의 반감기를 가진 '느린' 풀 (9 개 헤더) 입니다. 이는 명시적 아키텍처 설계 없이 두 시간 척도 메모리 시스템을 생성합니다.

의의와 한계
본 논문은 단일 얕은 재귀 레이어가 언어 모델링을 효과적으로 수행할 수 있으며, 매칭된 파라미터 수에서 깊은 베이스라인의 약 13% 이내까지 도달할 수 있다고 주장합니다. 그 의의는 '깊이'가 시간 재방문 재귀로 대체될 수 있음을 입증하고, 단일 벡터가 일반적으로 깊은 잔류 스트림에 숨겨져 있는 복잡한 등장적 메모리 구조 (빠른/느린 분할) 와 기하학적 속성 (지속적 앵커) 을 지원할 수 있음을 보여준 데 있습니다.

그러나 저자들은 현재 작업 상태에 대해 겸손합니다:

• 생산 준비 부족: 모델은 시퀀스 차원을 따라 병렬화되지 않고 Python 에서 순차적으로 실행되므로, 훈련 비용이 트랜스포머보다 훨씬 높습니다.
• 확장성: 결과는 1 억 3 천만 파라미터에서 단일 시드 (seed) 기반이며, 확장성 스윕 (scaling sweeps) 이 수행되지 않았습니다.
• 기술적 성격: 등장적 구조 (기본 토큰 방향 등) 에 대한 분석은 관찰적입니다. 논문은 제거 실험 (ablation) 을 통한 유용성에 대한 인과적 증명을 주장하지 않습니다.
• 컨텍스트 비용: 얕은 모델은 LAMBADA 와 같은 깊은 컨텍스트가 필요한 작업에서 측정 가능한 비용을 치르므로, 이 접근법이 유효함에도 불구하고 깊은 스택이 여전히 특정 컨텍스트가 많은 예측에 있어 장점을 보유하고 있음을 시사합니다.

이 작업은 '그라운딩 예측' 루프와 단일 벡터 상태 모델의 유효성에 대한 개념 증명 (proof of concept) 으로, 언어 모델이 정보를 조직화하는 방식을 관찰하는 새로운 렌즈를 제공합니다.