dkl_training module

Optimization of invertible flows in the weighted dataset problem using the DKL loss

Reminder: we have a dataset (x,p(x),f(x)) such that - x ~ p(x) - we want to learn a model that draws points according to f(x), which is positive, and known up to normalization

We can estimate the expected un-normalized log-likelihood of a random datapoint x~f(x) under the pdf of our flow model q(x) as

lL(f,q) = ∫ dx f(x) łog(q(x)) = E(f(x)/p(x) log(q(x)), x~p(x))

This is the same as minimizing the DKL between φ(x)=f(x)/∫ dx f(x) and q(x):

D_{KL}(φ|q) = ∫ dx φ(x) łog(φ(x)/q(x)) = - lL(f,q)*λ + η

where λ>0 and η are constants that are independent of the parameters of our flow q so maximizing the estimated log-likelihood is the same as minimizing the DKL. In either case, we can only do it up to a constant and use the loss lL(f,q) as a proxy

NB: An important point for importance sampling. In importance sampling, a natural loss is the variance of the integrand. One can however note that the optimum for both the variance loss and the DKL/ML is when the flow reproduces the target distribution.

class BasicDKLTrainer(*args, **kwargs)[source]

Bases: zunis.training.weighted_dataset.weighted_dataset_trainer.BasicTrainer

class BasicStatefulDKLTrainer(*args, **kwargs)[source]

Bases: zunis.training.weighted_dataset.weighted_dataset_trainer.BasicStatefulTrainer

weighted_dkl_loss(fx, px, logqx)[source]