Scaling and Distilling Text Embeddings for Better Diffusibility
Zekai Zhang, Yunjie Tian, Yanjin He, Xiaoyan Zhang, Dongdi Zhao, Qing Qu, Di Fu
Scaling and distilling text embeddings improves the diffusibility of continuous diffusion language models.
How can we improve the generation quality of continuous diffusion language models by optimizing the underlying text embeddings?
Continuous diffusion language models (DLMs) struggle to generate valid text because their latent embedding spaces are often too discriminative, causing generated trajectories to land on "invalid" embeddings that decode to incorrect or uncertain words. The authors find that scaling to stronger embedding models improves performance, but they solve the remaining generation failures by distilling the teacher's decoder probabilities into a student encoder. This forces the student to place plausible alternative words closer together in the latent space, creating a more connected and "diffusible" manifold. This distilled student model outperforms the teacher on generative perplexity and sampling stability, achieving a Gen. PPL of 31.2 on OpenWebText.
Paper Primer
The core challenge is that high-quality, discriminative embeddings—designed to distinguish between words for classification—are inherently "un-diffusible." Because these embeddings push semantically similar candidates apart, a continuous diffusion trajectory often gets stuck in the empty space between them, resulting in invalid embeddings that the decoder cannot map to a clear token.
The authors' method is a distillation-based re-shaping of the latent space: they train a student encoder to match the teacher's decoded probability distributions as soft labels. This acts like a gravitational pull: the student learns to cluster plausible word candidates together, ensuring that even if a diffusion trajectory is slightly imperfect, it lands in a region that decodes to a valid, high-confidence token.
Distilled embeddings significantly improve generative stability and quality over raw, high-performance teacher embeddings.
The student model achieves a Gen. PPL of 31.2 on OWT-1024, compared to 38.5 for the T5Gemma-2 teacher. The distilled student reduces the number of "invalid" embedding positions by nearly 90% compared to the teacher during generation.
Scaling embedding models within the same family provides a consistent baseline improvement for continuous DLMs.
Replacing T5-small with T5Gemma-2-270M reduces generative perplexity by approximately 40% at matched entropy.
Why does distillation improve generation if it technically degrades the encoder's discriminative power?
The authors observe a trade-off: while the student loses some downstream classification accuracy (e.g., 89.4 to 78.1 on SST-2), it gains "diffusibility." By sacrificing the strict separation of candidates, the student creates a smoother, more connected latent space that is easier for the diffusion process to navigate without collapsing.
Is this approach limited to specific model architectures?
The authors focus on encoder-decoder models (T5 family) because they are well-suited for this distillation setup. They note that encoder-only models like ModernBERT, while strong at discriminative tasks, performed poorly in their diffusion experiments, suggesting the choice of embedding architecture is critical for continuous DLMs.
Introduction: The Diffusibility Bottleneck
We expose why standard text embeddings hinder diffusion and propose scaling plus distillation to make them more diffusible.
Continuous diffusion language models (DLMs) generate text by diffusing over latent embeddings, but the quality of those embeddings determines how well the diffusion process can reach valid, decodable points. Standard embeddings, even strong ones like T5Gemma‑2, are so discriminative that sampling often lands in regions that decode to no real word, exposing a fundamental gap between embedding quality and diffusion performance.
Diffusibility measures how easily a diffusion process can wander through an embedding space and still land on a point that corresponds to a real token.
**Figure 1.** Overview of our performance and mechanism. *Left*: the same DLM trained on our embeddings (distilled from T5Gemma-2) reaches real-text entropy with the best quality. *Right*: T5Gemma-2 maps plausible candidates apart in the embedding space, so a generated embedding often reaches none of those. Our distilled embedding pulls the candidates closer, forming a continuous region where an embedding at least decodes to one of them.
The gap between embedding quality and diffusion performance drives the need for more diffusible latent spaces.
Preliminaries and Related Work
Foundations of continuous diffusion language models and a brief survey of related approaches.
We begin by defining the objects that the rest of the paper manipulates. Input token sequences $s \in \mathcal{V}^L$ are projected to embeddings $x = \text{Emb}(s) \in \mathbb{R}^{L \times d}$, which serve as the latent space for a diffusion process that adds Gaussian noise $\epsilon \in \mathbb{R}^{L \times d}$ at each timestep $t$.
A Continuous DLM treats a sequence of token embeddings as a continuous latent that is gradually corrupted by Gaussian noise and then restored by a learned denoiser.
ELF provides a recipe for turning any pretrained embedding model into a diffusion‑compatible latent by fixing the embedding extractor and adding a few diffusion‑specific components.
Related work can be split into three strands. First, early continuous DLMs used co‑trained embeddings, which tend to collapse to a single point to minimize mean‑squared error; later studies showed that pretrained embeddings avoid this collapse and yield higher quality generations. Second, discrete DLMs apply categorical corruption to tokens and have been adapted from autoregressive foundation models, demonstrating that continuous and discrete diffusion can co‑evolve. Third, the notion of “diffusibility” originates from image diffusion, where replacing VAE latents with pretrained representations such as DINOv2 improves sampling; analogously, scaling the T5‑Gemma‑2 encoder and distilling it for diffusibility is the textual counterpart.
Scaling Text Embeddings
Scaling embeddings reveals richer latents but also exposes generation challenges.
Even though larger encoders expose more information for diffusion, the resulting embeddings are often hard to generate faithfully. This tension motivates the scaling mechanism we introduce.
Scaling embeddings by training stronger encoders supplies the diffusion process with richer latent signals, but it also widens the space the sampler must explore.
Compute the relative perplexity reduction: (27 − 16) / 27 ≈ 0.41 → ≈ 40 % lower.
Entropy increases from 4.5 to 5.05, indicating richer information content.
Both encoders are normalized, so the distributional scale is comparable.
Scaling the encoder improves the diffusion’s signal‑to‑noise ratio, but the larger latent space also creates more “dead zones” where sampled embeddings have no nearby real token.
How does embedding scaling differ from simply increasing the diffusion model’s capacity?
Scaling changes the *input* distribution the diffusion sees, giving it richer semantic cues, whereas enlarging the diffusion model only adds more parameters to process the same latent space. The former can reduce perplexity even with the same diffusion architecture, while the latter may still struggle if the latent space is poorly aligned with real tokens.
T5Gemma‑2 is a 270 M‑parameter encoder that, after normalization, reaches the entropy of real text and yields the lowest generation perplexity among the evaluated models.
Why is T5Gemma‑2 more diffusible than ModernBERT even though both are strong encoders?
ModernBERT’s encoder‑only architecture yields embeddings that are highly discriminative but not organized for the diffusion denoiser’s averaging step; T5Gemma‑2’s encoder‑decoder design produces embeddings that better respect the smoothness assumptions of diffusion, leading to lower perplexity at comparable entropy.
**Figure 2.** Performance of the same ELF-B trained with different encoders on OWT-512. T5Gemma-2 reaches the entropy of real text, and its Gen. PPL is about 40% lower than T5-small's at the same entropy. Within the T5 family, T5Gemma-2 > T5Gemma-1 > T5-small.
**Figure 3:** Generated embeddings deviate from real T5Gemma-2's.
**Figure 4.** Visualizing the embedding error in 2D. Top: a sample generated at NFE 32; uncertain tokens (top-1 below 0.6) in purple coincide with the grammar mistakes. Bottom: four positions in the plane of their candidate words and the decoded probabilities. The disk marks the 0.91 nn range of a candidate, and an invalid one lands outside every disk and is decoded problematically. Distances are computed in the original 640-dimensional space; the 2D layout is a distance-preserving map for display. See details in App. B.4.
Generated embeddings can be invalid: an embedding is flagged invalid if its distance to every top‑k candidate exceeds 0.91 nn or its top‑1 probability falls below 0.6. In a 1024‑token sample we observed 142 uncertain positions, of which 109 were invalid.
Distilling for Diffusibility
Distillation aligns student embeddings with teacher soft labels to improve diffusibility.
Standard one‑hot training forces the encoder to push each token’s embedding away from all alternatives, which fragments the space and harms diffusion. Distillation replaces that pressure with soft teacher probabilities, encouraging nearby embeddings for plausible words.
Teach a student encoder to mimic the teacher decoder’s soft probability distribution, pulling plausible word embeddings together so the diffusion process lands on valid tokens more often.
How does this distillation differ from conventional knowledge distillation?
Standard KD aligns logits on a single hard target; here we align the full soft distribution over many plausible tokens, explicitly encouraging the student’s embedding space to group semantically related words together rather than merely copying class probabilities.
Compute KL for each token: $0.5\log\frac{0.5}{0.3}=0.255$, $0.2\log\frac{0.2}{0.3}=-0.058$, $0.2\log\frac{0.2}{0.2}=0$, $0.1\log\frac{0.1}{0.1}=0$, $0.0\log\frac{0.0}{0.1}=0$.
Sum the contributions: KL $=0.255-0.058=0.197$ (positive because the student under‑allocates to the most probable token).
Back‑propagate this loss; the gradient pushes the student’s probability for A up and reduces the spurious mass on E, tightening the embedding cluster around A, B, C.
KL on soft labels directly penalizes mismatches in the probability spread, forcing the student to bring the embeddings of the top candidates closer together while suppressing irrelevant ones.
Encode the input sequence with the teacher encoder.
Decode with the frozen teacher decoder to obtain soft label distribution $p_T$.
Encode the same input with the student encoder (parameters $\theta$).
Decode the student hidden states with the same teacher decoder to obtain $p_S$.
Compute $\mathcal{L}_{KD}(\theta)$ and update $\theta$ via gradient descent.
Repeat until convergence, optionally re‑initializing the student from every second teacher layer.
**Figure 5.** Illustration of the distillation pipeline and the effect on the embeddings. *Left*: The student encoder is trained to match the teacher's decoded probabilities (soft labels). *Right*: The student pulls the embeddings of plausible candidates closer while keeping the irrelevant word far. The right panel is an illustration; measurements on real embeddings are in App. A.4.
**Figure 6.** Distilled embeddings are more diffusible with lower embedding error. Same layout as Fig. 4. The teacher has more uncertain embeddings (in purple) which are mostly invalid, while the student's generated embeddings lie in the region formed by candidates that are close to each other. Over 64 sequences, the teacher has 142 uncertain positions per sequence, 109 of which are also invalid; the student 42 and 12; see App. A.1.
Experimental Results
Our distilled embeddings achieve lower perplexity and faster encoding than the T5Gemma‑2 baseline.
Recall that continuous diffusion language models (DLMs) generate text by diffusing over latent embeddings, and that scaling embeddings improves diffusibility but still benefits from distillation. This section reports the empirical outcomes of those scaled and distilled embeddings.
Student (ours) + ELF‑B beats the T5Gemma‑2 + ELF‑B teacher on OWT‑1024 by achieving lower generation perplexity at real‑text entropy.
Gen. PPL 31.2 ± 0.2 versus 38.5 ± 0.5 for the teacher (Table 1).
On LM1B, the distilled student also outperforms the T5Gemma‑2 teacher, delivering lower perplexity while matching entropy.
Gen. PPL 50.5 ± 0.5 versus 60.6 ± 0.3 for the teacher (Table 2).
The distilled student halves the encoding cost compared to the T5Gemma‑2 encoder.
Encoding latency drops from 0.21 T to 0.10 T per 1024‑token sequence (Table 3).
**Figure 7.** Round-trip denoising from the same sentence. *Left*: distance from each denoised embedding to the real embedding of the word it decodes to. *Right*: one sequence denoised from $t_0 = 0.3$; invalid tokens in purple.
**Figure 8.** Sampling ablation of ELF-B (left) and ELF-M (right) models. Each curve stands for different NFEs (64–512) under the same self-conditioning. The dotted curve on the right varies the sampler's noise scale. Our embedding consistently outperforms raw T5Gemma-2.
**Figure 9.** Student's and teacher's generated embeddings. Positions with top-1 below 0.6 from 64 sequences per model (NFE 32). Curve: the decoder's median top-1 per distance bin.
**Figure 10.** Ablation on the distillation objective. Gen. PPL–entropy curves on OWT-1024 with ELF-B. The teacher, the KL student and the MSE student are sampled under the same settings (256 samples, one seed), which vary the self-conditioning scale (1–6), the sampler's noise scale (1.0–3.0) and the NFE (16–256) but do not form a full grid; the CE student is sampled on the grid of Table 1. Each curve is the best PPL at the same entropy. These settings do not contain the teacher's setting in Table 1 (sc= 1, NFE= 64), so in this plot the teacher stops at entropy 5.38.
**Figure 11.** Effect of student depth. Gen. PPL–entropy curves for students with 3 to 18 layers and the T5Gemma-2 teacher on the 512-token OWT setting. All students use the same distillation and evaluation protocol. The vertical dotted line marks the entropy of real text, and the black star shows the real-text reference. Deeper students generally provide a better quality–diversity trade-off, although the improvement is not monotonic beyond 12 layers.
**Figure 12:** Distance between a word and its candidate on real text, when a share of the context tokens is replaced. Line: median over 317 positions; band: quartiles.
Qualitative Examples
Student-distilled embeddings dramatically lower perplexity compared to raw T5Gemma‑2 across sampling settings.
Student-distilled embeddings reduce sequence perplexity by 5.2 points at NFE = 64 compared to raw T5Gemma‑2.
Student + ELF‑M (sc 3/NFE 64) achieves seq. PPL 14.2 versus T5Gemma‑2 + ELF‑M (sc 2/NFE 64) with seq. PPL 19.4 on OWT‑1024.
Across both OWT‑1024 and LM1B, the distilled student embeddings consistently achieve lower perplexity than the original T5Gemma‑2, especially at low NFE where the raw model often collapses into multilingual gibberish while the student remains fluent.
Ethics and Data
Appendix A provides ethics, reproducibility details, and additional experimental analyses.
We use only publicly available corpora (OpenWebText, LM1B, WikiText‑103) and publicly released models (T5Gemma‑2, T5‑small, ModernBERT). No private data or human subjects are involved.
Our continuous DLM code builds on the public ELF‑pytorch implementation, altering only the embedding model while preserving the training and sampling pipeline (see App. B.1). Distillation details (App. B.2) and evaluation protocols (App. B.3) are provided in full.
We employ generative AI tools for implementation, dataset cleaning, literature search, and figure creation, but all core research ideas, scaling decisions, and final writing are authored by the human team. AI‑assisted code is inspected and re‑checked by the authors before inclusion.
The bibliography below enumerates all works cited throughout the paper.
In A.1 we generate $64$ sequences of $1024$ tokens per model with ELF‑B at $NFE=32$, retaining positions where the top‑1 decoded probability falls below $0.6$. The teacher yields $9064$ tokens (mean $1.39$ nn, median $1.38$) while the student yields $2664$ tokens (mean $0.82$ nn, median $0.81$).
A.2 tests decoder dependence by running the pretrained T5Gemma‑2 decoder with teacher forcing on the same sequences. On real text the top‑3 overlap is $99.3\%$ (average $1.5$ shared words), and on generated text the entropy correlation between decoders is $0.72$.
A.3 replaces the soft‑label KL distillation objective with (i) an MSE loss between student and teacher embeddings (no decoder) and (ii) a hard‑label cross‑entropy loss through the frozen decoder. The MSE student reaches an entropy of $5.26$, the CE student $4.18$, both below real‑text entropy.
We also vary student depth from $3$ to $18$ layers while keeping all other settings fixed. Near real‑text entropy ($5.05$) deeper students ($12$–$18$ layers) achieve generative PPL $31$–$34$, outperforming the $9$‑layer student ($49$) and the teacher ($46$); shallower students fall behind.
A.4 probes the effect of distillation on embeddings without diffusion by swapping each token with its top‑2 candidate in $317$ held‑out OWT positions and re‑encoding. On clean text the teacher and student distances are $0.96$ vs. $0.93$ nn; with $10\%$ corrupted context the teacher’s distance rises to $1.19$ nn while the student remains $0.97$ nn.
Downstream evaluation shows top‑1 token accuracy of $99.3\%$ for both teacher and student on WikiText‑103 and $99.4\%$ on held‑out arXiv; discriminative probes (SST‑2, STS‑B, MRPC) are modestly lower for the student.
**Table 4.** Distillation trades discriminative power for diffusibility. Reconstruction and probing results are reported on a 0–100 scale. Gray values indicate generation results below the real-text entropy.
Experimental Details
Experimental details for training, distillation, evaluation, and embedding visualizations.
We build on the public PyTorch ELF implementation, following its training and sampling configuration unless otherwise noted.
For the main experiments we pack OWT sequences up to 1024 tokens, yielding 9.0 M sequences (9.2 B tokens) from 7.8 M documents, and train ELF‑B and ELF‑M for 5 epochs with batch size 512, 2 k warm‑up steps, EMA 0.9999, and the Muon optimizer at learning rate $2\times10^{-4}$.
Encoder comparison uses a 25 % OWT subset with 512‑token sequences; we train ELF‑B on this subset for 15 epochs (58.6 k steps) to evaluate T5‑small, T5‑base, T5Gemma‑1, T5Gemma‑2, and ModernBERT under identical settings.
Distillation initializes a 9‑layer student encoder from evenly spaced layers of the 18‑layer T5Gemma‑2 teacher, freezes the teacher components, and trains the student to match the teacher’s top‑256 logits via token‑wise KL divergence for 50 k steps using AdamW (lr $3\times10^{-4}$) on 8 GPUs.
Evaluation re‑tokenizes all samples with the GPT‑2‑Large tokenizer, computes generation perplexity and unigram entropy, and measures MAUVE using 256‑token features on 1024 generated samples; sampling follows ELF’s SDE sampler with churn $\gamma=1.5$, noise scale 2.0, and varies self‑conditioning scale $sc\in\{1,2,3\}$ and $NFE\in\{64,128,256,512\}$.
We select the configuration with lowest generation perplexity that reaches the real‑text entropy of 5.43; if none does, we report the highest‑entropy point and shade the row gray, then evaluate the chosen points on three seeds with 1024 samples each.
Tokenization differences across encoders are mitigated by re‑tokenizing with GPT‑2‑Large; median generated lengths are 1021–1027 tokens versus 1008 for real text, with negligible impact on unigram entropy.
Embedding‑space visualizations (Figs. 1, 4, 6) plot a generated token embedding against the top‑3 candidate embeddings, normalizing distances by the median nearest‑neighbor distance (13.5 for teacher, 12.4 for student) to express similarity in “nn” units.
**Table.** MAUVE scores for various models using 256-token and 1024-token features.