Continue diffusie in token-tijd voor taalmodellering

Samenvatting

In dit artikel introduceren we token time continuous diffusion (TTCD), een nieuw diffusie taalmodel dat (a) opereert in continue ruimte, waarbij Gaussische ruis deterministisch wordt omgezet naar een uiteindelijk token canvas zonder verdere monstername, en cruciaal (b) een nieuw concept van per-token tijden omvat, waarbij sommige tokens sneller van ruis naar token evolueren dan andere. Modellering in continue ruimte helpt TTCD om het parallel samplen van meerdere tokens te vermijden, een belangrijke bron van onnauwkeurigheid bij hoge versnellingen voor modellen die uitsluitend in discrete ruimte itereren. Het concept van per-token tijden helpt TTCD om conditionele generatie beter te modelleren, maakt het mogelijk dat meer zekere tokens sneller verlopen, en zorgt voor gedifferentieerde invloeden tussen tokens tijdens verfijning. TTCD presteert beter dan discrete modellen bij hoge versnellingen. We trainen een TTCD-model met 160M parameters op OpenWebText en distilleren het vervolgens zelf; we vinden dat bij hoge versnellingen we vergelijkbaar zijn in kwaliteit van onvoorwaardelijke generatie, en beter presteren in conditionele generatie, vergeleken met verschillende bestaande modellen van vergelijkbare grootte die op dezelfde data zijn getraind en zelf-gedistilleerd. Ook bij het oplossen van Sudoku behalen we vergelijkbare winsten.

English

In this paper we introduce token time continuous diffusion (TTCD), a new diffusion language model which (a) operates in continuous space, deterministically mapping Gaussian noise to a final token canvas with no further sampling, and crucially (b) incorporates a new notion of per-token times, with some tokens proceeding from noise to token at a faster rate than others. Continuous space modeling helps TTCD avoid the parallel sampling of multiple tokens, which is a key source of inaccuracy at high speedups for models that iterate purely in discrete space. The notion of per-token times helps TTCD to better model conditional generation, allows for more sure tokens to proceed at a faster rate, and allows for differentiated inter-token influences during refinement. TTCD outperforms discrete models at high speedups. We train a 160M parameter TTCD model on OpenWebText, and then self-distill it; we find that at high speedups we are comparable in unconditional generation quality, and outperform in conditional generation, several existing models of similar size trained, on the same data, and self-distilled. We achieve similar gains in Sudoku solving as well.