기술 요약: Hidden Decoding의 스케일링 (Hidden Decoding at Scale)

문제 정의

대규모 언어 모델(LLM)의 스케일링은 전통적으로 트랜스포머 백본(backbone)을 확장하는 방식에 의존해 왔으나, 이는 막대한 비용이 드는 사전 학습(pretraining) 과정을 요구하며 서빙 비용을 증가시킵니다. 이미 강력한 성능을 가진 파운데이션 모델의 경우, 백본을 추가로 확장하는 것은 비현리적인 경우가 많습니다. 최근 연구들은 고정된 파라미터 내에서도 토큰당 연산량을 할당함으로써 추론 및 정확도를 개선할 수 있다고 제안하지만, 기존 방식들은 다음과 같은 심각한 확장성 문제를 안고 있습니다:

• 심도-재귀형(Depth-Recurrent/Looped) 트랜스포머: 동일한 트랜스포머 블록을 토큰당 여러 번 재사용하여 연산량을 늘립니다. 그러나 루프형 연산의 반복적 특성은 가장 큰 모델들을 학습시키는 데 사용되는 파이프라인 병렬 처리(pipeline parallelism) 가정을 깨뜨려 GPU 유휴 시간(idle time)을 발생시키기 때문에, 프런티어 모델(예: 100B+ 파라미터)로 확장하기 어렵습니다.
• 시퀀스 길이 확장(Sequence-Length Expansion): 단순히 더 긴 입력을 제공하는 방식은 표준 학습 파이프라인과 호환되지만, 동일한 토큰에 대해 컨텍스트 윈도우를 늘리는 것 이상의 연산을 확장하는 방법이 필요합니다.

본 논문은 기존 트랜스포머 백본의 깊이나 너비를 변경하지 않으면서도, 파이프라인 병렬 처리 및 대규모 학습 인프라와 호환성을 유지하며 토큰당 연산량을 높이는 과제를 해결하고자 합니다.

방법론: Hidden Decoding

저자들은 계속적 사전 학습(Continued Pretraining, CPT) 단계에서 적용되는 시퀀스 길이 확장 방식인 Hidden Decoding을 제안합니다. 핵심 메커니즘은 각 입력 토큰을 단일 순방향 패스(forward pass) 내에서 처리되는 $n$개의 병렬 스트림으로 확장하는 것입니다.

1. 멀티 스트림 토큰 확장 (Multi-Stream Token Expansion)

• 임베딩 확장: 단일 어휘 임베딩 테이블 대신, 모델은 $n$개의 독립적인 임베딩 테이블($E_1, \dots, E_n$)을 활용합니다. 입력 토큰 $x_i$는 $n$개의 표현식 시퀀스 $(E_1(x_i), E_2(x_i), \dots, E_n(x_i))$로 확장됩니다.
• 시퀀스 구성: 이 표현식들은 인터리빙(interleaved)되어 길이 $nL$($L$은 원래 시퀀스 길이)의 확장된 시퀀스를 형성합니다.
• 학습 목적 함수: 오직 마지막 스트림($E_n$)만이 지도 학습(supervised)됩니다. 모델은 마지막 스트림 위치의 은닉 상태(hidden state)를 기반으로 다음 토큰 $x_{i+1}$을 예측합니다. 처음 $n-1$개의 스트림은 직접적인 손실(loss)이 없는 "잠재 연산 상태(latent computation states)"로서 작용하며, 이를 통해 최종 예측 전 내부적으로 표현을 정교화할 수 있게 합니다.

2. 스트림 분리 어텐션 (Stream-Factorized Attention)

확장된 시퀀스로 인해 어텐션 비용이 이차 함수적으로($O(n^2 L^2)$) 증가하는 것을 방지하기 위해, 저자들은 스트림 분리 어텐션을 도입했습니다:

• 내부 스트림 레이어 (Intra-Stream Layers): 대부분의 레이어는 동일한 스트림 내의 위치로 어텐션을 제한합니다. 이를 통해 복잡도를 약 $O(nL^2)$로 줄입니다.
• 교차 스트림 레이어 (Cross-Stream Layers): 소수의 레이어(예: 베이스 모델의 풀 어텐션 또는 슬라이딩 윈도우 레이어)는 스트림 간의 혼합을 허용합니다.
• 효율성: 이 설계는 추가되는 어텐션 비용을 $n$에 대해 이차 함수가 아닌 선형에 가깝게 유지하여, 100B 이상의 파라미터를 가진 모델에서도 학습이 가능하게 합니다.

3. 점진적 확장 및 초기화 (Progressive Expansion and Initialization)

큰 확장 계수(expansion factor)를 학습하기 위해 저자들은 점진적 확장(예: $n=1 \to 2 \to 4 \to 8$) 기법을 사용합니다. 또한 **순환 복제 초기화(Cyclic Replication Initialization)**를 사용하는데, 여기서 새로운 임베딩 테이블은 기존의 것을 복사하여 초기화합니다(예: $E_{n+k} \leftarrow E_k$). 이를 통해 모델이 초기에는 작은 설정의 동작을 재현하면서도 이후 부드럽게 적응할 수 있도록 합니다.

4. 아키텍처 최적화 (WeLM 특화 사항)

WeLM 백본의 경우, 저자들은 KV-Mirror 설계를 활용합니다. 중간 스트림은 손실(loss)을 가지지 않으며 오직 Key와 Value를 통해서만 컨텍스트에 기여하므로, 미러링된 레이어(이전 레이어로부터 KV를 계산하는 레이어)는 중간 스트림에 대한 처리를 완전히 건너뛸 수 있습니다. 이는 80B 모델 설정에서 학습 시간을 약 20% 추가로 단축시킵니다.

주요 기여

1. 프런티어 규모의 고정 백본 스케일링: 본 논문은 100B+ MoE 규모에서 시퀀스 길이 확장의 첫 성공적인 적용 사례를 보여줍니다. 저자들은 WeLM-HD4-80B 및 WeLM-HD4-617B ($n=4$)를 학습시켜, 트랜스포머 백본을 키우지 않고도 비-HD 모델 대비 성능 향상을 달성했습니다.
2. 효율적인 시퀀스 길이 확장: 스트림 분리 어텐션의 도입은 확장된 시퀀스를 사용하여 학습할 때 발생하는 비용을 $n$에 대해 선형에 가깝게(예: 4배 확장은 밀집 어텐션의 16배 비용보다 훨씬 낮은 약 4.4배~5.1배의 비용 발생) 유지할 수 있게 합니다.
3. 확장 계수 스케일링: 실험 결과, 확장 계수 $n$을 높임에 따라(1에서 8로) 언어 모델링 손실과 다운스트림 정확도가 일관되게 개선됨을 확인하였으며, 이는 $n$이 고정된 백본을 위한 실질적인 스케일링 조절 장치(knob)임을 입증합니다.

실험 결과

저자들은 프런티어 규모의 MoE 모델과 밀집(dense) 베이스라인 모델에 대해 Hidden Decoding을 평가했습니다.

• 프런티어 규모 성능 (WeLM-HD4-80B & 617B):

  • WeLM-HD4-80B는 9개의 공유 하드 벤치마크 모두에서 WeLM-80B를 능가했으며, 특히 SciCode (+4.2), PHYBench (+4.0), FrontierMath (+3.2)에서 유의미한 이득을 보였습니다.
  • WeLM-HD4-617B는 GPQA Diamond (+2.1), HLE (+1.8), FrontierMath (+2.0)를 포함한 9개 벤치마크 전체에서 WeLM-617B보다 향상된 성능을 보였습니다.
  • 이 모델들은 효과를 격리하여 관찰하기 위해 초기 SFT-only 포스트 트레이닝 레시피를 사용하여 학습되었습니다.

• 확장 계수($n$)에 따른 스케일링:

  • 밀집 모델인 Qwen3-8B-Base에서 $n$을 1에서 8로 높였을 때, 평균 정확도가 단조 증가하는 결과(예: MMLU가 78.2에서 81.6으로 개선)를 보였습니다.
  • 80B MoE의 경우, $n$을 1에서 8로 높였을 때 MMLU는 85.0에서 87.5로 상승했고, Pile-test BPB는 0.386에서 0.378로 낮아졌습니다.

• 절제 연구 (Ablation Studies):

  • 지도 설계: 마지막 스트림만을 감독하는 것이 모든 스트림에 손실을 적용하거나 스트림 출력을 합치는 것보다 더 나은 결과를 냈습니다.
  • KV 유지: 별도의 KV 캐시를 유지하는 것(스트림을 공유하는 것보다)이 정확도를 높였으며, 이는 중간 스트림이 독특하고 유용한 컨텍스트를 제공함을 확인시켜 줍니다.
  • 어텐션 구성: 소수의 풀 교차 스트림 레이어(예: 27개 중 4개)만 가진 모델이 모든 레이어에서 풀 교차 스트림 혼합을 수행하는 모델과 거의 대등한 성능을 보였으며, 이는 스트림 분리 어텐션의 효율성을 입증합니다.

• 추론 처리량 (Inference Throughput):

  • Hidden Decoding은 지연 시간에 민감한 작은 배치 시나리오에서 높은 처리량을 유지합니다(배치 크기 1에서 베이스라인 처리량의 83~88% 유지).
  • 긴 컨텍스트 및 높은 배치 시나리오에서는 처리량이 감소합니다(배치 크기 32, 32k 입력 시 약 27%로 하락). 이는 토큰당 연산량이 증가했음을 반영합니다.

의의 및 주장

본 논문은 Hidden-Decoding이 트랜스포머 백본을 추가로 확장하는 것이 너무 비용이 많이 들 때, 프런티어 LLM을 스케일링할 수 있는 실용적이고 엔지니어링적으로 현실적인 경로를 제공한다고 주장합니다.

• 호환성: 루프형 트랜스포머와 달리, Hidden Decoding은 추가 연량이 반복적 깊이가 아닌 더 긴 시퀀스로 나타나기 때문에 파이프라인 병렬 처리 및 표준 대규모 모델 학습 스택과 자연스럽게 호환됩니다.
• 효과성: 100B+ MoE 체급에서 시퀀스 길이 확장을 성공적으로 스케일링한 첫 번째 사례로, 추론 및 지식 벤치마크에서 측정 가능한 개선을 이끌어냈습니다.
• 확장성: 이 방법은 확장 계수 $n$이 유효한 스케일링 축임을 증명하며, 새로운 모델 아키텍처나 대규모의 백본 재학습 없이도 $n$이 커짐에 따라 성능 이득을 얻을 수 있음을 보여줍니다.

저자들은 이 연구를 강력한 파운데이션 모델이 직면한 "스케일링 벽(scaling wall)"에 대한 해결책으로 제시하며, 효율적인 잠재 연산을 통해 기존 파라미터로부터 더 많은 능력을 추출하는 방법을 제안합니다.