D2F 논문 리뷰: 확산 LLM이 자기회귀 모델보다 빠른 추론을 처음으로 달성한 방식

핵심 메시지

• 자기회귀 LLM은 토큰을 순차 생성해 느리지만 KV 캐시로 연산을 줄일 수 있고, 확산 LLM은 병렬 생성이 가능하지만 양방향 어텐션 때문에 KV 캐시를 쓸 수 없어 실제 속도에서 밀렸다.
• D2F는 시퀀스를 일정 크기의 블록으로 나눠 블록 사이는 순차적으로, 블록 안에서는 병렬로 처리한다. 이로써 KV 캐시와 병렬 디코딩을 동시에 활용한다.
• 앞 블록이 완전히 복원되기 전, 일부만 복원된 상태에서 뒤 블록 예측을 시작하도록 학습시켜 한 단계에서 더 많은 토큰을 확정한다.
• 학습은 비대칭 증류로 이뤄진다. 전체 문맥을 보는 기존 확산 LLM을 교사로, 제한된 앞 블록만 보는 D2F 모델을 학생으로 두고 두 예측 분포의 차이를 줄여 나간다.
• 추론에서는 활성 블록 윈도우를 유지하며 진행률이 임계값을 넘으면 새 블록을 추가하고, 블록의 활성 단계에 따라 보수적·공격적 토큰 확정을 나눠 적용한다.

쉽게 이해하기

우리가 흔히 쓰는 GPT나 제미나이 같은 언어 모델은 자기회귀 방식으로, 토큰을 하나씩 순서대로 만들기 때문에 느리다는 약점이 있다. 대안으로 주목받는 확산 기반 언어 모델은 마스킹된 시퀀스를 점진적으로 복원하며 여러 위치의 토큰을 한 번에 만들 수 있어 원리상 훨씬 빠를 여지가 있는데, 실제로 재 보면 오히려 기존 모델보다 느렸다. 원인은 KV 캐시다. 자기회귀 모델은 이전 단계에서 계산해 둔 키와 밸류 행렬을 저장해 두고 새로 들어온 토큰의 쿼리만 계산하면 되지만, 확산 모델은 양방향 어텐션을 쓰기 때문에 앞선 계산 결과를 그대로 재활용할 수 없다. 병렬 생성이라는 이점을 얻는 대가로 연산 재사용이라는 이점을 잃은 셈이며, 이 논문이 제시하는 D2F는 두 가지를 동시에 갖는 것을 목표로 한다.

핵심은 시퀀스를 일정 크기의 블록으로 나누는 것이다. 블록과 블록은 순차적으로 처리하므로 앞 블록의 KV 캐시를 활용할 수 있고, 블록 내부에서는 확산 모델답게 여러 토큰을 병렬로 만든다. 여기에 장치가 하나 더 붙는다. 앞 블록이 완전히 복원될 때까지 기다리지 않고 일부만 복원된 상태를 조건으로 뒤 블록 예측을 시작하도록 학습시켜, 한 단계에서 확정되는 토큰 수를 늘린다.

학습 방법은 비대칭 증류다. 교사는 전체 노이즈 블록을 모두 보고 예측하는 기존 확산 LLM이고 학생은 자기 블록과 그 앞 블록까지만 보는 D2F 모델인데, 교사의 가중치로 학생을 초기화한 뒤 두 예측 분포 사이의 KL 발산을 줄이는 방향으로 학습한다. 보는 문맥의 범위가 서로 달라 비대칭이라는 이름이 붙었고, 학생은 앞 블록이 완전히 복원되지 않아도 뒤 블록을 잘 맞히는 능력을 이 과정에서 얻는다. 학습할 때는 뒤 블록으로 갈수록 마스크가 많아지는 노이즈 스케줄을 적용한다.

추론 단계에는 파이프라인 병렬 디코딩이 쓰인다. 활성 블록 윈도우를 유지하고 마지막 블록의 진행률이 임계값을 넘으면 새 블록을 뒤에 추가하는데, 윈도우 크기가 일정하게 유지되므로 매 단계마다 전체 시퀀스를 다시 다룰 필요가 없어 연산량이 줄어든다. 새로 들어온 블록은 처음에는 보수적으로 동작해 확신도가 임계값을 넘는 토큰만 확정하고, 앞 블록이 충분히 진행되면 공격적인 단계로 전환해 확신도가 아주 높지 않아도 상대적으로 가장 높은 토큰을 확정한다. 확정되지 않은 토큰은 다시 마스킹해 다음 단계에서 예측한다.

실험은 확산 LLM 쪽에서 LLaDA 계열 8B 지시 모델과 Dream 계열 7B 베이스 모델을 D2F로 증류해 진행했고, 자기회귀 기준선으로는 LLaMA3 8B 지시 모델과 Qwen2.5 7B 베이스 모델을 뒀다. 캐싱이나 샘플링 최적화로 속도를 올리려 한 기존 확산 LLM 연구들도 비교에 포함됐지만, 이들은 자기회귀 모델의 속도를 넘어서지 못했다. 초당 처리 토큰 수는 전반적으로 개선됐고 점수는 기존 방식과 비슷한 수준을 유지했다. 기존 확산 LLM 대비 최대 50배 개선을 주장하며, 자기회귀 기준선과 비교해도 GSM8K에서 2.5배, HumanEval 제로샷에서 1.6배 빨랐고 Dream 기반 모델에서는 비슷한 점수를 유지하며 3.1배 높은 속도를 기록했다.

주요 인사이트

• 확산 LLM의 병렬 생성은 이론적 이점이었을 뿐, 연산 재사용을 포기한 대가를 치르고 있었다. 실제 속도는 아키텍처 특성만으로 결정되지 않는다.
• 블록 단위 분할은 순차성과 병렬성 중 하나를 고르는 문제를 두 층위로 나눠 둘 다 취하는 방식으로 바꿔 놓았다. 문제를 다른 단위로 자르는 것이 해법이 된 사례다.
• 앞 단계가 완전히 끝나기를 기다리지 않고 부분 결과를 조건으로 다음 단계를 시작하는 설계는 생성 파이프라인 전반에 적용 가능한 아이디어다.
• 교사와 학생이 보는 문맥 범위를 일부러 다르게 두는 비대칭 증류는, 학생에게 '덜 보고도 맞히는' 능력을 목표로 학습시키는 방법이다.
• 속도와 품질은 임계값 설정으로 조절되는 연속적인 선택지다. 하나의 최적값이 아니라 용도에 맞는 지점을 고르는 문제로 봐야 한다.

자주 묻는 질문

관련 AI 소식