Steering Training Trajectories in Diffusion Representation Learning

원문 초록

Across training runs of a diffusion autoencoder, we observe similar image reconstruction quality, but substantially different levels of disentanglement in the learned representations. In this work, we study how training dynamics contribute to this behaviour. By analysing training trajectories in a controlled setting - curves of reconstruction quality against disentanglement during training - we identify two recurring trajectory patterns, which we call learning regimes. These regimes differ in how reconstruction and disentanglement evolve, despite similar reconstruction at the end of training. In the reconstruction regime, models improve image fidelity before disentanglement, while in the disentanglement regime, models improve both reconstruction and disentanglement more gradually. We hypothesise that modulating shortcut pathways in the diffusion U-Net and controlling early noise-level exposure can shape trajectories. To that end, we introduce SteeringDRL, combining gated residual U-Nets with a noise curriculum to steer models towards better disentanglement. Across attribute disentanglement benchmarks, SteeringDRL has competitive disentanglement, with reduced seed sensitivity in many settings. Our method extends to spatial disentanglement with object-centric learning, improving unsupervised segmentation on synthetic and real-world datasets.