Steering Optimisation Trajectories in Diffusion Representation Learning

2026-07-06Computer Vision and Pattern Recognition

Computer Vision and Pattern RecognitionArtificial Intelligence
AI summary

The authors investigate why diffusion autoencoders can produce similar high-quality images but learn very different ways to represent image information inside the model. They find that early training can follow two paths: one focusing quickly on making images look good, and another gradually improving both image quality and how well the model separates different features (disentanglement). To guide the model toward better, more structured representations, the authors propose a method called SteeringDRL, which adjusts the network architecture and controls noise exposure during training. This approach improves how well the model understands and separates image parts, helping in tasks like object segmentation in images.

diffusion autoencoderslatent representationdisentanglementU-Netoptimization dynamicsnoise-level exposureSteeringDRLobject-centric learningimage segmentationresidual networks
Authors
Rajat Rasal, Avinash Kori, Tian Xia, Ben Glocker
Abstract
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.