PaperlayerView exact revision on arXiv

Source-linked AI summary

Coevolving Representations in Joint Image-Feature Diffusion

Theodoros Kouzelis, Spyros Gidaris, Nikos Komodakis

arXiv:2604.17492v1cs.CV

TL;DR

Existing joint image-feature diffusion methods keep the semantic representation space fixed, despite its role in guiding generation. CoReDi learns a lightweight projection jointly with diffusion and stabilizes it with stop-gradient targets, normalization, and anti-collapse regularization. The resulting adaptive representations improve generative modeling and convergence in both VAE latent and pixel-space diffusion.

  • Problem

    Existing joint diffusion approaches use semantic representation spaces constructed independently of the generative objective and fixed during training.

  • Method

    CoReDi jointly learns a lightweight projection of frozen visual features with the diffusion model, using stop-gradient targets, batch normalization, and explicit regularization against feature collapse.

  • Results

    CoReDi improves generative modeling and convergence in both VAE latent space and pixel space compared with fixed-representation joint diffusion models.

  • Takeaways & Limitations

    Adaptive semantic representations can develop stronger spatial structure and provide image-synthesis guidance beyond fixed representation spaces and VAE latent diffusion.

Abstract

from arXiv · show

Joint image-feature generative modeling has recently emerged as an effective strategy for improving diffusion training by coupling low-level VAE latents with high-level semantic features extracted from pre-trained visual encoders. However, existing approaches rely on a fixed representation space, constructed independently of the generative objective and kept unchanged during training. We argue that the representation space guiding diffusion should itself adapt to the generative task. To this end, we propose Coevolving Representation Diffusion (CoReDi), a framework in which the semantic representation space evolves during training by learning a lightweight linear projection jointly with the diffusion model. While naively optimizing this projection leads to degenerate solutions, we show that stable coevolution can be achieved through a combination of stop-gradient targets, normalization, and targeted regularization that prevents feature collapse. This formulation enables the semantic space to progressively specialize to the needs of image synthesis, improving its complementarity with image latents. We apply CoReDi to both VAE latent diffusion and pixel-space diffusion, demonstrating that adaptive semantic representations improve generative modeling across both settings. Experiments show that CoReDi achieves faster convergence and higher sample quality compared to joint diffusion models operating in fixed representation spaces.

1 Introduction

Existing joint diffusion methods use fixed semantic representation spaces, whereas CoReDi learns the projection jointly with diffusion so the space can adapt to image synthesis. Stable coevolution requires stop-gradient targets, batch normalization, and explicit anti-collapse regularization, and improves convergence across latent- and pixel-space diffusion.

  • Motivation: Existing joint approaches use predefined PCA or autoencoder projections that remain fixed while diffusion learns their distribution alongside VAE latents.These spaces complement low-level detail modeling with high-level semantic structure but are constructed independently of the generative objective.
  • CoReDi: CoReDi jointly learns a projection of frozen visual features with the diffusion model, making the semantic space a trainable component optimized for image synthesis.The learned representation evolves throughout training rather than serving as an externally imposed target.
  • Stable coevolution: Naively optimizing the projection is degenerate because both the representation input and its denoising target become trainable.This failure motivates explicit stabilization mechanisms for coevolution.
  • Stable coevolution: Stable coevolution requires stop-gradient targets, batch normalization, and explicit regularization against feature collapse.These components prevent trivial target minimization, stabilize feature scale and the noise schedule, and preserve expressive, complementary representations.
  • Results: CoReDi improves joint diffusion in both VAE latent space and pixel space, extending adaptive semantic guidance beyond compressed representations.The framework is reported to improve generative modeling across both settings.
  • Results: CoReDi converges ∼13× faster than REPA in latent space and improves convergence by ×2 over DeCo in pixel space.The figure also reports that CoReDi outperforms ReDi in latent space.

2 Related Work

Related work spans latent diffusion, pixel-space diffusion, semantic guidance, and methods for preventing representation collapse. These lines of work motivate combining semantic representations with generation while addressing pixel-space cost and collapse risks.

  • Latent and pixel diffusion: VAE latent diffusion reduces spatial dimensionality, computational cost, and learning difficulty compared with pixel-space diffusion.Latent models operate in the compressed space produced by a variational autoencoder.
  • Latent and pixel diffusion: Pixel-space diffusion avoids the VAE reconstruction bottleneck but requires strategies to manage high dimensionality and training complexity.Earlier approaches used progressively increasing resolutions and more complex training and inference procedures.
  • Semantic representations: Integrating visual representations into pixel diffusion to enhance generative performance remains largely unexplored.This identifies a gap at the intersection of semantic guidance and pixel-space generation.
  • Semantic representations: Prior semantic-guidance methods align pretrained representations with diffusion features, jointly optimize VAE and diffusion models, or replace VAE latents with pretrained representations.Examples include REPA, REPA-E, iREPA, and other representation-based generative models.
  • Collapse prevention: Self-supervised learning uses redundancy reduction and batch normalization to prevent representational collapse, including methods such as Barlow Twins, VICReg, and SimCLR.These mechanisms provide precedent for regulating learned representations against degenerate solutions.

3 Method

CoReDi jointly learns a projection of frozen visual features with the diffusion model, while stabilization and regularization let the semantic representation evolve without collapse.

  • Joint Image-Feature Synthesis: The model jointly processes noisy image and representation tokens and predicts their velocities through separate modality-specific outputs.Merged token fusion preserves the original token count and adds no computational overhead over a standard diffusion transformer.
  • Coevolving Representation Diffusion: CoReDi replaces fixed dimensionality reduction with a learnable linear projection jointly optimized with the diffusion model.The projection is gϕ(z0) = z0Wϕ, where Wϕ maps feature dimension D to d.
  • Stabilization: Batch normalization stabilizes projected feature scale and implicitly discourages sample collapse across images or spatial locations.The method omits trainable affine parameters so normalization controls scale and collapse rather than rescaling or shifting features.
  • Stabilization: Stop-gradient through the clean projected representation target prevents the projection from trivially modifying its target to reduce representation loss.This allows the representation space to coevolve while the diffusion model learns to denoise both modalities.
  • Regularization Methods: Explicit regularization addresses remaining channel collapse by promoting variation, orthogonality, or reduced redundancy in the learned representation.Feature variance regularization enforces a minimum desired standard deviation γ = 1, while orthogonality regularization encourages distinct projection directions.

4 Experiments

Experiments evaluate CoReDi across latent and pixel-space diffusion, visual encoders, stabilization choices, regularization strategies, and representation spatial structure. CoReDi generally improves quality and convergence, while ablations identify components needed for stable training.

  • Latent Space Diffusion: CoReDi improves latent-diffusion quality across model scales and settings, outperforming fixed or alternative baselines at comparable training budgets.For B/2, it reaches 16.4 FID at 400K iterations versus ReDi’s 21.4 and SiT’s 33.0; for XL/2, it matches 3.3 FID in half the iterations.
  • Latent Space Diffusion: 1.58 FID after 400 epochs outperforms REPA’s 1.80 and ReDi’s 1.72 after 800 epochs with classifier-free guidance.CoReDi achieves this result while cutting the reported training steps by half.
  • Visual Encoder Variation: CoReDi consistently improves over ReDi’s fixed PCA projection across four visual encoders, with the largest reported gain for DINOv2 from 30.9 to 24.7 FID.The comparison uses models evaluated at 200K steps.
  • Pixel Space Diffusion: 31.5 FID at 100K iterations matches DeCo-L/16 at 200K, indicating 2× faster convergence in pixel-space diffusion; further training reaches 21.5 FID.These results extend the reported benefit beyond VAE latent generation.
  • Stabilization Ablations: Removing stop-gradient worsens FID from 24.7 to 50.8, while removing batch normalization causes training collapse with FID 223.9.The reported failure modes are trivial target minimization for stop-gradient removal and near-zero recalibration for batch-normalization removal.
  • Regularization Ablations: 24.7 FID is achieved with feature-variance regularization, compared with 25.6 for orthogonality, 25.9 for covariance, and 37.2 without regularization.The no-regularization projection collapses features and performs worse than ReDi’s static PCA projection.
  • Spatial Structure: LDS, CDS, and RMSC all improve during training, and learned projections achieve higher spatial-structure scores than fixed PCA projections.The metrics measure complementary aspects of spatial self-similarity and diversity.

5 Conclusion

CoReDi makes the semantic representation space coevolve with the generative model through a learned projection. The framework requires stabilization and anti-collapse mechanisms, improves spatial structure, and extends benefits from latent to pixel-space diffusion.

  • Contribution: CoReDi jointly learns an adaptive semantic projection with the diffusion objective instead of relying on predetermined representations.The learned space can develop structure directly beneficial for image synthesis.
  • Stable Coevolution: Stop-gradient stabilization, batch normalization, and explicit regularization against feature collapse are identified as essential for stable coevolution.The conclusion reports that all three ingredients are empirically necessary.
  • Representation Structure: Coevolving representations develop stronger spatial structure during training, providing a potential explanation for the observed generative improvements.The conclusion presents this as a potential explanation rather than a definitive causal finding.
  • Scope: CoReDi extends beyond VAE latent spaces to pixel-space diffusion and yields consistent improvements in both settings.The reported scope includes latent-based and pixel-space generation.

A Additional Results and Ablations

Additional analyses report faster convergence, clarify classifier-free-guidance behavior, and show that coevolving representations develop increasingly structured spatial organization in pixel diffusion.

  • Detailed Quantitative Comparison: CoReDi matches ReDi’s converged generative performance at 2M iterations instead of 4M, indicating significantly faster convergence.The detailed comparison includes SiT-XL/2, CoReDi-XL/2, ReDi-XL/2, and REPA-XL/2 without classifier-free guidance.
  • Classifier-Free Guidance: A CFG weight of 1.8 achieves optimal performance when classifier-free guidance is applied exclusively to VAE latents.The guidance ablation is reported in Fig. 6.
  • Pixel-Space Spatial Structure: LDS, CDS, and RMSC improve consistently throughout pixel-space CoReDi training, showing increasingly structured spatial organization.The same phenomenon is reported for pixel diffusion rather than only latent-space diffusion.

B.1 Architecture settings

The experiments use distinct latent- and pixel-space architectures and optimization configurations. Latent experiments precompute VAE latents, while pixel experiments use DeCo-L/16 with a 16-channel projection.

  • Latent Space: Latent diffusion follows SiT configurations with B/2 and XL/2 transformer models using SD-VAE-FT-EMA latents and 2×2 patches.The listed models contain 130M and 675M parameters, respectively.
  • Optimization: A cosine-decay schedule is applied only to the learnable projection in XL experiments, reducing its learning rate to zero by 400K iterations.The DiT backbone retains a constant learning rate of 1e-4.
  • Pixel Space: Pixel diffusion uses DeCo-L/16 with 22 encoder blocks, 3 pixel-decoder MLP blocks, embedding dimension 1024, and 16 attention heads.The pixel-space projection maps to 16 channels.
  • Optimization: Models use AdamW with a constant learning rate of 1 × 10^-4, β1 = 0.9, β2 = 0.999, and batch size 256.Image latents are precomputed to accelerate latent-diffusion training.

C Evaluation Metrics

The evaluation follows ADM’s setup, generating 50K samples and measuring distributional distance, spatial structure, classifiability, diversity, realism, and coverage.

  • 50K samples are generated using ADM’s evaluation setup and scored with the official TensorFlow evaluation toolkit.Evaluations are conducted on NVIDIA A100 GPUs.
  • FID measures the distance between real and generated image feature distributions using Inception-v3 under a multivariate Gaussian assumption.
  • sFID applies FID to intermediate Inception-v3 spatial feature maps to better capture generated-image spatial structure.
  • IS rewards outputs that are confidently classifiable while remaining diverse across categories.
  • Precision and recall measure realism and coverage of the real data distribution in feature space, respectively.

C.1 Spatial Structure Evaluation Metrics

The spatial metrics assess how patch-token similarity changes with distance and how much spatial diversity representations preserve.

  • Local vs. Distant Similarity (LDS): LDS contrasts average cosine similarity for spatially near versus distant patch pairs, with larger values indicating stronger spatial organization.Near and distant pairs are separated using Manhattan distance thresholds.
  • Correlation Decay Slope (CDS): CDS measures the rate at which patch similarity decays with spatial distance.It fits a least-squares line to the spatial correlogram; larger values indicate faster decay and stronger spatial organization.
  • RMS Spatial Contrast (RMSC): RMSC measures spatial diversity using normalized patch features and their mean.Higher values reflect greater spatial diversity, while lower values indicate more uniform, spatially uninformative representations.

D Additional Qualitative Results

The paper provides qualitative samples of jointly generated images and visual representations, alongside comparisons between learned CoReDi projections and static PCA visualizations.

  • The qualitative results cover both generated images and their visual representations.
  • Figure 8 presents selected CoReDi-XL/2 samples trained for 1M steps on ImageNet at 256 × 256 resolution.Images and visual representations are jointly generated with classifier-free guidance using w = 4.0.
  • Figure 9 compares PCA visualizations of DINOv2, MOCOv3, SigLIPv2, and MAE features with visualizations from CoReDi’s learned projection.
Loading 2604.17492v1…