🤖 machine learning

Elucidating Representation Degradation Problem in Diffusion Model Training

This paper identifies "Representation Degradation" as a key training bottleneck in diffusion models caused by mismatched target recoverability and proposes Elucidated Representation Diffusion (ERD), a plug-and-play framework that dynamically reallocates optimization effort to stabilize learning and accelerate convergence.

Original authors: Zhipeng Yao, Dazhou Li, Zitong Zhang, Durude Mahee, Fan Zhu, Wenbin Zhang, Xinwei He, Yeying Jin, Rui Yu

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhipeng Yao, Dazhou Li, Zitong Zhang, Durude Mahee, Fan Zhu, Wenbin Zhang, Xinwei He, Yeying Jin, Rui Yu

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 Picture: The "Noisy Classroom" Problem

Imagine you are trying to teach a student (the AI model) how to draw a perfect picture of a cat. You don't just show them the cat; you start with a blank page and slowly add more and more static noise until the page is just a blur of white and gray.

The student's job is to look at this noisy page and guess what the original cat looked like. They do this by practicing on thousands of different "noise levels"—from a little bit of static to a lot of static.

The Problem: The authors discovered that while the student gets really good at fixing the "little bit of static" (low noise), they completely fall apart when trying to fix the "heavy static" (high noise). In fact, when the noise is heavy, the student starts hallucinating and drawing gibberish. This causes the whole training process to become unstable and inefficient.

The authors call this "Representation Degradation." It's like the student's brain starts to "melt" or lose its shape when the task gets too messy.

Why Does This Happen? (The "Broken Compass" Analogy)

To understand why the student fails, the authors looked at the math behind the learning process using something called the Neural Tangent Kernel (NTK). Think of the NTK as a compass that tells the student which direction to move to get closer to the correct answer.

The Signal vs. The Noise:
- When the noise is low, the "signal" (the actual cat) is strong. The compass points clearly, and the student learns fast.
- When the noise is high, the "signal" is drowned out. The compass becomes weak and shaky.
The Mismatch:
- The problem is that the student is forced to spend just as much time trying to fix the "heavy static" as they do the "light static."
- But here's the catch: In the "heavy static" zone, the information is so corrupted that it's mathematically impossible to recover the perfect cat. The student is trying to solve a puzzle where half the pieces are missing.
- Because the student keeps trying to force an answer in these impossible zones, they get confused. The "compass" (the math) starts pointing in random directions, dominated by pure noise rather than the actual image.
The Contamination:
- Because the student uses the same brain (parameters) for all noise levels, the confusion from the "heavy static" zone leaks over into the "light static" zone.
- It's like if a student gets frustrated and confused by a math problem they can't solve, and that frustration makes them forget how to do the easy problems they already knew. The whole learning process gets "contaminated" by the noise.

The Solution: "Elucidated Representation Diffusion" (ERD)

The authors propose a new training method called ERD. Think of this as a smart study schedule for the student.

Instead of forcing the student to spend equal time on every type of noise, ERD looks at how much of the "cat" is actually visible in the noise.

The Old Way: "Spend 1 hour on light noise, 1 hour on medium noise, 1 hour on heavy noise." (This wastes time on the heavy noise where the student can't learn anything useful).
The ERD Way: "Spend 1 hour on light noise, 1 hour on medium noise, but only 10 minutes on heavy noise."

ERD dynamically adjusts the "weight" of the training. It tells the model:

"Hey, in this heavy noise zone, the signal is too weak to learn from effectively. Let's stop wasting energy here."
"Focus your energy where the signal is actually recoverable."

By ignoring the zones where the model is likely to get confused by pure noise, the model learns faster, stays more stable, and produces better pictures.

What Did They Prove?

The paper doesn't just guess; they proved this with math and experiments:

Visual Proof: They showed that as noise increases, the model's internal "map" of the world collapses. It's like a 3D sculpture of a cat that slowly flattens into a 2D blob and then disappears.
Math Proof: They showed that the "compass" (NTK) loses its strength in high-noise areas, making it impossible for the model to learn the right direction.
Real Results: They tested this on famous AI models (like DiT and U-ViT) using ImageNet (a huge database of photos).
- Result: Their method (ERD) made the models train faster and produce higher-quality images (lower FID scores) compared to standard methods, without needing any extra hardware or complex changes to the model's architecture.

Summary

The Issue: Diffusion models get confused and "break" when trying to remove heavy amounts of noise, which ruins their ability to learn from the easy parts too.
The Cause: The model is forced to try to learn from information that is too corrupted to be recovered, causing "noise contamination."
The Fix: A new training rule (ERD) that tells the model to focus on the noise levels where learning is actually possible and ignore the levels where it's just guessing in the dark.
The Outcome: Faster training, more stable models, and better images.

Technical Summary: Elucidating Representation Degradation Problem in Diffusion Model Training

1. Problem Statement: Representation Degradation

Despite the remarkable success of diffusion models, their training remains inefficient due to a severe optimization bottleneck the authors term Representation Degradation. Empirical analysis reveals that as noise levels increase during the forward diffusion process, the outputs of trained models exhibit progressive structural distortion, eventually leading to a collapse of meaningful semantic structure.

The authors identify that this instability is not merely an architectural limitation but an intrinsic optimization issue driven by mismatched target recoverability. In high-noise regimes, corrupted inputs lose informative structure, weakening the local Neural Tangent Kernel (NTK) spectrum and reducing learnability along semantically meaningful directions. Consequently, standard training allocates substantial optimization effort to regions where the target is only weakly recoverable. This mismatch causes the model to be dominated by irreducible Bayes noise (the "Bayes floor") rather than learnable signal, leading to spectral degeneration and effective low-rank behavior in the learned representations.

2. Methodology: Elucidated Representation Diffusion (ERD)

To address this bottleneck, the authors propose Elucidated Representation Diffusion (ERD), a plug-and-play training framework that dynamically reallocates optimization effort based on effective recoverability without requiring external supervision or architectural modifications.

Theoretical Foundation

The analysis is grounded in the Neural Tangent Kernel (NTK) regime, treating diffusion training as a continuous-time, noise-conditioned regression process. The authors derive that the optimization dynamics at any noise level $\lambda$ are governed by three factors:

Allocation of optimization effort: How loss is weighted across noise levels.
Local kernel geometry: Induced by corrupted inputs.
Target recoverability: The extent to which the target signal can be recovered from the noisy input.

They formalize that when the recoverable signal is small relative to the irreducible Bayes noise, the gradient updates become dominated by noise contamination. This "Bayes-noise forcing" prevents the contraction of the excess error and degrades the learned representations.

The ERD Framework

ERD introduces a target-adaptive weighting rule derived from the effective amplitude of the recoverable signal at each noise level.

Effective Target Definition: For a generalized target $y_\lambda = c_x(\lambda)x_0 + c_\epsilon(\lambda)\epsilon$ , the authors define an effective target $\tilde{y}_\lambda$ that scales the source components ( $x_0$ and $\epsilon$ ) by their respective amplitudes in the corrupted input ( $\alpha(\lambda)$ and $\sigma(\lambda)$ ).
Recoverability Score: The recoverability score $\omega_y(\lambda)$ is defined as the expected root-mean-square (RMS) amplitude of this effective target:
$\omega_y(\lambda) = \sqrt{(c_x(\lambda)\alpha(\lambda))^2 + (c_\epsilon(\lambda)\sigma(\lambda))^2}$
Weighting Strategy: The loss weight $w(\lambda)$ is set proportional to $\omega_y(\lambda)$ . This strategy emphasizes noise levels with stronger recoverable signals and downweights regions where the target is dominated by irreducible noise, thereby mitigating the mismatch between optimization effort and learnability.

3. Key Contributions

Identification of Representation Degradation: The paper empirically and theoretically identifies that standard diffusion training suffers from structural distortion and collapse at high noise levels due to the model's inability to preserve semantically meaningful structure when optimization effort is misaligned with target recoverability.
Optimization Analysis via NTK: The authors develop an NTK-based framework demonstrating how the interplay of optimization allocation, local kernel geometry, and target recoverability induces spectral degeneration and effective low-rank collapse in high-noise regimes. They provide non-asymptotic bounds showing how Bayes-noise forcing controls the residual error.
ERD Framework: They propose an architecture-agnostic training method that derives a target-adaptive weighting rule from the recoverable signal. This approach improves optimization stability and accelerates convergence without relying on external priors (e.g., pre-trained alignment networks) or auxiliary objectives.

4. Experimental Results

The authors evaluated ERD on ImageNet (256×256) and CelebA (64×64) using various backbones, including DiT, U-ViT, and UNet, across different prediction targets ( $\epsilon_\theta$ , $x_\theta$ , $v_\theta$ , $u_\theta$ ).

Performance Gains: ERD consistently outperformed baselines and existing weighting strategies (such as Min-SNR and P2) across all model scales and architectures.
- On ImageNet 256×256 with DiT-XL/2 (675M parameters), ERD achieved an FID of 15.0 (without CFG), improving upon the baseline SiT-XL/2 (17.2) and DiT-XL/2 (19.5).
- With Classifier-Free Guidance (CFG), the U-ViT-H/2 model achieved an FID of 1.45, significantly narrowing the gap with state-of-the-art methods like REPA (1.42) and outperforming Min-SNR (1.57).
Ablation Studies: Experiments confirmed that ERD provides consistent improvements regardless of the prediction target or network architecture (UNet vs. Transformer).
Toy Experiments: Controlled experiments on a 2D Gaussian Mixture Model (GMM) visualized the "Representation Degradation" phenomenon, showing that global shared-parameter models suffer from gradient contamination and spectral collapse in high-noise regimes, which is mitigated when optimization is isolated or reweighted.

5. Significance and Claims

The paper claims that Representation Degradation is a fundamental, intrinsic bottleneck in standard diffusion training, distinct from architectural limitations. By shifting the focus from heuristic noise-level allocation to target recoverability, ERD addresses the root cause of training inefficiency: the misallocation of optimization mass to irreducible noise.

The significance of this work lies in:

Theoretical Insight: It bridges the gap between diffusion optimization and NTK dynamics, explaining why extreme-noise training leads to degraded representations (spectral weakening and Bayes-noise domination).
Practical Efficiency: It offers a simple, plug-and-play weighting rule that accelerates convergence and improves generation quality across diverse backbones and tasks without the computational overhead of external alignment networks or complex latent modeling.
Universality: The method is shown to be effective across different diffusion parameterizations (noise, data, velocity, flow matching) and architectures, suggesting a unified solution to the optimization instability inherent in the diffusion process.

The authors remain modest regarding limitations, noting that their analysis relies on continuous-time optimization and NTK regime assumptions, which may not fully capture finite-width, highly nonlinear dynamics. Additionally, ERD uses an effective-amplitude score as a tractable proxy for recoverability rather than estimating the exact posterior recoverable signal.