💬 NLP

Where Should Diffusion Enter a Language Model? Geometry-Guided Hidden-State Replacement

The paper introduces DiHAL, a geometry-guided hybrid model that improves upon continuous diffusion language models by identifying optimal layers within pretrained transformers to replace with a diffusion bridge for hidden-state reconstruction, thereby avoiding direct continuous-to-discrete token recovery and achieving superior performance.

Original authors: Injin Kong, Hyoungjoon Lee, Yohan Jo

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Injin Kong, Hyoungjoon Lee, Yohan Jo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Trying to Fix a Broken Radio

Imagine you have a very smart radio (a Large Language Model) that is great at predicting the next word in a sentence. It works by listening to a sequence of words and guessing what comes next, one by one. This is called an "autoregressive" model.

Recently, scientists tried to use a different kind of technology called Diffusion (the same tech that creates amazing AI images) to make these radios work better. In image diffusion, the AI starts with a picture full of static noise and slowly cleans it up until a clear image appears.

However, when researchers tried to do this with text, it didn't work well. Text is made of distinct blocks (like Lego bricks), while diffusion works best with smooth, continuous water. Trying to turn "noise" directly into "Lego bricks" is messy. The AI often gets confused, and the final words come out wrong.

The New Idea: Don't Fix the Bricks, Fix the Blueprint

The authors of this paper, Injin Kong and colleagues, asked a different question: "Where inside the radio should we let the diffusion technology in?"

Instead of trying to clean up the final Lego bricks (the words), they decided to let diffusion clean up the blueprint (the hidden internal thoughts) that the radio uses before it decides on the final words.

They call their new system DiHAL (Diffusion-Transformer Hybrid Architecture for Language). Think of it as a "Locate-and-Replace" strategy.

Step 1: The "Geometry" Detective

The paper argues that not all parts of the radio's brain are the same. Some parts are chaotic and hard to clean; others are organized and easy to fix.

To find the best spot, the authors created a "Geometry Score." Imagine the radio's brain as a mountain range:

Curvature (The Slope): Some areas are steep and smooth (easy to slide down to the right answer). Others are flat or jagged (hard to navigate).
Stiffness (The Structure): Some areas are rigid and hold their shape well. Others are wobbly.
Dimension (The Complexity): Some areas are simple, like a straight hallway. Others are a massive, confusing maze with too many dead ends.

The authors built a tool to measure these "geometric" features at every layer of the radio. They found that the early layers (near the beginning of the processing) are like a smooth, organized hallway. The later layers are like a complex, tangled maze.

The Discovery: Diffusion works best in the smooth, organized hallways near the start. It struggles in the tangled mazes at the end.

Step 2: The "Locate-and-Replace" Surgery

Once they found the perfect spot (usually the 2nd or 3rd layer of the radio), they performed a surgical swap:

Locate: They identified the specific layer where the "blueprint" is easiest to clean.
Replace: They removed the original early layers of the radio and replaced them with a Diffusion Bridge.
Keep: They kept the rest of the radio (the upper layers and the final speaker) exactly as it was.

How it works in practice:

The Diffusion Bridge takes the noisy input and slowly cleans it up to recreate the "blueprint" (the hidden state) that the original radio would have had at that specific layer.
Once the blueprint is clean, it is handed back to the original, untouched upper layers of the radio.
The original radio then finishes the job, turning that clean blueprint into the final words.

Why This is Better

The paper tested this on two massive 8-billion-parameter models (Llama-3 and Qwen3).

The Old Way: Trying to clean up the final words directly is like trying to fix a broken watch by hammering the gears. It's hard and often breaks the time.
The DiHAL Way: It's like taking the watch apart, cleaning the main spring (the hidden blueprint) with a gentle, precise tool, and then putting the watch back together. Because the upper layers of the radio are already experts at reading that specific blueprint, they can decode the clean signal perfectly.

The Results

The experiments showed that:

The Geometry Score Works: Their "detective tool" successfully predicted which layers were best for this surgery without needing to try every single one.
Better Quality: The new hybrid model produced text that was more accurate (lower "perplexity") and more diverse than other diffusion methods that tried to fix the words directly.
It's a Hybrid: It's important to note that DiHAL isn't a pure diffusion model. It's a mix. It uses diffusion to fix the early part of the process, but it still relies on the original, powerful transformer layers to do the final thinking and word selection.

The Bottom Line

The paper concludes that the reason diffusion struggles with text isn't just because text is "discrete" (like Lego bricks). It's because we were trying to apply diffusion in the wrong "room" of the AI's brain. By finding the room with the right geometry (smooth, organized, and simple) and letting diffusion clean up the internal blueprint there, we can get much better results without rebuilding the whole AI from scratch.

Technical Summary: DiHAL – Geometry-Guided Hidden-State Replacement

Problem Statement

Continuous diffusion language models have historically lagged behind autoregressive transformers. While diffusion models excel in image generation through iterative denoising, applying them to text is challenging due to the discrete nature of tokens. Prior approaches attempting to adapt diffusion to language—such as discrete token corruption, masked diffusion, or continuous diffusion over token embeddings and learned latents—often struggle with the "continuous-to-discrete" recovery bottleneck. Small errors in the continuous representation space can lead to incorrect token recovery, degrading generation quality.

The authors hypothesize that the performance gap is not solely due to discreteness but stems from applying diffusion in continuous spaces with unsuitable geometry. While Transformer language models operate on continuous hidden states, these states vary in their geometric properties across layers. The central question posed is: Where inside a pretrained Transformer does the representation space possess the geometry required for effective diffusion-based denoising and recovery?

Methodology: DiHAL

The paper proposes DiHAL (Diffusion-Transformer Hybrid Architecture for Language Generation), a "Locate-and-Replace" framework that integrates a diffusion bridge into a pretrained Transformer rather than training a standalone diffusion model.

1. Theoretical Motivation

The authors ground their approach in Langevin dynamics and concentration theory. They argue that a "diffusion-friendly" representation space must satisfy three geometric properties:

Fast Contraction: The space should allow noisy samples to contract quickly toward the target distribution (governed by strong log-concavity/curvature).
Stability: The denoising process should be robust to score-estimation errors (also governed by curvature).
Low Effective Complexity: The variation in the representation should be concentrated in a low-dimensional manifold, making it easier for the diffusion model to learn.

2. Locate: Geometry-Guided Layer Selection

Since the true density and Hessian of hidden states are unavailable, DiHAL uses empirical geometric proxies to score each Transformer layer:

Local Curvature ( $\hat{m}_{curv}$ ): Estimated via the inverse of the maximum eigenvalue of local neighborhood covariances. Higher values indicate compact neighborhoods.
Global Monotonicity ( $\hat{m}_{mono}$ ): Estimated via the Rayleigh quotient of the regularized precision matrix. Higher values indicate global directional stiffness.
Effective Rank ( $\hat{k}$ ): Estimated as the trace of the covariance divided by its spectral norm ( $\text{tr}(\Sigma)/\|\Sigma\|$ ). Lower values indicate concentration on a low-dimensional manifold.

A selection score is computed by combining these proxies (rewarding curvature and penalizing effective rank) to identify the optimal insertion layer $\ell^*$ .

3. Replace: Hidden-State Diffusion Bridge

Once the optimal layer $\ell^*$ is identified:

The lower Transformer prefix (layers $1$ to $\ell^*$ ) is replaced by a conditional diffusion bridge.
The bridge is trained to reconstruct the hidden state $h_{\ell^*}$ given an embedding-based condition derived from the source model's input.
The upper Transformer layers ( $\ell^*+1$ to $N$ ) and the original Language Model (LM) head are retained and frozen (or fine-tuned minimally).
At inference, the bridge denoises the condition to produce $\hat{h}_{\ell^*}$ , which is then passed to the retained upper layers for token prediction. This avoids direct token recovery, reframing the problem as hidden-state reconstruction.

Key Contributions

Formulation of Diffusion Insertion: The paper formulates the insertion of diffusion into pretrained Transformers as a geometry-guided interface-selection problem, proposing practical layer-wise proxies (local compactness, global stiffness, effective rank) to identify diffusion-friendly spaces.
Fixed Geometry Score: The authors introduce a fixed geometry score that correlates strongly with bridgeability (reconstruction quality) across 8B-scale backbones. This allows for the identification of effective insertion layers without exhaustive training of bridges for every candidate layer.
DiHAL Architecture: A hybrid architecture that replaces lower Transformer layers with a conditional diffusion bridge while reusing the upper layers and LM head. Under a diagnostic matched-budget setting, DiHAL demonstrates that hidden-state recovery outperforms embedding-based and latent-based continuous diffusion baselines.

Experimental Results

The authors evaluated DiHAL on Llama-3.1-8B-Instruct and Qwen3-8B backbones using the Dolma v1.7 corpus.

Layer Selection: The geometry score successfully predicted effective insertion layers. For Llama-3.1-8B, the score selected Layer 3; for Qwen3-8B, Layer 2. These layers are near the embedding interface, suggesting that early hidden states retain embedding-like geometric structure suitable for diffusion.
Correlation with Bridgeability: The fixed geometry score showed strong Spearman correlation ( $\rho \approx 0.91$ –$0.93$) with validation bridge loss, confirming that geometric proxies can cheaply identify layers where diffusion bridges can be successfully trained.
Diagnostic Comparison: In a matched-budget comparison against baselines (Diffusion-LM, SED, LD4LG, CoDAR), DiHAL achieved the lowest generative perplexity (Gen.PPL) and highest diversity. Specifically, on Qwen3-8B, DiHAL improved Gen.PPL from 144.83 (CoDAR) to 136.02 and diversity from 0.4777 to 0.5913.
Full Training: When fully trained, the geometry-selected layers remained competitive with validation-loss oracles (layers selected by explicitly training all bridges) and significantly outperformed "worst-layer" controls and the CoDAR baseline.

Significance and Claims

The paper claims that hidden-state geometry is a critical factor in determining where diffusion can effectively operate within a language model. The results suggest that the failure of continuous diffusion in language is not merely a consequence of discreteness but also a mismatch between denoising requirements and representation-space geometry.

Key takeaways include:

Reframing the Problem: Continuous diffusion for language need not generate text by directly recovering tokens; it can reconstruct an internal representation that the pretrained Transformer already knows how to decode.
Feasibility of Hybrid Models: It is feasible to replace parts of a pretrained Transformer with a diffusion bridge if the insertion point is chosen based on geometric criteria.
Limitations: The authors modestly note that DiHAL is not yet a standalone diffusion language model; it relies on the retained Transformer suffix for token prediction. Furthermore, inference latency is currently higher than the original backbone due to the denoising steps. The work is presented as a step toward understanding where diffusion can enter existing models, rather than a complete replacement for autoregressive generation at this scale.