기술 요약: K-Forcing

문제 정의

자기회귀(Autoregressive, AR) 언어 모델링은 텍-생성의 지배적인 패러다임이지만, 근본적인 효율성 병목 현상을 겪고 있습니다: 바로 엄격하게 순차적인 토큰 단위 디코딩입니다. 이 과정은 메모리 바운드(memory-bound)이며, $L$개의 토큰 시퀀스를 생성하기 위해 $L$번의 순방향 패스(forward pass)가 필요합니다. 이는 하드웨어 가속기의 활용도를 낮추고, 특히 높은 부하의 배치 서빙(batch serving) 시나리오에서 높은 지연 시간을 초고하게 됩니다.

기존의 가속 방법들은 산업 규모의 배포에서 이 병목 현상을 효과적으로 해결하는 데 어려움을 겪습니다:

• 추측적 디코딩 (Speculative Decoding, Draft-then-Verify): 단일 요청의 지연 시간을 줄이는 데는 효과적이지만, 이 방법들은 가변 길이의 수락된 출력을 생성합니다. 이는 정규적인 배칭(regular batching)을 방라도하며, 동기화 문제(예: ragged tensors)와 배치 크기가 커짐에 따라 증가하는 오버헤드를 유발하여, 높은 부하 상황에서 처리량(throughput) 이득을 상쇄하곤 합니다.
• 확산 언어 모델 (Diffusion Language Models, DLMs): 이 모델들은 반복적인 노이즈 제거를 통해 여러 토큰을 디코드합니다. 그러나 표준 정식화(예: Masked Diffusion Language Models)는 위치별 주변 예측기(per-position marginal predictors)를 학습합니다. 이론적 분석에 따르면, 결합 조건부 분포(joint conditional distribution)가 아닌 독립적인 주변 분포(marginals)로부터 샘플링하는 것은 분포적 오류를 유발합니다. 품질을 유지하기 위해 이들은 종종 많은 정제 단계(refinement steps)를 요구하며, AR에 비해 순방향 평가 횟수(NFE)를 유의미하게 줄이지 못합니다.

방법론: K-Forcing

저자들은 **결합 차기-$k$-토큰 디코딩(joint next-$k$-token decoding)**을 위해 설계된 푸시포워드(push-forward) 언어 모델링 패러다임인 K-Forcing을 제안합니다. 핵심 목표는 독립적인 균등 노이즈 변수를 단 한 번의 순방향 패스로 $k$개의 미래 토큰의 결합 샘플로 변환하는 조건부 매핑을 학습함으로써, 고정 길이 출력과 정규적인 배칭을 보존하는 것입니다.

1. 푸시포워드 정식화 (Push-Forward Formulation)

K-Forcing은 결합 분포 $p(x_{t+1}, \dots, x_{t+k} | x_{\le t})$를 명시적으로 모델링하는 대신(대규모 어휘집에서는 계산적으로 불가능함), 결정론적 맵 $G_\theta: \mathcal{V}^t \times [0, 1]^k \to \mathcal{V}^k$를 학습합니다.

• 입력: 컨텍스트 $x_{\le t}$ 및 $k$개의 i.i.d. 노이즈 변수 $z \sim \text{Uniform}(0, 1)^k$.
• 출력: $k$개의 미래 토큰의 결합 샘플 $(\hat{x}_{t+1}, \dots, \hat{x}_{t+k})$.
• 이론적 기초: 논문은 AR 오라클의 역 누적 분포 함수(inverse CDF)를 언롤링(unrolling)함으로써 폐쇄형 푸시포워드 매핑이 존재함을 증명합니다. 이 매핑은 결합 확률 테이블을 명시적으로 매개변수화하지 않고도 이론적으로 AR 교사(teacher)의 정확한 결합 조건부 분포를 재현할 수 있습니다.

2. 점진적 셀프 포싱 증류 (Progressive Self-Forcing Distillation)

이 매핑을 사전 학습된 AR 교사로부터 학습시키기 위해, 저자들은 기존 노이즈 역전(noise-inversion) 베이스라인의 두 가지 주요 한계(학습-추론 불일치 및 수치적 취약성)를 해결하는 점진적 셀프 포싱 증류를 도입합니다.

• 1단계 (Forward Distillation): $k=1$인 PFLM을 부트스트랩합니다. AR 교사는 노이즈 $z$를 샘플링하고 역-CDF 샘플러를 적용하여 타겟을 생성합니다. 학생(student)은 $(x_{\le t}, z)$를 토큰으로 매핑하는 법을 배웁니다. 이는 노이즈 분포가 추론 시와 동일(Uniform)하도록 보장하며, 좁은 CDF 빈(bin)을 역전하는 문제를 피하게 합니다.
• 2단계 (Self-Forcing Distillation): 더 큰 $k$(예: $k=2 \to 4$)로 확장하기 위해, 학습된 PFLM($k$)이 PFLM($2k$) 학생의 교사 역할을 수행합니다. 교사는 $z_{1:k}$를 사용하여 첫 $k$개 토큰을 생성한 다음, 생성된 토큰들에 조건화되어 $z_{k+1:2k}$를 사용하여 다음 $k$개 토큰을 생성합니다. 학생은 전체 노이즈 벡터가 주어졌을 때 전체 $2k$ 시퀀스를 한 번에 예측하도록 학습됩니다. 이 과정은 예측 윈도우를 재귀적으로 두 배로 늘립니다 ($1 \to 2 \to 4$).

3. 아키텍처 설계

저자들은 AR 백본을 재사용하는 **완전 인과적 아키텍처(Fully Causal Architecture)**를 제안합니다.

• 각 노이즈 변수는 시퀀스의 별도 토큰으로 취급됩니다.
• 공유된 예측 헤드가 미래 토큰을 디코딩합니다.
• 인과적 어텐션(Causal attention)은 $j$번째 미래 토큰이 컨텍스트와 접두사 노이즈 변수 $(z_1, \dots, z_j)$에만 의존하도록 보장하며, 이는 이론적인 역-CDF 구조를 반영합니다.
• 이 설계는 모든 확률적 결정이 단일 공유 잠재 토큰을 통해 라우팅되는 정보 병목 현상(표준 다중 토큰 예측 아키텍처의 경우)을 방지합니다.

주요 기여

1. 기존 한계 분석: 논문은 드래프트-후-검증(draft-then-verify) 방식이 가변 길이 출력 때문에 배치 서빙에서 실패하며, 확산 모델은 결합 분포가 아닌 주변 분포로부터 샘플링하기 때문에 NFE를 효과적으로 줄이지 못한다는 점을 식별했습니다.
2. 푸시포워드 정식화: K-Forcing을 고정 길이 출력과 결합 샘플링의 요구 사항을 충족하는, 균등 노이즈를 토큰 블록으로 매핑하는 암시적 생성 모델로 정식화했습니다.
3. 학습 전략: 노이즈 역전과 관련된 학습-추론 불일치 및 수치적 취약성을 제거하는 점진적 셀프 포싱 증류를 도입했습니다.
4. 경험적 검증: 표준 인과적 트랜스포머 백본을 사용하여 LM1B 및 OpenWebText에서 통계적 모델링 변경만으로 AR 디코딩에 비해 상당한 배치 서빙 처리량 이득을 얻을 수 있음을 입증했습니다.

실험 결과

LM1B 및 OpenWebText에서 표준 인과적 트랜스포머 백본을 사용하여 평가되었습니다:

• 처리량 (Throughput): $k=4$ 토큰을 한 번의 순방향 패스당 생성하도록 구성했을 때, K-Forcing은 다양한 배치 크기(4, 16, 128)에 대해 AR 대비 2.4~3.5배의 속도 향상을 달-성했습니다.
• 품질 (Quality): AR 교사에 비해 생성 품질의 약간의 저하가 발생합니다.
  • LM1B의 경우, $k=4$는 쌍체 비교 LLM-as-a-Judge 평가에서 약 42.9%의 승률을 기록했습니다.
  • OpenWebText의 경우, $k=4$는 약 39.4%의 승률을 기록했습니다.
  • 더 낮은 $k$ 값(예: $k=2$)은 더 매끄러운 트레이드오프를 제공합니다. $k=2$는 AR과 거의 구별할 수 없는 수준(승률 ~50%)이면서 약 1.6배의 속도 향상을 제공합니다.
• SOTA와의 비교: K-Forcing은 (여러 토큰을 언마스킹할 때 급격히 저하되는) MDLM 및 (Medusa, PTP와 같은) 추측적 디코딩 베이스라인과 비교하여 품질-NFE 트레이드오프 측면에서 우위를 점합니다. K-Forcing은 동일한 NFE에서 MDLM보다 낮은 Gen-PPL과 높은 승률을 달성하며, 유사한 품질을 달성하기 위해 추측적 방법보다 더 적은 순방향 패스를 사용합니다.

의의 및 주장

본 논문은 실제 고부하 배포 환경에서 AR 생성을 가속화하기 위한 유망한 경로를 제시한다고 주장합니다.

• 기존 방법들이 불규칙한 배칭으로 인해 실패하는 고부하 배치 서빙의 특정 병목 현상을 해결합니다.
• 푸시포워드 매핑을 통한 결합 다중 토큰 샘플링이 순차적 디코딩의 실행 가능한 대안임을 입형하며, 효율적인 KV-캐시 관리와 연속 배칭에 필요한 고정 길이 구조를 보존함을 보여줍니다.
• 저자들은 이 작업을 푸시포워드 언어 모델이라는 더 넓은 연구 방향을 향한 초기 단계로 규정합니다. 특히 더 큰 $k$ 값에서 AR 교사와의 비사소한 품질 격차가 존재함을 인정하며, 학습 비용을 줄이기 위한 커스텀 블록 희소 어텐션 커널 개발 및 사전 학습된 AR 교사에 의존하지 않는 학습 패러다임 탐색을 향후 연구 방향으로 제시합니다.