ALoDLM이 토큰 난이도별로 확산 언어모델의 연산량을 배분한다

무엇인가

확산 언어모델(DLM)은 마스킹된 위치를 병렬로 복원해 빠르게 생성하지만, 같은 크기의 자기회귀(AR) 모델과 비교하면 품질 격차가 남아 실무 도입이 막혀 왔다. 이 논문은 그 원인을 '연산-난이도 불일치(computation-difficulty mismatch)'로 진단한다. 부분적으로 관측된 시퀀스 안에서 어떤 미지 토큰은 쉽게 맞힐 수 있지만, 다른 토큰은 훨씬 많은 연산이 필요하다. 그런데 기존 DLM은 각 디노이징 스텝마다 모든 미지 위치에 동일한 깊이의 디노이저를 적용한다. 쉬운 예측에 연산을 낭비하고 어려운 예측에는 연산을 굶기는 셈이다. 신뢰도 기반 디코딩도 저신뢰 토큰을 뒤 스텝으로 미룰 뿐, 미해결 토큰의 잠재 연산을 버리고 다음 스텝에서 마스크 임베딩부터 다시 계산하기 때문에 결국 같은 고정 연산을 준다.

어떻게 동작하나

ALoDLM의 핵심은 토큰별 적응형 잠재 반복(token-adaptive latent recurrence)이다. 모델은 프렐류드(토큰 임베딩과 선택적 앞부분 트랜스포머 블록), 리커런트 코어(중간 블록), 코다(나머지 뒷부분 블록)로 나뉜다. 각 디노이징 스텝(바깥 루프) 안에서 최대 K번의 안쪽 반복 패스가 돌고, 매 패스마다 h~(s)=RecurrentCore(h(s-1)), r(s)=Coda(h~(s))로 상태를 갱신한다. 읽기 상태 r(s)는 두 개의 헤드로 동시에 들어간다. LMHead는 어휘 분포를, ExitGate는 스칼라 로짓을 내놓아 시그모이드로 정지 확률 λ_i(s)를 만든다. 이때 ExitGate로 가는 그래디언트는 백본으로 흐르지 않도록 stop-gradient가 걸리고, 마지막 패스에서는 λ_i(K)=1로 강제해 정지 확률의 합이 1이 되게 한다.

무엇과 다른가

토큰 확정 절차는 다음과 같다. 패스 s에서 예측 엔트로피가 임계값 τ 이하인 미해결 위치들이 확정 후보 집합 B(s)가 되고, 각 위치에서 ŷ_i를 샘플링해 미해결 집합 U에서 제거한다. 확정된 토큰은 다음 반복 입력에서 h_i(s)=Emb(ŷ_i)로 이산 문맥이 되고, 미확정 위치는 자신의 누적 잠재 상태 h~_i(s)를 그대로 이어받아 정제를 계속한다. 안쪽 루프 종료는 누적 정지 확률 a_i(s)=a_i(s-1)+(1-a_i(s-1))λ_i(s)로 결정되며, U가 비거나 U 전체의 평균 누적 정지 확률이 임계값 q에 도달하면 멈춘다. 만약 한 번도 확정된 토큰이 없으면 엔트로피가 가장 낮은 위치를 강제로 확정해 진행을 보장한다.

어떻게 쓰나

학습은 토큰별 연산 배분을 잠재 변수로 다루는 방식이다. 각 토큰이 몇 번째 반복 패스에서 확정되는지를 나타내는 종료 스케줄 z를 정의하고, 확정된 토큰이 이후 문맥을 바꾸기 때문에 스케줄 전체에 대한 주변화가 지수적으로 커진다는 문제를 지적한다. 대신 ExitGate로 파라미터화한 변분 분포 q_φ(z|x_t, y)를 도입하고, 잘린 기하분포를 사전분포 π(z)로 두어 조건부 NELBO를 유도한다. 목적함수는 궤적 손실과 KL 항의 합이며, 이산 스케줄 때문에 통상적 역전파가 불가능하므로 스코어 함수 추정량(REINFORCE) 기반의 비편향 그래디언트 추정기를 쓴다. 분산을 줄이려고 샘플링된 종료 깊이 이전 패스들의 예측을 재사용해 가중 감독을 주고, 1차 패스 예측 비용을 대조 변량으로 뺀다. 이 중간 감독을 제거하면 디노이저 정규화 파라미터 4,096개 기준 조건부 그래디언트 분산이 학습 스텝 1,000·6,500·17,000에서 각각 1.76배, 1.49배, 1.42배로 커진다고 보고한다.

전제와 한계

실험은 Qwen3-1.7B와 Qwen3-8B를 WeDLM의 스트리밍 블록 확산 프레임워크로 변환해 수행했고, 최대 반복 깊이는 K=4다. 리커런트 코어는 1.7B에서 28개 레이어 전부, 8B에서 중간 16개 레이어다. 기존 변환 파이프라인과 달리 지속 사전학습(CPT)을 건너뛰고 5B 토큰 코퍼스에 AdamW(학습률 1e-5, 가중치 감쇠 0.01)로 바로 SFT했다. 평가는 ARC-Easy/Challenge, MMLU, MMLU-Pro, GPQA-Diamond, GSM8K, MATH-500, HumanEval, MBPP-sanitized와 EvalPlus 변형까지 11개 벤치마크에서 OpenCompass 프로토콜, 그리디 디코딩, 4,096 토큰 제한으로 이뤄졌다. 비교 대상은 Qwen3, Dream, LLaDA, Fast-dLLM-v2, SDAR, WeDLM이며 각자 최적화된 추론 엔진(dInfer, JetEngine, vLLM)을 쓴다. ALoDLM은 깊이 인식 KV 캐싱을 적용한 vLLM으로 서빙된다.

결과는 1.7B와 8B에서 평균 65.5와 80.3으로, 평가한 모든 DLM과 대응 Qwen3 AR 기준선을 앞섰다. ALoDLM-8B는 WeDLM-8B를 평균 5.2점 차로 이기며 11개 과제 중 10개에서 앞선다. Qwen3 AR 대비로는 1.7B가 63.8 대 65.5, 8B가 78.5 대 80.3이다. 코드 생성에서 특히 강해 4개 코딩 벤치마크를 모두 석권했고, MBPP+에서만 AR 모델에 근소하게 뒤진다. 효율에서는 GSM8K 기준 ALoDLM-8B가 vLLM으로 서빙한 Qwen3-8B의 약 2.7배 처리량을 비슷하거나 더 높은 정확도로 낸다. q=0.5에서 τ를 0.1에서 0.6으로 올리면 처리량이 278.7에서 508.3 토큰/초로 늘고 정확도는 93.8%에서 92.3%로 떨어진다(τ=0.9면 89.9%). τ=0.2에서 q를 0.1에서 0.9로 올리면 정확도가 93.3%에서 93.8%로 오르고 처리량은 455.3에서 309.0 토큰/초로 준다. B200 한 장에서 WeDLM-8B와 비교한 정확도-처리량 프런티어는 약 650 토큰/초 부근에서 교차하고, 정확도 93.25%에서 ALoDLM이 612.4 토큰/초로 WeDLM의 564.2보다 8.5% 빠르며, 토큰당 GFLOPs도 133.5 대 154.6으로 13.6% 적다. q를 올려 토큰당 평균 루프 수를 1.6에서 2.34로 늘리면 11개 벤치마크 평균이 77.9%에서 79.1%로 오르는 테스트타임 스케일링도 확인된다. 별도 난이도 라벨 없이도 숫자 토큰의 평균 1차 정지 확률이 0.369로 전체 기준 평균 0.417보다 약 11.5% 낮아, 수치 예측에 더 많은 잠재 정제를 배분하는 경향이 학습된다.

개발자 관점에서 이 논문은 확산 언어모델을 서빙 파이프라인에 넣을지 판단하는 근거가 된다. τ와 q 두 개의 디코딩 노브로 정확도와 처리량을 연속적으로 조절할 수 있고, vLLM과 깊이 인식 KV 캐싱을 전제로 한 수치가 제시된다는 점이 실무적으로 중요하다. 다만 처리량 이득은 프롬프트와 도메인에 따라 크게 달라지므로, 실제 워크로드에서 τ·q 스윕을 직접 돌려 프런티어를 확인해야 한다. 또한 리커런트 코어를 중간 레이어에 두는 배치가 마지막 레이어 반복보다 학습 후반에 더 좋았다는 점, K=4가 K=8과 비슷한 성능을 절반 깊이로 낸다는 점은 자체 모델을 만들 때 참고할 만하다.

저자들이 밝힌 한계는 두 가지다. 첫째, 프리필 단계에서 전체 반복 깊이에 대한 KV 캐시를 구성하기 때문에 AR 모델보다 첫 토큰까지 걸리는 시간(TTFT)이 길어질 수 있고, 이는 짧은 응답이나 지연에 민감한 애플리케이션에서 병렬 디코딩의 이점을 상쇄한다. 둘째, 생성 속도가 입력에 의존한다. 반복 연산량과 동시 확정 토큰 수가 토큰 신뢰도와 적응적 정지 결정에 달려 있어, 그리디 디코딩에서도 프롬프트·데이터셋·도메인에 따라 처리량이 크게 변한다. 학습 데이터에 적게 등장한 도메인에서는 추가 정제가 필요해 병렬 확정이 줄고, 최적화된 AR 기준선 대비 속도 우위가 줄거나 아예 역전될 수 있다고 명시한다.