Jacobi 포싱으로 구현한 고속 인과 병렬 디코딩

초록

본 논문은 기존 자동회귀(AR) 모델의 인과적 사전학습 특성을 유지하면서, 자체 생성한 병렬 디코딩 궤적을 이용해 점진적으로 왜곡을 줄이는 ‘Jacobi Forcing’이라는 새로운 증류 방식을 제안한다. 이를 통해 AR 모델을 효율적인 병렬 디코더로 전환하고, 다중 블록 디코딩과 재활용(rejection‑recycling) 기법을 결합해 코딩·수학 벤치마크에서 최대 4배에 달하는 속도 향상을 달성한다.

상세 분석

이 논문은 대규모 언어 모델(LLM)의 추론 지연을 줄이기 위해 ‘다중 토큰 생성’이라는 패러다임을 재조명한다. 기존의 diffusion‑based LLM(dLLM) 접근법은 양방향 어텐션을 도입해 전체 시퀀스를 한 번에 디노이징하지만, 사전 학습 시 사용된 인과적 마스크와 충돌해 KV‑cache 재사용이 불가능하고, 마스크된 데이터 분포가 실제 데이터와 크게 달라 품질 저하가 발생한다. 저자들은 이러한 ‘pre‑train to post‑train mismatch’를 해소하기 위해 Jacobi Decoding을 기반으로 한 자체 생성 궤적을 활용한다. Jacobi Decoding은 초기 무작위 토큰 블록을 인과적 마스크만으로 반복 업데이트해 고정점에 수렴하는 과정이며, 이 고정점은 전통적인 AR greedy 디코딩 결과와 동일함을 이론적으로 보장한다.

‘Jacobi Forcing’은 두 단계로 구성된다. 첫 번째는 진행적 노이즈 스케줄이다. 큰 블록을 작은 서브‑블록으로 나누고, 각 서브‑블록에 점진적으로 노이즈 비율을 0→1 로 증가시키는 순환 스케줄을 적용한다. 이는 각 토큰이 너무 긴 노이즈 연쇄에 의존하는 것을 방지해 학습 안정성을 크게 높인다. 두 번째는 진행적 일관성 손실(progressive consistency loss)이다. 각 서브‑블록의 특정 노이즈 비율에 해당하는 궤적 지점을 선택해, 깨끗한 컨텍스트와 노이즈 컨텍스트를 동시에 입력으로 하는 KL‑다이버전스 손실을 계산한다. 이때 기존 CLLM이 필요로 하던 O(N)개의 포워드 패스를 O(1)로 줄이기 위해 시퀀스 패킹과 블록‑와이즈 희소 어텐션 마스크를 도입, 전체 배치에 걸쳐 깨끗·노이즈 토큰을 한 번에 처리한다.

학습 단계에서는 일관성 손실에 전통적인 AR 손실을 가중치 w 로 결합해 품질을 보전한다. 또한, 학습이 진행될수록 더 큰 블록 크기의 궤적을 다시 수집해 재증류함으로써 ‘속도 포화’를 깨뜨리고 추가 20% 정도의 가속을 얻는다.

추론 최적화 측면에서는 두 가지 혁신이 있다. Rejection‑Recycling은 미수렴 단계에서 생성된 고품질 n‑gram을 풀에 저장하고, 이후 반복에서 첫 토큰이 일치하면 해당 n‑gram 전체를 후보로 검증한다. 이는 한 번의 포워드 패스로 다수 토큰을 빠르게 확정시킬 수 있게 한다. Multi‑Block Decoding은 여러 블록을 동시에 유지·업데이트하면서, 앞선 블록이 아직 수렴하지 않아도 뒤쪽 블록에서 이미 확정된 토큰을 활용해 전체 진행 속도를 높인다. 두 기법을 결합하면 평균 3.8×, 최악의 경우 4.0×에 달하는 실시간 속도 향상을 달성한다.

실험에서는 코딩(HumanEval, MBPP)과 수학(MATH) 벤치마크에서 기존 AR 모델 대비 3.8배4.5배의 토큰 수용량 증가와 3.5배4.0배의 벽시계 시간 감소를 보고한다. 품질 측면에서는 Pass@1, 정확도 등에서 0.2% 이하의 미미한 손실만을 보이며, 대규모 모델(7B~13B)에서도 동일한 경향을 확인한다. 전체적으로 이 접근법은 인과적 사전학습 특성을 보존하면서도 병렬 디코딩 효율을 크게 끌어올리는 실용적인 솔루션으로 평가된다.