TRADE: Transfer of Distributions between External Conditions with Normalizing Flows
This addresses a common but challenging scenario in fields like molecular simulations, offering a more stable and flexible solution compared to existing methods.
The paper tackles the problem of modeling distributions that depend on external control parameters, such as temperature in molecular simulations, by introducing TRADE, which formulates learning as a boundary value problem to efficiently learn parameter-dependent distributions without restrictive assumptions, achieving excellent results across applications like Bayesian inference and molecular simulations.
Modeling distributions that depend on external control parameters is a common scenario in diverse applications like molecular simulations, where system properties like temperature affect molecular configurations. Despite the relevance of these applications, existing solutions are unsatisfactory as they require severely restricted model architectures or rely on energy-based training, which is prone to instability. We introduce TRADE, which overcomes these limitations by formulating the learning process as a boundary value problem. By initially training the model for a specific condition using either i.i.d.~samples or backward KL training, we establish a boundary distribution. We then propagate this information across other conditions using the gradient of the unnormalized density with respect to the external parameter. This formulation, akin to the principles of physics-informed neural networks, allows us to efficiently learn parameter-dependent distributions without restrictive assumptions. Experimentally, we demonstrate that TRADE achieves excellent results in a wide range of applications, ranging from Bayesian inference and molecular simulations to physical lattice models.