# Augmented KRnet for density estimation and approximation

arXiv: https://arxiv.org/abs/2105.12866 · HTML: https://mlsanity.com/paper/2105.12866/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 33 references that are in the map.
- First posted on arXiv 26 May 2021.

## Builds on (most cited first, 23 of 23)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [Conditional Image Generation with PixelCNN Decoders](https://mlsanity.com/paper/1606.05328/) (2016)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [Wasserstein GAN](https://mlsanity.com/paper/1701.07875/) (2017)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [Augmented Neural ODEs](https://mlsanity.com/paper/1904.01681/) (2019)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Beyond Finite Layer Neural Networks: Bridging Deep Architectures and Numerical Differential Equations](https://mlsanity.com/paper/1710.10121/) (2017)
- [Sylvester Normalizing Flows for Variational Inference](https://mlsanity.com/paper/1803.05649/) (2018)
- [Inference via low-dimensional couplings](https://mlsanity.com/paper/1703.06131/) (2017)
- [Monge-Amp`ere Flow for Generative Modeling](https://mlsanity.com/paper/1809.10188/) (2018)
- [Potential Flow Generator with L2 Optimal Transport Regularity for Generative Models](https://mlsanity.com/paper/1908.11462/) (2019)
- [Adaptive deep density approximation for Fokker-Planck equations](https://mlsanity.com/paper/2103.11181/) (2021)
- [Coupling the reduced-order model and the generative model for an importance sampling estimator](https://mlsanity.com/paper/1901.07977/) (2019)

## Built on by (most cited first, 4 of 4)

- [DAS-PINNs: A deep adaptive sampling method for solving high-dimensional partial differential equations](https://mlsanity.com/paper/2112.14038/) (2021)
- [Bounded KRnet and its applications to density estimation and approximation](https://mlsanity.com/paper/2305.09063/) (2023)
- [Deep adaptive sampling for surrogate modeling without labeled data](https://mlsanity.com/paper/2402.11283/) (2024)
- [Moving Sampling Physics-informed Neural Networks induced by Moving Mesh PDE](https://mlsanity.com/paper/2311.16167/) (2023)

## Abstract (arXiv)

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.