Augmented KRnet for density estimation and approximation

TL;DR

This paper introduces augmented KRnets, combining discrete and continuous models, to improve invertibility and efficiency in flow-based generative modeling, achieving full nonlinear updates with fewer iterations and maintaining exact invertibility.

Contribution

The paper proposes augmented KRnets that incorporate augmented dimensions, enabling full nonlinear updates in fewer iterations while preserving exact invertibility and linking to neural ODEs.

Findings

Augmented KRnet achieves full nonlinear updates in two iterations.

The model maintains exact invertibility through neural ODE reformulation.

Numerical experiments demonstrate improved effectiveness of the proposed models.

Abstract

In this work, we have proposed augmented KRnets including both discrete and continuous models. One difficulty in flow-based generative modeling is to maintain the invertibility of the transport map, which is often a trade-off between effectiveness and robustness. The exact invertibility has been achieved in the real NVP using a specific pattern to exchange information between two separated groups of dimensions. KRnet has been developed to enhance the information exchange among data dimensions by incorporating the Knothe-Rosenblatt rearrangement into the structure of the transport map. Due to the maintenance of exact invertibility, a full nonlinear update of all data dimensions needs three iterations in KRnet. To alleviate this issue, we will add augmented dimensions that act as a channel for communications among the data dimensions. In the augmented KRnet, a fully nonlinear update is achieved in two iterations. We also show that the augmented KRnet can be reformulated as the discretization of a neural ODE, where the exact invertibility is kept such that the adjoint method can be formulated with respect to the discretized ODE to obtain the exact gradient. Numerical experiments have been implemented to demonstrate the effectiveness of our models.