Scaling and Distilling Text Embeddings for Better Diffusibility

Scaling and Distilling Text Embeddings for Better Diffusibility. By Zekai Zhang and six others, from the University of Michigan. Published on the first of October. Its code repository is public and real: eleven source files.

Diffusion language models are an alternative to generating text autoregressively, one token after another. A recent line of them, continuous diffusion language models, applies latent diffusion to continuous text embeddings rather than to discrete tokens. That raises a practical question which sounds simple and is not. Which embedding makes the best latent space to diffuse in? The authors call the property they are after diffusibility, and the answer depends on more than how good an embedding is at telling words apart.

A search through different embeddings produces two observations. First, scaling the embedding model to stronger members of the same family greatly improves generative performance: T5, then T5Gemma-1, then T5Gemma-2. Second, even the strongest of those is not optimal when used raw. T5Gemma-2 embeddings are so discriminative that even the embeddings of plausible alternative words are separated from one another. That leaves generation vulnerable to imperfect sampling. A continuous diffusion process that does not land cleanly often fails to reach any of the valid alternatives, and ends up at an invalid embedding instead.

The remedy is distillation. T5Gemma-2 acts as the teacher, and a student encoder is trained to learn the teacher's decoded probabilities as soft labels. A soft label spreads weight across the plausible alternatives instead of naming one word. Learning from such labels, the student pulls the alternative embeddings closer together, while keeping the encoding and decoding mechanism intact. The distilled embeddings form a latent space that is more connected, and in the authors' term more diffusible, than the vanilla T5Gemma-2 embeddings the student learned from.

The headline number comes from a medium-sized diffusion language model trained on OpenWebText. Sampled at real-text entropy, it reaches a generative perplexity of seventeen point eight. Real text itself scores fifteen point four. On that generative perplexity measure, the model outperforms GPT-2-M. The distilled embeddings are also reported to improve over the vanilla T5Gemma-2 embeddings they were distilled from.

Two pressures are at work here, and they point different ways. A stronger encoder gives a better space to generate in, which is why scaling helps. Yet the sharper an encoder separates words, the less room an imperfect sampler has to land somewhere valid. Distilling with soft labels is how the authors keep the strength of the larger model while bringing plausible alternatives back within reach of each other.