TRADE: Transfer of Distributions between External Conditions with Normalizing Flows

TL;DR

TRADE introduces a novel method for modeling parameter-dependent distributions using normalizing flows, formulated as a boundary value problem, enabling efficient learning across external conditions without restrictive assumptions.

Contribution

It presents TRADE, a boundary value problem approach for transfer learning of distributions dependent on external parameters, overcoming limitations of previous energy-based and restricted models.

Findings

Achieves excellent results in Bayesian inference, molecular simulations, and lattice models.

Efficiently propagates distribution information across conditions.

Does not rely on energy-based training or restrictive model architectures.

Abstract

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.