Dateneffiziente Autoregressiv-zu-Diffusions-Sprachmodelle mittels On-Policy-Distillation

Zusammenfassung

Wir untersuchen die Transformation von autoregressiven Modellen (ARLMs) in Diffusions-Sprachmodelle (DLMs). Anstatt von Grund auf neu zu trainieren, ersetzen bisherige Arbeiten die kausale Aufmerksamkeit in ARLMs durch bidirektionale Aufmerksamkeit und trainieren das resultierende Modell anschließend mit einem DLM-Ziel. Diese Ansätze führen jedoch zu zwei Verteilungsverschiebungen. Erstens kann der Übergang von einem Nächster-Token-Vorhersageziel zu einem DLM-Ziel Wissen verwerfen, das das ARLM während des Trainings erworben hat. Zweitens leiden Standard-DLMs unter einer Diskrepanz zwischen Training und Inferenz, da der Trainingsverlust auf zufällig maskierten Sequenzen definiert ist, nicht auf den im Inferenzprozess auftretenden Trajektorien, die durch konfidenzbasiertes Decodieren erzeugt werden. Um beide Herausforderungen zu bewältigen, führen wir ein On-Policy-Diffusions-Sprachmodell (OPDLM) ein, bei dem On-Policy-Distillation (OPD) zur ARLM-zu-DLM-Transformation eingesetzt wird. Konkret wird OPDLM durch Selbst-OPD trainiert, wobei der Schüler, ein ARLM mit bidirektionaler Aufmerksamkeit, seine eigenen Trajektorien generiert, und der Lehrer, das ursprüngliche eingefrorene ARLM, sein Wissen destilliert, indem er Ziel-Logits für diese Trajektorien bereitstellt. Durch das direkte On-Policy-Training beseitigt OPDLM die Diskrepanz zwischen Training und Inferenz in DLMs, während die Destillation aus dem ursprünglichen Modell die Wissensbewahrung aus dem ARLM verbessert. Empirische Ergebnisse zeigen, dass OPDLM 15- bis 7.000-mal weniger Trainings-Tokens benötigt und dabei über eine Vielzahl von Aufgaben hinweg eine starke Leistung erbringt. OPDLM vermeidet die prohibitiv hohen Kosten des DLM-Pretrainings und positioniert die DLM-Transformation als eine Form des ARLM-Post-Trainings.

English

We study the transformation of autoregressive models (ARLMs) into diffusion language models (DLMs). Rather than pretraining from scratch, prior work replaces the causal attention in ARLMs with bidirectional attention and then trains the resulting model using a DLM objective. However, these approaches incur two distribution shifts. First, transitioning from a next-token prediction objective to a DLM objective can discard knowledge acquired by the ARLM during training. Second, standard DLMs suffer from a train-inference mismatch, as the training loss is defined on randomly masked sequences rather than the trajectories encountered at inference produced by confidence-based decoding. To address both challenges, we introduce an On-Policy Diffusion Language Model (OPDLM) in which On-Policy Distillation (OPD) is employed for ARLM-to-DLM transformation. Specifically, OPDLM is trained via self-OPD, where the student, an ARLM with bidirectional attention, generates its own trajectories, and the teacher, the original frozen ARLM, distills its knowledge by providing target logits on these trajectories. By training directly in an on-policy manner, OPDLM eliminates the train-inference mismatch in DLMs, while distillation from the original model enhances knowledge retention from the ARLM. Empirical results demonstrate that OPDLM requires 15x to 7,000x fewer training tokens with strong performance across a wide variety of tasks. OPDLM avoids the prohibitive cost of DLM pretraining and positions DLM transformation as a form of ARLM post-training.