Dynamics Hessian Penalty for Disentangled Representation Learning

Disentangled representation learning seeks to represent distinct factors of variation in separate latent dimensions. Some methods encourage disentanglement by restricting interactions between latent dimensions during generation. When imposed on the complete generative mapping, these restrictions can favor additive structure and limit the nonlinear composition of latent effects in observation space. We introduce the Dynamics Hessian Penalty (DHP) for iterative generators, which applies this interaction prior to the individual conditional predictions, penalizing their mixed second-order derivatives with respect to the latent representation. Our analysis shows that the complete generative mapping need not collapse to an additive form, even when all such local mixed derivatives vanish. We implement DHP using stochastic finite differences and apply it to diffusion and rectified-flow models, without unrolling the sampler during training or adding inference cost. Experiments show strong disentanglement on synthetic benchmarks and improvements in both disentanglement and generation quality on real-world images. Qualitative results further illustrate controllable generation through latent manipulation.