기술 요약: 언어 모델을 위한 계층적 잠재 예측 (Hierarchical Latent Prediction)

문제 정의

교사 강요(teacher-forced) 학습을 통한 표준 차기 토큰 예측(Next-Token Prediction, NTP)은 매우 유능한 트랜스포머 언어 모델의 발전을 가능하게 했습니다. 그러나 이 패러다임은 **노출 편향(exposure bias)**이라 불리는 훈련-추론 불일치를 초래합니다. 추론 과정에서 자기 회귀적 생성은 실제 정답(ground truth)이 아닌 모델 자신의 이전 출력값에 의존하며, 이는 오류가 누적되어 장기 생성 품질을 저하시키는 결과를 낳습니다.

이를 완테하기 위한 최근의 시도들인 다중 토큰 예측(Multi-Token Prediction, MTP) 및 차기 잠재 예측(Next-Latent Prediction, NextLat)은 미래의 토큰이나 잠재 상태를 예측하는 보조 목적 함수를 도입합니다. 하지만 이러한 접근 방식들은 다음과 같은 한계에 직면해 있습니다:

• MTP는 다단계 롤아웃(multi-step rollouts) 과정에서 발생하는 누적 오류 문제로 어려움을 겪는 경우가 많습니다.
• NextLat는 잠재 공간 내에서 자기 예측을 강제하지만, 일반적으로 단 한 단계 앞만을 예측합니다. 이는 국소적인 전이 일관성(local transition consistency)을 포착하는 데는 효과적이지만, 긴 시간적 척도에서의 구조를 포착하기 위한 간접적인 그래디언트 압력을 제공하는 데는 한계가 있으며, 순차적 역학(sequential dynamics)이 전개됨에 따라 장기 의존성이 약화됩니다.

핵리적인 과제는 아키텍처 오버헤드나 추론 지연을 발생시키지 않으면서, 어떻게 사전 학습 중에 언어 모델의 표현(representation)이 다중 스케일 예측 구조를 인코딩하도록 형성할 것인가 하는 점입니다.

방법론: 계층적 잠재 예측 (Hierarchical Latent Prediction, HiLP)

저자들은 잠재 공간 내 자기 예측 목적 함수에 시간적 계층 구조를 도입하는 학습 방법인 **계별적 잠재 예측(HiLP)**을 제안합니다. 이 아키텍처는 네 개의 공동 학습된 파라미터 맵으로 구성됩니다: 출력 헤드($p_\theta$), 하위 레벨 역학($p_\psi$), 상위 레벨 역학($p_\phi$), 그리고 결합된 NTP 헤드($p_\rho$).

아키텍처 구성 요소

1. 하위 레벨 잠재 상태 ($h_t$): 백본 트랜스포머 $G_\theta$가 토큰 접두사 $X_{1:t}$로부터 하위 레벨 잠재 상태 $h_t$를 생성합니다.
2. 상위 레벨 추상 잠재 상태 ($u_t$): 윈도우 크기 $k$를 가진 슬라이딩 윈도우 어텐션(Sliding-Window Attention, SWA) 모듈이 마지막 $k$개의 하위 레벨 상태를 집계하여 상위 레벨 추상 잠재 상태 $u_t = \text{SWA}(h_{t-k+1:t})$를 만듭니다.
3. 상위 레벨 역학: 모델 $p_\phi$는 토큰 입력 없이 추상 공간에서 $k$ 스텝 앞을 예측($\hat{u}_{t+k}$)함으로써, 단일 단계 전이를 넘어선 룩어헤드 플래닝(lookahead planning)을 장려합니다.
4. 결합된 NTP 헤드: 차기 토큰 예측은 하위 레벨 잠재 상태 $h_t$와 상위 레벨 잠재 상태 $\tilde{u}_t$ 모두에 조건화됩니다 (여기서 $\tilde{u}_t$는 결합된 헤드가 SWA 윈도우를 통해 하위 잠재 상태를 업데이트하는 것을 방지하기 위해 stop-gradient를 포함합니다).

학습 목적 함수

전체 손실(loss)은 다섯 가지 항의 가중 합으로 구성됩니다:

1. 표준 NTP 손실 ($L_{ntp}$): 차기 토큰 예측을 위한 표준 교차 엔트로피입니다.
2. 하위 레벨 전이 일관성 ($L_h$): 하위 역학 모델 $p_\psi$가 다음 하위 레벨 잠재 상태를 정확하게 예측하도록 보장하는 SmoothL1 손실입니다.
3. KL 발산 ($L_{KL}$): 예측된 하위 레벨 잠재 상태의 분포가 타겟 분포와 일치하도록 보장합니다.
4. 상위 레벨 전이 일관성 ($L_u$): 상위 레벨 역학 모델 $p_\phi$가 $k$ 스텝 앞의 추상 잠재 상태를 정확하게 예측하도록 보장하는 SmoothL1 손실입니다.
5. 결합된 NTP 손실 ($L_{cntp}$): $h_t$와 $\tilde{u}_t$를 모두 활용하는 결합 헤드를 위한 교차 엔트로피입니다.

추론 전략: HiLP의 핵심 설계 특징은 이 계층적 장치 전체(SWA 모듈, 역학 모델, 결합 헤드)가 추론 시에는 폐기된다는 점입니다. 오직 표준 NTP 헤드만이 사용되므로, 아키텍처 오버헤드나 추론 지연 증가가 발생하지 않습니다.

주요 기여

1. 계층적 표현 학습: 슬라이딩 윈도우 어텐션을 통한 잠재 상태 기반의 다중 스케일 자기 예측 학습을 트랜스포머 사전 학습에 통합하는 HiLP를 도입했습니다.
2. 평가: 코딩 및 다단계 추론 성능에 대한 다운스트림 정확도와 추측 디코딩(speculative decoding) 효율성을 기준으로 MTP 및 NextLat과 비교하여 HiLP의 벤치마킹을 수행했습니다.
3. 순수 학습 시 개입 (Pure Training-Time Intervention): 복잡한 계층 구조가 순수하게 보조적인 학습 신호로서 기능함을 입증하여, 배포 시 표준 자기 회귀 디코딩의 효율성을 유지할 수 있음을 보여주었습니다.

실험 결과

저자들은 8개의 NVIDIA B200 GPU를 사용하여 100B 토큰에 대해 1B 파라미터 모델을 학습시켰습니다.

벤치마크 성능

• 코딩 (HumanEval): HiLP는 11.33의 pass@1 점수를 달성하여, 표준 NTP (8.77), MTP (9.21), NextLat (10.58)을 능가했습니다.
• 다단계 추론: 심볼릭 및 다단계 벤치마크(LLM을 위한 DataComp)에서 HiLP는 다중 스케일 잠재 예측 덕분에 향상된 성능을 보였습니다.

추측 디코딩 (Speculative Decoding)

HiLP는 추측 디코딩 작업에서 우수한 효율성을 보여주었습니다:

• 코드 데이터: HiLP는 드래프팅 단계당 평균 3.41개의 수락된 토큰을 달성하여, MTP(2.57) 및 NextLat(3.26)보다 높은 성능을 보였습니다.
• Nemotron-ClimbMix 데이터: HiLP는 2.59개의 수락된 토큰을 달성하여, NextLat(2.57)과 MTP(2.11)를 약간 상회했습니다.
• 잠재 롤아웃 안정성: 잠재 교차 엔트로피 곡선 분석 결과, HiLP는 NextLat과 유사한 근거리 예측 손실을 유지하면서도 더 긴 호라이즌(최대 8 스텝)에서 오류를 유의미하게 감소시켜, 더 안정적인 다단계 잠재 롤아웃을 보여주었습니다.

효율성

HiLP는 보조 파라미터를 추가하지만(백본 1.06B 대비 총 1.27B), 추론 경로는 기본 모델과 동일합니다. 학습 처리량(throughput)은 HiLP가 약 **65%**로, NTP 베이스라인보다는 낮지만 MTP(49%)보다는 높았습니다. 이는 추가적인 계층적 손실 그래프로 인한 비용을 반영합니다.

의의 및 주장

본 논문은 HiLP가 차기 토큰 예측의 "근시안적 격차(myopic gap)"를 잠재 공간 내 자기 예측 목적 함수에 직접적인 시간적 계층 구조를 도입함으로써 성공적으로 해결했다고 주장합니다. 추론 시점에 복잡성을 소비하는 기존의 계층적 모델(예: Large Concept Models 또는 Byte Latent Transformers)과 달리, HiLP는 사전 학습 중에 다중 스케일 예측 구조를 인코딩하도록 표현을 형성한 뒤, 추론 시에는 이를 완전히 제거합니다.

저자들은 이 접근 방식이 표준 토큰 레벨 자기 회귀 디코딩의 속도나 단순성을 희생하지 않으면서도, 모델이 더 풍부한 잠재 전이 역학을 학습하고 장기 의존성을 더 잘 포착할 수 있게 한다고 상정합니다. 이 연구는 NTP를 넘어서는 사전 학습 목적 함수, 잠재 공간 언어 모델링, 그리고 계층적 모델의 접점에 위치하며, 언어 모델의 장기 추론 및 계획 능력을 향상시키기 위한 참신한 방법을 제시합니다.