루프형 확산 언어 모델 (Looped Diffusion Language Models)

마스크 확산 모델 (Masked Diffusion Models, MDMs)은 언어 모델링을 위한 자기회귀 모델 (Autoregressive Models)의 유망한 대안으로 부상했으나, MDM을 위한 트랜스포머 (Transformer) 아키텍처의 효과적인 설계는 여전히 미개척 분야로 남아 있습니다. 본 논문에서 우리는 트랜스포머의 초기-중기 레이어 (Early-middle layers)를 선택적으로 루프 (Looping) 시키는 것이 MDM의 훈련 효율성 (Training efficiency)과 모델 성능 (Model performance)을 모두 크게 향상시킨다는 것을 보여줍니다. 우리는 이 접근 방식을 LoopMDM (Looped Masked Diffusion Model)이라 명명하며, 이는 두 가지 핵심적인 이점을 제공합니다. 즉, 훈련 시 레이어를 루프 시키면 파라미터 (Parameters)를 추가하지 않고도 깊이 확장 (Depth-scaling) 효과를 얻을 수 있으며, 추론 시 (Inference-time) 루프 횟수를 조절함으로써 유연한 연산 확장 (Compute scaling)이 가능해집니다. 이러한 단순함에도 불구하고 결과는 놀랍습니다. 여러 사전 훈련 코퍼스 (Pre-training corpora)에 걸쳐 LoopMDM은 동일한 크기의 MDM과 성능은 대등하면서도 훈련 FLOPs (Floating Point Operations)는 최대 3.3배 적게 사용하며, 최종 성능은 GSM8K에서 최대 8.5포인트 차이를 보이는 등 다양한 추론 벤치마크 (Reasoning benchmarks)에서 이를 능가합니다. 심지어 LoopMDM은 단계별 연산량 (Per-step compute)이 유사하게 훈련된 더 깊은 비루프형 (Non-looped) MDM보다도 뛰어난 성능을 보이며, 이는 선택적 루핑이 단순한 깊이 확장보다 더 효과적임을 나타냅니다. 또한, LoopMDM은 루프 횟수를 늘림으로써 추론 시 연산량을 확장할 수 있습니다. 샘플링 과정 (Sampling process) 전반에 걸쳐 루프 횟수를 적응적으로 조정하면 성능을 유지하면서도 연산 효율성 측면에서 추가적인 이득을 얻을 수 있습니다. 마지막으로, 어텐션 분석 (Attention analysis)을 통해 루핑이 마스크된 위치 (Masked positions) 간의 상호작용을 촉진함으로써 MDM에서 효과적이라는 증거를 제시합니다. 우리의 코드와 가중치 (Weights)는 공개될 예정입니다.