PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling and Distilling Text Embeddings for Better Diffusibility

Zekai Zhang, Yunjie Tian, Yanjin He, Xiaoyan Zhang, Dongdi Zhao, Qing Qu, Di Fu

arXiv:2610.01016v1cs.CL

TL;DR

Continuous DLMs use text embeddings as their latent space, but which embeddings are most diffusible remains an open question. The paper compares scaled embeddings and distills T5Gemma-2 using its decoder probabilities as soft labels. Scaling reduces generative perplexity, and the distilled embeddings further improve diffusibility and generation outcomes.

  • Problem

    For continuous DLMs, it remains unclear which text-embedding space is most diffusible, and prior work either changes embeddings alongside the diffusion model or uses relatively older embedding models.

    Source: §1 · p. 2

  • Method

    The paper holds the ELF diffusion framework fixed, scales embeddings within the T5 family, and distills T5Gemma-2 by training a student encoder on the teacher decoder’s probabilities.

    Source: §1 · pp. 2–3

  • Results

    Replacing ELF’s T5-small embeddings with T5Gemma-2-270M reduces generative perplexity by about 40% at the same entropy, and distillation further improves perplexity and MAUVE while stabilizing samples.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results suggest that more connected embeddings can be more suitable for continuous diffusion, trading some discriminative power for generation.

    Source: §6 · p. 11

    The work leaves few- or one-step generation and practical-scale models for real-world tasks such as QA for future study.

    Source: §6 · p. 11

Abstract

from arXiv · show

Diffusion language models (DLMs) offer a promising alternative to autoregressive (AR) language generation. Recent advances in continuous DLMs, which apply latent diffusion to continuous text embeddings, raise a practical question: which embedding makes the best latent space, i.e., the most diffusible? To answer this, we search through different embeddings and find that scaling the embedding model to stronger ones within the same family (T5 to T5Gemma-1 to T5Gemma-2) greatly improves generative performance. But the raw T5Gemma-2 embeddings are still not optimal. They are so discriminative that even the embeddings of plausible alternative words are separated, which makes the generation vulnerable to imperfect sampling. Consequently, continuous diffusion often fails to reach any of them and ends up at an invalid embedding instead. To address this, we distill T5Gemma-2 into a student encoder that learns the teacher's decoded probabilities as soft labels. Learning from such soft labels makes the student pull the alternative embeddings closer while maintaining the encoding-decoding mechanism. The distilled embeddings form a more connected and diffusible latent space, improving over the vanilla T5Gemma-2 embeddings. As a result, our medium-sized DLM achieves Gen. PPL 17.8 (against real-text PPL 15.4) at real-text entropy on OpenWebText, outperforming GPT-2-M on Gen. PPL.

1 Introduction

Continuous DLMs diffuse text embeddings, but which embedding space is most suitable remains an open question. This paper isolates the embedding choice, finding that scaling and distillation improve diffusibility while identifying invalid embeddings and practical-scale evaluation as open concerns.

1 Introduction

  • Existing text-embedding studies either jointly modify embeddings and the diffusion model or rely on relatively older embedding models, limiting isolation of the embedding-space contribution.The paper therefore studies the embedding space in isolation, holding the ELF diffusion framework fixed.Source: §1 · p. 2
  • Replacing ELF’s T5-small embeddings with T5Gemma-2-270M reduces generative perplexity by about 40% at the same entropy, and distillation further improves perplexity and MAUVE while stabilizing samples.The study fixes the diffusion framework and changes the embeddings; the student learns from T5Gemma-2 decoder probabilities.Source: §1 · p. 2
  • Soft-label distillation pulls plausible alternative-word embeddings closer, countering the separation that can make imperfect sampling end at invalid embeddings.T5Gemma-2’s information-dense, distinct representations can separate words that could fill the same position, making a continuous trajectory liable to settle between candidates.Source: §1 · p. 2
  • Whether continuous DLMs work at larger practical scales remains untested, and their generated embeddings can fail to represent any word.The authors frame scaling and distillation as a basis for further testing rather than evidence of performance at those scales.Source: §1 · p. 3

2 Preliminary and Related Work

The paper follows ELF’s latent-diffusion setup, changing the frozen text-embedding model while keeping the diffusion design largely fixed. It situates this approach among work on pretrained text embeddings, discrete DLMs, and embedding diffusibility.

2.1 Preliminaries

  • ELF maps token sequences to frozen text embeddings, trains a denoiser on sequence-wise Gaussian-corrupted embeddings, and decodes generated embeddings back to tokens.The decoder shares the denoiser and adds a trainable head; token-wise corruption and cross-entropy training improve robustness to imperfect embeddings.Source: §2.1 · p. 3
  • The optimal denoiser is the posterior mean over training embeddings, weighted by their proximity to the noisy input.The weights decrease with the distance between the noisy input and each candidate embedding.Source: §2.1 · p. 3

2.2 Related work

  • Prior continuous DLMs used simple co-trained embeddings, which can collapse under MSE training; later work found pretrained text embeddings perform better.This paper follows the pretrained-embedding line and tests its potential.Source: §2.2 · p. 4
  • Discrete DLMs diffuse corrupted tokens and can adapt autoregressive foundation models, while this paper explores continuous DLMs using embeddings adapted from such models.Related work also studies discrete DLMs enhanced with continuous embeddings or continuous relaxations.Source: §2.2 · p. 4
  • In text, scaled T5Gemma-2 provides a strong latent space, which this work enhances by distilling the encoder for diffusibility rather than explicitly redesigning the latent.The paper presents this as the text counterpart to findings that pretrained image representations can improve diffusibility over VAE latents.Source: §2.2 · p. 4

3 Scaling Text Embedding for Continuous DLMs

Scaling within the T5 family improves diffusion performance, with T5Gemma-2 reaching real-text entropy and lowering Gen. PPL at matched entropy. Yet its generated embeddings can miss plausible candidate words, producing invalid embeddings and uncertain or incorrect decoding.

3.1 Scaling improves the diffusibility of embeddings

  • At matched entropy, T5Gemma-2 cuts Gen. PPL about 40% versus T5-small and alone reaches real-text entropy.Within the T5 family, performance improves from T5-small through T5Gemma-1 to T5Gemma-2.Source: §3.1 · pp. 4–5, Figure 2 · p. 5
  • ModernBERT performs well on downstream tasks but is less suitable for diffusion and generates repetitive sentences.The paper defines scaling as overall pretraining scale (data, compute, and recipe), not parameter count alone.Source: §3.1 · pp. 4–5

3.2 Scaled embeddings are hard to generate perfectly

  • Across 1024-token generations, 142 positions are uncertain on average and 109 are invalid; invalid means no candidate is within 0.91 nn.The count is reported for generations at NFE 32.Source: §3.2 · pp. 5–6, Figure 3 · p. 5
  • Separated plausible candidates can leave imperfect sampling trajectories between them, where embeddings decode with low confidence or to a wrong word.The paper relates this failure to the denoiser's weighted mean of candidate embeddings.Source: §3.2 · pp. 5–6

4 Enhancing the Diffusibility via Distillation

Distillation trains a student encoder on T5Gemma-2's decoded probabilities, drawing plausible candidate embeddings closer. Student embeddings show fewer uncertain and invalid generated positions, and denoising returns embeddings closer to the original sentence.

4.1 Distilling T5Gemma-2’s Decoded Probabilities

  • The student learns the teacher decoder’s probabilities as soft labels, bringing plausible candidate embeddings closer while keeping irrelevant words far.The right panel of Figure 5 is an illustration; measurements on real embeddings are reported in Appendix A.4.Source: §4 · p. 6, §4.1 · p. 7, Figure 5 · p. 7
  • Replacing the soft-label KL objective with one-hot CE or embedding MSE leaves students below real-text entropy 5.43; a randomly initialized student cannot converge.Students initialized from all 18 teacher layers to the 9 used in the method match or outperform the teacher, so the gain is not due to reduced capacity.Source: §4.1 · p. 7
  • The student encodes and decodes WikiText, but distillation lowers SST-2 classification performance from 89.4 to 78.1.The student was distilled on OWT; the paper reports this discriminative-performance cost alongside increased diffusibility.Source: §4.1 · p. 7

4.2 Distilled embeddings are easier to generate

  • Across 64 sequences, teacher/student uncertain positions are 142/42 per sequence, with 109/12 also invalid.The student’s generated embeddings lie among close candidate embeddings; teacher samples contain more uncertain and invalid embeddings.Source: Figure 6 · p. 8, §4.2 · p. 8
  • In round-trip denoising, the student returns valid embeddings closer to the original sentence, while the teacher has larger embedding error and changes the sentence.The student’s returned sentence has meaningful fluctuations.Source: §4.2 · p. 8

5 Main Experiments

Across OWT-1024 and LM1B, ELF models with distilled embeddings improve over raw T5Gemma-2 embeddings; the section also evaluates sampling behavior and embedding-scaling costs.

5 Main Experiments

  • Evaluation trains ELF-B and ELF-M on OWT-1024 and LM1B, selecting the lowest Gen. PPL configuration that reaches real-text entropy.Samples are re-tokenized with the GPT-2 tokenizer; the authors report that tokenization effects are negligible and sequences are comparable.Source: §5 · pp. 8–9, §5.2 · p. 9

5.1 Generation Performance on OWT-1024

  • Across different sampling configurations, distilled embeddings consistently achieve lower Gen. PPL at higher entropy than the T5Gemma-2 teacher; it collapses at extremely low NFEs (16–32) from under-integration.Source: §5.1 · p. 9

5.2 Generation Performance on LM1B

  • On LM1B, the distilled student surpasses raw T5Gemma-2 on PPL, while the corpus serves as a check rather than a comparison.The authors note that a retrained autoregressive model has lower PPL than real text on this problematic corpus.Source: §5.2 · p. 9, Table 2 · p. 11

5.3 Efficiency

  • ELF-M generates a 1024-token sequence in 1.5 s at NFE 64 and 11.8 s at NFE 512, compared with GPT-2-M’s 10.7 s for 1024 serial steps.The student halves T5Gemma-2 encoding cost; embeddings are frozen and can be cached, and affect sampling only at final decoding.Source: §5.3 · p. 10, Table 3 · p. 11

6 Conclusion

The conclusion argues that scaling improves continuous-DLM diffusibility but that scaled embeddings may remain difficult to generate perfectly. Distillation further improves diffusibility, while practical-scale tasks and few-step generation remain open.

6 Conclusion

  • Scaling improves diffusibility, but strong embeddings can still be hard for continuous diffusion to generate perfectly.The paper defines diffusibility in terms of generating valid, decodable embeddings from the same DLM.Source: §6 · p. 11
  • Distillation trades embedding discriminative power for generation, favoring spaces where similar embeddings are more connected.The paper presents this connectivity as beneficial for diffusion.Source: §6 · p. 11
  • Few- or one-step generation, practical-scale tasks such as QA, and adapting AR models directly into DLM embeddings remain open.The paper notes that raw T5Gemma-2 under-integrates at low NFE and is itself adapted from Gemma-3.Source: §6 · p. 11

Ethics statement

The work uses public corpora and models and involves no human subjects or private data. Its unfiltered web-style generations may be false or biased and are used only for evaluation.

Ethics statement

  • The study uses public corpora and models, with no human subjects or private data.Generated unfiltered web-style text may be false or biased; samples are used only for evaluation.Source: Ethics statement · p. 11

Reproducibility statement

The authors adapt public ELF-pytorch code, changing only the embedding model while retaining the training and sampling recipe.

Reproducibility statement

  • The continuous-DLM implementation changes only the embedding model; training and sampling follow the public ELF-pytorch recipe.Distillation details, evaluation protocol, sampling grid, and table-row selection are described in appendices B.2 and B.3; baselines use authors’ code and checkpoints when available.Source: Reproducibility statement · p. 12

Disclosure of AI use

The authors disclose using generative AI across implementation, experimentation, data preparation, literature search, qualitative analysis, and figure and text preparation. They state that the research ideas originated with the human authors and that they checked AI-assisted code and results.

Disclosure of AI use

  • Generative AI supported implementation, experiment design feedback, dataset cleanup, literature search, and qualitative analysis.The authors also used it to create and edit figures and suggest parts of the text.Source: Disclosure of AI use · p. 12
  • The authors say the research ideas originated with them, and they checked AI-assisted code and re-checked results against external baselines.They state that they take responsibility for the final content and artifacts.Source: Disclosure of AI use · p. 12
  • They report not using generative AI to create synthetic datasets, formulate mathematical claims, write proofs, or assist with translation.Source: Disclosure of AI use · p. 12

A.1 Statistical results on the generated embeddings

The appendix examines generated-embedding errors, decoder robustness, distillation choices, and how distillation changes the embedding space. It also reports training and distillation details for the experiments.

A.1 Statistical results on the generated embeddings

  • The teacher’s uncertain generated embeddings were farther from candidate words than the student’s, and teacher sequences with greater median distance had worse Gen. PPL.Among uncertain positions, the teacher averaged 142 per sequence, including 109 invalid; the student averaged 42, including 12 invalid, though the analysis inspected only uncertain positions.Source: §A.1 · p. 18

A.2 Candidates are robust to different decoders

  • The pretrained decoder’s top-3 overlapped the deployed head’s top-3 on 99.3% of real-text positions, and their uncertainty correlated at 0.72 on generated text.The comparison covered 764 real-text and 3,702 generated-text positions.Source: §A.2 · p. 19

A.3 Ablation on distillation

  • At real-text entropy, the soft-label student reached Gen. PPL 31.2, while MSE and hard-label CE students did not reach that entropy.Students with 12–18 layers achieved 31–34 Gen. PPL at entropy 5.05, versus 46 for the teacher; the 18-layer student still beat it.Source: §A.3 · p. 19

A.4 How distillation changes the embeddings

  • With 10% of context tokens replaced, candidate embeddings moved to 1.19 nn for T5Gemma-2 but stayed at 0.97 nn for the student.On clean text, the two encoders were close at 0.96 and 0.93 nn, respectively.Source: §A.4 · pp. 20–21
  • On WikiText-103, teacher and student top-1 token accuracy was 99.3 for both; on held-out arXiv papers, it was 99.3 versus 99.4.The frozen student encoder scored lower on linear probes for SST-2, STS-B, and MRPC; Table 4 frames the trade-off between reconstruction and discriminative power.Source: §A.4 · p. 21, Table 4 · p. 21

B.1 Training Details

  • The main ELF-B and ELF-M experiments trained on 9.0M packed OWT sequences for five epochs, following ELF’s training design.The encoder comparison instead used a 25% OWT subset with 512-token sequences and trained ELF-B for 15 epochs.Source: §B.1 · p. 21

B.2 Distillation details

  • The student was initialized from nine evenly spaced T5Gemma-2 layers and trained against teacher-decoder soft targets while the teacher encoder and decoder stayed frozen.The targets retained and renormalized the teacher’s top-256 logits; the student was trained with token-wise KL divergence implemented as cross-entropy.Source: §B.2 · pp. 21–22

B.3 Evaluation Details

The evaluation defines a shared tokenization and sampling procedure for comparing generation quality and entropy. It also describes embedding visualizations and provides qualitative generation examples.

B.3 Evaluation Details

  • OWT reference text has Gen. PPL 15.4 and unigram entropy 5.43 under GPT-2-Large tokenization, which is also applied to generated samples.The operating point is the lowest Gen. PPL among configurations reaching entropy 5.43; if none does, the highest-entropy point is reported.Source: §B.3 · p. 22
  • The main evaluation excludes NFE below 64 because samples collapse in the cited model settings, although the teaser includes those points.At NFE 16–32, T5Gemma-2+ELF-B has Gen. PPL 77–335 and T5-small+ELF-M has 42–84.Source: §B.3 · p. 22
  • Truncating real-text windows by 5% changed unigram entropy by only 0.031; the observed 1–2% token-length difference contributed less than 0.01 to reported entropy.Both real and generated text were re-tokenized with GPT-2-Large before computing unigram entropy.Source: §B.3 · p. 23

B.4 Embedding-space Visualization Details

  • Embedding panels show generated embeddings alongside top-3 candidate embeddings, with distances normalized by each space’s median real-embedding nearest-neighbor distance.The qualitative examples include T5Gemma-2 samples that collapse into a language mix at low NFE, while distilled-student samples stay fluent.Source: §B.4 · p. 23, C Generated Examples · p. 24

C.1 OWT-1024

The provided OWT-1024 and LM1B excerpts show the kinds of text generated in the paper’s qualitative examples.

C.1 OWT-1024

  • OWT-1024 excerpts include mixed-language fragments and long passages with implausible or corrupted claims.Source: §C.1 · pp. 25–27

C.2 LM1B

  • LM1B excerpts combine partial statements with abrupt shifts among topics such as politics, economics, and news.Source: §C.2 · p. 27

Related papers

Loading 2610.01016v1…