Rotations on Latent Hyperspheres: a Geometry-Aware Guiding Framework for Diffusion Models

ICLR 2026 Conference Submission21692 Authors

Keywords: Diffusion Models, Reinforcement Learning, Latent Space Optimization, Physics-Informed Machine Learning, High-Dimensional Geometry, Constraint-Aware Generation

Abstract:

Diffusion models have emerged as a powerful tool across diverse domains. However, their purely data-driven nature can produce samples that deviate from domain-governing constraints. We introduce a plug-and-play, Reinforcement Learning framework that operates in the latent space of pre-trained diffusion models to optimize initial noise samples.
Our approach, motivated by the near-spherical geometry of high-dimensional Gaussian distributions, employs a novel rotation-matrix-based scheme for efficient latent space exploration. This steers the model toward more feature-preserving outputs, guided by task-specific rewards computed on the final samples. We evaluate our method on two diffusion models: one trained on solutions of the Darcy Flow PDE and another on a synthetic dataset with complex structural features. Across both settings, our framework yields significant improvements in sample quality, achieving a ${\\sim}26 \\%$ relative reduction in PDE residual and a ${\\sim}44 \\%$ relative improvement on the synthetic dataset's feature-alignment metric, compared to the vanilla diffusion models. Finally, we show that rotation-matrix-based exploration significantly outperforms unconstrained exploration, validating our geometry-aware approach and establishing a more effective method for latent space control.

Primary Area: generative models

Submission Number: 21692