💻 computer science

Coevolving Representations in Joint Image-Feature Diffusion

The paper proposes Coevolving Representation Diffusion (CoReDi), a framework that dynamically adapts semantic representation spaces during training via a learnable linear projection and stabilization techniques, thereby achieving faster convergence and higher sample quality in joint image-feature diffusion models compared to methods using fixed representations.

Original authors: Theodoros Kouzelis, Spyros Gidaris, Nikos Komodakis

Published 2026-04-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Theodoros Kouzelis, Spyros Gidaris, Nikos Komodakis

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

Imagine you are teaching a robot artist how to paint a masterpiece.

In the past, you'd give the robot two things:

The Canvas (The Image): A blank space where it learns to paint pixels (colors, textures, edges).
The Reference Book (The Features): A pre-written encyclopedia of "what things look like" (e.g., "a cat has pointy ears," "a sky is blue").

The Old Way (Fixed Representations):
Previously, researchers gave the robot a fixed, rigid dictionary. No matter what the robot was trying to paint, it had to use the exact same definitions from that dictionary. If the dictionary said "cat," the robot had to use that specific definition, even if it wasn't the best way to describe a fluffy kitten vs. a wild tiger. The dictionary never changed, even as the robot got better at painting. This was like trying to write a modern novel using a dictionary from 1800; it works, but it's clunky and limits your creativity.

The New Way (CoReDi):
The paper introduces CoReDi (Coevolving Representation Diffusion). Instead of giving the robot a static dictionary, they give it a living, breathing dictionary that learns while the robot paints.

Here is how it works, using simple analogies:

1. The Dance of Co-evolution

Imagine the robot (the painter) and the dictionary (the semantic guide) are dancing together.

In the old method, the dictionary stood still, and the robot tried to dance around it.
In CoReDi, the dictionary moves with the robot. As the robot learns to paint better, the dictionary updates its definitions to match the robot's specific style. If the robot starts painting in a "cyberpunk" style, the dictionary instantly learns what "cyberpunk" means in that context. They grow up together.

2. The Three Safety Nets (Why it didn't work before)

The authors realized that if you just let the dictionary change freely, it might go crazy. It might decide that "cat" and "dog" are the same word, or that "blue" and "red" are identical. This is called Feature Collapse (the dictionary becomes useless gibberish).

To stop this, they added three "safety nets":

The "Stop-Gradient" Brake: Imagine the robot is trying to cheat. It could just change the dictionary to say "The answer is always 0" so it wins the game easily. The "Stop-Gradient" is like a referee who says, "You can change the dictionary, but you can't change the target answer to make it easy." This forces the dictionary to actually learn meaningful things, not just cheat.
The "Batch Normalization" Thermostat: Sometimes, the dictionary gets too hot (numbers get huge) or too cold (numbers get tiny), which confuses the robot. The thermostat keeps the temperature just right, ensuring the dictionary stays stable and readable.
The "Anti-Collapse" Regularizer: This is like a strict teacher who says, "You must use all your words, and they must mean different things!" If the dictionary tries to make 10 words all mean "cat," the teacher penalizes it. This ensures the dictionary stays diverse and rich.

3. The Result: A Better Artist

Because the dictionary evolves with the painter:

Faster Learning: The robot doesn't waste time figuring out how to translate the old dictionary into its new style. It learns much faster (the paper says up to 13 times faster in some cases!).
Better Quality: The final paintings are sharper and more accurate because the "definitions" of what to paint are perfectly tuned to the painting style.
No Bottleneck: The method works even if you skip the "compressed" canvas and paint directly on the raw pixels, removing a layer of complexity that usually slows things down.

The Big Picture

Think of CoReDi as upgrading from a static map to a GPS that learns your driving habits.

Old Way: You drive with a paper map that never updates. If you take a new route, the map is useless.
CoReDi: You drive with a GPS that watches how you drive, learns your shortcuts, and updates the map in real-time to help you get to your destination faster and smoother.

The paper proves that when the "guide" (the semantic features) and the "artist" (the diffusion model) grow up together, they create better art, faster, than if they were stuck with a fixed, outdated guide.

1. Problem Statement

Recent advances in diffusion models have explored joint image-feature generative modeling, where low-level image latents (from a VAE) are coupled with high-level semantic features extracted from pre-trained visual encoders (e.g., DINOv2, MAE).

Current Limitation: Existing approaches (e.g., ReDi, REPA) rely on a fixed representation space. They typically use a static projection (like PCA) to compress high-dimensional semantic features into a lower-dimensional space before training begins. This projection is constructed independently of the generative objective and remains unchanged during training.
The Core Question: Should the representation space guiding diffusion remain fixed, or should it adapt jointly with the generative model to better serve the specific task of image synthesis?
The Challenge: Naïvely optimizing a learnable projection jointly with the diffusion model leads to degenerate solutions (feature collapse), where the model trivially minimizes loss by collapsing features rather than learning meaningful representations.

2. Methodology: CoReDi Framework

The authors propose Coevolving Representation Diffusion (CoReDi), a framework where the semantic representation space evolves dynamically alongside the diffusion model.

A. Learnable Adaptive Projection

Instead of a fixed PCA projection, CoReDi introduces a learnable linear projection $g_\phi(\cdot)$ (a trainable weight matrix $W_\phi$ ) that maps frozen visual encoder outputs to a lower-dimensional space. This projection is optimized jointly with the diffusion model parameters $\theta$ .

B. Stabilization Mechanisms

To prevent the optimization from collapsing into trivial solutions, the authors identify and implement three critical components:

Stop-Gradient (SG) on Targets: In the representation loss calculation, gradients are stopped from flowing through the clean target representation ( $\tilde{z}_0$ ). This prevents the projection from trivially modifying the target to minimize the loss, forcing the diffusion model to learn to denoise the evolving representation.
Batch Normalization (BN): Applied after the learnable projection. Diffusion models are sensitive to input scale; BN stabilizes feature statistics, preserves the intended noise schedule, and acts as an implicit regularizer against sample collapse.
Explicit Regularization: Even with SG and BN, feature collapse (redundant channels) can occur. The authors propose explicit regularizers to enforce diversity:
- Feature Variance Regularization: Penalizes channels with low standard deviation across samples, ensuring each channel remains active.
- Orthogonality Regularization: Penalizes the deviation of the projection weight matrix from an orthogonal matrix ( $W^T W \approx I$ ), ensuring distinct projection directions.
- Covariance Regularization: Penalizes off-diagonal entries in the channel covariance matrix to decorrelate features.
- Finding: Feature Variance Regularization yielded the best results.

C. Extension to Pixel Space

While most joint models rely on VAE latents, CoReDi extends to pixel-space diffusion (building on the DeCo architecture). This removes the reconstruction bottleneck imposed by VAE compression, allowing the co-evolving semantic features to guide raw pixel generation directly.

3. Key Contributions

CoReDi Framework: A novel method for jointly modeling images and semantic representations where the representation space itself is a learnable, co-evolving component.
Stability Analysis: Identification of the three necessary ingredients (Stop-Gradient, Batch Norm, Explicit Regularization) required to achieve stable co-evolution without feature collapse.
Dual-Setting Validation: Demonstration that adaptive representations improve generative performance in both VAE latent space and pixel space.
Spatial Structure Discovery: Empirical evidence showing that co-evolving representations naturally develop stronger, more structured spatial organization over time compared to fixed projections.

4. Experimental Results

The paper evaluates CoReDi on ImageNet (256x256) using various model scales (SiT-B/2, SiT-XL/2) and visual encoders (DINOv2, MOCOv3, SigLIPv2, MAE).

Latent Space Performance:
- Faster Convergence: CoReDi-XL/2 achieves an FID of 3.3 in 2M iterations, matching the performance of ReDi-XL/2 which requires 4M iterations (2x speedup).
- Higher Quality: At 400K iterations, CoReDi-B/2 achieves an FID of 16.4, significantly outperforming ReDi (21.4) and SiT (33.0).
- With CFG: CoReDi achieves an FID of 1.58 in 400 epochs, outperforming REPA (1.80) and ReDi (1.72) which were trained for 800 epochs.
Pixel Space Performance:
- CoReDi-L/16 achieves comparable performance to DeCo-L/16 at 200K steps but reaches that level in only 100K steps (2x acceleration).
- At 200K steps, CoReDi reaches an FID of 21.5, a substantial improvement over the baseline.
Ablation Studies:
- Removing Stop-Gradient causes FID to degrade from 24.7 to 50.8.
- Removing Batch Normalization causes total training collapse (FID > 200).
- Removing Regularization leads to feature collapse (FID 37.2), performing worse than static PCA.
Spatial Structure: Metrics (LDS, CDS, RMSC) show that CoReDi representations develop significantly stronger spatial organization during training compared to fixed PCA projections, correlating with improved generation quality.

5. Significance

Paradigm Shift: CoReDi challenges the assumption that semantic priors must be static. It demonstrates that the "language" of the semantic features should adapt to the specific needs of the generative model.
Efficiency: By accelerating convergence (up to 2x faster) and improving sample quality, CoReDi offers a more computationally efficient path to high-fidelity image synthesis.
Generalizability: The method works across different visual encoders and extends beyond latent spaces to raw pixels, suggesting that adaptive semantic guidance is a universal improvement for diffusion models.
Theoretical Insight: The paper provides a deeper understanding of why joint diffusion works, linking performance gains to the emergence of structured spatial organization in the learned representation space.