균일 확산 언어모델의 루프가 학습 연산을 최대 3.3배 줄인다

무엇인가

이 논문은 균일 확산 언어모델(UDLM)의 학습과 추론에서 '루프(looping)'가 얼마나 유용한지, 그리고 학습 때 루프 횟수를 어떤 분포로 뽑아야 하는지를 다룬다. UDLM은 마스크 대신 어휘의 임의 토큰으로 시퀀스를 오염시키고, 어느 위치가 오염됐는지 표시하지 않는다. 그래서 디노이저는 이미 오염됐을 수 있는 문맥을 보고 깨끗한 토큰을 예측해야 하며, 이 목표는 마스크 확산 언어모델(MDLM)보다 학습 FLOPs를 더 많이 쓰지만 하류 추론 성능은 더 강하다고 저자들은 설명한다. 기존 연구에서 MDLM에는 중간 블록을 반복하는 선택적 루프가 학습 효율을 높였지만, 오염 위치를 알 수 없는 UDLM에서도 같은 효과가 있는지는 열린 질문이었다.

어떻게 동작하나

방법은 Duo++ 학습 목표와 Ψ-샘플러를 그대로 두고 디노이저 계산만 바꾼다. 170M Duo++ small DiT를 헤드 블록, 중간 블록 슬라이스, 테일 블록으로 나눈 뒤 입력을 한 번 임베딩하고 헤드에 통과시킨다. 이어 중간 슬라이스를 공유 가중치로 S번 반복하고 테일과 softmax로 깨끗한 토큰 예측을 낸다. S=1이 원래 네트워크이고, 루프는 파라미터를 늘리지 않고 깊이만 늘린다. 루프 동안 noisy 입력은 고정되고 diffusion step 사이에서만 바뀐다. 학습 때는 매 옵티마이저 스텝마다 1부터 Smax까지의 루프 횟수 S를 분포 p(S)에서 뽑아 S번 언롤하고, Duo++ 목적함수를 모든 반복에 역전파한다. 그래서 하나의 체크포인트를 추론 때 여러 S로 실행할 수 있다.

무엇과 다른가

루프 횟수 분포는 세 가지다. 기본은 Smax=12의 균등 샘플링이고, FLOP-비례 샘플링은 p_fp(S)=F(S)^-1/ΣF(r)^-1로 각 루프 횟수에 같은 기대 학습 FLOPs를 배분해 작은 S를 더 자주 뽑는다. 노이즈-비례 샘플링은 더 noisy한 입력에 더 큰 S를 배정하지만 주변 분포는 균등하게 유지한다. 추론에서는 모든 디노이징 스텝에 같은 S를 쓰는 고정 루프와, 연속 반복 사이 은닉 상태의 상대 변화가 0.1 미만이면 멈추는 적응형 루프를 비교한다. 적응형 루프는 최대 12회를 넘지 않는다.

어떻게 쓰나

논문은 루프가 왜 도움이 되는지도 UDLM의 오염 커널로 설명한다. UDLM은 q_t(z|x)=α_t 1{z=x}+(1-α_t)ν(z)로 토큰을 유지하거나 균등 분포 ν로 교체하는데, 관측 토큰만 보고는 유지됐는지 알 수 없다. 관측 토큰이 e_i일 때 유지 확률 γ는 α_t π(i)/(α_t π(i)+(1-α_t)ν(e_i))이고, UDLM에서는 γ/(1-γ)=Kα_t/(1-α_t)π(i)가 된다. 깨끗한 토큰의 사후분포는 관측 토큰에 놓인 점질량 γ와 문맥 예측 π를 (1-γ)로 섞은 형태다. Duo++ 디노이저는 이 문맥 예측 π를 추정하므로, 루프로 x_hat을 정교화하면 유지 여부 판단과 교체된 토큰의 복원이 함께 좋아질 수 있다. 실제 분석에서 S=1에서 이미 유지/교체 구분 AUROC는 0.96 이상이었고, S=1에서 12로 늘릴 때 퍼플렉시티 감소의 93~94%는 '어떤 토큰을 유지할지'가 아니라 '교체된 위치의 문맥 예측'에서 나왔다. 이득은 t=0.5~0.7의 중간 노이즈에서 가장 컸고, 아주 noisy한 입력에서는 작았다.

전제와 한계

실험은 170M Duo++ small DiT(12블록, 폭 768, 시퀀스 길이 1,024, GPT-2 토크나이저)로 진행했다. 블록 1~2를 Smax=12로 루프했고, OpenWebText와 FineWeb-Edu sample-10BT에서 비루프 베이스라인(1M 옵티마이저 스텝, 배치 512)과 동일 학습 FLOPs로 비교했다. 루프 모델은 비루프 UDLM이 약 0.5조 토큰으로 학습했을 때의 검증 NLL에 OpenWebText에서 최대 2.8배, FineWeb-Edu에서 3.3배 적은 학습 FLOPs로 도달했다. 본문 요약은 두 데이터셋에서 2.7~3.3배 절감으로 제시한다. FLOP-비례 샘플링은 최종 NLL에서 균등 샘플링과 거의 같았고(OpenWebText 3.159 대 3.154, FineWeb-Edu 3.122 대 3.119), S=1에서는 더 좋았다(3.289 대 3.306). 노이즈-비례 샘플링은 이득이 없었고, 최고 결과가 FLOP-비례보다 NLL +0.029 뒤졌으며, 다른 두 분포가 약 30% 적은 연산으로 그 최고 NLL에 도달했다. 루프 위치는 블록 1~5 어디서 시작해도 비슷했고, 늦은 슬라이스는 나빴으며 Smax를 12보다 늘려도 추가 이득이 없었다.

추론에서는 같은 NFE(디노이저 호출 수)에서 루프 UDLM이 비루프 UDLM보다 생성 퍼플렉시티를 최대 9.8% 낮췄고, 대부분의 이득은 S=6~8의 중간 루프 횟수에서 나왔다. 다만 NFE를 맞춘 비교이지 FLOPs를 맞춘 비교는 아니다. S=6은 S=1 대비 추론 FLOPs가 1.48배, S=12는 2.05배다. 적응형 루프는 평균 약 5.5회를 돌아 고정 S=6보다 조금 저렴하면서 대부분의 샘플러-NFE 조합에서 같거나 더 좋았다. 7개 out-of-domain 데이터셋에서는 추론 시 루프가 비루프 대비 zero-shot 퍼플렉시티를 평균 11.5% 낮췄다. 중요한 단서는 루프로 학습한 체크포인트를 S=1로 실행하면 모든 데이터셋에서 비루프 베이스라인보다 나빴고, 추론 때 루프를 돌려야 이득이 나타난다는 점이다. Zero-shot likelihood는 깊이가 늘수록 계속 좋아져 균등 체크포인트에서는 S=12가 모든 데이터셋에서 최고였고, 적응형 루프는 S=12 이득의 88~98%를 유지하면서 루프 반복의 10~28%를 건너뛰었다.

10개 zero-shot 다중선택 과제에서는 적응형 루프가 고정 S=12와 같은 정확도를 유지하면서 균등 샘플링 체크포인트에서 26.0%, FLOP-비례 체크포인트에서 34.8% 적은 루프를 돌았다. FLOP-비례 체크포인트는 모든 동일 루프 정책에서 균등 체크포인트를 이겼고, 적응형 루프로 고정 S=12와 비슷한 성능을 약 3분의 1 적은 반복으로 냈다. 매크로 정확도는 35.38%였고, 균등 체크포인트보다 12% 적은 루프(7.82 대 8.89)를 쓰면서 더 정확했다. 이득은 HellaSwag, OpenBookQA, PIQA, BoolQ에서 컸고 MathQA, WinoGrande에서는 거의 없었다. 대신 FLOP-비례 체크포인트는 동일 루프 횟수에서 생성 퍼플렉시티가 균등 체크포인트보다 평균 1.55% 높았다. 즉 깊이 배분은 과제에 따라 맞바꿈이 생긴다.

개발자 관점에서 이 논문은 루프를 파라미터 추가 없는 '계산 깊이' 노브로 제시한다. 학습 FLOPs를 줄이고, 추론에서는 NFE뿐 아니라 실제 루프 FLOPs와 적응형 중단 정책을 함께 봐야 한다. 다중선택·분류처럼 정확도가 중요한 과제에는 FLOP-비례 샘플링과 적응형 루프가 유리할 수 있지만, 생성 품질이 중요한 경우에는 균등 샘플링이 더 나을 수 있다. 논문이 명시한 범위와 전제는 170M Duo++ small DiT, OpenWebText와 FineWeb-Edu, 블록 1~2 루프, Smax=12이며, 추론 결과는 OpenWebText 체크포인트로 평가했다. 학습 목표와 Ψ-샘플러를 바꾸지 않았으므로 결론은 Duo++ 계열에 묶여 있고, 저자들은 마스크 확산과 균일 확산의 루프 효과를 reasoning 과제까지 비교하는 후속 연구를 계획한다고 밝힌다. 루프 위치 선택 실험은 50k 스텝, 단일 시드였다는 점도 전제로 남는다.