Data-efficiënte autoregressief-naar-diffusie taalmodellen via on-policy distillatie

Samenvatting

Wij bestuderen de transformatie van autoregressieve modellen (ARLMs) naar diffusie-taalmodellen (DLMs). In plaats van vanaf nul te pretrainen, vervangt eerder werk de causale aandacht in ARLMs door bidirectionele aandacht en traint vervolgens het resulterende model met een DLM-doelstelling. Deze benaderingen leiden echter tot twee distributieverschuivingen. Ten eerste kan de overgang van een volgende-token voorspellingsdoelstelling naar een DLM-doelstelling kennis weggooien die door de ARLM tijdens de training is verworven. Ten tweede hebben standaard DLMs te lijden onder een train-inferentie mismatch, omdat het trainingsverlies wordt gedefinieerd op willekeurig gemaskeerde sequenties in plaats van de trajecten die tijdens de inferentie worden aangetroffen en die worden geproduceerd door vertrouwensgebaseerd decoderen. Om beide uitdagingen aan te pakken, introduceren wij een On-Policy Diffusion Language Model (OPDLM), waarbij On-Policy Distillation (OPD) wordt gebruikt voor de ARLM-naar-DLM-transformatie. Specifiek wordt OPDLM getraind via zelf-OPD, waarbij de student, een ARLM met bidirectionele aandacht, zijn eigen trajecten genereert, en de leraar, het oorspronkelijke bevroren ARLM, zijn kennis distilleert door doellogits op deze trajecten te leveren. Door direct op een on-policy manier te trainen, elimineert OPDLM de train-inferentie mismatch in DLMs, terwijl distillatie van het oorspronkelijke model de kennisbehoud van de ARLM verbetert. Empirische resultaten tonen aan dat OPDLM 15 tot 7000 keer minder trainingssymbolen nodig heeft met sterke prestaties over een breed scala aan taken. OPDLM vermijdt de buitensporige kosten van DLM-pretraining en positioneert DLM-transformatie als een vorm van ARLM-natraining.

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.