IDLM: 역증류 디퓨전 언어 모델

IDLM: Inverse-distilled Diffusion Language Models

디퓨전 언어 모델(DLM)은 최근 텍스트 생성 분야에서 우수한 성능을 달성했다. 그러나 다단계 샘플링으로 인해 추론 속도가 느려져 실제 활용에 한계가 있다. 이를 해결하기 위해 우리는 본래 연속형 디퓨전 모델의 가속을 위해 개발된 기법인 역증류(Inverse Distillation)를 이산(discrete) 환경으로 확장한다. 그럼에도 불구하고 이 확장은 이론적 및 실용적인 과제를 동시에 수반한다. 이론적 관점에서 역증류 목적 함수는 해의 유일성을 보장하지 않아 최적화되지 않은 결과(suboptimal solutions)를 초래할 수 있다. 실용적 관점에서 이산 공간에서의 역전파는 까다롭고 종종 불안정하다. 이러한 문제를 극복하기 위해, 우리는 먼저 제안하는 역 공식화(inverse formulation)가 유일한 해를 가짐을 증명하는 이론적 결과를 제시하여 유효한 최적화를 보장한다. 그런 다음 효과적인 학습을 지원하기 위해 그래디언트가 안정적인 완화(gradient-stable relaxations) 기법을 도입한다. 결과적으로 여러 DLM에 대한 실험을 통해, 우리의 방법론인 역증류 디퓨전 언어 모델(IDLM)이 교사 모델의 엔트로피와 생성 퍼플렉서티(perplexity)를 유지하면서도 추론 단계 수를 4배에서 64배까지 줄일 수 있음을 입증하였다.

Diffusion Language Models (DLMs) have recently achieved strong results in text generation. However, their multi-step sampling leads to slow inference, limiting practical use. To address this, we extend Inverse Distillation, a technique originally developed to accelerate continuous diffusion models, to the discrete setting. Nonetheless, this extension introduces both theoretical and practical challenges. From a theoretical perspective, the inverse distillation objective lacks uniqueness guarantees, which may lead to suboptimal solutions. From a practical standpoint, backpropagation in the discrete space is non-trivial and often unstable. To overcome these challenges, we first provide a theoretical result demonstrating that our inverse formulation admits a unique solution, thereby ensuring valid optimization. We then introduce gradient-stable relaxations to support effective training. As a result, experiments on multiple DLMs show that our method, Inverse-distilled Diffusion Language Models (IDLM), reduces the number of inference steps by 4x-64x, while preserving the teacher model's entropy and generative perplexity.