Steering Optimisation Trajectories in Diffusion Representation Learning
For researchers working on representation learning with diffusion models, this work provides a method to improve latent quality and disentanglement, though the improvements are incremental.
The paper identifies two distinct optimization regimes in diffusion autoencoders that affect latent representation quality, and introduces SteeringDRL, which uses gated residual U-Nets and a noise-level curriculum to steer training toward better representations. SteeringDRL improves disentanglement metrics and reduces seed sensitivity, and extends to object-centric learning with improved segmentation on synthetic and real-world datasets.
We study why diffusion autoencoders can achieve similar image quality while learning substantially different latent structures. We trace this behaviour to optimisation dynamics; we analyse curves of image reconstruction against latent representation quality, revealing trajectories that organise around two distinct regimes early in training. Models in the reconstruction regime prioritise image fidelity early, whereas those in the disentanglement regime improve reconstruction and disentanglement more gradually. We hypothesise that this behaviour can be influenced by targeting shortcut pathways in the diffusion U-Net and controlling early noise-level exposure, thereby shaping the reconstruction-disentanglement trade-off during training. To steer optimisation toward stronger representations, we introduce SteeringDRL, combining gated residual U-Nets with a simple noise-level exposure curriculum for training. Across disentanglement benchmarks, SteeringDRL improves representation quality and reduces seed sensitivity. Our method further extends to spatial disentanglement in object-centric learning, improving segmentation quality on synthetic and real-world datasets.