One Latent, Many Tokens: Jointly Learning Compressed Embeddings for Efficient Language Diffusion

Yulin Yuan; Ying Zhang; Xiangming Meng

cs.AISep 27, 2026

One Latent, Many Tokens: Jointly Learning Compressed Embeddings for Efficient Language Diffusion

Authors: Yulin Yuan, Ying Zhang, Xiangming Meng

Organizations: Zhejiang University · University of Cambridge · ZJU-UIUC Institute, Zhejiang University

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.

Figures & tables

Figure 1: Schematic comparison of fixed and jointly learned compressed spaces, showing more structured embeddings and broader decodable regions with joint training.

Figure 2: Overview of JPEG-DLM . The frozen T5 encoder maps tokens to features. The context compressor maps a corrupted view to $\boldsymbol{z}_{c}$ , while the target compressor provides the stop-gradient clean target embedding $\boldsymbol{z}_{\mathrm{tgt}}$ and is updated by an exponential moving average. The shared-weight network predicts the clean target embedding from $\boldsymbol{z}_{c}$ for JEP and from the noisy embedding $\boldsymbol{z}_{t}$ for flow matching. The decoding module maps the predicted clean embedding back to features and tokens using $\mathcal{L}_{\mathrm{feat}}$ and $\mathcal{L}_{\mathrm{ce}}$ .

Figure 3

Method	Param. (M)	Gen-PPL $\downarrow$	Ent.	MAUVE $\uparrow$	Speed $\uparrow$	Gen-PPL $\downarrow$	Ent.	MAUVE $\uparrow$	Speed $\uparrow$
		LM1B ( $L=128$ )				OWT ( $L=1024$ )
AR ( Radford et al., 2019 )	130	151.36	4.397	0.947	$1.00\times$ ( $1.62\times$ w/ KV cache)	41.06	5.583	0.930	$1.00\times$ ( $7.14\times$ w/ KV cache)
MDLM ( Sahoo et al., 2024 )	130	192.80	4.410	0.923	$0.52\times$	41.46	5.287	0.755	$0.32\times$
SEDD-Absorb ( Lou et al., 2024 )	170	182.57	4.405	0.941	$0.40\times$	41.15	5.260	0.728	$0.25\times$
Duo ( Sahoo et al., 2025 )	130	161.80	4.382	0.941	$0.41\times$	75.88	5.536	0.903	$0.15\times$
FLM ( Lee et al., 2026 )	179	176.74	4.402	0.863	$0.66\times$	61.37	5.335	0.373	$0.21\times$

Table 1: Unconditional generation on LM1B and OWT. Speed is relative to AR without KV cache. Among diffusion and flow models, bold and underlined mark the best and second-best Gen-PPL, MAUVE, and speed. JPEG-DLM has 105M DiT and token projection parameters, or 154M including the compressor and decompressor.

Figure 3: Generation quality at different latent lengths (a) and latent widths (b).

Panel of Figure 4 — Figure 4: Joint and two-stage training. (a) Gen-PPL and MAUVE across compression rates; arrows show relative changes from COSMOS to JPEG-DLM , and diamonds mark ELF-B. COSMOS uses final-layer T5-small features at every rate. (b) Correlation between latent positions at $r=0.5$ .

Figure 5: Gen-PPL over interpolations between 1,024 pairs of clean target embeddings at a compression rate of

r=0.5

, without added noise and at three noise levels. Here,

t

denotes the COSMOS forward time, with

t=0

indicating no added noise. For the three noisy settings, JPEG-DLM uses noisy embeddings at the same SNRs along its linear path.

Variant	$\mathcal{L}_{\mathrm{jep}}$	$\mathcal{L}_{\mathrm{feat}}$	$\sigma$	$\mathcal{L}_{\mathrm{ce}}$ timesteps	Gen-PPL $\downarrow$	Ent.	MAUVE $\uparrow$
Base	$\checkmark$	$\checkmark$	1.0	$t=1$	48.00	4.064	0.912
sep. nets	$\checkmark$	$\checkmark$	1.0	$t=1$	92.45	4.117	0.786
w/o $\mathcal{L}_{\mathrm{jep}}$	$\times$	$\checkmark$	1.0	$t=1$	51.65	4.016	0.849
w/o $\mathcal{L}_{\mathrm{jep}}$ + sep. nets	$\times$	$\checkmark$	1.0	$t=1$	131.94	4.188	0.632
w/o $\mathcal{L}_{\mathrm{feat}}$	$\checkmark$	$\times$	1.0	$t=1$	51.57	4.069	0.825
$\sigma=0.5$	$\checkmark$	$\checkmark$	0.5	$t=1$	61.43	3.993	0.817

Table 2: Ablation study on OWT-128. Base uses a shared-weight network.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

SIGReg	VICReg	Gen-PPL $\downarrow$	MAUVE $\uparrow$	Entropy
$\times$	$\times$	63.11	0.856	4.090
$\checkmark$	$\times$	48.60	0.865	4.062
$\times$	$\checkmark$	40.77	0.868	3.965
$\checkmark$	$\checkmark$	48.00	0.912	4.064

Appendix

Table 3: Regularizer ablation on OWT-128 with

r=0.5

d=512

, 64 sampling steps, and self-conditioning at scale

3.0

. The variance and covariance terms are enabled or disabled as a pair. The shaded row reproduces Base from Table 2 .

Figure 6: Gen-PPL on LM1B across sampling steps. Error bars show standard deviations over five seeds. Standard deviations are unavailable for the 8-step baselines.

Steps	Gen-PPL $\downarrow$	Ent.	MAUVE $\uparrow$
8	$225.24$	$4.162$	$0.432$
16	$137.86$	$4.270$	$0.896$
32	$96.76$	$4.263$	$\mathbf{0.951}$
64	$78.64$	$4.248$	$0.940$
128	$69.45$	$4.238$	$0.941$
256	$\mathbf{64.07}$	$4.230$	$0.933$

Appendix

Table 4: Generation quality of JPEG-DLM on LM1B with

r=0.25

at different sampling budgets. Results are means over five seeds. The shaded row is the 32-step configuration used in the main comparison.

Figure 7: GPT-2 token lengths of ELF generations and two OWT reference sets. The reference constructed from fixed 1,024-token blocks concentrates near the block length, whereas the reference tokenized and decoded with T5 and the ELF samples vary in length. After discretization, the intersection between the ELF and fixed-length reference histograms is only 0.9%.

Parameter	LM1B	OWT-1024
Data and tokenization
Dataset	LM1B ( Chelba et al., 2014 )	OWT
Token sequence length	128	1024
Sequence construction	Blocks packed across sentence boundaries	Blocks packed across sentence boundaries (no padding)
Model architecture
Frozen text encoder	T5-small	T5-small

Appendix

Table 5: Model and training configurations of JPEG-DLM on LM1B and OWT-1024. Sampling compute counts one multiply–accumulate as one operation. For

K

solver steps, JPEG-DLM uses

K+1

DiT evaluations, including the final clean embedding prediction, followed by one decoding pass. The estimate at full latent length uses the same

K+1

evaluation budget.

			LM1B		OWT-1024
Method	SC	Samples	Tokenizer	Steps	Tokenizer	Steps
AR	$\times$	1,024	BERT	–	GPT-2	–
MDLM	$\times$	1,024	BERT	128	GPT-2 + mask	1024
SEDD-Absorb	$\times$	1,024	GPT-2	128	GPT-2	1024
Duo	$\times$	1,024	BERT	128	GPT-2	1024
FLM	$\times$	1,024	BERT	128	GPT-2	1024

Appendix

Table 6: Tokenization and sampling configurations for the main unconditional generation comparison. BERT ( Devlin et al., 2019 ) denotes bert-base-uncased , and GPT-2 ( Radford et al., 2019 ) denotes GPT-2 BPE. T5 ( Raffel et al., 2020 ) uses SentencePiece ( Kudo and Richardson, 2018 ) . A checkmark in the SC column indicates self-conditioning.

Parameter	OWT-128
Dataset	linluqiu/openwebtext-len128-packing_padding
Sequence length	128
Text encoder / tokenizer	Frozen T5-small / T5
Encoder feature	Final hidden layer
Latent length / compression rate $r$	128 / 1.0, 64 / 0.5, 32 / 0.25, 16 / 0.125
Latent width	512

Appendix

Table 7: Configuration of COSMOS on OWT-128.

Layer index	Gen-PPL $\downarrow$	Ent.	MAUVE $\uparrow$
$-1$	$240.811$	$3.9033$	$0.1167$
$-2$	$154.551$	$3.9736$	$0.3009$
$-3$	$163.196$	$3.9717$	$0.2677$

Appendix

Table 8: COSMOS performance at

r=1.0

with T5-small layer indices

-1

-2

, and

-3

. Results are means over five seeds using 1,024 samples per seed.

Explore similar work

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space
May 8, 2026Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4Diffusion Language Models Latent Diffusion Model
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
Scaling and Distilling Text Embeddings for Better Diffusibility
Oct 1, 2026Zekai Zhang, Yunjie Tian, Yanjin He +4Diffusion Language Models Text Embeddings
University of Michigan
ELF: Embedded Language Flows
May 11, 2026Keya Hu, Linlu Qiu, Yiyang Lu +5Diffusion Language Models Latent Flow
MIT

May 8, 2026, cs.CLJ/K move · Enter open · S save

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov, Nikita Gushchin+3

Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia+2

Diffusion Language Models Latent Diffusion Model Diffusion Models Latent Space

Open paper AskPDF arXiv

Latent diffusion models offer an attractive alternative to discrete diffusion for non-autoregressive text generation by operating on continuous text representations and denoising entire sequences in parallel. The major challenge in latent diffusion modeling is constructing a suitable latent space. In this work, we present the Latent Diffusion Language Model (LDLM), in which the latent encoder, diffusion model, and decoder are trained jointly. LDLM builds its latent space by reshaping the representations of a pre-trained language model with a trainable encoder, yielding latents that are easy to both denoise and decode into tokens. We show that naive joint training produces a low-quality diffusion model, and propose a simple training recipe consisting of an MSE decoder loss, diffusion-to-encoder warmup, adaptive timestep sampling, and decoder-input noise. Ablations show that each component substantially impacts generation performance. On OpenWebText and LM1B, LDLM achieves better generation performance than existing discrete and continuous diffusion language models while being

2{\text -}13\times

faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.

Abstract

Explore similar work

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Scaling and Distilling Text Embeddings for Better Diffusibility

ELF: Embedded Language Flows