Steering Optimisation Trajectories in Diffusion Representation Learning

2026-07-06

AI Summary

Although diffusion autoencoders achieve high-quality image reconstruction, the optimization dynamics governing their latent representations remain poorly understood. This work uncovers, for the first time, a bimodal optimization trajectory in early training stages—characterized by competing reconstruction-dominant and disentanglement-dominant modes—and proposes to steer this trajectory by modulating shortcut connections in the U-Net architecture and scheduling noise-level exposure. To this end, we introduce a novel paradigm that integrates a gated residual U-Net with a curriculum-based noise exposure strategy. Our approach substantially enhances representation quality and reduces sensitivity to random seeds across multiple disentanglement benchmarks, while significantly improving segmentation performance in object-centric learning tasks on both synthetic and real-world datasets.

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.

Problem

• diffusion autoencoders
• latent representation
• disentanglement
• optimization dynamics
• reconstruction-disentanglement trade-off

Innovation

• diffusion autoencoders
• optimization trajectories
• disentanglement
• noise-level curriculum
• gated residual U-Net