Technical Summary: SMOOTHIE (Smoothing Diffusion on Token Embeddings)

Problem Statement

Diffusion models have achieved state-of-the-art results in continuous domains like images and audio but face significant challenges when applied to text generation due to the discrete nature of language. Existing approaches generally fall into two categories, each with distinct limitations:

1. Gaussian Diffusion in Latent Spaces: These methods apply Gaussian noise to continuous token embeddings. While they preserve semantic structure (based on the Euclidean semantic space hypothesis), they struggle with the discrete decoding step, often requiring an explicit mapping from continuous vectors back to tokens which can be unstable or lossy.
2. Categorical/Simplex Diffusion: These methods operate directly on the vocabulary probability simplex or categorical distributions. While they respect the discrete nature of text and eliminate the need for explicit decoding, they typically disregard semantic relationships between tokens during the noising process. This leads to a more erratic and less meaningful degradation of information.

The core problem is the lack of a diffusion framework that simultaneously respects the discrete nature of text and progressively removes semantic information in a way that maintains semantic coherence during the forward process.

Methodology: SMOOTHIE

The authors propose SMOOTHIE (Smoothing Diffusion on Token Embeddings), a novel framework that bridges the gap between continuous and discrete diffusion by smoothing token embeddings based on semantic similarity.

1. Forward Diffusion Process (Noising)

Instead of adding noise directly to embeddings or categorical probabilities, SMOOTHIE represents each token $w_i$ in a target sequence as a vector of negative squared Euclidean distances between its embedding and the embeddings of all tokens in the vocabulary:
$$D_0 = \left( -\|E_{w_i} - E_j\|^2 \right)_{i,j=1}^{m,V}$$
where $E$ is a fixed embedding matrix (e.g., from BERT), $m$ is sequence length, and $V$ is vocabulary size.

The forward process gradually perturbs these distance vectors by adding Gaussian noise:
$$q(D_t | D_0) = \mathcal{N}\left(D_t \mid \frac{1}{\sigma_t^2}D_0, \delta^2 I\right)$$
Here, $\sigma_t$ is a noise scheduler controlling the bandwidth, and $\delta$ controls stochasticity. The noisy vector $D_t$ is converted into a probability distribution over the vocabulary using a softmax function: $p_t = \text{softmax}(D_t)$.

Key Insight: As noise increases (and $\sigma_t$ grows), the probability mass initially centered on the correct token gradually distributes to other tokens. Crucially, because the input is based on Euclidean distances, the probability mass first spreads to semantically similar tokens (those with smaller Euclidean distances) before reaching distant tokens. This mimics a "heat dissipation" principle, ensuring semantic information is removed progressively rather than chaotically.

2. Reverse Diffusion Process (Denoising)

The reverse process is modeled as a Markov chain where a neural network $f_\theta$ predicts the original distance vector $D_0$ (or a scaled version) from the noisy input $p_t$. The model is trained to minimize the Mean Squared Error (MSE) between the predicted and actual distance vectors:
$$L(\theta) = \mathbb{E}[\|D_0 - g_\theta(p_t, t)\|^2]$$
The authors prove that predicting the distance vector is equivalent to predicting the expected embedding, allowing the model to output a clean embedding estimate which is then converted back to a token via nearest-neighbor search or argmax.

3. Sampling and Solver

The authors introduce a specific sampling scheme (Eq. 9) that updates the latent state by applying the forward process directly to the model's prediction, rather than relying on the previous latent state as in standard DDPM solvers.
$$D_{t-1} = \frac{1}{\sigma_{t-1}^2}D_0(f_\theta(p_t, t)) + \tilde{\delta}\epsilon$$
This approach ignores the previous hidden state during the update step. The authors argue this is beneficial for the sparse data distribution of text embeddings, reducing the risk of the latent representation drifting off the data manifold.

4. Architecture and Conditioning

• Architecture: Based on Transformer decoder layers with UNet-style skip connections.
• Conditioning: For sequence-to-sequence tasks, an additional Transformer encoder processes the source sequence, integrated via cross-attention.
• Self-Conditioning: The model utilizes self-conditioning, feeding its own previous prediction as an input to subsequent steps to improve stability and quality.

Key Contributions

1. Novel Diffusion Framework: SMOOTHIE is the first text diffusion method that simultaneously respects the discrete nature of text and progressively removes semantic information by smoothing token embeddings based on Euclidean distances.
2. Generalization of Simplex Diffusion: The authors demonstrate theoretically that their distance-based latent space generalizes the simplex-based latent space (used in SSD-LM and TESS). Simplex diffusion is shown to be a special case of SMOOTHIE where the distance metric is trivial (0 for same token, 1 otherwise).
3. Improved Sampling Solver: The paper proposes a solver that ignores the previous hidden state during updates, which empirically outperforms standard DDPM and DDIM solvers on sequence-to-sequence tasks.
4. Empirical Superiority: Extensive experiments show SMOOTHIE outperforms existing diffusion-based models (including embedding-based and simplex-based variants) and achieves quality comparable to autoregressive models (like BART and GPT-2) across multiple tasks.

Experimental Results

The authors evaluated SMOOTHIE on one unconditional generation task (ROCStories) and four sequence-to-sequence tasks (Quora Question Pairs, Quasar-T, XSum, ParaDetox).

• Performance: SMOOTHIE consistently outperformed other diffusion-based baselines (DiffuSeq, TESS, SSD-LM, TEncDM, etc.) and alternative latent space implementations (Embedding-based and Simplex-based diffusion trained within the same framework).
• Metrics:
  • XSum: Achieved a BERTScore of 68.8 and ROUGE-L of 26.0, outperforming other diffusion methods and approaching autoregressive baselines.
  • ROCStories: Achieved the highest MAUVE score (80.7) among diffusion models, nearly matching GPT-2 (78.9), while maintaining high diversity.
  • ParaDetox: Achieved the highest J-Score (51.7) for text detoxification.
• Ablation Studies:
  • $\tilde{\delta}$ (Stochasticity): Lower values of $\tilde{\delta}$ (e.g., 0.1) improved perplexity and coherence for sequence-to-sequence tasks, while slightly higher values were optimal for unconditional generation to maintain diversity.
  • Solver: The proposed solver (Eq. 9) consistently outperformed DDPM and DDIM across all tasks.
  • Embeddings: The model relies on pre-trained BERT embeddings. Ablation showed that while GPT-2 embeddings performed similarly, GloVe embeddings resulted in lower quality, highlighting the importance of embedding quality.

Significance and Claims

The paper claims that SMOOTHIE represents a significant step forward in text diffusion by resolving the trade-off between semantic preservation and discrete structure. By incorporating semantic relationships directly into the noise addition process via Euclidean distances, the method enables a "natural decoding process" where information is removed gradually and meaningfully.

The authors modestly note that while the method is applicable to other categorical domains with inherent similarities (e.g., graphs, protein sequences), it currently relies on pre-trained embeddings rather than end-to-end training, which remains a challenge due to stability issues with MSE loss. They also acknowledge the computational bottleneck of calculating pairwise distances at every step but argue that the performance gains justify the cost, and that the method offers a robust alternative to autoregressive models for tasks requiring high diversity and controllability.

The work suggests that the choice of latent space and the specific diffusion solver are as critical as the model architecture itself, with SMOOTHIE providing a unified framework that leverages the strengths of both continuous and discrete diffusion approaches.