Generative models regain lost variability in cosmology simulations
Recovering Missing Degrees of Freedom in Simulation-Trained Field-Level Generative Models
Machine Learning
Summary
Simulation-trained models sometimes miss some natural randomness found in the real world, making their predictions less varied than they should be. The authors treat this loss of randomness as missing parts of the system called degrees of freedom. They add new, physics-based features inside the model to recover much of this missing variation without needing more training data. Testing this on models predicting cosmic weak-lensing patterns shows the method helps models better handle differences between training and real data. This approach also makes the models less sensitive to unexpected changes during inference.
What this means in practice
- •For cosmology simulation teams: Improve simulation-trained generative models by restoring lost variability, enhancing accuracy in strongly nonlinear cosmological field inference.
- •For machine learning engineers: Enhance robustness of physics-informed generative models to distribution shifts by integrating latent corrections without additional data.
Authors
Masato Shirasaki, Shiro Ikeda
Abstract
Simulation-trained generative models are emerging as a powerful framework for field-level inference in cosmology, particularly in strongly nonlinear regimes where analytical approaches become intractable. However, the finite size of training datasets often leads to variance shrinkage, indicating that part of the realization-level variability of the underlying physical process is not captured by the learned generative dynamics. In this Letter, we interpret variance shrinkage as a diagnostic of missing effective degrees of freedom and introduce physics-informed latent modifications to the generative dynamics of a continuous normalizing flow. We show that these latent corrections recover a substantial fraction of the missing variability without requiring additional training simulations. Using weak-lensing field-level inference as a representative cosmological application, we further demonstrate through posterior analysis that the learned latent directions are actively exploited during inference and behave as effective nuisance parameters that absorb part of the mismatch between the training and target distributions. As a result, the model exhibits improved robustness to out-of-distribution shifts in field-level inference. Our work establishes variance shrinkage as a practical diagnostic for identifying missing degrees of freedom in simulation-trained generative models and provides a framework for recovering them through physics-informed modifications of generative dynamics.