- The paper shows that diffusion autoencoders with identical architectures and losses can follow distinct reconstruction or disentanglement trajectories because of early optimisation dynamics, shortcut connections, and noise-level exposure.
- SteeringDRL combines near-zero-initialised gated residual connections with an importance-corrected SNR curriculum to limit reconstruction shortcuts and guide learning toward disentangled representations while preserving image quality.
- The method achieves strong benchmark results, including FactorVAE 1.000 and DCI 0.954 on Shapes3D, DCI 0.624 on Cars3D, and improved ClevrTex object segmentation with 68.31 mIoU and 69.69 mBO.
Diffusion autoencoders can produce images of comparable perceptual quality while learning latent representations that differ dramatically in their alignment with the underlying generative factors. This paper investigates why such divergent representational outcomes arise under identical architectures and objectives, and shows that the divergence traces to optimisation dynamics rather than to the loss or model capacity alone. The authors introduce SteeringDRL, a combination of architectural and curriculum-based interventions that steers training trajectories toward stronger disentangled representations, and validate the approach on attribute disentanglement and object-centric learning (OCL) benchmarks.
Optimisation regimes in diffusion autoencoders
The core empirical observation comes from plotting optimisation trajectories—reconstruction quality (LPIPS) against representation quality (DCI, FactorVAE, MIG) over the course of training—for a 3M-parameter diffusion autoencoder on Shapes3D across 25 runs. Rather than a continuum of outcomes, trajectories organise around two regimes that emerge early in training. Models in the disentanglement regime improve reconstruction and disentanglement gradually and jointly, while models in the reconstruction regime reach near-optimal reconstruction quickly but lag substantially in disentanglement; by roughly 576K images seen, the regimes are already clearly separated. Notably, regime membership does not correlate strongly with weight initialisation or noise-sampling seed, and the pattern persists under alternative loss-weighting functions such as w(λ)=sech(λ), suggesting it reflects a general optimisation property rather than a specific configuration choice.
The two regimes also correspond to different internal mechanisms. Cross-attention maps in disentanglement-regime models are spatially localised and aligned with ground-truth factors, whereas reconstruction-regime models exhibit diffuse attention. This points to a pathway that bypasses cross-attention conditioning: skip connections, which prior work has shown provide a high-bandwidth route around the U-Net trunk.
Two hypotheses
The paper formalises its explanation as two hypotheses. H1 (shortcut pathways): since disentanglement is learned without direct supervision, the U-Net can exploit skip connections to optimise reconstruction without organising information in the latent z. H2 (noise-level exposure): diffusion training imposes an information bottleneck over noise levels—high-noise levels force reliance on the semantic conditioning signal, while low-noise levels allow fine-detail recovery—and the exposure to specific noise levels early in training determines the relative pace of reconstruction versus disentanglement learning.
Both hypotheses are supported by controlled interventions. Zeroing skip connections improves disentanglement but sharply degrades reconstruction (LPIPS rises from roughly 0.002 to 0.157), confirming that skips act as reconstruction shortcuts at the expense of the representation pathway. Fixed-band experiments show that restricting the sampled logSNR band to high-noise values ([−5,0]) yields the strongest early disentanglement but poor reconstruction (LPIPS $0.2929$), while the full band [−5,12] accelerates reconstruction without collapsing disentanglement. A further finding with practical significance: the maturity of the frozen VQ-VAE modulates regime visibility. With a longer-trained VQ-VAE, reconstruction improves and regime separation weakens, plausibly because finer visual detail is encoded directly in the latent space, reducing pressure on z to organise information.
SteeringDRL
The method combines two interventions targeting H1 and H2. First, a gated residual U-Net replaces concatenative skip connections with gated residuals, G(h)=Softplus(ξ)⊙h, initialised near zero so that shortcut pathways are suppressed precisely during the early-training window in which regime commitment occurs. Residual connections inside spatial transformer blocks are similarly gated, and cross-attention queries are normalised via AdaLN-Zero modulation by λ. Second, an SNR curriculum trains on a bounded, linearly widening band [λL(m),λU(m)] that expands to the full objective over mmax steps, with an importance-weighting correction to keep the estimator unbiased. The formulation is built on the VDM++ framework in SNR space, so it applies uniformly to z0-prediction and z1-prediction parameterisations via an equivalent weighting conversion.
The ablation sequence on Shapes3D shows that the components interact non-trivially. The z2 loss-weighting alone improves disentanglement over EncDiff (DCI z3 vs. z4) but increases variance, and adding the curriculum to the standard U-Net does not consistently help. Only when combined with the gated residual U-Net does the full method deliver strong disentanglement (DCI z5, FactorVAE z6) with substantially reduced seed variance. Curriculum ablations further demonstrate that the initial band selects the regime—starting from high-noise biases toward reconstruction, starting from low-noise steers toward disentanglement—while z7 governs convergence speed, with z8 over 32K steps performing best.
Attribute disentanglement results
On standard benchmarks with z9 tokens, SteeringDRL achieves state-of-the-art results on Cars3D and Shapes3D: FactorVAE [−5,0]0 and DCI [−5,0]1 on Cars3D (versus [−5,0]2/[−5,0]3 for the strongest prior method), and FactorVAE [−5,0]4 with DCI [−5,0]5 on Shapes3D, with lower variance across seeds than competing approaches. On MPI3D-toy it attains the best DCI ([−5,0]6) while remaining competitive on FactorVAE. Gains are larger at [−5,0]7, close to the true number of generative factors, where the DCI improvement over EncDiff on Cars3D is substantial ([−5,0]8 vs. [−5,0]9). Latent-editing experiments attribute the Cars3D gains to better isolation of the difficult Style/Object Type factor, which EncDiff spreads across multiple latents. The implication is that the same generative fidelity can be obtained with latents that are considerably better aligned with causal factors, provided optimisation is steered appropriately.
Spatial disentanglement with object-centric learning
Extending to OCL, the authors instantiate the inductive bias with SlotAttention and scale to a $0.2929$0138M-parameter U-Net. On ClevrTex, SteeringDRL outperforms SlotDiffusion and methods with additional object-centric inductive biases on mIoU ($0.2929$1) and mBO ($0.2929$2), while also improving reconstruction (LPIPS $0.2929$3) and composition (FID $0.2929$4). On PascalVOC with pretrained DINO ViT features, it improves over the SlotDiffusion baseline on all reported metrics. A notable difference emerges between the two settings: with features learned from scratch, representations and reconstructions improve progressively together, whereas with pretrained features, learning shifts from feature discovery to slot binding, and SteeringDRL converges faster than the baseline. This suggests the steering mechanism is complementary to, and amplified by, semantic features already present in the encoder.
Limitations and open questions
The paper is explicit that regime separation is clearest in low-capacity settings; larger models tend to align with the reconstruction regime, masking the phenomenon, and the interventions are effective in that regime rather than the separation being naturally visible. The VQ-VAE checkpoint materially affects results—longer VQ-VAE training reduces regime separation—and the interaction between latent autoencoder quality and representation learning is left as an open question. Curriculum design is acknowledged to be configuration-sensitive: different initial bands and widening durations change behaviour across architectures, capacities, and datasets, and the authors call for a more systematic study of SNR curricula. Zeroing skips entirely is not a viable solution since it sharply degrades reconstruction, which is why the softer gated residual is preferred, but the optimal balance between shortcut suppression and reconstruction fidelity is not characterised. Finally, the applicability to transformer architectures, where shortcut pathways are structurally constrained, remains untested.
Conclusion
This paper reframes the variability of diffusion autoencoder representations as a property of optimisation dynamics: early commitment to reconstruction- or disentanglement-oriented trajectories, driven by skip-connection shortcuts and early noise-level exposure. SteeringDRL demonstrates that simple, targeted interventions—gated residual bottlenecks and an SNR curriculum with importance correction—can reliably steer training toward the disentanglement regime, yielding state-of-the-art attribute disentanglement on Cars3D and Shapes3D and improved object-centric segmentation on synthetic and real-world data. The central design principles, restricting reconstruction shortcuts and controlling when the model encounters specific noise levels, offer a practical route to stronger representations without sacrificing generative fidelity.
