PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Injin Kong, Hyoungjoon Lee, Yohan Jo

arXiv:2605.14368v1cs.CLcs.AI

TL;DR

Continuous diffusion language models face a representation-space problem in addition to discrete-token recovery. DiHAL selects diffusion-friendly hidden-state interfaces geometrically and replaces a lower transformer prefix with a bridge that reconstructs hidden states for the retained suffix. Across 8B-scale backbones, the geometry score predicted effective shallow insertion layers under fixed-budget training, while the paper reports improved hidden-state recovery in its diagnostic comparison.

  • Problem

    Continuous diffusion language models lag behind autoregressive Transformers, and diffusion must handle unsuitable continuous spaces and recovery to discrete language.

  • Method

    DiHAL scores transformer layers with geometric proxies, selects a hidden-state interface, and replaces lower layers with a conditional diffusion bridge while retaining upper layers and the original LM head.

  • Results

    Across two 8B-scale backbones, geometry-selected insertion points predict bridgeability under fixed-budget training and remain competitive with validation-loss oracles after full training.

  • Takeaways & Limitations

    Hidden-state geometry can help identify feasible locations for diffusion-based replacement inside pretrained language models.

  • Takeaways & Limitations

    DiHAL is not a standalone diffusion language model because token prediction still depends on the retained transformer suffix and LM head.

Abstract

from arXiv · show

Continuous diffusion language models lag behind autoregressive transformers, partly because diffusion is applied in spaces poorly suited to language denoising and token recovery. We propose DiHAL, a geometry-guided diffusion-transformer hybrid that asks where diffusion should enter a pretrained transformer. DiHAL scores layers with geometry-based proxies, selects a diffusion-friendly hidden-state interface, and replaces the lower transformer prefix with a diffusion bridge while retaining the upper layers and original LM head. By reconstructing the selected-layer hidden state rather than tokens, DiHAL avoids direct continuous-to-discrete recovery. Experiments on 8B-scale backbones show that the geometry score predicts effective shallow insertion layers under a fixed bridge-training protocol and that hidden-state recovery improves over continuous diffusion baselines in a diagnostic comparison matching the diffusion/recovery training budget. These results suggest that hidden-state geometry helps identify where diffusion-based replacement is feasible inside pretrained language models.

1 Introduction

DiHAL argues that diffusion’s difficulty in language may reflect unsuitable continuous-space geometry rather than continuity alone. It selects diffusion-friendly hidden-state interfaces and replaces lower transformer layers with a diffusion bridge while retaining upper layers and the LM head.

  • Diffusion language models still lag behind autoregressive Transformers, especially in continuous diffusion settings.
  • The paper hypothesizes that unsuitable geometry in continuous representation spaces contributes to diffusion’s difficulty in language.
  • A diffusion-friendly space should be easy to denoise, stable under imperfect score estimates, and simple enough for diffusion to learn.
  • DiHAL scores transformer layers with geometric proxies and replaces the lower transformer prefix with a diffusion bridge at a selected hidden-state interface.
  • The bridge reconstructs selected-layer hidden states while the retained upper layers and original LM head produce token predictions.

2 Background

Language models are discrete at their token interfaces but perform most computation in continuous hidden spaces. Continuous token-level diffusion can suffer recovery errors, motivating diffusion over internal hidden states instead of direct token decoding.

  • Autoregressive language models map discrete tokens into embeddings, transform continuous hidden states, and project final states to vocabulary logits.
  • Diffusion generates samples by adding noise and learning a reverse denoising process, which requires choosing a representation to denoise in language.
  • Continuous diffusion language models denoise token embeddings or learned latent vectors, while discrete diffusion models corrupt tokens directly.
  • Small continuous denoising deviations can flip recovered tokens, creating recovery errors at the token-level interface.
  • Targeting internal transformer hidden states changes the objective from direct token decoding to hidden-state reconstruction.

3 Method

DiHAL motivates diffusion-friendly interfaces with geometric proxies and uses a conditional bridge to reconstruct a selected hidden state. The method treats its theoretical conditions as qualitative guidance rather than guarantees for transformer activations.

  • 3 Method: DiHAL replaces part of a pretrained transformer with a conditional diffusion bridge while retaining the upper layers and original LM head.
  • 3.1 Geometric Principles for Diffusion-Friendly Layer Selection: A diffusion-friendly representation should support rapid denoising contraction, stability under score-estimation error, and low effective complexity.
  • 3.1 Geometric Principles for Diffusion-Friendly Layer Selection: Under strong log-concavity, larger curvature yields faster convergence and greater stability when score estimates are imperfect.
  • 3.1 Geometric Principles for Diffusion-Friendly Layer Selection: Effective rank measures the effective number of active variance directions and serves as a proxy for intrinsic dimensionality.
  • 3.1 Geometric Principles for Diffusion-Friendly Layer Selection: The theoretical results motivate geometric desiderata but are not guarantees that transformer hidden states satisfy the idealized assumptions.
  • 3.2 Locate: The selection score rewards curvature and monotonicity proxies while penalizing effective rank, and bridgeability is measured by validation loss under matched training.
  • 3.3 Replace: Hidden-State Diffusion Module: The bridge reconstructs the selected-layer hidden state from an embedding-derived condition, after which retained upper layers and the LM head produce token probabilities.

4 Experiments

Experiments on two 8B-scale backbones test whether geometry identifies bridgeable hidden-state interfaces and whether hidden-state diffusion improves diagnostic generation quality. The fixed score predicts effective shallow layers under matched bridge-training budgets, while DiHAL remains competitive in final evaluation.

  • Experimental Setup: Experiments use Llama-3.1-8B-Instruct and Qwen3-8B, estimating layer-wise geometry from hidden states and testing candidate layers with one-epoch bridge training.The sweep trains one bridge per candidate on 150K examples and measures validation bridge loss without fitting the geometry score.
  • Layer-Wise Geometry: Layer geometry varies systematically with depth: input-adjacent layers show larger local curvature, while effective rank and global monotonicity follow different trends.Selection therefore balances curvature, global stiffness, and representation complexity rather than maximizing one proxy.
  • Layer-Wise Geometry: The geometry score selects layer 3 for Llama-3.1-8B and layer 2 for Qwen3-8B, placing diffusion near the embedding interface.These layers are treated as candidate hidden spaces that preserve embedding-like structure while remaining easier to denoise than token embeddings.
  • Fixed-Budget Layer Sweep: Spearman ρ = 0.9143 ± 0.0069 on Llama-3.1-8B and ρ = 0.9267 ± 0.0157 on Qwen3-8B, with rank gaps of 2 and 1, respectively.The correlations use repeated score-estimation runs and indicate that the fixed score identifies layers close to the best observed bridgeability.
  • Diagnostic Matched-Budget Comparison: DiHAL improves diagnostic generation over CoDAR, reducing Gen.PPL from 144.83 to 136.02 and increasing diversity from 0.4777 to 0.5913.The comparison uses the same diffusion/recovery training budget, while DiHAL reconstructs an internal hidden state and retains the pretrained suffix and LM head.
  • Top-Layer Full Training and Evaluation: In final evaluation, the geometry-selected layer improves over the worst-layer control and remains competitive with the validation-loss oracle, but its relationship to the autoregressive teacher defines DiHAL’s scope.It outperforms the oracle in NLL and PPL on Llama-3.1-8B; on Qwen3-8B, the oracle is slightly better while the geometry-selected layer remains comparable and improves over CoDAR.

5 Related Work

Related work adapts diffusion to text through discrete token corruption, masked refinement, and continuous diffusion over embeddings or learned latents. DiHAL instead targets internal transformer hidden states, connecting diffusion geometry with hybrid transformer replacement.

  • Diffusion Language Modeling: Diffusion language modeling includes discrete token diffusion, masked iterative refinement, and continuous diffusion over embeddings or learned latents.These approaches adapt diffusion to text’s discrete generation problem through different corruption and representation choices.
  • Diffusion Language Modeling: Continuous methods can introduce projection or decoding errors when recovering tokens from denoised vectors, whereas DiHAL studies transformer hidden states as the denoising space.The hidden-state target changes recovery from direct token decoding to reconstruction consumed by the remaining transformer.
  • Geometry and Efficient Generation: DiHAL connects representation-geometry analysis with efficient-generation methods by identifying internal transformer representations suitable for diffusion-based replacement.Its position combines geometry-dependent diffusion behavior with hybrid modules and transformer reduction strategies.

6 Limitations

DiHAL is a hybrid replacement inside a pretrained transformer, not a standalone diffusion language model, and its evaluation is constrained by compute-limited exploration of larger or deeper replacements.

  • Scope: DiHAL still depends on the retained transformer suffix and LM head for token prediction rather than operating as a standalone diffusion language model.The method moves diffusion to an internal hidden-state interface but does not replace the full language model.
  • Scope: Compute constraints limit exploration of larger bridges, longer training, and deeper transformer-prefix replacement.The paper identifies incorporating geometric proxies into bridge training as a possible route toward deeper replacement.

7 Conclusion

DiHAL reframes continuous diffusion as internal hidden-state reconstruction: it locates diffusion-friendly layers, replaces the lower transformer prefix with a bridge, and preserves upper layers and the original LM head. Experiments indicate that interface geometry matters, while the method remains dependent on retained pretrained components.

  • Conclusion: DiHAL locates diffusion-friendly layers with geometry proxies, replaces the lower transformer prefix with a conditional diffusion bridge, and preserves the upper layers and original LM head.The bridge reconstructs a selected hidden representation that the pretrained transformer can decode.
  • Conclusion: Across two 8B-scale backbones, geometry-selected insertion points are embedding-adjacent and predict bridgeability under fixed-budget training.They remain competitive with validation-loss oracles after full training.
  • Conclusion: Middle and late hidden states are much harder to reconstruct, linking diffusion failures to representation-space geometry as well as discreteness.
  • Conclusion: DiHAL is not yet a standalone diffusion language model and still relies on retained pretrained layers.This bounds the current architecture’s scope while motivating future diffusion–transformer hybrids.

A Proofs of Theorems

The proofs establish invariant-measure existence, contraction, uniqueness, and stability under score perturbations using Langevin dynamics, couplings, strong convexity, and Lyapunov arguments.

  • Proofs of Theorems: Strong convexity supplies gradient monotonicity, which supports contraction estimates for synchronously coupled Langevin processes.
  • Proofs of Theorems: The Gibbs measure is invariant under the stated smoothness and convexity assumptions.
  • Proofs of Theorems: W2(ν, µ) ≤e−mtW2(ν, µ) for all t ≥0 implies ν = µ for any invariant ν when m > 0.
  • Proofs of Theorems: The unique invariant distribution is established by applying the contraction estimate to an arbitrary invariant distribution.
  • Proofs of Theorems: A uniformly bounded score perturbation is controlled through synchronous coupling, strong convexity, Lyapunov bounds, and concentration arguments.The proof uses truncation to justify the unbounded Lyapunov function.
  • Proofs of Theorems: Under bounded approximation terms and a covariance norm bounded below, the effective-rank and concentration bounds remain controlled.

B Interpretation of the Geometric Proxies

The geometric proxies operationalize theoretical ideas about curvature and intrinsic dimension using observable layer activations. They are theory-motivated surrogates rather than exact estimates, computed from pooled or tokenwise representations with optional projection.

  • Interpretation: The empirical proxies target the functional roles of curvature, stability, and intrinsic dimension rather than recovering exact theoretical constants.
  • Curvature proxy: ˆmmono summarizes typical directional stiffness using covariance precision and sampled representation pairs.
  • Curvature proxy: ˆmcurv measures local compactness and sharp curvature, with larger values indicating smaller spread along the most variable direction.
  • Intrinsic dimension: Effective rank reff(Σ) measures active variance directions and exactly recovers intrinsic dimension in the ideal isotropic subspace case.
  • Intrinsic dimension: When representations concentrate near a low-dimensional manifold, effective rank measures meaningful variation directions rather than nominal ambient dimension.
  • Implementation: Proxy estimation uses mean, last-token, or tokenwise representations, excludes padding, and may apply random projection for stability and efficiency.The main experiments use mean pooling and retain a bounded representation set.

C.7 Implementation Details of the Diffusion Bridge

The diffusion bridge converts selected hidden states into a latent-denoising problem and reconstructs them for the retained transformer layers. Its image-like layout is an interface to the denoiser, not a claim that language representations have image semantics.

  • Bridge conditioning: The bridge conditions on the source model’s pre-transformer hidden activation through a learned projection into the UNet conditioning interface.
  • Bridge pipeline: The target hidden state is reshaped, VAE-encoded, noised, denoised by a UNet, decoded, and projected back into hidden space.
  • Sensitivity analysis: Table 5 examines layer-selection sensitivity to coefficient perturbations using the final linear curvature–dimension trade-off.
  • Hidden-to-image layout: For S = 1024 and H = 4096, the hidden state uses a 32 × 32 grid with C = 3 channels for the SD-v1.5 VAE interface.
  • Design qualification: The 2D layout is an architectural interface rather than a semantic image representation, and it is compared with sequence-native backbones.
  • Training: Training updates the UNet and bridge projections while keeping the VAE frozen and the diffusion scheduler fixed.

D Experimental Setup Details

The experiments use sampled activation corpora, geometry estimation, and fixed-budget bridge training to evaluate candidate transformer insertion layers.

  • Data and representations: The activation corpus samples up to 300,000 Dolma v1.7 sequences from 700 source files using stratified round-robin sampling and seed 42.Sequences are tokenized with the source-model tokenizer, and embeddings, decoder-layer outputs, input IDs, and attention masks are saved.
  • Data and representations: Special tokens are retained in stored input IDs and activation tensors, while padding positions are excluded from masked averaging.Decoded prompt reconstructions remove special tokens.
  • Data and representations: Mean pooling is the default representation for geometry estimation, with last-token pooling and token-level sampling also supported.
  • Geometry estimation: Geometry estimation uses dproj = 0, ridge coefficient 10^-3, k = 64 nearest neighbors, 512 anchor points, 200,000 sampled pairs, and 95% bootstrap confidence intervals.Layerwise proxies are estimated from repeated subsamples.
  • Bridge training: The fixed-budget layer sweep trains each candidate for one epoch on up to 150,000 examples with batch size 4, learning rate 3 × 10^-5, AdamW, FP16, and 37,500 maximum steps.Validation uses a deterministic shard-level split with ratio 0.1 and seed 42; bridgeability denotes reconstruction ease under this budget.
  • Compute: The layer sweep uses NVIDIA H100 GPUs with a matched 40 GPU-hour budget, while final full-training runs use NVIDIA B200 GPUs for 40 hours per main run.Final runs train the selected bridge for four epochs on 300,000 examples.

D.1 Diffusion Bridge Architecture and Backbone Choices

The study evaluates a Stable-Diffusion-style UNet as DiHAL’s conditional hidden-state denoising bridge and measures its parameter and inference-cost trade-offs.

  • Bridge backbone: The Stable-Diffusion-style UNet bridge is chosen to denoise high-dimensional conditional hidden-state representations at the selected internal boundary.The target is an internal hidden state h_l* conditioned on the embedding-derived representation of the same input.
  • Bridge backbone: 1163.59 mean validation loss is achieved by the Stable-Diffusion-style UNet on layer 16, 17.6% below the MLP-based hidden DDPM bridge’s 1411.36.The comparison uses S=512 and N=500 under the same validation setting.
  • Bridge backbone: 58.4% lower validation loss than ddpm_hidden_transformer and 68.3% lower than ddpm_hidden_conv1d are reported for the Stable-Diffusion-style UNet.
  • Bridge backbone: The authors characterize the UNet as practical in this setting rather than generally optimal, reusing its latent denoising structure without image supervision.
  • Parameter accounting: DiHAL’s parameter accounting distinguishes trainable parameters from active bridge or module parameters, while retained upper layers and the original LM head remain part of the active interface.
  • Inference cost: End-to-end cost measurements include bridge denoising, retained upper transformer layers, and the original LM head across insertion depths and denoising steps.
  • Inference cost: Geometry-selected shallow insertions are slower than the original backbone, and latency rises substantially with NFEs because bridge denoising dominates current overhead.Deeper insertions reduce latency at NFE=1 but are less reliable reconstruction targets; peak memory also increases.
  • Inference cost: The current implementation is not an end-to-end acceleration method because denoising overhead and hidden-state reconstruction difficulty limit practical acceleration.Deeper replacement trades retained-transformer computation against these costs.
Loading 2605.14368v1…