Efficient language generation with jointly learned compressed embeddings
One Latent, Many Tokens: Jointly Learning Compressed Embeddings for Efficient Language Diffusion
Artificial Intelligence
Summary
Generating text with diffusion language models can be slow because each word is processed one piece at a time. To speed this up, some methods group the words into shorter codes, but these codes are fixed before training and can be tricky to work with. The authors introduced JPEG-DLM, which trains the code compressor, model, and decoder all together so the codes become easier to predict and turn back into words. This results in faster and better text generation compared to previous models.
What this means in practice
- •For language model engineers: Build faster text generation models by integrating jointly trained compressed embeddings with diffusion-based approaches.
- •For natural language processing teams: Improve throughput of large-scale language deployment by adopting joint embedding compression methods to reduce computational costs.
Authors
Yulin Yuan, Ying Zhang, Xiangming Meng
Abstract
Most continuous diffusion language models process one latent position per token at each sampling step, making generation expensive. Two-stage methods lower the cost by reducing the latent length, but they fix the compressed embedding space before training the diffusion model. Embeddings from the fixed space can be difficult to model with diffusion and decode reliably into tokens, which limits generation quality after compression. To address this problem, we introduce JPEG-DLM (Joint-embedding Prediction for Efficient Generation with Diffusion Language Model), which jointly trains a compressor, a flow matching model and a decoding module. With joint-embedding prediction, JPEG-DLM learns compressed embeddings that are more structured, easier to model with diffusion and reliably decodable into tokens. JPEG-DLM achieves the lowest mean Gen-PPL and highest throughput among recent diffusion and flow models on LM1B and OWT. At a compression rate of 0.5 on OWT, it reaches a Gen-PPL of 34.52 and approximately 2.3 times ELF's throughput. These results suggest that jointly learning compressed embeddings offers a promising path toward efficient diffusion language modeling. Code will be released soon.