Paper deep dive
VINA: Variational Invertible Neural Architectures
Shubhanshu Shekhar, Mohammad Javad Khojasteh, Ananya Acharya, Tony Tohme, Kamal Youcef-Toumi
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 89%
Last extracted: 7/20/2026, 2:40:19 PM
Summary
This paper introduces Variational Invertible Neural Architectures (VINA), a unified framework for Invertible Neural Networks (INNs) and Normalizing Flows (NFs) based on variational unsupervised loss functions. It provides theoretical performance guarantees for posterior accuracy in INNs and distributional accuracy in NFs under realistic assumptions, and demonstrates effectiveness on ocean-acoustic inversion problems.
Entities (8)
Relation Signals (7)
VINA → unifies → Normalizing Flows
confidence 95% · We introduce a unified framework for INNs and NFs based on variational unsupervised loss functions
VINA → unifies → Invertible Neural Networks
confidence 95% · We introduce a unified framework for INNs and NFs based on variational unsupervised loss functions
VINA → appliedto → Ocean-Acoustic Inversion
confidence 90% · demonstrating the effectiveness of our approach on a realistic ocean-acoustic inversion problem
VINA → usesmetric → Wasserstein Metric
confidence 90% · quantifying posterior accuracy... in terms of the W 1 metric
Ocean-Acoustic Inversion → usesdatafrom → SWellEx-96
confidence 85% · SWellEx-96 experiment conducted off the coast of San Diego near Point Loma
VINA → instantiatedby → iResNet
confidence 80% · concrete instantiation of a practically relevant model class (iResNet models)
VINA → instantiatedby → Coupling-based Architectures
confidence 80% · coupling-based architectures demonstrate improved efficiency
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The distinctive architectural features of normalizing flows (NFs), notably bijectivity and tractable Jacobians, make them well-suited for generative modeling. Invertible neural networks (INNs) build on these principles to address supervised inverse problems, enabling direct modeling of both forward and inverse mappings. In this paper, we revisit these architectures from both theoretical and practical perspectives and address a key gap in the literature: the lack of theoretical guarantees on approximation quality under realistic assumptions, whether for posterior inference in INNs or for generative modeling with NFs. We introduce a unified framework for INNs and NFs based on variational unsupervised loss functions, inspired by analogous formulations in related areas such as generative adversarial networks (GANs) and the Precision-Recall divergence for training normalizing flows. Within this framework, we derive theoretical performance guarantees, quantifying posterior accuracy for INNs and distributional accuracy for NFs, under assumptions that are weaker and more practically realistic than those used in prior work. Building on these theoretical results, we conduct extensive case studies to distill general design principles and practical guidelines. We conclude by demonstrating the effectiveness of our approach on a realistic ocean-acoustic inversion problem.
Tags
Links
- Source: https://arxiv.org/abs/2602.20480v1
- Canonical: https://arxiv.org/abs/2602.20480v1
Trouble viewing inline? Open PDF directly →
Full Text
171,055 characters extracted from source content.
Expand or collapse full text
VINA: Variational Invertible Neural Architectures Shubhanshu Shekhar ∗ EECS Department University of Michigan, Ann Arbor shubhan@umich.edu Mohammad Javad Khojasteh ∗ EME Department Rochester Institute of Technology mjkem@rit.edu Ananya Acharya EME Department Rochester Institute of Technology a2334@rit.edu Tony Tohme Department of Mechanical Engineering Massachusetts Institute of Technology tohme@mit.edu Kamal Youcef-Toumi Department of Mechanical Engineering Massachusetts Institute of Technology youcef@mit.edu Abstract The distinctive architectural features of normalizing flows (NFs), notably bijectivity and tractable Jacobians, make them well-suited for generative modeling. Invertible neural networks (INNs) build on these principles to address supervised inverse problems, enabling direct modeling of both forward and inverse mappings. In this paper, we revisit these architectures from both theoretical and practical perspectives and address a key gap in the literature: the lack of theoretical guarantees on approximation quality under realistic assumptions, whether for posterior inference in INNs or for generative modeling with NFs. We introduce a unified framework for INNs and NFs based on variational unsupervised loss functions, inspired by analogous formulations in related areas such as generative adversarial networks (GANs) and the Precision-Recall divergence for training normalizing flows. Within this framework, we derive theoretical performance guarantees, quantifying posterior accuracy for INNs and distributional accuracy for NFs, under assumptions that are weaker and more practically realistic than those used in prior work. Building on these theoretical results, we conduct extensive case studies to distill general design prin- ciples and practical guidelines. We conclude by demonstrating the effectiveness of our approach on a realistic ocean-acoustic inversion problem. 1 Introduction Recently, machine learning approaches, particularly those based on deep neural networks, have emerged as effective alternatives to conventional inverse problem solvers. Among these, invertible neural networks (INNs) [Ardizzone et al., 2018] stand out due to their ability to model complex, non-linear relationships ∗ Equal contribution. 1 arXiv:2602.20480v1 [cs.LG] 24 Feb 2026 while ensuring invertibility between input and output spaces. INNs offer advantages such as bijectivity and computationally tractable Jacobians, which make them particularly suitable for solving inverse problems, where both forward and inverse mappings need to be computed accurately and efficiently. INNs have been applied for solving inverse problems in various fields, including epidemiology [Radev et al., 2021], astrophysics [Ardizzone et al., 2018], optics [Luce et al., 2023], geophysics [Zhang and Curtis, 2021, Wu et al., 2023], and reservoir engineering [Padmanabha and Zabaras, 2021]. X Y Z Invertible Neural Network forward modeling inverse prediction Figure 1: INNs can be represented by an invertible map T that approximates the relation from the input X ∈ R d x to the output Y ∈ R d y and a latent variable Z ∈ R d x −d y . The invertibility of T means that for any y ∈ R d y , we also have an approximate posterior sampling distribution via T −1 (y,Z). The name INN in the literature has been used to refer to supervised models trained to solve inverse problems, where the training loss involves a super- vised loss (to model observed outputs) as well as an unsupervised component [Ardizzone et al., 2019, Guan et al., 2024, Ardizzone et al., 2018, Hagemann and Neumayer, 2021]. There exists a related class of models based on similar architectural principles, referred to as normalizing flows (NFs), which are primarily used for unsupervised density estimation and generative modeling [Papamakarios et al., 2021, 2017, Gomez et al., 2017, Kingma and Dhariwal, 2018, Dinh et al., 2014]. For ease of exposition, we also use this convention and distinguish between INNs and NFs based on the existence of the super- vised loss in the training process. A typical inverse problem can be described as follows: We are given a datasetD =(X i ,Y i )∈ R d x +d y n i=1 consisting of n input-output pairs, where X i ∈ R d x is the input feature vector and Y i ∈ R d y is the target vector, from a joint distribution (X,Y ) ∼ P X,Y . Usually we have d x ≥ d y representing some inherent information loss in a measurement process (forward map). This joint distribution can be decomposed into a prior P X and a possibly randomized transformation P Y|X representing the forward process. The general goal of the Bayesian approach to inverse problems [Stuart, 2010, Dashti et al., 2012] is to learn or approximate the posterior distribution P X|Y given the dataD. As depicted in Figure 1, the approach taken by INNs to achieve this goal is to incorporate an additional latent output variable, denoted by Z ∈ R d z where d z = d x − d y . This latent variable is designed to represent information related to X that is not contained in Y , and its distribution P Z is selected by the data analyst [Hagemann and Neumayer, 2021]. An INN consists of a map T ≡ (T y ,T z ) : R d x → R d y ×d z that is continuous and invertible with a continuous inverse, and usually they are also assumed to be continuously differentiable to enable gradient based training procedures. Let T denote an appropriate parametrized collection of such functions (we discuss some architectural details in Section A.1). The model is trained by minimizing an objective of the form T ∗ INN ∈ arg min T∈T E XY [∥T y (X)− Y∥ 2 2 ] + λD(P Y,Z ,P Y,T z (X) ) ,(1) where the two terms correspond to the supervised and the unsupervised (population) losses, respectively, and λ > 0 is a regularizing constant. In this display, we use D to denote some divergence or distance measure (such as f -divergence, kernel-MMD, Wasserstein metric, etc.) between the joint distributions of P Y,T z (X) and P Y,Z = P Y P Z . The supervised term encourages the component T y to approximate the forward map P Y|X , while the unsupervised term matches the joint distribution of P Y,T z (X) with the output-latent variable pair P Y,Z . Crucially, due to the invertibility of T , it is possible to show that under certain conditions, making these two terms small induces an approximate posterior sampling distribution: for a given Y = y, the distribution of T −1 (y,Z) for Z ∼ P Z is close (in the same divergence D) to the true posterior P X|Y =y . Note that unlike Hagemann and Neumayer [2021], we use P Y,T z (X) instead of P T y (X),T z (X) as the second argument of D(·,·) to simplify some of the technical arguments. We expect that our results can be extended to the unsupervised loss with P T y (X),T z (X) , and we consider this choice in some of our experiments. In the purely unsupervised setting, the same invertible models reduce to NFs. In particular, we have access to a datasetD =X i ∈ R d n i=1 , and our goal is to estimate the distribution P X of these observations. As in the case of INNs, the idea is to introduce a latent variable Z ∼ P Z , also in R d , and learn an invertible map T : R d → R d that minimizes some notion of divergence between P Z and the forward map T (X); that 2 is, T ∗ NF ∈ arg min T∈T D(P Z ,P T (X) ). By making the divergence between the distributions of Z and T (X) small, the invertibility of T can be used to establish closeness (in the same divergence) between the distributions of the inverse map X, and T −1 (Z), as desired. Within this formulation, in this paper, we view NFs as a purely unsupervised variant of the more general class of INNs. The choice of divergence D in the training objective critically shapes the behavior of INNs (and NFs) as different choices encourage different behavior in the trained model. For instance, forward relative entropy (or KL divergence) induces mode coverage behavior, while reverse relative entropy encourages mode seeking. Jensen-Shannon (JS) divergence provides a natural compromise between these two extremes [Polyanskiy and Wu, 2025]. Other popular choices include the Wasserstein metrics that are known to provide informative gradients, and kernel maximum mean discrepancy (MMD) which is an important instance of the family of integral probability metrics (IPMs) [Sriperumbudur et al., 2012]. In practice, these choices influence not only the final model’s eventual behavior but also the optimization dynamics, sample quality, and robustness to constraints such as limited data and architectural constraints. 1.1 Overview of our contributions We now present an overview of our main contributions, which can be divided into the proposal of unifying variational framework, theoretical characterization of the quality of approximation achieved by the empirical risk minimization (ERM)-based model, experimental case studies that verify and augment the theory, and a real-world ocean acoustic inversion application. A unifying framework. The starting point of our work is the observation that a large class of divergence measures used to define the unsupervised loss term in the INN and NF training objective admit a variational (or dual) representation. This allows an interpretation of the training process as a saddle-point optimization problem. While we present the formal details in Section 3.1, we illustrate the idea using an f -divergence between two distributions on R d , denoted by D f (P ∥ Q), which admits the following classical Donsker- Varadhan (DV) type variational representation [Polyanskiy and Wu, 2025, Theorem 7.26]: D f (P ∥ Q) =sup g:E Q [f ∗ (g(X))]<∞ E P [g(X)]− E Q [f ∗ (g(X))],where f ∗ (b) = sup a∈R ab− f (a) denotes the convex conjugate of f (also known as Fenchel conjugate [Rockafellar, 1970]). With this formu- lation, the training of an INN given the dataset (augmented with samples from the latent distribution P Z ) D = (X i ,Y i ,Z i ) ∼ P X,Y,Z n i=1 , a model class T , and a “critic class” G ⊂ g : E Q [f ∗ (g(X))] < ∞, can be represented as ˆ T n ∈ arg min T∈T 1 n n X i=1 ∥T y (X i )− Y i ∥ 2 2 + λ sup g∈G 1 n n X i=1 g(Y i ,Z i )− 1 n n X j=1 f ∗ g(Y,T z (X)) . Similar expressions can be obtained for the general family of IPMs, Wasserstein-1 (W 1 ) metric using Kantorovich-Rubinstein duality [Villani, 2008]. We refer to this class of models as Variational Invertible Neural Architectures (VINA). Theoretical analysis. In Section 3.2 (INNs) and Section B (NFs), we analyze the quality of the ERM model ˆ T n under the unified variational framework. In both cases, we present our results at three “levels”. At the first level (Theorem 3.1 and Theorem B.1), we state how certain assumptions, such as realizability, uniform learnability, moment bounds, and representation ability of the critic class can allow us to convert the empirical loss into a high probability guarantee on the approximation quality of ˆ T n in terms of the W 1 metric (see Remark 3.3 for justification of this choice). At the next level (Section 3.2.1 and Section B.1), we identify sufficient and more naturally verifiable conditions under which the requirements of the previous level are satisfied. In particular, we present conditions on the Rademacher complexity, Lipschitz regularity of 3 the model class, and certain tightness properties associated with the latent distribution and model class. At the third and final level (Section 3.2.2 and Section B.2), we present a concrete instantiation of a practically relevant model class (iResNet models), and critic class (norm constrained reproducing kernel Hilbert spaces (RKHSs)), and verify that they satisfy all the conditions obtained at the previous level. Unlike existing theoretical guarantees, our results incorporate some practical aspects of training INNs and NFs. Ardizzone et al. [2018] showed that if a model T achieves exactly zero population loss (supervised squared error + unsupervised kernel-MMD loss), then sampling via T −1 (y,Z) exactly recovers the poste- rior P X|y . In practice, however, zero population loss is seldom achieved. Hagemann and Neumayer [2021] addressed this issue by bounding the W 1 distance between the true posterior and the estimated posterior in terms of the population loss (combination of supervised and unsupervised loss). However, their analysis places a rather strong bounded support assumption on the input and output of the network, which excludes common choices such as Gaussian latent variables. In this work, we relax this assumption by only requiring a finite-moment assumption which significantly broadens the scope of our result to more practical scenarios. In fact, we obtain the first formal quantification of the quality of the posterior distribution represented by INNs trained via empirical risk minimization (ERM) in terms of the Wasserstein (W 1 ) metric, and also discuss how it can be extended to other metrics such as kernel-MMD. Empirical Results. We then perform a thorough empirical study to deliver practically useful insights and general design principles that complement the theoretical results mentioned above. More specifically, we study the impact of the different design choices involved in our framework, such as the effect of the choice of divergence / metric, the choice of the latent distribution, and the prior loss. The key insights obtained from these empirical evaluations are summarized below. • In Section 4.1.1, we study the significance of integrating the knowledge of prior distribution on the input space (X = R d x ) into the training of INNs. Our empirical findings show that employing a well-specified prior can significantly improve the performance of INNs, but a misspecified can lead to a degradation in the quality of the generated samples. • In Section 4.1.2, we compare coupling-based and iResNet architectures using f -divergence-based un- supervised losses in both forward and backward directions. We observe that training performance is influenced by multiple factors, including the network architecture, the training procedure, and the capacity of the critic class, with backward f -divergence losses tending to yield lower inference error. In line with prior work in the literature (e.g., Behrmann et al. [2019]), coupling-based architectures demonstrate improved efficiency relative to iResNet-based models. • In Section 4.1.3, we observe an approximately local U-shaped dependence between latent dimension and inference performance of INNs. More specifically, as we change the latent dimension size d z (and appropriately modify d x by padding), we observe that increasing dimension initially improves results, and then starts degrading, and the overall trend is multimodal. • In Section 4.1.4, we investigate the practical aspects of solving inverse problems using Wasserstein Distance [Villani, 2008]. Calculating the Wasserstein distance is computationally challenging and sus- ceptible to the curse of dimensionality [Fournier and Guillin, 2015]. Here, we deploy the entropic estimate of the Wasserstein distance (W 2 ), both Sinkhorn approximation[Cuturi, 2013] and Sinkhorn divergence [Peyr ́e et al., 2019], as a metric within our training process. We observe that lower en- tropic regularization improves sample quality but increases training time for both Sinkhorn Divergence and Sinkhorn approximation of Wasserstein Distance, while Sinkhorn divergence demonstrates greater stability than the Sinkhorn approximation. • In Section 4.1.5, we study the effect of the support of the latent distribution when training with the KL divergence and the Sinkhorn divergence. It is well known that f -divergences perform poorly when the two distributions are mutually singular[Zhang et al., 2019, 2020], while Wasserstein distances can provide informative gradients even when supports do not overlap because it reflects the underlying geometric cost. Our empirical observations in this section validate this hypothesis, and the models trained with the Sinkhorn divergence were more robust to mismatch between the supports of the latent and prior distributions. 4 • In Section 4.1.6, we study the relation between the number of finite moments of X, and the performance of the NF model. As predicted by our theoretical results, we find that as the number of finite moment of X increases, the W 1 distance between the true and estimated distributions decreases. Overall, our exploratory experiments in Section 4.1 provide empirical validation of some of the theoretical predictions of Section 3, and also provide some general practical insights in training invertible architectures. Ocean-Acoustic Application. To demonstrate the practical relevance of our approach, we employ the insights gained from the exploratory experiments in Section 4.1 to the ocean–acoustic inversion setting of the SWellEx-96 experiment conducted off the coast of San Diego near Point Loma [Yardim et al., 2010, Meyer and Gemba, 2021].Geoacoustic inversion (GI), as demonstrated in the SWellEx-96 study, is a challenging and computationally intensive problem [Dosso and Dettmer, 2011, Huang et al., 2006, Chapman and Shang, 2021]. This paper utilizes synthetic data to simplify the task for the initial application of an INN- based framework. These preliminary studies revealed that the latent-space dimensionality plays a critical role in estimating the posterior, and that physically meaningful priors, such as uniform priors for uncertain quantities like sound speed, improve stability. We demonstrate that invertible architectures can offer a favorable computational trade-off relative to likelihood-based MCMC sampling for GI, in which likelihood evaluations rely on repeated calls to the forward model KRAKEN [Porter, 1992]. Although no method is universally best, moving this computational burden offline into training enables a pre-trained INN to support rapid, near-real-time posterior inference at test time. To summarize, our work advances both the theory and practice of invertible neural architectures. We provide a first unified ERM analysis of invertible models that provides explicit bounds (in W 1 metric) on the quality of approximation achieved by the trained model under realistic moment and capacity assumptions. Through concrete instantiations we show that these conditions are satisfied by practically useful models, and finally we implement a series of empirical case studies and apply our ideas to a real-world ocean-acoustic inversion problem. 1.2 Organization of the paper The remainder of the paper is organized as follows. We present a thorough review of the related work in Section 2, and then recall some background on the architecture and training details of existing invertible models. We present our main results in Section 3, and in particular, propose our variational training strategy in Section 3.1 and derive theoretical results for INNs in Section 3.2. Also, we extend our theoretical analysis to variational NFs in Section B. We then move on to the empirical part of the paper in Section 4. In particular, in Section 4.1 we present a series of observations about the effect of various design choices on the performance of invertible models. Some of these observations verify the theoretical predictions from the previous section (such as the effect of the complexity of critic class, and the effect of the number of finite moments of the latent distributions), while others concern important aspects such as the latent dimension, choice of divergence measure etc. Finally, in Section 4.2, we apply our ideas to a real-world task of ocean- acoustic inversion. 2 Related Works Invertible architectures (NFs and INNs): There are numerous ways to achieve invertibility in neural architectures [Tabak and Turner, 2013, Kobyzev et al., 2020, Keller et al., 2021], such as by using residual connections [Gomez et al., 2017, Jacobsen et al., 2018, Behrmann et al., 2019], triangularization [Bogachev et al., 2005, Marzouk et al., 2016, Parno et al., 2016, JM and Welling, 2017], and coupling-based normalizing flows (e.g., RealNVP Dinh et al. [2016] and Glow Kingma and Dhariwal [2018]). In practice, coupling-based architecture strikes the right balance between computational efficiency and expressivity. In fact, under appropriate assumptions, they have been shown to be universal diffeomorphism approximators [Teshima et al., 2020, Jin et al., 2024]. While NFs were developed largely for tractable density modeling, their supervised variants INNs have become increasingly popular for solving inverse problems. We explore this particular direction in Section 4.2. However, INNs (as defined in our paper) are not the only way to employ invertible neural networks to solve 5 inverse problems. There exist other approaches, such as using conditional INNs [Ardizzone et al., 2019, Winkler et al., 2019]. Another approach is to use conditional NFs to learn the likelihoods [Papamakarios et al., 2019], which can be used to estimate posterior by combining it with a prior distribution using Bayes rule. NFs have also been integrated with sampling-based Bayesian inference, to combine the advantages of both approaches [Song et al., 2017, Winter et al., 2023, Kruse et al., 2025], as will be discussed further below. Our research focuses on INNs, which provide a probabilistic framework for addressing inverse problems. However, it’s important to recognize the growing popularity of neural networks in solving inverse prob- lem. For instance, In non-probabilistic contexts, works Ying [2022], Fan and Ying [2019], Khoo and Ying [2019] have designed specialized neural architectures that embed physical formulation to recover unknown parameters while reducing the reliance on large amounts of data. INNs vs. sampling-based Bayesian methods: Conventional Markov Chain Monte Carlo (MCMC) methods [MacKay, 2003, Brooks et al., 2011, Andrieu et al., 2003, Doucet and Wang, 2005, Korattikara et al., 2014, Kungurtsev et al., 2023, Atchad ́e and Rosenthal, 2005] are widely used for sampling from complex probability distributions, including posterior distributions arising in inverse problems [Geweke, 1989]. From a pushforward viewpoint, MCMC defines a Markov transition kernel with the target posterior as its stationary distribution; under standard ergodicity conditions, the chain converges asymptotically to the exact posterior. However, standard likelihood-evaluation-based MCMC methods often require substantial computational time to achieve adequate sampling [Roy, 2020, Jones and Qin, 2022], which becomes a bottleneck, particularly in inverse problems where the forward process is computationally expensive to evaluate. Likelihood-free MCMC variants exist but can be substantially more computationally demanding in practice [Beaumont, 2019]. On the other hand, invertible generative models parameterize an invertible map that pushes a simple base distribution to an induced (learned) distribution that approximates the target distribution. When the Jacobian determinant is tractable, this learned distribution can be evaluated exactly via the change-of- variables formula. In context of inverse problem, once an INN is trained, approximating the posterior involves running the network in reverse, which is computationally inexpensive. We demonstrate that invertible architectures can offer a favorable computational trade-off relative to repeated likelihood-evaluation-based MCMC sampling for a practical inverse problem. This speedup comes from shifting computation offline into training. In fact, under appropriate assumptions, INNs can even perform in amortized settings, where a model trained extensively can generalize effectively to new instances without retraining [Radev et al., 2023, 2021]. Since practically relevant performance guarantees on invertible architecture are scarce in literature (we discuss some existing results later in this section), Gabri ́e et al. [2021, 2022], Brofos et al. [2022], Sch ̈onle and Gabri ́e [2023] have incorporated NFs to characterize the proposal distributions of MCMC methods, thereby accelerating convergence in practical tasks, while retaining theoretical guarantees for MCMC convergence. In this work, we establish error bounds for our proposed class of invertible models under appropriate as- sumptions, providing theoretical characterizations of posterior accuracy for INNs and generative accuracy for NFs. Taken together, our results clarify when offline-trained invertible maps with controlled approximation error can replace repeated online likelihood-based Monte Carlo sampling. Posterior sampling in high dimensions is often computationally challenging [Montanari and Wu, 2023], and there is no universally best configuration for these problems. As demonstrated in our empirical studies, INNs exhibit sensitivity to several design and optimization choices such as the learning rate and architectural configuration. To address this, in our cases studies, we employed tree-structured Parzen estimator [Bergstra et al., 2011] to search for promising hyperparameters including: the size of the subnetworks, the latent dimensionality, and the number of coupling layers. Unsupervised loss: The choice of training objective critically shapes the behavior of INNs as these models are trained through optimization. INNs were initially trained using MMD [Ardizzone et al., 2018], which only requires samples and avoids the need for explicit density modeling. However, due to challenges such as poor scaling in high dimensions and sensitivity to kernel choice [Ramdas et al., 2015], negative log- likelihood (NLL) [Dinh et al., 2016] has become more prevalent in recent INN architectures [Ren et al., 2020, Kruse et al., 2021]. The use of NLL can be justified as an approximation of training by minimizing the KL- divergence [Kingma et al., 2019, Papamakarios et al., 2021]. Training an INN using NLL in the literature is 6 Table 1: Definitions of important acronyms. NFnormalizing flowINNinvertible neural network IKinverse kinematicsBNNBayesian neural network NLLnegative log-likelihoodMMDmaximum mean discrepancy ERMempirical risk minimizationKLKullback-Leibler JSJensen–ShannonIPMintegral probability metric SLsupervised lossUSLunsupervised loss GANgenerative adversarial networkGIGeoacoustic inversion based on making assumptions about the distribution of the output samples, for instance Gaussian distribution has been used in prior work [Ren et al., 2020, Kruse et al., 2021]. In this paper, we eliminate the need for such distributional assumptions by relying solely on samples and leveraging the variational formulation of f-divergences, as introduced by Nguyen et al. [2010] and follow-up works [Nowozin et al., 2016, Ruderman et al., 2012, Ke et al., 2021]. Our approach is inspired by analogous formulations in related domains, such as generative adversarial networks (GANs) [Goodfellow et al., 2014, Nowozin et al., 2016], Flow-GAN [Grover et al., 2018] and the Precision-Recall divergence for training NFs [Verine et al., 2023]. To train invertible architectures, we take a unified look at the usage of a class of variational distance metrics over the space of probability measures, such as IPMs [Sriperumbudur et al., 2012], which include the MMD [Li et al., 2017] and Wasserstein metrics [Arjovsky et al., 2017, Coeurdoux et al., 2022], and the f -divergence family [Csisz ́ar, 1967, Arjovsky and Bottou, 2017], and we study its implication in supervised inverse problems (variational INN) and in unsupervised generative modeling tasks (Variational NF). Training based on the variational representation of f-divergence enables stronger theoretical guarantees. By extending and unifying the results of Nguyen et al. [2010] and Hagemann and Neumayer [2021], we derive new convergence bounds for INN that are both tighter and hold under weaker assumptions than prior work. In particular, Hagemann and Neumayer [2021, Theorem 2] establish an approximation guarantee for the posterior distribution of an INN represented by a homeomorphism T in terms of the error bounds on the supervised and unsupervised losses. Our results extend their result in two main ways: first we relax a strong requirement of bounded observations imposed by Hagemann and Neumayer [2021] that omits usual choices like Gaussian latent distributions, and second, we analyze the performance of a data-driven model ˆ T n trained using an ERM strategy, unlike Hagemann and Neumayer [2021] who work with a fixed T satisfying certain approximation guarantees. Latent distribution: It is well established that the structure and topology of the latent space play a crucial role in the accuracy of generative models [Gurumurthy et al., 2017, Bevins et al., 2023, Stimper et al., 2022, Laszkiewicz et al., 2021, Hickling and Prangle, 2024, Fadel et al., 2021]. While the most commonly used latent distributions for INNs is the standard Gaussian distribution [Ardizzone et al., 2018], the works in Behrmann et al. [2021], Hagemann and Neumayer [2021] shows that a simple Gaussian prior is often insufficient for multimodal problems and choosing a suitable latent distribution can improve robustness. In this paper, we study the effect of size of latent variable and support set of this variable in performance of INN. A selective discussion of the existing architectural choices and design of INNs is provided in Section A.1. 3 Main Results We begin this section by presenting a general unsupervised training cost for variational invertible neural architecture (VINA) in Section 3.1. As mentioned earlier, this formulation unifies several commonly used cost functions in the literature. For simplicity we focus our presentation in Section 3.1 for the case of INNs, but we then illustrate how it naturally applies to NFs in Section B. 7 3.1 Variational unsupervised loss (USL) Our main methodological contribution in this paper begins with the observation, also used in priors works such as Nowozin et al. [2016], Zhang et al. [2019], Grover et al. [2018], that several commonly used statistical divergence measures used in generative modeling admit variational representations. More specifically, with T : R d x → R d y × R d z denoting an INN, we wish to minimize some divergence of the form D φ (P Y,Z ,P Y,T z (X) ) = sup g∈G φ E [φ(g,X,Y,Z,T )], emp. loss =⇒ ˆ L = sup g∈G φ 1 n n X i=1 φ (g,X i ,Y i ,Z i ,T ).(2) Here G φ denotes some class of functions, usually from X = R d x to R, and we will refer to it as the “critic class” following the convention used in the generative adversarial network (GAN) literature [Goodfellow et al., 2014]. In practice, this class can be represented by machine learning models such as neural networks, or more analytically tractable classes such as RKHSs. Our theoretical results in the next two subsections explore the trade-offs involved in choosing more expressive G φ and approximation guarantees. For different choices of (φ,G φ ), the above formulation recovers various popular loss functions such as relative entropy (and more generally, the f -divergence family), kernel-MMD (and more generally, the IPM family), energy distance, and the Wasserstein metric. Besides the conceptual unification, this approach also allows us to obtain theoretical guarantees on the approximation performance of INNs trained via empirical risk minimization strategy, as we discuss in Section 3.2 for INNs (and in Section B for NFs). Before proceeding further, we recall two important family of distance or divergence measures that are realizations of (2). The first class of distance metrics are the IPMs, which for some function class G that is closed under negation (i.e., if g ∈G, then so does −g), are defined as IPM(P Y,Z ,P Y,T z (X) ) = sup g∈G E [g(Y,Z)− g(Y,T z (X))] which corresponds to φ IPM (g,x,y,z,T ) = g(y,z)− g(y,T z (x)). Perhaps the most important element of the IPM family is the kernel-MMD metric, where the critic class (also known as the witness class) is the unit norm ball RKHS associated with a positive-definite kernel k : X ×X → R. An immediate consequence of the reproducing property of such G is that for square integrable kernels, we have MMD(P Y,Z ,P Y,T z (X) ) =∥μ P Y,Z − μ P Y,T z (X) ∥ k ,where μ P = Z k(x,·)dP (x), is referred to as the kernel mean-embedding of a distribution P ,∥·∥ k denotes the RKHS norm, and R above is a Bochner integral. This leads to a very natural interpretation of MMD metric between two distributions P and Q: it is the distance (as measured by the RKHS norm) between the two representative elements μ P and μ Q associated with the two distributions. Other important IPMs include the total variation distance, and the Wasserstein 1 metric (also known as the earth-mover distance), which by using Kantorovich-Rubinstein duality (Theorem 1.14 Villani [2021], and Theorem 11.8.2. Dudley [2018]) can be represented as an IPM associated with the class G of 1-Lipschitz functions. The second important class of divergence measures captured by (2) is the family of f -divergences. For any convex, lower semicontinuous f : (0,∞) → (0,∞) with f (1) = 0, the f -divergence between P Y,Z and P Y,T z (X) , assuming that P Y,Z ≪ P Y,T z (X) , is defined as D f (P Y,Z ∥ P Y,T z (X) ) = Z q T (y,T z (x))f p Y,Z (y,z) p T (y,T z (x)) ν(x,y,z)dxdydz, where we have assumed that both distributions admit densities p T and p Y,Z , with respect to some common dominating measure ν, for all choices of T in the model class. These divergence measures are also known to admit the following variational representation D f (P Y,Z ∥ P Y,T z (X) ) = sup g (E[g(Y,Z)− f ∗ (g(Y,T z (X)))]),(3) which corresponds to φ f (g,x,y,z,T ) = g(y,z)− f ∗ (g(y,T z (x)). Here, the supremum is over all measurable g : Y ×Z → R make the right-hand side (RHS) well-defined (by avoiding ∞−∞). This implies that if 8 we restrict our attention to any smaller class of functions G, get a lower bound on the divergence, and this gap (that we later refer to as the variational gap) decreases by enlarging the critic class G. This fact will play a central role in our analysis in the next two subsections. The canonical member of the f -divergence family is the (forward) relative entropy (or Kullback- Leibler (KL) divergence), with generator function f (u) = u logu. The corresponding φ function in (2) is φ KL (g,x,y,z,T ) = logg(y,z)−g(T (x)) + 1. Other important f -divergences are the squared Hellinger diver- gence with f (u) = ( √ u− 1) 2 and φ H (g,x,y,z,T ) = g(y,z)−g(T (x))/(1−g(T (x))), and the Jensen Shannon divergence with f (u) =−u(+1) log((1 + u)/2) + u logu and φ JS (g,x,y,z,T ) = g(y,z) + log 2− e g(T (x)) . We now proceed to a discussion of INNs defined using a variational objective as introduced in (2). 3.2 Variational invertible neural networks (V-INN) Let X = R d x , Y = R d y and Z = R d z , with d x = d y + d z , and suppose we have a dataset (X i ,Y i ,Z i ) ∈ X ×Y ×Z : 1 ≤ i ≤ n drawn i.i.d. from the joint distribution P XY Z . Our goal is to learn an invertible neural network (INN) from a family of homeomorphisms (i.e., each T :X →Y×Z is a continuous bijective map with a continuous inverse) using this training dataset. Each candidate INN T can be decomposed into T y :X →Y and T z :X →Z, and consequently, the INN training process involves population loss functions: L y (T ) = E P XY ∥T y (X)− Y∥ 2 2 ,and L z (T )≤ D(P Y,Z ,P Y,T z (X) ), where D(P,Q) is an appropriate notion of distance or divergence between probability distributions P,Q ∈ P(Y ×Z). In this section, we will focus on the case of D being an f -divergence and employ its variational definition stated in (3) within the empirical risk minimization (ERM) framework. The use of an “≤” instead of “=” when introducing L z (T ) above is a consequence of the variational definition with a restricted critic class as we make precise in (INN4) in Assumption 3.1. More explicitly, we work with a function class G n (the “critic class”), and define L z (T ) as L z (T ) = sup g∈G n E P YZ [g(Y,Z)]− E P XY [f ∗ (g(Y,T z (X)))]≤ D f (P Y,Z ,P Y,T z (X) ).(4) To train an INN from a family T (i.e., T might represent all INNs with a fixed architecture), we use the empirical analogs of the population loss terms, b L y,n (T ) = 1 n n X i=1 ∥T y (X i )− Y i ∥ 2 2 , b L z,n (T ) = sup g∈G n ( 1 n n X i=1 g(Y i ,Z i )− 1 n n X i=1 f ∗ ◦ g(Y i ,T z (X i )) ) , where f ∗ denotes the convex conjugate of f . The model learned by ERM can then be defined as ˆ T n ∈ arg min T∈T b L y,n (T ) + λ b L z,n (T ),(5) for some regularization parameter λ > 0. The first loss term ( b L y,n ) forces the learned map ˆ T n to approximate Y from an input X, while the second loss term ( b L z,n ) ensures that the joint law of (Y, ˆ T n,z (X)) matches the true joint distribution (Y,Z). Thus, intuitively, if the sample-size n is large enough, we should expect ˆ T n to be a good proxy for T ∗ , an element of T that minimizes the population risk: T ∗ ∈ arg min T∈T L y (T ) + λL z (T ). The performance of ˆ T n defined above is governed by two effects: (i) the finite sample effect caused by the uniform deviations between the empirical objectives b L y,n and b L z,n from their population counterparts; and (i) the variational or approximation error induced by restricting the divergence representation to a classG n . To make the variational loss smaller, we need to increase the capacity of G n . However, that also causes the uniform deviations between the empirical and population losses to increase. Thus, the crucial challenge is to find the right trade-offs between these two effects to simultaneously drive the overall error to zero. For the rest of this section, we follow a three-step roadmap. First, we state a general result (Theorem 3.1) that obtains an upper bound on the quality of the posterior approximation provided by ˆ T n under a set of high-level assumptions (stated in Assumption 3.1). Second, we present more concrete, verifiable, sufficient 9 conditions for satisfying the requirements of Assumption 3.1 in Section 3.2.1. Finally, we specialize the general theorem to a specific choice of INN architecture (iRes-Nets in Definition 3.1) and critic class (RKHS over truncated domain in Definition 3.2) using Jensen-Shannon divergence in Section 3.2.2. We now present the assumptions required for stating the main result of this section. Assumption 3.1. To analyze the performance of the INN model ˆ T n defined in (5), we place the following assumptions: • (INN1): Realizability and Bi-Lipschitz. There exists a T ∗ ∈ T such that Y = T ∗ y (X) almost surely (a.s.) and P Y,Z = P Y,T ∗ z (X) , and furthermore sup T∈T Lip(T ), Lip(T −1 ) ≤ J <∞. • (INN2): Uniform Convergence. For any confidence level δ ∈ (0, 1), there exist deterministic sequences r n ≡ r n (δ) : n ≥ 1 and u n ≡ u n (δ) : n ≥ 1, with r n ,u n → 0, such that the following conditions hold with probability at least 1− δ: sup T∈T | b L y,n (T )− L y (T )|≤ u n ,andsup T∈T | b L z,n (T )− L z (T )|≤ r n . • (INN3): Moment Bounds. There exists an R <∞, such that the following holds for some a > 0: max E[∥(Y,Z)∥ 1+a ], sup T∈T E[∥(Y,T z (X)∥ 1+a ] ≤ R. Here ∥·∥ to denotes the ℓ 2 norm, and we use ∥(y,z)∥ as a shorthand for p ∥y∥ 2 +∥z∥ 2 . • (INN4): Variational Approximation Gap. The function class G n contains the 0 function for all n≥ 1, and there exists a vanishing deterministic sequence η n → 0, such that sup T∈T D f (P Y,Z ∥ P Y,T z (X) )− L z (T )≤ η n . In other words, we assume that the capacity of the critic class G n grows with n to approximate the likelihood ratios of all distributions modeled by elements of T and the true data distribution P Y,Z . Theorem 3.1. Suppose Assumption 3.1 holds, and D f satisfies the Pinsker-type inequality c f p D f (P ∥ Q)≥ TV (P,Q) for all distributions P,Q, and for some constant c f > 0. Then, for every measurable A⊂Y with P Y (A) > 0, on the (1− δ) probability event of (INN2), we have W 1 P (A) X ,P (A) ˆ T −1 n (Y,Z) ≲ (r n + u n + η n ) a 2(1+a) , where W 1 denotes the 1-Wasserstein metric, P (A) X = P X|Y∈A , P (A) ˆ T −1 n (Y,Z) = P ˆ T −1 n (Y,Z)|Y∈A , and ≲ suppresses the constant factors and lower order terms. The exact expression of the upper bound is in (32). The proof of this result is in Section C.1.Note that the Pinsker-type inequality TV (P,Q) ≤ c f p D f (P ∥ Q) is satisfied by several divergences, such as relative entropy, Jensen Shannon, Hellinger, and chi-squared divergence. Remark 3.1. The result of Theorem 3.1 is obtained by combining two ingredients: (i) using the Pinsker-type inequality, we convert the control over the given f -divergence into control over total variation, and (i) a truncation argument (Lemma D.1) that converts the total variation into a bound on W 1 under just finite (1 + a) moment condition. One crucial advantage of our result over Hagemann and Neumayer [2021, Theorem 2] is that we do not impose the bounded support condition on the random variables (X,Y,Z). The mild (1 + a) moment requirement significantly broadens the applicability of our result, and in particular, allows us to consider more realistic models used in practice. Additionally, the dependence on the parameter a quantifies how the heavy tails affect the approximation quality. We empirically verify this insight in Section 4.1.6. 10 Remark 3.2. Although Theorem 3.1 characterizes the distance between the true (and unknown) posterior and the INN posterior in terms of W 1 metric, we can use existing inequalities to translate it into other dis- tances. For example, by Sriperumbudur et al. [2010, Theorem 21], we know that MMD(P (A) X ,P (A) ˆ T −1 n (Y,Z) ) ≲ W 1 (P (A) X ,P (A) ˆ T −1 n (Y,Z) ) for commonly used kernels, such as the Gaussian and Mat ́ern kernels, and thus Theo- rem 3.1 also implies a bound on the kernel-MMD distance between the true and estimated posteriors. Remark 3.3. A closer look at the proof of Theorem 3.1 reveals several reasons why W 1 metric is the appropriate choice to present the result: (i) By Kantorovich-Rubinstein duality, the W 1 metric is defined via Lipschitz critic functions. This fact coupled with the assumed bi-Lipschitz structure (INN1) allows us to transfer bounds from posterior distributions to the forward distributions. (i) Since W 1 is a metric, it satisfies the triangle inequality which leads to a natural decomposition in the proof into terms that correspond directly to the supervised and unsupervised training losses. (i) The Lipschitz formulation of W 1 is exactly what makes our truncation argument(Lemma D.1) effective for handling heavy tailed distributions under a mild (1 + a) moment assumption. Remark 3.4. While stating Theorem 3.1, we only needed to explicitly place a (1 + a) moment requirement, for some a > 0. However, since L y consists of squared losses, the existence of the uniform convergence rates in (INN2) implicitly places a stronger (2 + β) moment conditions, for some β > 0. We will make this explicit when identifying verifiable sufficient conditions of these assumptions in Proposition 3.1. 3.2.1 Verifiable Sufficient Conditions for Assumption 3.1 Theorem 3.1 gives us a general result that translates uniform convergence and the variational approximation gap guarantees on the two losses, L y ( ˆ T n ) and L z ( ˆ T n ) into a bound on the W 1 metric between the true posterior, and the INN posterior. We now present more easily verifiable sufficient conditions for the assump- tions (INN2)-(INN4) required by Theorem 3.1, working under the realizability part of the assumption (INN1). We begin with the usual characterization of uniform convergence in terms of the Rademacher complexities (see Definition A.1 in Section A) of the associated function classes. Proposition 3.1. Assume that the critic class G n consists of uniformly bounded functions; that is, sup g∈G n ∥g∥ ∞ ≤ b n , for some b n < ∞. Let F y = x 7→ ⟨u,T y (x)⟩ : T ∈ T , ∥u∥ ≤ 1, and define its Rademacher complexity as R n (F y ) = E ε n ,X n " sup f∈F y n X i=1 ε i f (X i ) # ,with ε n i.i.d. ∼ Rademacher(−1, +1). Introduce the two critic class complexities R (1) n (G n ) = E ε n ,Y n ,Z n " sup g∈G n 1 n n X i=1 ε i g(Y i ,Z i ) # ,R (2) n (G n ,T ) = E ε n ,X n ,Y n " sup T,g 1 n n X i=1 ε i g(Y i ,T z (X i )) # . Suppose maxE[∥Y∥ 2+β ], sup T∈T E[∥T y (X)∥ 2+β ]≤ R y <∞ for some β > 0, and Π K n : R d y →y ∈ R d y : ∥y∥ ≤ K n denoting the projection on a ball of radius K n in Y = R d y , define the projected empirical and population supervised losses: b L K n y,n (T ) = 1 n n X i=1 ∥Π K n (T y (X i ))− Π K n (Y i )∥ 2 ,and L K n y (T ) = E ∥Π K n (T y (X))− Π K n (Y )∥ 2 . Assume that the function f (in fact, its convex conjugate f ∗ ) satisfies A 1,n : = sup |u|≤b n |(f ∗ ) ′ (u)| <∞,and A 2,n : = sup |u|≤b n |f ∗ (u)| <∞.(6) Then, the following two conditions hold with probability at least (1− δ), for a given δ in(0, 1): sup T∈T | b L z,n (T )− L z (T )| ≲R (1) n + A 1,n R (2) n + (b n + A 2,n )n −1/2 ,(7) sup T∈T | b L y,n (T )− L y (T )| ≲ K n R n (F y ) + K 2 n n −1/2 + 4R y K −β n ,(8) 11 where ≲ suppresses constants that depend on δ and R y . Thus, a sufficient condition to satisfy (INN2) with vanishing sequences of r n ,u n : n≥ 1, is to select an appropriate sequence of K n ,b n , and G n to drive both these terms to 0. The proof of (7) relies on some standard arguments from empirical process theory [Shorack and Wellner, 2009], while the justification of (8) requires a combination of symmetrization and vector contraction along with a truncation argument, and we present the details in Section C.2. Next, we present a simple result stating that the uniform moment bound required in (INN3) can be satisfied if the model class T is uniformly Lipschitz. As we will observe in Section 3.2.2, this condition holds for an important class of INNs. Proposition 3.2. Suppose there exist constants J 0 ,J < ∞, such that with ∥·∥ denoting the ℓ 2 norm, we have ∥T (x)∥≤ J 0 + J∥x∥, for all x∈X , and for all T ∈T . Then, assuming that E[∥X∥ 1+a ] <∞ for some a > 0, the following uniform moment bounds hold under the realizability assumption: max E ∥(Y,Z)∥ 1+a , sup T∈T E ∥(Y,T z (X)∥ 1+a ≤ R,with R = 2 2a+1 J 1+a 0 + J 1+a E[∥X∥ 1+a ] . Proof. Since we are working under the realizability assumption, we have (Y,Z) d = T ∗ (X). Using the fact that (x + y) p ≤ 2 p−1 (x p + y p ), with p = 1 + a, and the condition that ∥T (x)∥ ≤ J 0 + J∥x∥, we have E [∥(Y,Z)∥ p ]≤ 2 p−1 (J p 0 + J p E[∥X∥ p ]). Next, we look at the term (Y,T z (X)) and observe that ∥(Y,T z (X))∥ p ≤ 2 p−1 (∥Y∥ p +∥T z (X)∥ p )≤ 2 p−1 (∥(Y,Z)∥ p +∥T (X)∥ p ) which implies that E [∥(Y,T z (X))∥ p ]≤ 2 p−1 E[∥(Y,Z)∥ p ] + 2 2(p−1) (J p 0 + J p E[∥X∥ p ]) ≤ 2× 2 2(p−1) (J p 0 + J p E[∥X∥ p ]). Taking the maximum of the two bounds gives us the required expression for R. Proposition 3.3. Introduce the notation P ≡ P Y Z and Q T ≡ P Y,T z (X) , and assume that for all T ∈T , we have P ≪ Q T . For each T ∈T , let g ∗ T denote a maximizer in the definition of the population L z (T ) in (4), and let ̄g T = clip(g ∗ T ,−b n ,b n ) denote its clipped version for b n > 0. For some K n > 0, let B K n denote the ball u∈ R d y +d z :∥u∥≤ K n , and define the following terms: • An approximation error term δ n ≡ δ n (b n ,K n ) defined as δ n ≡ δ n (b n ,K n ) : = sup T∈T inf g∈G n E P [| ̄g T − g|1 B K n ] + E Q T [| ̄g T − g|1 B K n ] . • The tightness term τ 1 ≡ τ 1 (K n ), and the clipping error term τ 2 ≡ τ 2 (b n ), as τ 1 (K n ) = sup T∈T P (B c K n ) + Q T (B c K n ), τ 2 (b n ) = sup T∈T E P [(|g ∗ T |− b n ) + ] + A 1,n E Q T (|g ∗ T |− b n ) + , where we use (a) + to denote max0,a for any a∈ R. Suppose there exists a sequence of K n ,b n : n≥ 1 such that lim n→∞ τ 2 + A 1,n δ n + (b n + A 2,n )τ 1 = 0,(9) where A 1,n and A 2,n were defined in (6). Then the variational gap η n also converges to zero; that is, if (9) is true,=⇒ η n = sup T∈T D f (P X ∥ Q T )− L z (T ) n→∞ −→ 0. 12 The proof of this result is in Section C.3. The structure of (9) is worth further discussion. The term δ n is a purely approximation property of the critic class G n , and it enforces the requirement that every optimal critic function can be approximated in L 1 by a bounded element in G n with vanishing error. The terms τ 1 and τ 2 can be interpreted as tail approximation terms: τ 1 is controlled by the moment bounds from (INN3), while τ 2 depends on how quickly the tails of the optimal critic function decays with the clipping level b n . To summarize, the results of this section provide a concrete recipe to apply the more abstract Theorem 3.1 to a particular instantiation of the variational INN pipeline. In particular, for a chosen model and critic class pair, it suffices to show that the Rademacher complexities of the associated function classes decay sufficiently quickly to 0 (Proposition 3.1), to establish a uniform bi-Lipschitz property, and through it the moment bounds (Proposition 3.2), and show that the critic class and well-approximate the optimal witness function on high probability regions (Proposition 3.3). In the next subsection, we illustrate these steps for a specific variational INN architecture with Gaussian RKHS critic class. 3.2.2 INN Example In this section, we instantiate our abstract assumptions and conditions from Section 3.2.1 and Assumption 3.1 with a concrete class of INN models and critic function classes. The main goal of this section is demonstrate the the conditions derived in the previous sections are satisfied by a non-trivial, reasonably expressive, and practically relevant pair of model and critic classes. In particular, we show that a residual INN with bounded weights and biases, trained with a Jensen-Shannon based loss and a Gaussian RKHS critic class, fits into our framework developed in the previous sections. We begin with a formal description of the model class. Definition 3.1. As before, we assume d x = d y + d z , and for an integer M ≥ 1 and s ∈ (0, 1), define T : R d x → R d x as follows, with I d x denoting the d x × d x identity matrix: T = (I d x + F M )◦ (I d x + F M−1 )◦·◦ (I d x + F 1 ), F j (x) : = W j,2 tanh(W j,1 x + b j,1 ) + b j,2 , with the following constraints for all j ∈ [M ] =1,...,M: ∥W j,1 ∥ op ≤ s, ∥W j,2 ∥ op ≤ s,max∥b j,1 ∥,∥b j,2 ∥≤ B. For every j ∈ [M ], the matrices W j,1 and W j,2 have dimensions d j × d x and d x × d j respectively, and let H : = max j∈[M ] d j . We split the output coordinates to write T (x) = (T y (x),T z (x)), with T y (x) ∈ Y = R d y and T z (x) ∈ Z = R d z . It is easy to verify that this function class is bi-Lipschitz and there exist constants J 0 ,J (depending on M,B,s,H ) such that for all x ∈ X , we have ∥T (x)∥ ≤ J 0 + J∥x∥ and ∥T −1 (y,z)∥ ≤ J 0 + J∥(y,z)∥. This choice of the INN model classT above is driven by the need for a global bi-Lipschitz bounds on both the forward and inverse maps. The residual architecture with tanh activations and the norm constraints on the weight and bias terms provide a convenient way of controlling global Lipschitz constants. In contrast the popular coupling-based architectures such as RealNVP of Dinh et al. [2014] generally do not admit such global control of Lipschitz constants. We now present the details of our critic class construction. Definition 3.2. Let k n ≡ k γ n denote the Gaussian kernel on R d y +d z , k n ((y,z), (y ′ ,z ′ )) = exp −γ n ∥(y,z)− (y ′ ,z ′ )∥ 2 = exp −γ n ∥y− y ′ ∥ 2 +∥z− z ′ ∥ 2 , and let H k n denote its RKHS. Fix a radius parameter K n > 0 and a uniform upper bound b n > 0 and define the critic class as G n = g = h1 C K n : h∈H k n , ∥h∥ k n ≤ b n ,with C K n = [−K n ,K n ] d x . Note that 0∈G n , and each g ∈G n also satisfies ∥g∥ ∞ ≤ b n due to the reproducing property and the fact that sup x p k n (x,x) = 1. This particular choice of the critic classG n is chosen as a compromise between analytical tractability and representation power. The RKHS associated with Gaussian kernels admit clean bounds on their Rademacher complexity, and are also known to be universal in the sense that they are approximate any continuous 13 bounded function over compact domains arbitrarily well in sup norm. These are the exact properties that we need in Proposition 3.1 to ensure uniform convergence, and in Proposition 3.3 for controlling the variational approximation gap. The additional truncation on cube C K n by the indicator function is a purely technical choice: it ensures the uniform boundedness of the critic class while also decoupling the approximation on C K n from tail events, which can be handled separately via moment bounds. Furthermore, since the ℓ 2 -ball B K n used by Proposition 3.3 is contained in C K n , any approximation bound C K n is also valid on B K n . We now present the choice of the f -divergence to be used for the unsupervised loss. Definition 3.3. We will employ a Jensen-Shannon divergence based unsupervised loss in our INN training. This corresponds to f (t) = x log 2x x+1 + log 2 x+1 , and f ∗ (t) = log(1 + e t )− log 2 with A 1,n = sup |t|≤b n |(f ∗ ) ′ (t)|≤ 1,and A 2,n = sup |t|≤b n |f ∗ (t)|≤ log((1 + e b n )/2)≤ b n . Having introduced all the main components, we can now present the main result of this section. Theorem 3.2. Let (T ,G n ) denote a pair of model and critic classes as introduced in Definition 3.1 and Def- inition 3.2 respectively, and define the parameters K n = M (s √ H + B) + p d x + p 2 logn, γ n = K 2+ε n for some fixed ε > 0, b n = C b γ d x /4 n K 2+d x /2 n = C b K 2+d x + d x 4 ε n . Then, under the realizability assumption (INN1), and with these choices of the parameters, the ERM model ˆ T n defined in (5) satisfies W 1 P X|Y∈A , P ˆ T −1 n (Y,Z)|Y∈A = o(1),w.p. at least 1− δ. In other words, with these parameters, the models introduced above satisfy the sufficient conditions for (INN2)-(INN4) to hold, as derived in Section 3.2.1. The proof of this result is in Section C.4. As mentioned earlier, the purpose of this theorem is to illustrate that our abstract theoretical results are applicable to a non-trivial, reasonably realistic variational INN pipeline. In particular, Theorem 3.2 demonstrates that one can work with architectures that are expressive enough to model high-dimensional data and still verify the technical assumptions required for our posterior accuracy guarantees. Extending our arguments to other classes of invertible architectures, such as coupling-based flows equipped with appropriate spectral or Jacobian regularization, is an interesting direction for future work. 4 Empirical Results In this section, we first present a series of small-scale probing experiments to understand the effects of various design choices involved in training invertible models (§ 4.1), and then apply the insights gained to a practically relevant ocean-acoustic inversion problem (§ 4.2). Our objective in Section 4.1 is two-fold: • The first objective is to provide empirical support for several theoretical claims derived from the results of the previous section on INNs and from Section B on NFs. In particular, we investigate the use of the variational USL framework Section 3.1 for training invertible architectures, and we compare the performance of different f -divergences with that of IPMs. We further analyze the behavior of the forward and backward training losses when variational f -divergences are employed. Additionally, we empirically examine the main theoretical result for NFs presented in Section B by studying the relationship between the number of finite moments of X and the W 1 distance between the true and the estimated distributions. The observed polynomial convergence rates are consistent with our theoretical predictions and analogous to the behavior is established for INNs in Theorem 3.1. 14 • The other direction is empirically explore the effects of certain design choices in practical performance of INNs. This includes the prior loss or reconstruction loss, the effect of dimensionality of the latent space, and role of the support set, and the effect of entropy regularization in Wasserstein metric. This empirical analysis complements the existing literature and providing insights for applied implementations. The software accompanying this paper is available on GitHub, https://github.com/ananya-ac/INN-Project. 4.1 Effect of design choices on INN training As mentioned earlier, in this section, we study the effects of various design choices (such as prior distribution, latent dimension, entropic regularization parameter, etc.) on the training of invertible neural networks. Before presenting the empirical studies in detail, we first clarify two general aspects. First, unless otherwise stated, all case studies were trained for 10 epochs using a batch size of 512. The model parameters were optimized with the Adam optimizer with a learning rate of 10 −4 and ℓ 2 regularization. Second, since the NLL is a widely adopted loss function for trainingINNs [Dinh et al., 2016, Ren et al., 2020, Kruse et al., 2021], we incorporate the NLL in our case studies, complementing the USL described in Section 3.1 (the relevant definitions are reviewed in Section A.2). Also, as discussed in Section 2, while the Gaussian assumption underlying the computation of the NLL can be restrictive, we avoid such distributional assumptions by leveraging the variational formulation of f-divergences. 4.1.1 Effect of prior on training In this case study, we use the inverse kinematics (IK) framework described next. Inverse kinematics example: IK is a core problem in robotics, where the objective is to determine the joint configurations required to achieve a desired end-effector position in space [Niku, 2020]. Efficiently solving this problem is essential for tasks like motion control, path planning, and real-time manipulation. Traditional methods such as analytical solutions and numerical optimization often face challenges as the robot configurations become complicated and number of degrees of freedom (DOF) increases. Machine learning algorithms are increasingly being employed for solving IK problems [Toquica et al., 2021], with the 4 DOF robot example becoming a standard benchmark in analyzing the analyzing the practical aspect of training invertible architecture INN [Ardizzone et al., 2018, Kruse et al., 2021, Hagemann and Neumayer, 2021]. We consider an articulated robotic arm that moves vertically along a rail and rotates at three joints. These four degrees of freedom constitute the parameter vector X = [X (1) ,X (2) ,X (3) ,X (4) ] ⊤ . The dataset is generated using Gaussian priors defined as X (i) ∼ N (0,σ i ), where the standard deviations are specified as σ 1 = 0.25 and σ 2 = σ 3 = σ 4 = 0.5, and the forward kinematic process of the robotics arm is modeled as Y (1) = X (1) + l 1 sin(X (2) ) + l 2 sin(X (3) − X (2) ) + l 3 sin(X (4) − X (2) − X (3) ); Y (2) = l 1 cos(X (2) ) + l 2 cos(X (3) − X (2) ) + l 3 cos(X (4) − X (2) − X (3) ), where the arm segment lengths are given by l 1 = 0.5, l 2 = 0.5, and l 3 = 1.0. We trained an INN to model the relationship between the joint angles of a robotic arm and a point in its 2D workspace. All models were trained using the NLL loss formulation for the INN. The input consisted of 4 variables corresponding to the joint angles of the robotic arm, and the output was a 2-dimensional vector representing a point in the 2D workspace. For this case study, we use two different regularization priors: a Gaussian prior and a uniform prior. The Gaussian prior represents the correct prior (Normal prior with μ = [0, 0, 0, 0] and σ = [0.25, 0.5, 0.5, 0.5]) as the data-generation process follows the same Gaussian distribution. The corresponding prior loss can be computed using the negative log-likelihood, up to a constant, as follows p L x = 1 n n X i=1 (T −1 (Y i ,Z i )− ̃ X i ) 2 . Here, ̃ X i N i=1 is the data generated from the input. A uniform prior [Andrle et al., 2021] also represents an inconsistent with the true distribution. If the input variable ranges from a to b (b > a), then the loss term 15 is as follows: p L x = 1 n n X i=1 max(0,T −1 (Y i ,Z i )− b) + max(0,a− T −1 (Y i ,Z i )) . Here, a = 0 and b = 1. (a) No prior(b) Correct Prior, λ ′ = 1 (c) Misspecified Prior, λ ′ = 1 (d) Correct Prior, λ ′ = 100 Figure 2: The four plots illustrate the results of the inverse kinematics case study with different prior regularizations. (a) shows samples from a model trained without any prior loss (i.e., with prior weight λ ′ = 0). (b) and (c) display samples from models trained with a prior weight of λ ′ = 1. In (b), the model assumes a Gaussian prior over the joint configuration, resulting in a scaled L 2 loss between the ground truth and the generated samples. In contrast, (c) uses a uniform prior U (0, 1), which is inconsistent with the true distribution. (d) illustrates results from a model trained with a Gaussian prior, using λ ′ = 100. The cross indicates the true end-effector position. All models were trained using the NLL loss formulation for the INN. To demonstrate the effect of incorporating a prior objective during training, the prior weight λ ′ was set to values of 0, 1, and 100. Figure 2 illustrates the effect of incorporating the prior loss p L x into the training objective. We observe that using a suitable prior, along with a well-chosen prior weight λ ′ , leads to more plausible samples. In the context of the inverse kinematics case study, this corresponds to configurations that keep joint angles close to zero, resulting in “straighter” arm trajectories, as shown in Figure 2b. Since the prior had a mean of 0 and a small variance, it led to samples that favored configurations where the arm positions were straight, with all joint angles close to 0 degrees. However, when the assumed prior is incorrect, incorporating it into the training objective leads to inaccurate joint configurations, as seen in Figure 2c. Finally, Figure 2d demonstrates that when the prior weight λ ′ is set too high, even with a suitable prior, the prior loss dominates the training, resulting in degraded performance and unrealistic outputs. Increasing the coefficient for p L x highlights this effect further. We conclude that using prior distribution knowledge in training INNs acts as a form of regularization by constraining the estimated posterior to align with known distributions. By incorporating such priors, the model is guided toward more plausible solutions, preventing it from learning overly complex or irrelevant patterns. Essentially, the prior acts as a form of bias that encourages the model to explore a more reasonable subset of possible solutions, leading to robustness in generation tasks. The designers need to navigate the fine line between Bayesian wisdom and risky assumptions. A well-chosen prior can improve the INN performance by effectively regularizing the model, and poorly chosen priors can dominate the results, leading to biased or misleading conclusions. We conclude this section by discussing the reason behind the importance of prior cost on training by focusing on the Gaussian case. Fig. 3 compares the training procedure of Bayesian neural network (BNN), INN, and INN with prior for solving an inverse problem. The standard direct formulation in BNN depends on a discriminative supervised loss term on X, which is one motivation for introducing INN for inverse problems [Ardizzone et al., 2018]. In the case of Gaussian distributions, we have demonstrated that p L x can serve as an effective regularizer by imposing a supervised cost on X. In this case, as demonstrated in Fig. 3c, training with prior propose a hybrid approach applying both supervised and unsupervised costs on X. This approach leverages the regularization with prior to enhance the performance of INN. 16 SC x y Bayesian Neural Network forward modeling inverse prediction (a) USC x y z SC USC Invertible Neural Network forward modeling inverse prediction (b) USC + SC x y z SC USC Invertible Neural Network forward modeling inverse prediction (c) Figure 3: Schematic illustration of how different training paradigms are used to solve inverse problems in the Gaussian setting. (a) Bayesian neural network for solving inverse problem. There is supervised cost (SC) between predicted and true x. (b) INN Trained without p L x . There is unsupervised cost (USC) between predicted and true x. (c) Shows an INN model trained with p L x , which serves as a reconstruction loss. 4.1.2 Comparison of architectures This observation compares the performance of INNs on an inverse problem using coupling-based and iResNet architectures, trained with different variational f-divergences as loss functions, considering both forward and backward losses. In general, it is challenging to optimize over all measurable function in the the critic classG, defined in Section 3.1. Hence, in practice, finite-parameter family of functions are used to approximate the optimal critic function. We adopt the method proposed in Nowozin et al. [2016], to approximate the critic function by a neural network. In this case, we replace the function g in equation (3) with V ω (x) = g f (V ′ ω (x)), where V ′ ω : X → R denotes a neural network with parameters ω, and g f is a monotonic activation function (mapping into the domain of f ∗ ). Here, V ′ ω is implemented as a fully connected network with two hidden layers of size 64 with ReLU activations [Goodfellow et al., 2016]. The function g f is determined by the specific choice of the f -divergence [Nowozin et al., 2016]. For the Kullback–Leibler (KL) divergence, g f employs an identity activation, i.e., g f (v) = v. In the case of the Reverse KL divergence, the activation is given by g f (v) =− exp(−v), whereas for the Jensen–Shannon (JS) divergence, the activation takes the form g f (v) = log(2)− log(1 + exp(−v)). In this numerical study, we consider the USL in both the backward and forward directions. In the backward direction, we wish to minimize D f (P T −1 (Y,Z) ∥ P X ) using the following joint objective function F (θ,ω) = E[V ω (T −1 θ (Y,Z)]− E[f ∗ (V ω (X))], where T −1 and V are parametrized by θ and ω, respectively. In particular, in this case study, we refer to the following regularized empirical estimation as the backward loss L backward (θ,ω) = 1 N N X i=1 V ω (T −1 θ (Y i ,Z i ))− 1 N N X i=1 f ∗ (V ω (X i )) + 1 N N X i=1 ∥T −1 θ (Y i ,Z i )− X i ∥ 2 (10) Similarly, in the forward direction, we wish to minimize D f (P Y,Z ∥ P T (X) ), using the the objective function F (θ,ω) = E[V ω (Y,Z)]− E[f ∗ (V ω (T θ (X)))].(11) Here, the following empirical estimation is called the forward loss L forward (θ,ω) = 1 N N X i=1 V ω (Y i ,Z i )− 1 N N X i=1 f ∗ (V ω (T θ (X i ))) + 1 N N X i=1 ∥T θ (X i )− (Y i ,Z i )∥ 2 . 17 0.140.150.150.160.16 0 2 4 6 Training Time (mins) ℓ 2 -Distance Forward Loss JSD KLD RKL 0.160.180.20.220.240.260.280.30.32 0.2 0.4 0.6 0.8 Training Time (mins) ℓ 2 -Distance Backward Loss JSD KLD RKL 0.680.70.720.740.760.780.8 0 2 4 6 Training Time (mins) ℓ 2 -Distance JSD KLD RKL 0.811.21.41.6 0.2 0.4 0.6 0.8 Training Time (mins) ℓ 2 -Distance JSD KLD RKL Figure 4: Comparison of f -divergence performance for coupling-based and iResNet architectures in the inverse kinematics case study. The x-axis shows training time and the y-axis shows ℓ 2 -distance between true and generated end-effector positions. Ellipses are centered at the mean values, with horizontal and vertical diameters indicating the standard deviations in training time and ℓ 2 -distance. The right and left columns correspond to forward loss (11) and backward losses (10), respectively. The top row corresponds to the coupling-based architecture whereas the figures on the bottom row generated based on iResNet layers. The training scheme, which incorporates f-divergences, differs from the configuration described in the previous section. Unlike models trained with the other cost functions, this setup includes both a discriminator network and a generator network (cf. Nowozin et al. [2016]). Both backward and forward losses are obtained by finding the saddle point of F (θ,ω): θ ∗ ,ω ∗ = arg min θ max ω F (θ,ω). Each training epoch includes gradient updates for both the discriminator and the generator, and both networks were trained using the Adam optimizer with a learning rate of 2× 10 −4 . To evaluate the impact of different f-divergences we consider the inverse kinematics problem, describe in Section 4.1.1, aiming to recover joint parameters from the end-effector position. We train INN models using each of the previously discussed f-divergence objectives to investigate their influence on the quality of generated samples. Training performance depends on multiple factors, including the network architecture, and, in this case study, we consider both coupling-based and residual invertible architectures. Both coupling- based and iResNet architecture comprise four layers, each with affine transformations parameterized by 18 261014182226303438 0 0.5 1 1.5 2 2.5 3 3.5 Latent Dimension KLD, ℓ 2 -Distance Effect of latent dimension 10 −1.6 10 −1.4 10 −1.2 10 −1 10 −0.8 10 −0.6 10 −0.4 10 −0.2 Sinkhorn, ℓ 2 -Distance Figure 5: This diagram illustrates the impact of varying latent dimensions on the posterior samples. The circular dots represent the variational formula- tion of the forward KL divergence and the hollow squares correspond to the Sinkhorn divergence. The x-axis shows latent dimension and the y-axis shows ℓ 2 -distance between true and generated end-effector positions. 0.350.40.450.5 0.2 0.4 0.6 0.8 Training Time (mins) ℓ 2 -Distance Effect of Entropy Regularization parameter 0.010.1 110 Figure 6: This diagram illustrates the effect of en- tropy regularization on the Wasserstein distance. The circular dots represent the debiased Sinkhorn divergence, while the hollow squares indicate the Sinkhorn approximation of the W 2 distance. The x-axis shows training time and the y-axis shows the ℓ 2 -distance between true and generated end-effector positions. feedforward neural networks with a hidden size of 128. The initial weights of the network were initialized by sampling values from a standard normal distribution. The latent dimension was chosen to be 14 for both configurations. Figure 4 compares the mean square error and the computation times of the f-divergence losses. We take the mean error over 20 runs of the f-divergences under the given task. As depicted in this figure, training performance is influenced by many factors, including the network architecture, forward and backward training, as well as the capacity of the critic class. Also, different variational formulations of f-divergences may be better suited to different settings. As shown in Figure 4, the backward loss in equation (10) yields better performance, and coupling-based designs are more efficient. 4.1.3 Effect of latent dimension in invertible neural networks To investigate the impact of latent dimension size on training and inference, we train models with different latent dimensions (which can be tuned by appropriately changing d x by padding). The loss functions employed are the Sinkhorn divergence and the variational formulation of the forward KL divergence. Also, similar to previous examples, the IK problem described in Section 4.1.1 is used for this case study. As illustrated in Figure 5, we observe an approximately local U -shaped relationship between latent dimension and inference performance, where the overall trend is multimodal. Additionally, training with the Sinkhorn divergence exhibits greater stability compared to the forward KL divergence. While in this case study we focus on the dimension of latent variable, the work in Hagemann and Neumayer [2021] investigated the number of modes of P Z , demonstrating that parameterizing P Z in dependence on the labels Y is helpful for obtaining stable INNs with reasonable Lipschitz constants. 4.1.4 Effect of entropic regularization on Sinkhorn divergence and Wasserstein distance ap- proximation To demonstrate the effect of entropic regularization on the debiased Sinkhorn divergence and the Sinkhorn approximation of the Wasserstein distance, we vary the entropic regularization parameter across experi- 19 ments. The model is trained for the inverse kinematics problem same training configuration described in Section 4.1.1. As illustrated in Figure 6, a decrease in entropic regularization tends to improves sample quality but increases training time. Additionally, training with the Sinkhorn divergence exhibits greater stability compared to the Sinkhorn approximation of the Wasserstein distance. 4.1.5 Effect of support set of the latent distribution and the data distribution We investigate the effect of the support set of the latent distribution when training with the variational formulation of the KL divergence and the Sinkhorn divergence. In this experiment, the data distribution is fixed as U (0, 1). The latent distributions considered for comparison are U (0, 1), U (3, 4), U (5, 6), U (10, 11) andU (15, 16). As illustrated in Figure 7 and Table 2, the quality of the generated samples deteriorates as the support of the latent distribution becomes increasingly disjoint from that of the true distribution. This effect is observed for both cost functions; however, the model trained using the f -divergence formulation of the KL divergence exhibits a more pronounced degradation compared to the model trained with the Sinkhorn divergence. Thus, our empirical results support the hypothesis that models trained with the Sinkhorn divergence are more robust to mismatches between the supports of the latent and prior distributions. -0.5-0.250.00.250.500.751.01.251.5 0 20 40 60 80 100 120 Frequency -0.5-0.250.00.250.500.751.01.251.5 0 20 40 60 80 100 120 -0.5-0.250.00.250.500.751.01.251.5 0 20 40 60 80 100 120 -0.5-0.250.00.250.500.751.01.251.5 0 20 40 60 80 100 120 Value Frequency -0.5-0.250.00.250.500.751.01.251.5 0 50 100 Value -0.5-0.250.00.250.500.751.01.251.5 0 50 100 Value Figure 7: Top Row: Samples from NF models trained using the Sinkhorn divergence. The left image corresponds to a latent distribution U (0, 1) used to simulate the data, the middle image uses U (3, 4), and the right image uses U (5, 6). Bottom Row: Samples from NF models trained using an f -divergence approximation of the forward KL divergence. The latent distributions mirror those used in the Sinkhorn divergence experiments. 4.1.6 Effect of moments of X on the performance of NF In this experiment, we investigate how the number of finite moments of the P X affects the performance of the NF. To ensure that the moment of X can be adjusted, we assumed X has a Pareto distribution [Arnold, 2014]. That is, the dataset was generated by sampling from a Pareto distribution characterized by a shape 20 P Z Divergence KL DivergenceSinkhorn Divergence U (0, 1)0.0030.001 U (3, 4)1.220.24 U (5, 6)1.360.91 U (10, 11)3.071.28 U (15, 16)1.431.33 Table 2: This table shows the MMD between samples from the true distributionU (0, 1) and those generated by the trained model. parameter α and a scale parameter x m . In this case the probability density function of X is given by f (X;x m ,α) = αx α m X α+1 , X ≥ x m , 0,X < x m , where α > 0 and x m > 0. Also, the NF trained with the variational formulation of KL-Divergence. The latent space was defined by a 10-dimensional random vector drawn from a standard normal distribution. Here, the parameter α is varied from 1 to 10, and, as illustrated in Fig. 8, the Wasserstein-1 distance between the true and generated distributions decreased. This is consistent with our theoretical results in Section B that as the number of finite moment of X increases, the W 1 distance between the true and estimated distributions decreases. 12345678910 10 −1 10 0 10 1 α W 1 Distance config 1 config 2 config 3 Figure 8: This figure illustrates the Wasserstein-1 distance for increasing values of α in the Pareto distribu- tion, which corresponds to increasing number of finite moments. The y-axis is set to a logarithmic (base 10) scale. Configurations 1–3 specify INNs with 2, 6, 8 coupling layers and subnetworks of hidden sizes 128, 64, 128, respectively. 4.2 Geoacoustic Inversion To demonstrate the practical relevance of our approach, we employ the insights gained from the previous exploratory experiments to the GI setting of the SWellEx-96 experiment [Yardim et al., 2010, Meyer and Gemba, 2021] conducted off the coast of San Diego near Point Loma. This experimental setting is one of the most used, documented, and understood studies in the undersea acoustics community. ∗ Here, synthetic ∗ see http://swellex96.ucsd.edu/ 21 Figure 9: The SWellEx-96 experiment environment. The acoustic source is towed by a research vessel and transmits signals at various frequencies. The acoustic sensor consists of a vertical line array (VLA). Based on the measurements collected at the VLA, the objective is to estimate posterior distributions over parameters of interest (e.g. water depth, sound speed at the water-sediment interface, source range and depth, etc.). data corresponding to this experiment is used to simplify the task for the initial application of an INN-based framework in ocean acoustics. As depicted in Figure 9, the data is collected via a vertical line array (VLA). The specification of the 21 hydrophones of the VLA and sound speed profile (SSP) in the water column is provided in the SWellEx-96 documentation. The SSP and sediment parameters are considered to be range-independent. Water depth refers to the depth of the water at the array. The source is towed by a research vessel which consists of a comb signal comprising frequencies of 49, 79, 112, 148, 201, 283, and 388 Hz. While in the SWellEx-96 experiment the position of the source changes with time, for this task we consider the instant when the source depth is 60 m and the distance (or range) between the source and the VLA is 3 km. The sediment layer is modeled with the following properties. The seabed consists initially of a sediment layer that is 23.5 meters thick, with a density of 1.76 g/cm 3 , and an attenuation of 0.2 dB/kmHz. The sound speed at the bottom of this layer is assumed to be 1593 m/s. The second layer is mudstone that is 800 meters thick, possessing a density of 2.06 g/cm 3 , and an attenuation of 0.06 dB/kmHz. The top and bottom sound speeds of this layer are 1881 m/s and 3245 m/s respectively. The description of the geoacoustic model of the SWellEx-96 experiment is complemented by a half-space featuring a density of 2.66 g/cm 3 , an attenuation of 0.020 dB/kmHz, and a sound speed of 5200 m/s. Based on the measurements at the VLA, the objective of this task is to infer the posterior distribu- tion over the water depth as well as the sound speed at the water-sediment interface. For this task, we assume all the quantities above to be known. The unknown parameters m 1 (the water depth) and m 2 (the sound speed at the water-sediment interface) follow a uniform prior in [200.5, 236.5] m and [1532, 1592] m/s, i.e. m 1 ∼U ([200.5, 236.5]) and m 2 ∼ U ([1532, 1592]), where U (Ω) denotes a uniform distribution in the domain Ω. The received pressure y on each hydrophone and for each frequency is a function of unknown parameters m (e.g. water depth, sound speed at the water-sediment interface, etc.) and additive noiseε as follows y = s(m,ε) = F (m) +ε,ε∼N (0, Σ) where Σ is the covariance matrix of data noise. Here, s(m,ε) is a known forward model that, assuming an additive noise model, can be rewritten as F (m) +ε, where F (m) represents the undersea acoustic model [Jensen et al., 2011]. The SWellEx-96 experiment setup involves a complicated environment and no closed-form analytical solution is available for F (m). In this case, F (m) can only be evaluated numerically, and we use the normal-modes program KRAKEN [Porter, 1992] for this purpose. The signal-to-noise ratio is 15 dB. Inspired by Zhang and Curtis [2021], for the invertible architectures, we include data noiseε as additional model parameters to be learned. In this context, the input of the network is obtained by augmenting the unknown parameters m with additive noiseε. 22 The pressures received on the hydrophones are considered in the frequency domain, and hence they can be complex numbers. While the invertible architectures can be constructed to address complex numbers, in this case study, we stack the real and imaginary parts of the pressure field at the network’s output. That is, the pressure y = Rey + i Imy will be represented as Rey, Imy at the network’s output. This case study is analyzed noise and signal data collected from five hydrophones. Each hydrophone captured complex measurements across seven frequencies, which we processed by separating into real and imaginary components and then concatenating into a single vector. This processing approach resulted in input data tensors containing 70 signal measurements. We combined the input parameters with noise measurements, creating tensors of sizes 72. 200210220230 0 0.05 0.1 0.15 0.2 Water Depth Probability Density NLL 200210220230 0 0.05 0.1 0.15 0.2 0.25 Water Depth Probability Density KLD 200210220230 0 0.05 0.1 0.15 0.2 0.25 Water Depth Probability Density JSD 200210220230 0 0.1 0.2 0.3 Water Depth Probability Density NLL 200210220230 0 0.05 0.1 0.15 0.2 Water Depth Probability Density KLD 200210220230 0 0.1 0.2 0.3 Water Depth Probability Density JSD Figure 10: Posterior comparisons between models trained with different loss functions for the GI dataset, using signal measurements from five hydrophones. The inferred parameter is water depth, with the red dotted line indicating the true value. The INN configuration for the upper row uses Configuration 1 whereas the configuration used in the bottom row is Configuration 2. The y-axis shows un-normalized probability scores. For the GI dataset, training is conducted using a batch size of 128 over 10 epochs with an 80/20 train– validation split. The optimization settings vary by objective: the INN trained with the NLL loss uses a learning rate of 3× 10 −4 , while the INN and critic network V ω trained with the f-divergence objective use a learning rate of 2× 10 −4 together with L 2 regularization of 2× 10 −5 . All models employ the Adam optimizer with betas (0.8, 0.9). Two different coupling-based configuration are considered. Configuration 1. corresponds to a model with sub-networks of hidden size 64, a latent dimension of 18, and 3 coupling layers. Configuration 2. corresponds to a model with sub-networks of hidden size 403, a latent dimension of 2, and 4 coupling layers. The results are presented in Figs 10 and 11, and in Tables 3, 4 and 5. The histograms, in Figs 10 and 11, 23 are based on 1,000 samples drawn from the trained INN. Also, while training times are compared in Table 5 the inference time is also provided in its caption. 1,5001,5201,5401,5601,580 0 0.05 0.1 0.15 0.2 0.25 Sound Speed Probability Density NLL 1,5001,5201,5401,5601,580 0 0.05 0.1 0.15 0.2 0.25 Sound Speed Probability Density KLD 1,5001,5201,5401,5601,580 0 0.05 0.1 0.15 0.2 Sound Speed Probability Density JSD 1,5001,5201,5401,5601,580 0 0.05 0.1 0.15 0.2 0.25 Sound Speed Probability Density NLL 1,5001,5201,5401,5601,580 0 0.1 0.2 0.3 0.4 Sound Speed Probability Density KLD 200210220230 0 0.1 0.2 0.3 Water Depth Probability Density JSD Figure 11: Posterior comparisons between models trained with different loss functions for the GI dataset, using signal measurements from five hydrophones. The inferred parameter is sound speed, with the red dotted line indicating the true value. The INN configuration for the upper row uses Configuration 1 whereas the configuration used in the bottom row is Configuration 2. Method Epochs 1510 NLL6.1922.9022.48 KL Divergence7.063.463.27 JS Divergence6.633.213.26 (a) Configuration 1 Method Epochs 1510 NLL25.162.320.75 KL Divergence10.5610.6611.56 JS Divergence11.2011.0412.26 (b) Configuration 2 Table 3: Comparison of the L 2 loss for water depth using two different model configurations. The parameter being recovered is water depth. For GI, standard likelihood-evaluation-based MCMC is computationally intensive, as each likelihood evaluation necessitates solving the forward model with an acoustic propagation simulator KRAKEN [Porter, 1992]. Alternately, a pre-trained INN facilitates rapid, near-real-time posterior inference by eliminating the need for repeated forward model evaluations. This speedup (inference time of 0.0060 seconds per 1000 samples) comes from shifting computation offline into training. Posterior sampling is often computationally challenging, especially in high dimensions [Montanari and Wu, 2023], and there is no universally best con- figuration for these problems. As demonstrated in our case studies, INNs can be sensitive to architectural and optimization choices; therefore, automated hyperparameter search methods such as the tree-structured 24 Method Epochs 1510 NLL365.26211.71193.99 KL Divergence31.360.922.24 JS Divergence27.400.617.26 (a) Configuration 1 Method Epochs 1510 NLL9.330.731.91 KL Divergence9.261.256.31 JS Divergence8.732.2910.24 (b) Configuration 2 Table 4: Comparison of the L 2 loss for sound speed using two different model configurations introduced. The parameter being recovered is sound speed. Method Epochs 1510 NLL4.8716.7830.50 KL Divergence1.375.4410.91 JS Divergence1.345.5811.08 (a) Configuration 1 Method Epochs 1510 NLL4.2214.1926.02 KL Divergence1.716.7613.36 JS Divergence1.866.6412.45 (b) Configuration 2 Table 5: Comparison of training times for two model configurations. The configuration used for the left table achieves an inference time of 0.0060 seconds per 1000 samples, while the configuration used for the right table achieves 0.0036 seconds for the same number of samples. Parzen estimator [Bergstra et al., 2011] can be useful in practice. 5 Conclusion In this work, we have presented a unified theoretical and practical framework that connects INNs and NFs through a common variational formulation, motivated by related works of Nowozin et al. [2016], Zhang et al. [2019], Grover et al. [2018]. By analyzing these architectures under a shared perspective, we derived new theoretical guarantees on the approximation quality of both posterior and generative distributions under weaker, more realistic assumptions than those used in prior studies. Our findings bridge a key gap in the literature by providing a principled understanding of when and why these models succeed in both generative and inverse problem settings. Beyond the theoretical results, our empirical investigations yield general design principles that can guide practitioners in the effective implementation of INN- and NF-based systems. Finally, the application of our framework to a ocean-acoustic inversion task demonstrates its practical value and robustness. Acknowledgements M. J. Khojasteh acknowledges support from the Gleason Endowment at RIT and also thanks Prof. Florian Meyer and Dr. Mohsen Sadr for their constructive discussions and suggestions. K. Youcef-Toumi acknowl- edges the support from the Center for Complex Engineering Systems (CCES) at King Abdulaziz City for Science and Technology (KACST) and Massachusetts Institute of Technology (MIT). 25 References C. Andrieu, N. De Freitas, A. Doucet, and M. I. Jordan. An introduction to MCMC for machine learning. Machine learning, 50:5–43, 2003. A. Andrle, N. Farchmin, P. Hagemann, S. Heidenreich, V. Soltwisch, and G. Steidl. Invertible Neural Networks versus MCMC for Posterior Reconstruction in Grazing Incidence X-Ray Fluorescence . In International Conference on Scale Space and Variational Methods in Computer Vision, pages 528–539. Springer, 2021. L. Ardizzone, J. Kruse, S. Wirkert, D. Rahner, E. W. Pellegrini, R. S. Klessen, L. Maier-Hein, C. Rother, and U. K ̈othe. Analyzing Inverse Problems with Invertible Neural Networks. arXiv preprint arXiv:1808.04730, 2018. L. Ardizzone, C. L ̈uth, J. Kruse, C. Rother, and U. K ̈othe. Guided Image Generation with conditional Invertible Neural Networks. arXiv preprint arXiv:1907.02392, 2019. M. Arjovsky and L. Bottou. Towards Principled Methods for Training Generative Adversarial Networks. In International Conference on Learning Representations, 2017. M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein Generative Adversarial Networks. In International conference on machine learning, pages 214–223. PMLR, 2017. B. C. Arnold. Pareto distribution. Wiley StatsRef: Statistics Reference Online, pages 1–10, 2014. Y. F. Atchad ́e and J. S. Rosenthal. On Adaptive Markov Chain Monte Carlo algorithms. Bernoulli, 11(5): 815–828, 2005. M. A. Beaumont. Approximate bayesian computation. Annual review of statistics and its application, 6(1): 379–403, 2019. J. Behrmann, W. Grathwohl, R. T. Chen, D. Duvenaud, and J.-H. Jacobsen. Invertible Residual Networks. In International conference on machine learning, pages 573–582. PMLR, 2019. J. Behrmann, P. Vicol, K.-C. Wang, R. Grosse, and J.-H. Jacobsen. Understanding and mitigating exploding inverses in Invertible Neural Networks. In International Conference on Artificial Intelligence and Statistics, pages 1792–1800. PMLR, 2021. J. Bergstra, R. Bardenet, Y. Bengio, and B. K ́egl. Algorithms for hyper-parameter optimization. Advances in neural information processing systems, 24, 2011. H. Bevins, W. Handley, and T. Gessey-Jones. Piecewise Normalizing Flows. arXiv preprint arXiv:2305.02930, 2023. V. I. Bogachev, A. V. Kolesnikov, and K. V. Medvedev. Triangular transformations of measures. Sbornik: Mathematics, 196(3):309, 2005. J. Brofos, M. Gabri ́e, M. A. Brubaker, and R. R. Lederman. Adaptation of the Independent Metropolis- Hastings Sampler with Normalizing Flow Proposals . In International Conference on Artificial Intelligence and Statistics, pages 5949–5986. PMLR, 2022. S. Brooks, A. Gelman, G. Jones, and X.-L. Meng. Handbook of Markov chain Monte Carlo. CRC press, 2011. N. R. Chapman and E. C. Shang. Review of geoacoustic inversion in underwater acoustics. Journal of Theoretical and Computational Acoustics, 29(03):2130004, 2021. R. T. Chen, J. Behrmann, D. K. Duvenaud, and J.-H. Jacobsen. Residual flows for Invertible Generative Modeling. Advances in neural information processing systems, 32, 2019. 26 F. Coeurdoux, N. Dobigeon, and P. Chainais. Sliced-Wasserstein normalizing flows: beyond maximum likelihood training. arXiv preprint arXiv:2207.05468, 2022. I. Csisz ́ar. On Information-Type measure of difference of probability distributions and indirect observations. Studia Sci. Math. Hungar., 2:299–318, 1967. M. Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in neural infor- mation processing systems, 26, 2013. M. Dashti, S. Harris, and A. Stuart. Besov priors for Bayesian inverse problems. Inverse Problems and Imaging, 6(2):183–200, 2012. L. Dinh, D. Krueger, and Y. Bengio. NICE: Non-linear Independent Components Estimation. arXiv preprint arXiv:1410.8516, 2014. L. Dinh, J. Sohl-Dickstein, and S. Bengio.Density estimation using Real NVP.arXiv preprint arXiv:1605.08803, 2016. S. E. Dosso and J. Dettmer. Bayesian matched-field geoacoustic inversion. Inverse Problems, 27(5):055009, 2011. A. Doucet and X. Wang. Monte Carlo methods for signal processing: a review in the statistical signal processing context. IEEE Signal Processing Magazine, 22(6):152–170, 2005. R. M. Dudley. Real Analysis and Probability. Chapman and Hall/CRC, 2018. S. G. Fadel, S. Mair, R. da S. Torres, and U. Brefeld. Principled Interpolation in Normalizing Flows. In Machine Learning and Knowledge Discovery in Databases. Research Track: European Conference, ECML PKDD 2021, Bilbao, Spain, September 13–17, 2021, Proceedings, Part I 21, pages 116–131. Springer, 2021. Y. Fan and L. Ying. Solving Inverse Wave Scattering with Deep Learning . arXiv preprint arXiv:1911.13202, 2019. N. Fournier and A. Guillin. On the rate of convergence in Wasserstein distance of the empirical measure. Probability theory and related fields, 162(3):707–738, 2015. M. Gabri ́e, G. M. Rotskoff, and E. Vanden-Eijnden. Efficient bayesian sampling using normalizing flows to assist markov chain monte carlo methods. arXiv preprint arXiv:2107.08001, 2021. M. Gabri ́e, G. M. Rotskoff, and E. Vanden-Eijnden. Adaptive Monte Carlo augmented with normalizing flows. Proceedings of the National Academy of Sciences, 119(10):e2109420119, 2022. J. Geweke. Bayesian inference in econometric models using monte carlo integration. Econometrica: Journal of the Econometric Society, pages 1317–1339, 1989. A. N. Gomez, M. Ren, R. Urtasun, and R. B. Grosse. The Reversible Residual Network: Backpropagation without Storing Activations. Advances in neural information processing systems, 30, 2017. I. Goodfellow, Y. Bengio, and A. Courville. Deep learning. MIT press, 2016. I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative Adversarial Networks . Advances in neural information processing systems, 27, 2014. A. Grover, M. Dhar, and S. Ermon. Flow-GAN: Combining Maximum Likelihood and Adversarial Learning in Generative Models. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018. X. Guan, X. Wang, H. Wu, Z. Yang, and P. Yu. Efficient Bayesian Inference using Physics-Informed Invertible Neural Networks for Inverse Problems. Machine Learning: Science and Technology, 5(3):035026, 2024. 27 S. Gurumurthy, R. Kiran Sarvadevabhatla, and R. Venkatesh Babu. DeLiGAN : Generative Adversarial Networks for Diverse and Limited Data . In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 166–174, 2017. P. Hagemann and S. Neumayer. Stabilizing Invertible Neural Networks using Mixture Models. Inverse Problems, 37(8):085002, 2021. T. Hickling and D. Prangle. Flexible Tails for Normalizing Flows. arXiv preprint arXiv:2406.16971, 2024. C.-F. Huang, P. Gerstoft, and W. S. Hodgkiss. Uncertainty analysis in matched-field geoacoustic inversions. The Journal of the Acoustical Society of America, 119(1):197–207, 2006. J.-H. Jacobsen, A. Smeulders, and E. Oyallon. i-RevNet: Deep Invertible Networks.arXiv preprint arXiv:1802.07088, 2018. F. B. Jensen, W. A. Kuperman, M. B. Porter, H. Schmidt, and A. Tolstoy. Computational ocean acoustics, volume 2011. Springer, 2011. B. Jin, Z. Zhou, and J. Zou. On the Approximation of Bi-Lipschitz Maps by Invertible Neural Networks . Neural Networks, 174:106214, 2024. J. M. T. JM and M. Welling. Improving Variational Auto-Encoders using convex combination linear Inverse Autoregressive Flow . page 162, 2017. G. L. Jones and Q. Qin. Markov chain Monte Carlo in practice. Annual Review of Statistics and Its Application, 9(1):557–578, 2022. L. Ke, S. Choudhury, M. Barnes, W. Sun, G. Lee, and S. Srinivasa. Imitation Learning as f-divergence minimization. In Algorithmic Foundations of Robotics XIV: Proceedings of the Fourteenth Workshop on the Algorithmic Foundations of Robotics 14, pages 313–329. Springer, 2021. T. A. Keller, J. W. Peters, P. Jaini, E. Hoogeboom, P. Forr ́e, and M. Welling. Self Normalizing Flows. In International Conference on Machine Learning, pages 5378–5387. PMLR, 2021. Y. Khoo and L. Ying. SwitchNet: a neural network model for forward and inverse scattering problems. SIAM Journal on Scientific Computing, 41(5):A3182–A3201, 2019. D. P. Kingma and P. Dhariwal. GLOW: Generative Flow with Invertible 1x1 Convolutions. Advances in neural information processing systems, 31, 2018. D. P. Kingma, M. Welling, et al. An Introduction to Variational Autoencoders . Foundations and Trends® in Machine Learning, 12(4):307–392, 2019. I. Kobyzev, S. J. Prince, and M. A. Brubaker. Normalizing Flows: An Introduction and Review of Current Methods. IEEE transactions on pattern analysis and machine intelligence, 43(11):3964–3979, 2020. A. Korattikara, Y. Chen, and M. Welling. Austerity in MCMC land: Cutting the Metropolis-Hastings budget. In International conference on machine learning, pages 181–189. PMLR, 2014. J. Kruse, L. Ardizzone, C. Rother, and U. K ̈othe. Benchmarking Invertible Architectures on Inverse Problems . arXiv preprint arXiv:2101.10763, 2021. L. A. Kruse, A. Tzikas, H. Delecki, M. Arief, and M. J. Kochenderfer. Enhanced Importance Sampling Through Latent Space Exploration in Normalizing Flows. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 17983–17989, 2025. V. Kungurtsev, A. Cobb, T. Javidi, and B. Jalaian. Decentralized Bayesian Learning with Metropolis- Adjusted Hamiltonian Monte Carlo. Machine Learning, 112(8):2791–2819, 2023. M. Laszkiewicz, J. Lederer, and A. Fischer.Copula-Based Normalizing Flows.arXiv preprint arXiv:2107.07352, 2021. 28 C.-L. Li, W.-C. Chang, Y. Cheng, Y. Yang, and B. P ́oczos. MMD GAN: Towards Deeper Understanding of Moment Matching Network . Advances in neural information processing systems, 30, 2017. A. Luce, A. Mahdavi, H. Wankerl, and F. Marquardt. Investigation of Inverse Design of Multilayer thin-films with conditional Invertible Neural Networks. Machine Learning: Science and Technology, 4(1):015014, feb 2023. doi: 10.1088/2632-2153/acb48d. URL https://dx.doi.org/10.1088/2632-2153/acb48d. D. J. MacKay. Information Theory, Inference and Learning Algorithms. Cambridge university press, 2003. Y. Marzouk, T. Moselhy, M. Parno, and A. Spantini. An introduction to sampling via measure transport. arXiv preprint arXiv:1602.05023, 2016. A. Maurer. A vector-contraction inequality for Rademacher complexities . In International Conference on Algorithmic Learning Theory, pages 3–17. Springer, 2016. F. Meyer and K. L. Gemba. Probabilistic focalization for shallow water localization. J. Acoust. Soc. Am., 150(2):1057–1066, 08 2021. A. Montanari and Y. Wu. Posterior sampling in high dimension via diffusion processes. arXiv preprint arXiv:2304.11449, 2023. X. Nguyen, M. J. Wainwright, and M. I. Jordan. Estimating divergence functionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information Theory, 56(11):5847–5861, 2010. S. B. Niku. Introduction to Robotics: Analysis, Control, Applications . John Wiley & Sons, 2020. S. Nowozin, B. Cseke, and R. Tomioka. f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization . Advances in neural information processing systems, 29, 2016. G. A. Padmanabha and N. Zabaras. Solving Inverse Problems using conditional Invertible Neural Networks. Journal of Computational Physics, 433:110194, 2021. G. Papamakarios, T. Pavlakou, and I. Murray. Masked Autoregressive Flow for Density Estimation. Advances in neural information processing systems, 30, 2017. G. Papamakarios, D. Sterratt, and I. Murray. Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows. In The 22nd international conference on artificial intelligence and statistics, pages 837–848. PMLR, 2019. G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. Journal of Machine Learning Research, 22(57):1–64, 2021. M. Parno, T. Moselhy, and Y. Marzouk. A multiscale strategy for Bayesian Inference using transport maps. SIAM/ASA Journal on Uncertainty Quantification, 4(1):1160–1190, 2016. G. Peyr ́e, M. Cuturi, et al. Computational Optimal Transport: With applications to Data Science. Founda- tions and Trends® in Machine Learning, 11(5-6):355–607, 2019. Y. Polyanskiy and Y. Wu. Information Theory: From Coding to Learning. Cambridge university press, 2025. M. B. Porter. The KRAKEN normal mode program . Naval Research Laboratory, Washington DC, 1992. S. T. Radev, F. Graw, S. Chen, N. T. Mutters, V. M. Eichel, T. B ̈arnighausen, and U. K ̈othe. OutbreakFlow: Model-based Bayesian Inference of disease outbreak dynamics with Invertible Neural Networks and its application to the COVID-19 pandemics in Germany. PLoS computational biology, 17(10):e1009472, 2021. S. T. Radev, M. Schmitt, L. Schumacher, L. Elsem ̈uller, V. Pratz, Y. Sch ̈alte, U. K ̈othe, and P.-C. B ̈urkner. BayesFlow: Amortized Bayesian workflows with neural networks. arXiv preprint arXiv:2306.16015, 2023. A. Ramdas, S. J. Reddi, B. P ́oczos, A. Singh, and L. Wasserman. On the Decreasing Power of Kernel and Distance based Nonparametric Hypothesis Tests in High Dimensions . In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015. 29 S. Ren, W. Padilla, and J. Malof. Benchmarking deep inverse models over time, and the neural-adjoint method. Advances in Neural Information Processing Systems, 33:38–48, 2020. R. T. Rockafellar. Convex Analysis. Princeton university press, 1970. V. Roy. Convergence diagnostics for Markov chain Monte Carlo. Annual Review of Statistics and Its Application, 7(1):387–412, 2020. A. Ruderman, M. Reid, D. Garc ́ıa-Garc ́ıa, and J. Petterson. Tighter Variational Representations of f- Divergences via Restriction to Probability Measures . arXiv preprint arXiv:1206.4664, 2012. C. Sch ̈onle and M. Gabri ́e. Optimizing Markov chain Monte Carlo convergence with normalizing flows and Gibbs sampling. In NeurIPS 2023 AI for Science Workshop, 2023. S. Shalev-Shwartz and S. Ben-David. Understanding Machine Learning: From Theory to Algorithms . Cambridge university press, 2014. G. R. Shorack and J. A. Wellner. Empirical processes with applications to statistics. SIAM, 2009. J. Song, S. Zhao, and S. Ermon. A-NICE-MC: Adversarial Training for MCMC. Advances in neural information processing systems, 30, 2017. B. K. Sriperumbudur, A. Gretton, K. Fukumizu, B. Sch ̈olkopf, and G. R. Lanckriet. Hilbert Space Embed- dings and Metrics on Probability Measures. The Journal of Machine Learning Research, 11:1517–1561, 2010. B. K. Sriperumbudur, K. Fukumizu, A. Gretton, B. Sch ̈olkopf, and G. R. Lanckriet. On the empirical estimation of integral probability metrics. Electronic Journal of Statistics, 6:1550–1599, 2012. V. Stimper, B. Sch ̈olkopf, and J. M. Hern ́andez-Lobato. Resampling Base Distributions of Normalizing Flows . In International Conference on Artificial Intelligence and Statistics, pages 4915–4936. PMLR, 2022. A. M. Stuart. Inverse Problems: a Bayesian Perspective. Acta numerica, 19:451–559, 2010. E. G. Tabak and C. V. Turner. A Family of Nonparametric Density Estimation Algorithms. Communications on Pure and Applied Mathematics, 66(2):145–164, 2013. T. Teshima, I. Ishikawa, K. Tojo, K. Oono, M. Ikeda, and M. Sugiyama. Coupling-based Invertible Neu- ral Networks Are Universal Diffeomorphism Approximators. Advances in Neural Information Processing Systems, 33:3362–3373, 2020. J. S. Toquica, P. S. Oliveira, W. S. Souza, J. M. S. Motta, and D. L. Borges. An Analytical and a Deep Learning model for solving the Inverse Kinematic problem of an Industrial parallel robot. Computers & Industrial Engineering, 151:106682, 2021. A. Verine, B. Negrevergne, M. S. Pydi, and Y. Chevaleyre. Precision-Recall Divergence Optimization for Generative Modeling with GANs and Normalizing Flows . Advances in Neural Information Processing Systems, 36:32539–32573, 2023. C. Villani. Optimal Transport: Old and New, volume 338. Springer, 2008. C. Villani. Topics in Optimal Transportation, volume 58. American Mathematical Soc., 2021. C. Winkler, D. Worrall, E. Hoogeboom, and M. Welling. Learning Likelihoods with Conditional Normalizing Flows. arXiv preprint arXiv:1912.00042, 2019. S. Winter, T. Campbell, L. Lin, S. Srivastava, and D. B. Dunson. Machine Learning and the Future of Bayesian Computation . arXiv preprint arXiv:2304.11251, 2023. S. Wu, Q. Huang, and L. Zhao. Fast Bayesian Inversion of Airborne Electromagnetic Data Based on the Invertible Neural Network. IEEE Transactions on Geoscience and Remote Sensing, 61:1–11, 2023. 30 C. Yardim, P. Gerstoft, and W. S. Hodgkiss. Geoacoustic and Source Tracking using Particle Filtering: Experimental results. The Journal of the Acoustical Society of America, 128(1):75–87, 2010. L. Ying. Solving Inverse Problems with Deep Learning. In Proc. Int. Cong. Math, volume 7, pages 5154–5175, 2022. M. Zhang, T. Bird, R. Habib, T. Xu, and D. Barber. Variational f-divergence Minimization. arXiv preprint arXiv:1907.11891, 2019. M. Zhang, P. Hayes, T. Bird, R. Habib, and D. Barber. Spread Divergence. In International Conference on Machine Learning, pages 11106–11116. PMLR, 2020. X. Zhang and A. Curtis. Bayesian Geophysical Inversion using Invertible Neural Networks. Journal of Geophysical Research: Solid Earth, 126(7):e2021JB022320, 2021. 31 A Additional Background Definition A.1 (Rademacher Complexity). Let F denote a class of real-valued functions over the domain R d , and let X n i.i.d. ∼ P X denote an i.i.d. dataset. With ε n denoting n i.i.d. Rademacher random vari- ables (i.e., ±1 w.p. 1/2 each), we define the Rademacher complexity of the function class F as R n (F,X n ) = E ε n " sup f∈F 1 n n X i=1 ε i f (X i ) X n # We will denote the expected Rademacher complexity (with respect to the data X n ) asR n (F ) = E X n [R n (F,X n )]. We now recall a result about vector contraction for Rademacher complexities from Maurer [2016]. Fact A.1. For some d ≥ 1, let φ i : R d → R denote a collection of J -Lipschitz functions in ℓ 2 norm with φ i (0) = 0. Then, for a given set of points x 1 ,...,x n in some domain X , and a function class H consisting of h :X → R d , we have the following: E ε n " sup h∈H 1 n n X i=1 ε i φ i (h(x i )) # ≤ 2JE σ n " sup h∈H 1 n n X i=1 ⟨σ i ,h(x i )⟩ # . Above, we assume that (ε 1 ,...,ε n ) are i.i.d. Rademacher, and each (σ 1 ,...,σ n ) represents i.i.d. −1, 1 d - valued random vectors with i.i.d. Rademacher coordinates. Finally, we now recall a version of a well-known concentration inequality that we will use to obtain high probability uniform deviation bounds. Fact A.2 (McDiarmid’s Inequality). Let X i : 1 ≤ i ≤ n denote a stream of independent X -valued observations, and let h :X n → R be a function satisfying a bounded difference property: |h(x 1 ,...,x i−1 ,x i ,x i+1 ,...,x n )− h(x 1 ,...,x i−1 ,x ′ i ,x i+1 ,...,x n )|≤ c, for all i∈ [n], x n ∈X n and x ′ i ∈X . Then we have the following for all ε > 0: P (h(X n )− E[h(X n )] > nε)≤ exp − 2nε 2 c 2 . A.1 Discussion of Invertible Neural Architectures This section outlines the key factors that ensure practical invertibility in INNs, which are formally defined in Section 1, including architectural structures, bidirectional training, and padding schemes. Model Architecture: A key aspect in designing INNs is constructing architectures that are both expressive and invertible. One of the most popular design is constructed by concatenating affine coupling blocks Kingma and Dhariwal [2018], Dinh et al. [2016, 2014]. In this case, invertible architecture T is defined by a sequence of reversible blocks, where each block consists of two complementary affine coupling layers. The the block’s input u∈ R d u is split into u 1 ∈ R d u 1 and u 2 ∈ R d u −d u 1 . Each transformation stage applies a learned mapping to a subset of the input features, while the remaining features are left unchanged, ensuring invertibility of the overall transformation. In other words, we have v 1 v 2 = u 1 ⊙ exp (s 1 (u 2 ))+ t 1 (u 2 ) u 2 , o 1 o 2 = v 1 v 2 ⊙ exp (s 2 (v 1 ))+ t 2 (v 1 ) The mappings s i and t i for i = 1, 2 are arbitrarily neural network. Each block can be inverted, that is, given the output o = [o 1 , o 2 ] the inverse can be calculated as u 2 = o 2 − t 2 (o 1 ) ⊘ exp (s 2 (o 1 )) 32 and u 1 = o 1 − t 1 (u 2 ) ⊘ exp(s 1 (u 2 )) Note that when the coupling block is inverted, the subnetworks s i and t i need not to be invertible they are only evaluated in the forward direction. Also, this structure yields a triangular Jacobian, making the log-determinant computationally tractable. Another popular approach to achieve invertibility is using residual connections (iResNet models) Gomez et al. [2017], Jacobsen et al. [2018], Behrmann et al. [2019]. In this way, the mapping T is constructed by the successive composition of residual blocks, T = (I +F M )◦ (I +F M−1 )◦·◦ (I +F 1 ), where “◦ ” denotes the composition operation. A sufficient condition for invertibility of each block is that the residual sub-network has a Lipschitz constant less than one. However, unlike coupling-based methods that offer exact solutions for the inverse and the log-determinant of the Jacobian, they need to be approximated in the case of iResNet models Chen et al. [2019]. In summary, coupling layers provide exact inversion and efficient Jacobian computation, but their fixed partitioning limits flexibility and can affect stability. Invertible residual networks support more flexible and expressive architectures, needing only a Lipschitz constraint for stability. However, they have less efficient inversion and Jacobian computation than coupling layers. Bi-Directional Training: Invertible networks allow applying losses in both the input and output domains, since the mapping can be evaluated in both directions. During training, one may alternate between forward and inverse passes, accumulating gradients from both directions. This bidirectional loss formulation can improve the effectiveness and stability of training Ardizzone et al. [2018].In the forward direction, as discussed in (1), the supervised loss (SL) L y (T ) and USL L z (T ) will be minimized. Similarly, in the backward direction an USL L x (T ) can be considered. For the SL we consider mean square error, and the USL is discussed in the next section. Padding: Padding adjusts input dimensionality to meet a model’s structural requirements Ardizzone et al. [2018]. For INNs, it resolves mismatches by adding entries enabling bijective mappings. Common strategies like zero-padding and repetition-padding keep forward and inverse passes dimensionally consistent and can support stable learning. For instance, if zero-padding is used on either side of the network, additional loss terms, called reconstruction loss, are required to ensure no information is encoded in the padding dimensions and forcing these dimensions to remain inactive. A.2 Negative Log-Likelihood (NLL) For a map T (X;θ) 7→ [Y,Z] parameterized by θ, and assuming Y and Z are independent, the change-of- variables formula implies that the density P T −1 (Y,Z) is P T −1 (Y,Z) = P T y (X) P T z (X) · det J X7→ [Y,Z] (X) where J X7→ [Y,Z] (X;θ) denotes the Jacobian of the map T parameterized by θ. As described next, this expression can be used to define an unsupervised training loss. In particular, we aim to minimize the forward KL divergence between the true posterior distribution P X|Y and P T −1 (Y,Z) , given by L(θ) = D KL P X|Y ,P T −1 (Y,Z) =−E P X|Y logP T −1 (Y,Z) + const. =−E P X|Y logP T y (X) + logP T z (X) + log det J X7→ [Z,Y ] (X) + const. The empirical approximation of the above loss is as follows. ˆ L(θ) =− 1 n n X i=1 logP T y (X i ) + logP T z (X i ) + log det J X7→ [Y,Z] (X i ) + const. (12) As we can see, minimizing the above Monte Carlo approximation of the KL divergence is equivalent to maximizing likelihood (or minimizing negative log-likelihood). Assuming P Z is standard Gaussian and P Y 33 is a multivariate normal distribution around y gt and a small standard deviation σ, NLL loss in Eq. (12) becomes L NLL (θ) = 1 n n X i=1 1 2 T y (X i )− y gt 2 σ 2 + 1 2 T z (X i ) 2 − log det J X7→ [Y,Z] (X i ) ! Note that we started with the backward loss between P X|Y and P T −1 (Y,Z) and end up with a supervised loss and unsupervised loss in the forward direction. Also, the Gaussian assumption on Y can be restrictive and as discussed in Section 2, we eliminate the need for such distributional assumptions by relying solely on samples and leveraging the variational formulation of f-divergences, which yields a lower bound on the true value which is discussed more in Section 4.1.2. B Variational normalizing flow (V-NF) In this section, we introduce a class of normalizing flows trained using a variational objective. To simplify our presentation, we focus on the concrete example of the variational form of relative entropy, although, similar arguments can be developed for a larger class of f -divergences. Our discussion in this section closely follows the structure of Section 3.2 for INNs, and in fact some of the results we present can be directly inferred from the results of Section 3.2. However, we include all the details to keep this section self-contained and independent of Section 3.2. Let X = R d denote the observation space and let Z = R d denote the latent space in which the latent variable Z ∼ P Z takes its values. Let T : X → Z denote a class of diffeomorphisms † , and for any T ∈ T , define the model distribution on X as Q T = (T −1 ) # P Z ⇐⇒ Q T (E) = P Z (z ∈Z : T −1 (z)∈ E), for any measurable E ⊂X. In other words, for Z ∼ P Z , the model T represents a random variable X T = T −1 (Z) ∼ Q T with density q T (x) = p Z (T (x))detJT (x), with JT denoting the associated Jacobian. Training procedure. As mentioned in Section 3.1, we consider a variational objective for training NFs, and in particular, in this section we focus on the concrete case of forward relative entropy loss by using a function class G n (critic or witness function class that we allow to grow with n). Formally, the population loss that will be used to train a model is L NF (T ) = sup g∈G n E X∼P X [logg(X)]− E X T ∼Q T [g(X T )] + 1. Here we have used the variational definition of relative entropy presented in Nguyen et al. [2010], and by definition, L NF (T )≤ D KL (P X ∥ Q T ) with equality if the true likelihood ratio dP X /dQ T is contained in G n . Given n i.i.d. draws X 1 ,...,X n from the data distribution P X , and with m n independent latent random variables Z 1 ,...,Z m n i.i.d. ∼ P Z , we define the empirical loss associated with T ∈T as b L NF n (T ) = sup g∈G n ( 1 n n X i=1 logg(X i )− 1 m m n X i=1 g(T −1 (Z i )) ) + 1. Observe that here we use n to denote the size of the true dataset that we are trying to model, while m n denotes the number of latent data points we use. In general m ≡ m n can be much larger than the n, as the latent random variables are much easier to generate (a common choice is to select P Z to be a known multivariate Gaussian distribution). We can now define the empirical NF model ˆ T n as ˆ T n ∈ arg min T∈T b L NF n (T ).(13) † Our theoretical guarantees in Theorem B.1 are also valid for homeomorphisms, but due to the specific requirements during the training of NF, here we state the definition based on diffeomorphisms. 34 The performance of the ERM estimator with a variational objective defined above is governed by two factors: the sampling error, and the alproximation error due to the restriction of the critic class to G n . Classical results from empirical process theory can allow us to control the difference between L NF and b L NF in terms of certain notions of capacity of the class G n . Informally, smaller G n will lead to much better finite sample approximation. On the other hand, reducing the difference D KL (P X ∥ Q T ), that we refer to as the variational gap, requires a larger G n . This represents the key challenge in our analysis: finding the right tradeoff for the size of the critic class G n that simultaneously controls these two terms. We first present a general result that assumes the existence of such a sequence of G n : n ≥ 1 that satisfies the conditions stated in Assumption B.1. Then, in Section B.1, we identify sufficient that are easier to verify for Assumption B.1 to hold, and finally in Section B.2, we discuss a concrete realistic example satisfying these conditions in Section B.2. Assumption B.1. To analyze the NF model ˆ T n defined in (13), we require the following assumptions: • (NF1): Realizability. For the P X ,P Z under consideration, there exists a T ∗ ∈ T with Q T ∗ = P X , which implies that L NF (T ∗ ) = 0 assuming that 1 ∈ G n (note that we allow the critic class G n to grow with n, keeping T fixed). • (NF2): Uniform Convergence. There exists a deterministic sequence r n → 0 and a confidence level δ ∈ (0, 1), such that with probability 1− δ, we have sup T∈T | b L NF (T )− L NF (T )|≤ r n . In other words, we assume that the function classes (G n ,T ) are small enough to ensure uniform learn- ability. • (NF3): Variational Approximation Gap. P X ≪ Q T for all T ∈T , and there exists a determin- istic sequence η n converging to 0 as n→∞, such that sup T∈T D KL (P X ∥ Q T )− L NF (T )≤ η n . In other words, we assume that the capacity of the critic class G n grows with n to approximate the likelihood ratios of all distributions modeled by elements of T and the true data distribution P X . • (NF4): Moment Bounds. There exists positive constants R and a, such that max E[∥X∥ 1+a ], sup T∈T E[∥T −1 (Z)∥ 1+a ] ≤ R. With these assumptions in place, we can now state the main result of this section. Theorem B.1. Under Assumption B.1, the empirical risk minimizer ˆ T n defined in (13) satisfies the following with probability at least 1− δ: W 1 (P X ,Q ˆ T n ) ≲ (r n + η n ) a 2(1+a) , where ≲ suppresses a multiplicative constant that depends on a, and W 1 is the 1-Wasserstein metric. The proof of this statement is given in Section B.3. This theorem obtains a H ̈older-type transfer from the original empirical variational objective to the Wasserstein error. In particular, it says that the quality of approximation achieved by the NF model ˆ T n depends on the sum of uniform empirical-process deviation r n and the variational gap η n . The exponent a/(1 + a) encodes the tail behavior from Assumption B.1: distributions with lighter tails (or larger a) yield a stronger, near- √ · relation between W 1 and r n +η n , while heavier tails cause a weaker transformation. 35 Remark B.1. The proof of Theorem B.1 proceeds in three steps.(i) On the high-probability event from (NF2), ERM and realizability assumption imply that L NF ( ˆ T n ) ≤ 2r n . Combining this with the vari- ational gap assumption (NF3) yields D KL (P X ∥ Q ˆ T n ) ≤ 2r n + η n . (i) The next step is to use Pinsker’s inequality to get ∥P X − Q ˆ T n ∥ TV ≤ p (1/2)(2r n + η n ). (i) Finally, we employ the truncation argument formalized in Lemma D.1 along with the bounded (1 + a)-moment assumption (NF4) to get the required bound in W 1 metric. As this outline suggests, this exact argument goes through for any distance or divergence that dominates the total variation metric via a Pinsker-type inequality. Hence, Theorem B.1 can be stated more generally for such a class divergence measures, which includes the Jensen Shannon, Hellinger, and chi-squared divergence. B.1 Verifiable Sufficient Conditions for Assumption B.1 In this section, we work under the realizability assumption (NF1), and derive sufficient conditions on (T ,G n ) for the other conditions to hold, such that the W 1 metric between the model distribution and the true distribution converges to zero as the sample size n increases. We begin with the uniform convergence condition (NF2), and identify sufficient conditions in terms of the Rademacher complexities of the function classes involved. Proposition B.1. Let (ε i ) i≥1 denote i.i.d. Rademacher random variables, and with H n denoting the func- tion class logg : g ∈G n , define the expected Rademacher complexities R n (H n ) = E ε n ,X n " sup h∈H n 1 n n X i=1 ε i h(X i ) # ,andR m n (H n ,T ) = E ε m n ,Z m n " sup h,T 1 m n n X i=1 ε i h(T −1 (Z i )) # . If ∥h∥ ∞ ≤ b n for all h∈H n , for any δ > 0, we have the following: sup T∈T b L NF n (T )− L NF (T ) ≤ 2R n (H n ) + 2e b n R m n (H n ,T ) + ( b n √ n + e b n √ m n ) p 2 log(4/δ) Thus, to verify (NF2), we need to identify uniformly bounded H n (equivalently G n ) whose expected Rademacher complexity can be appropriately controlled to drive the above term to 0 with n. This result follows from the standard symmetrization via Rademacher variables argument, and we present the details in Section B.4. We next observe that if the model class T satisfies a uniform Lipschitz property, then the uniform moment control assumption (NF4) follows easily. Proposition B.2. Suppose the function class T satisfies the following uniform affine growth property: sup T∈T ∥T −1 (z)∥≤ J 0 + J 1 ∥z∥ for all z ∈Z. If E P Z [∥Z∥ 1+a ] = M Z <∞ for an a > 0, we have the following under the realizability assumption (NF 1): max E[∥X∥ 1+a ], sup T∈T E ∥T −1 (Z)∥ 1+a ≤ 2 a J 1+a 0 + J 1+a 1 M Z = : R. The proof of this result is exactly the same as that of Proposition 3.2, and we omit it to avoid repetition. In other words, assuming realizability, a sufficient condition for (NF4) is if the latent variable Z has finite (1 +a) moments, and the model class is globally Lipschitz and uniformly bounded at the origin. Finally, we proceed to the condition of vanishing variational approximation gap (NF3). Proposition B.3. Suppose D KL (P X ∥ Q T ) < ∞ for all T ∈ T , and let v T = dP X /dQ T and ℓ T = logv T denote the likelihood ratio and log-likelihood ratio (resp.) of P X and Q T . Fix some sequence (K n ) n≥1 such that K n n→∞ −→ ∞, and define the following uniform tail expectation terms, with X ∼ P X , X ′ ∼ Q T , and H n =logg : g ∈G n : τ 1 ≡ τ 1 (K n ,T ,H n ) = sup T∈T sup h∈H n E P X [|ℓ T (X)− h(X)| 1 ∥X∥ ∞ >K n ], and τ 2 ≡ τ 2 (K n ,T ,H n ) = sup T∈T sup h∈H n E Q T h e h(X ′ ) − e ℓ T (X ′ ) 1 ∥X ′ ∥ ∞ >K n i . 36 Next, introduce a uniform approximation error term computed over a restricted domain: δ n ≡ δ n (K n ,T ,H n ) = sup T∈T inf h∈H n sup x:∥x∥ ∞ ≤K n |h(x)− ℓ T (x)|. Suppose there exists a sequence of K n : n≥ 1 with K n ↑∞, such that lim n→∞ δ n + τ 1 + τ 2 = 0.(14) Then the variational gap η n also converges to zero; that is, if (14) holds, then η n = sup T∈T D KL (P X ∥ Q T )− L NF (T ) n→∞ −→ 0. The proof of this result is in Section B.5. B.2 V-NF Example We now illustrate that the assumptions required by Theorem B.1 are satisfied by practically useful models, by constructing specific (T ,G n ), and verifying this assumptions. Our model class is an iResNet flow with uniformly controlled Jacobians that ensures certain useful global regularity properties. Our critic class is induced by a Gaussian RKHS restricted to a growing cube C K n = [−K n ,K n ] d , which allows a neat separation between (i) approximation on a compact set, and (i) tail control outside the compact set, as needed by Proposition B.3. Definition B.1. We will work with i-ResNets of the form T = (I d + F M )◦ (I d + F M−1 )◦·◦ (I d + F 1 ), where “◦ ” denotes the composition operation, and F j (x) = W j,2 tanh (W j,1 x + b j,1 ) + b j,2 , withmax∥W j,1 ∥ op ,∥W j,2 ∥ op ≤ s∈ (0, 1),andmax∥b j,1 ∥, ∥b j,2 ∥≤ B, ∀j ∈ [M ]. Each W j,1 lies in R d j × R d while W j,2 lies in R d × R d j , with d j denoting the dimension of the hidden layer in block j. Similarly, b j,1 ∈ R d j and b j,2 ∈ R d . Let H denote the maximum value of the hidden layer dimension; that is, H = max 1≤j≤M d j , ∥·∥ op denotes the operator norm, and ∥·∥ denotes the ℓ 2 norm. Next, we introduce the critic function class G n . Definition B.2. Let k n (x,y)≡ k γ n (x,y) = exp −γ n ∥x− y∥ 2 denote a Gaussian kernel with scale param- eter γ n > 0, and let H k n denote the reproducing kernel Hilbert space (RKHS) associated with this kernel. Then, we defile the critic class G n as G n =e h : h∈H n ,where H n =h1 [−K n ,K n ] d : h∈H k n , with∥h∥ k n ≤ b n for some b n ,K n > 0. By the reproducing property, it follows that ∥h∥ k γ n ≤ b n implies ∥h∥ ∞ = sup x∈X ⟨h,k γ n (x,·)⟩ k γ n ≤ sup x∈X p k γ n (x,x)∥h∥≤ b n . As a result, every g ∈G n satisfies 0 < e −b n ≤ g(x)≤ e b n . As mentioned above, we work with the i-ResNet model classT in Definition B.1 since this satisfies a global Lipschitz continuity condition that allows us to work with unbounded latent random variables Z ∼ N (0,I d ). This global Lipschitz property turns out to be useful in establishing that the (T ,G n ) pair of Definition B.1 and Definition B.2 satisfy the conditions required by Assumption B.1. We can also ensure these assumptions are satisfied for the coupling-based architecture Dinh et al. [2016] by restricting the latent distribution P Z to be supported on a compact domain. We now state the main result of this section. 37 Theorem B.2. Let (T ,G n ) be as in the two definitions above, and define the parameters K n = M (s √ H + B) + √ d + p 2 logn, γ n = K 2+ε n for some fixed ε > 0, b n = C b γ d/4 n K 2+d/2 n = C b K 2+d+ d 4 ε n , m n =⌈e 4b n ⌉. Then, under the realizability assumption (NF1), and with these choices of the parameters, the ERM model ˆ T n defined in (13) satisfies W 1 (P X ,Q ˆ T n ) = o(1),w.p. at least 1− δ. In other words, with these parameters, the models in Definition B.1 satisfy the sufficient conditions for (NF2)-(NF4) to hold, as derived in Section B.1. The proof of this result is in Section B.6. B.3 Proof of Theorem B.1 Introduce the uniform approximation event E n =sup T∈T | b L NF (T )− L NF (T )|≤ r n , and note that under the assumption (NF2), we have P(E n )≥ 1− δ. For the rest of this proof, we will work under the event E n . By definition of ERM model, we knlw that for any T ∈ T , we have b L NF ( ˆ T n ) ≤ b L NF (T ). This implies the following chain: L NF ( ˆ T n )≤ b L NF ( ˆ T n ) + r n (event E n ) ≤ b L NF (T ∗ ) + r n (ERM property) ≤ (L NF (T ∗ ) + r n ) + r n (event E n ) = 2r n (assumption (NF1)). Now, by the variational gap assumption (NF3), we know that D KL (P X ∥ Q ˆ T n )− L NF ( ˆ T n ) ≤ sup T∈T |D KL (P X ∥ Q T )− L NF (T )|≤ η n . Thus, combining the two displays above, we obtain the following bound on the relative entropy between the true and the NF distributions, under the 1− δ probability event E n : D KL (P X ∥ P ˆ T n )≤ η n + 2r n Pinsker’s =⇒ ∥P X − P ˆ T n ∥ TV ≤ r 1 2 (2r n + η n ).(15) We can now conclude the proof with the following chain: W 1 (P X ,P ˆ T n ) ≲ ∥P X − P ˆ T n ∥ TV a 1+a ≲ (2r n + η n ) a 2(1+a) . Here the first inequality (modulo constants) is due to Lemma D.1, and the second inequality is by (15). B.4 Proof of Proposition B.1 This follows from the standard uniform convergence results for bounded function classes. Observe that sup T∈T | b L NF (T )− L NF (T )|≤ sup h∈H n |(P (X) n − P (X) )h| + sup T∈T sup h∈H n |(P (Z) m n − P (Z) )e h ◦ T −1 | : = D X + D Z . 38 Note that above we have used the notation P (X) n h = 1 n P n i=1 h(X i ), P (X) h = E X∼P X [h(X)], P (Z) m n e h ◦ T −1 = 1 m n P m n i=1 e h(T −1 (Z i )) and P (Z) e h ◦ T −1 = E Z∼P Z [e h(T −1 (Z)) ]. Now observe that by the classical symmetrization technique, we have the following upper bound on the expected value of D X [Shalev-Shwartz and Ben-David, 2014, Lemma 26.2]: E[D X ]≤ 2R n (H n ) = 2E [R n (H n ,X n )] Here,R n denotes the Rademacher complexity whose definition is recalled in Definition A.1 in Section A. Since we have assumed that∥h∥ ∞ ≤ b b , for all h∈H n , an application of the bounded difference concentration inequality (recalled in Fact A.2 in Section A) immediately implies with some constant C > 0: P X n D X ≥ E[D X ] + Cb n r log(2/δ) n ! ≤ P X n D X ≥ 2R n (H n ) + Cb n r log(2/δ) n ! ≤ δ 2 .(16) A similar argument works for the term D Z . In particular, we know that the function class e h◦T −1 : h∈H n ,T ∈T is uniformly bounded from above by e b n . Thus, we can again get a high probability upper bound using the expected Rademacher complexity and the bounded difference inequality: P Z m n D Z ≥ E[D Z ] + Ce b n s log(2/δ) m n ≤ P Z n D Z ≥ 2R m n (G n ◦T −1 ) + Ce b n s log(2/δ) m n ≤ δ 2 . (17) Together, (17) and (16), along with the observation thatR m n (G n ◦T −1 )≤ e b n R m n (H n ,T ) by the contraction result, imply the required uniform convergence bound with probability at least 1− δ: D X + D Z ≤ 2R n (H n ) + 2R m n (H n ,T ) + C p log(2/δ) b n √ n + e b n √ m n = : r n . Thus, in order to establish a uniform convergence guarantee, it suffices to find a sequence of(n,m n ,b n ,H n ) : n≥ 1, such that the term r n in RHS above converges to 0. B.5 Proof of Proposition B.3 Throughout this proof, we assume that X ∼ P X and X ′ ∼ Q T for T ∈ T , where Q T represents the distribution of T −1 (Z) for some latent variable Z ∼ P Z . Our argument below does not rely on how X ′ is generated, or on any properties of P Z . Now, observe that for any T ∈T , the relative entropy D KL (P X ∥ Q T ) is equal to E P [ℓ T (X)], which implies that ∆ n (T ) : = D KL (P X ∥ Q T )− sup h∈H n n E P X [h(X)]− E Q T [e h(X ′ ) ] + 1 o = inf h∈H n n E P X [ℓ T (X)− h(X)] + E Q T [e h(X ′ ) − e ℓ T (X ′ ) ] o , where the second equality uses the fact that E Q T [e ℓ T ] = 1 (recall that ℓ T is the log-likelihood ratio between P X and Q T ). Now, for some K n → ∞, define the events E n = ∥X∥ ∞ ≤ K n and F n = ∥X ′ ∥ ∞ ≤ K n , and observe that ∆ n (T )≤ inf h∈H n n E P X [(ℓ T (X)− h(X)) 1 E n ] + E Q T [ e h(X ′ ) − e ℓ T (X ′ ) 1 F n ] o + τ 1 + τ 2 ,(18) where the two terms τ 1 and τ 2 are defined as τ 1 (K n ,T ,H n ) = sup T∈T sup h∈H n E P X [|ℓ T (X)− h(X)| 1 ∥X∥ ∞ >K n ], and τ 2 (K n ,T ,H n ) = sup T∈T sup h∈H n E Q T h e h(X ′ ) − e ℓ T (X ′ ) 1 ∥X ′ ∥ ∞ >K n i . 39 It remains to analyze the first term in the RHS of (18). We proceed by observing that e ℓ T − e h = e ℓ T e h−ℓ T − 1 ≤ e ℓ T |ℓ T − h| + e |ℓ T −h| 2 |ℓ T − h| 2 , where the inequality uses a second-order Taylor approximation of e x around 0. More specifically, we use |e u − 1|≤|u|≤ e |u| u 2 2 with u← ℓ T − h. Now, let us introduce the following terms: e b n : = sup h∈H n sup x:∥x∥ ∞ ≤K n e h(x) ,and δ n = sup T∈T inf h∈H n sup x:∥x∥ ∞ ≤K n |ℓ T (x)− h(x)|. Using these definitions in (18), we obtain ∆ n (T )≤ E P X [δ n 1 E n ] + E Q T e ℓ T (X ′ ) δ n + e δ n δ 2 n 2 1 F n + τ 1 + τ 2 ≤ δ n + δ n + e δ n δ 2 n 2 E Q T [e ℓ T ] + τ 1 + τ 2 = δ n 2 + e δ n δ n /2 + τ 1 + τ 2 . Since the RHS is independent of T , taking a supremum over the function class T , gives us the required η n = sup T∈T ∆ n (T ) ≤ 2 + e δ n δ n /2 δ n + τ 1 + τ 2 , which converges to 0 if δ n ,τ 1 ,τ 2 converge to 0. This concludes the proof. B.6 Proof of Theorem B.2 To prove this result, we need to verify that the sufficient conditions obtained in Proposition B.1 (uniform con- vergence), Proposition B.2 (moment bounds), and Proposition B.3 (variational gap), are satisfied by our spe- cific choices in Definition B.1 and Definition B.2. We verify these conditions in Section B.6.1, Section B.6.2, and Section B.6.3 respectively. Before proceeding to these steps, we first establish certain properties of the function class T that will be used often. Suppose P Z is a Gaussian distribution with identity covariance, and density p Z (z) ∝ exp(−∥z∥ 2 /2), which implies that ψ(z) = logp Z (z) = −∥z∥ 2 /2 + const. Under the realizability assumption, there exists T ∗ ∈T such that P X = Q T ∗ . Every T ∈T can be represented as T = G M ◦ G M−1 ◦ ...◦ G 1 ,where G j (x) = x + F j (x), F j (x) = W j,2 tanh(W j,1 x + b j,1 ) + b j,2 . By assumption, we have ∥W j,i ∥ op ≤ s ∈ (0, 1) and ∥b j,i ∥ ≤ B for all j,i, where ∥·∥ denotes the ℓ 2 norm, and ∥·∥ op denotes the operator norm induced by ∥·∥. B.6.1 Verification of the Uniform Convergence Assumption To simplify the notation, introduce the function classes S n = g ◦ T −1 : g ∈ G n ,T ∈ T and U n = (1/b n )h◦ T −1 : h ∈ H n ,T ∈ T. Recall that H n consists of functions of the form ̃ h1 [−K n ,K n ] d for all ̃ h lying in the RKHS of a Gaussian kernel k n (denoted by H k n ), with ∥ ̃ h∥ k n ≤ b n . We first consider the termR m n (S n ), since the bound on the other term follows similarly. Observe that S n = e b n u : u ∈ U n by definition, leading to the following inequality via the contraction lemma [Shalev-Shwartz and Ben-David, 2014, Lemma 26.9]: R m n (S n ,Z m n )≤ e b n R m n (U n ,Z m n ). It remains for us to control the Rademacher complexity of the composition class U n . For any z m n = (z 1 ,...,z m n ), introduce the terms I ≡ I(z n ,K n ) = i ∈ [m n ] : z i ∈ [−K n ,K n ] d and s n ≡ s n (z m n ,K n ) = 40 |I|. For z m n such that s n > 0, observe that b n R m n (U n ,z m n ) = E σ " sup T∈T sup h∈H n 1 m n m n X i=1 σ i h(T −1 (z i )) # = E σ " sup T∈T sup ̃ h:∥ ̃ h∥ k n ≤b n 1 m n X i∈I σ i ̃ h(T −1 (z i )) # = E σ sup T∈T sup ̃ h:∥ ̃ h∥ k n ≤b n * ̃ h, 1 m n X i∈I σ i k n (T −1 (z i ),·) + k n (reproducing property) = E σ sup T∈T b n 1 m n X i∈I σ i k n (T −1 (z i ),·) k n optimal h∝ b n m n X i k n (T −1 (z i ),·) ! = E σ b n m n s sup T∈T X i∈I X j∈I σ i σ j k(T −1 (z i ),T −1 (z j )) ≤ b n m n v u u u t E σ m n X i=1 m n X j=1 σ i σ j Jensen’s + sup x,x ′ k(x,x ′ )≤ 1 = b n √ s n m n . As a result, with an i.i.d. sample Z m n , the expected Rademacher complexity satisfies b n R m n (U n ) = b n E Z m n [R m n (U n ,Z m n )]≤ b n m n E Z m n v u u t m n X i=1 1 Z i ∈[−K n ,K n ] d ≤ b n s P Z ([−K n ,K n ] d ) m n ≤ b n √ m n . An exactly analogous argument impliesR n (H n ) ≤ b n / √ n. Combining these two bounds, we see that the uniform convergence rate obtained in Proposition B.1 reduces to the following (up to leading constants): r n ≲ b n √ n + b n e b n √ m n .(19) This can be made to converge to 0 with n, by selecting b n = o( √ n) and m n = Ω(e (2+ε)b n ) for any ε > 0. Observe that the choices of b n ,m n in Theorem B.2 satisfy these conditions, as we discuss further at the end of this section. B.6.2 Verification of the Moment Assumption Lipschitz property. Observe that tanh is 1-Lipschitz as its derivative is sech 2 ∈ [0, 1], which implies Lip(F j )≤∥W j,2 ∥ op ∥W j,1 ∥ op ≤ s 2 = : ρ∈ (0, 1).(20) Hence, for every block G j = I d + F j , we have (1− ρ)∥x− x ′ ∥≤∥G j (x)− G j (x ′ )∥≤ (1 + ρ)∥x− x ′ ∥ Since T is the composition of M such blocks, we immediately obtain (1− ρ) M ∥x− x ′ ∥≤∥T (x)− T (x ′ )∥≤ (1 + ρ) M ∥x− x ′ ∥. 41 This allows us to conclude that every T ∈T is bi-Lipschitz with Lip(T )≤ (1 + s 2 ) M ,and Lip(T −1 )≤ (1− s 2 ) −M . Since these constants are independent of T , they are also uniformly valid over the class T : sup T∈T max Lip(T ), Lip(T −1 ) ≤ max (1− s 2 ) −M , (1 + s 2 ) M = : J 1 .(21) Values at x = 0. Fix any T = G M ◦...◦G 1 in T , and consider x 0 = 0, x j = G j (x j−1 ) = x j−1 +F j (x j−1 ) for j ∈ [M ]. Observe that by (20), we have for any x and j: ∥F j (x)∥≤∥F j (0)∥ + s 2 ∥x∥. Now, ∥F j (0)∥ can be bounded as ∥F j (0)∥ =∥W j,2 tanh(b j,1 ) + b j,2 ∥≤∥W j,2 ∥ op (∥ tanh(b j,1 )∥ +∥b j,2 ∥). We know that ∥W j,2 ∥ op ≤ s and ∥b j,i ∥≤ B for i = 1, 2. Furthermore, let H denote the maximum dimension of the hidden layer in M blocks; that is, H = max j∈[M ] d j . Then, ∥ tanh(b j,1 )∥ ≤ p d j tanh(B) ≤ √ H. These facts imply that ∥F j (0)∥≤ C 0 : = s √ H + B. This implies that for any j ∈ [M ], due to the Lipschitz property, we have ∥x j ∥ =∥G j (x j−1 )∥ =∥x j−1 + F j (x j−1 )∥≤∥x j−1 ∥ +∥F j (x j−1 )∥≤∥x j−1 ∥ +∥F j (0)∥ + s 2 ∥x j−1 ∥ ≤ C 0 + (1 + s 2 )∥x j−1 ∥. Applying this inequality iteratively, with x 0 = 0 and x M = T (x 0 ), we have ∥T (0)∥≤ C 0 M−1 X i=0 (1 + s 2 ) i = C 0 (1 + s 2 ) M − 1 s 2 .(22) Finally, for any T , let y T denote T −1 (0), and observe that ∥y T ∥ =∥T −1 ◦ T (y T − 0)∥≤ Lip(T −1 )∥T (y T )− T (0)∥ = Lip(T −1 )∥T (0)∥.(23) Together, (22) and (23) imply sup T∈T max ∥T (0)∥, ∥T −1 (0)∥ ≤ J 1 (s √ H + B) (1 + s 2 ) M − 1 s 2 = : J 0 .(24) Since Z ∼ N (0,I d ) has finite moments of all orders and ∥T −1 (z)∥≤ J 0 + J 1 ∥z∥ uniformly in T , we have that sup T∈T E[∥T −1 (Z)∥ a ] < ∞ for all a > 0. Under the realizability assumption, X d = (T ∗ ) −1 (Z), hence this also implies E∥X∥ 1+a <∞. Thus, with an arbitrary a > 0, and M Z = 2 a/2 Γ((a + d)/2)/Γ(d/2), we get R = 2 a J 1+a 0 + J 1+a 1 2 a/2 Γ((a + d)/2)/Γ(d/2) . Recall that J 0 was defined in (24), and J 1 in (21). This completes the proof. B.6.3 Verification of the Variational Gap Assumption To verify this assumption for the model introduced in Definition B.1, we will first obtain the Jacobian of the functions in T , which will then allow us to obtain a closed form expression for the log likelihood ratio dP X /dQ T and characterize its behavior on a cube C K n = [−K n ,K n ] d . 42 Lemma B.1. Consider any T = G M ◦ G M−1 ◦·◦ G 1 ∈ T of Definition B.1, where each block G j (x) = (I d + F j )(x) = x + W j,2 tanh(W j,1 x + b j,1 ) + b j,2 . Then, we have the following: JG j (x) = I d + W j,2 D j (x)W j,1 ,where D j (x) = diag(sech 2 (W j,1 x + b j,1 )) ∥JF j (x)∥ op ≤ s 2 ,and ∥JF j (x)− JF j (y)∥ op ≤ 2s 3 ∥x− y∥. This result, proved in Section D.3, tells us that every G j is a diffeomorphism with singular values in [1− s 2 , 1 + s 2 ], and hence, det JG j (x) > 0 which implies that log det JG j is well defined. We now obtain a closed-form expression for the log-likelihood ratio ℓ T = logdP X /dQ T . Lemma B.2. The log-likelihood ratio function ℓ T (x) = logdP X /dQ T for any T ∈T is equal to ℓ T (x) = 1 2 ∥T (x)∥ 2 −∥T ∗ (x)∥ 2 + log det JT ∗ (x)− log det JT (x). Furthermore, when restricted to the domain [−K,K] d , each ℓ T is Lipschitz with constant L 1 = O(K) (see (51) for exact expression), which implies the following uniform bound sup T∈T sup x∈[−K,K] d |ℓ T (x)|≤ L 0 + L 1 √ dK, for a universal constant L 0 stated in (52) This result, proved in Section D.4, establishes that when restricted to a box [−K,K] d in the domain X = R d , the likelihood ratio function class ℓ T : T ∈ T is well behaved. With these results at hand, we can now proceed towards verifying the sufficient conditions for the variational gap to vanish as n goes to ∞. Approximation error δ n . We begin by obtaining a characterization of the approximation error δ n . Lemma B.3. Consider the given H n = ̃ h1 [−K n ,K n ] d : ̃ h ∈ H k n , ∥ ̃ h∥ k n ≤ b n , where H k n is the RKHS associated with the Gaussian kernel k n (x,y) = exp(−γ n ∥x− y∥ 2 ) for a scale parameter γ n . For parameters b n , γ n , and K n , we have b n ≳ γ d/4 n K 2+d/2 n =⇒ δ n (K n ) ≲ K n √ γ n .(25) Here a n ≲ b n represents that there exists a constant c that does not change with n, such that a n ≤ cb n for all n large enough. The proof of this statement is in Section D.5. Tail term τ 1 .Let us consider the first tail term τ 1 , whose definition is recalled below: τ 1 ≡ τ 1 (K n ,T ,H n ) = sup T∈T sup h∈H n E P X [|ℓ T (X)− h(X)| 1 ∥X∥ ∞ >K n ]. Fix an h and T , and observe that E P X [|ℓ T (X)− h(X)|1 ∥X∥ ∞ >K n ]≤ E P X [|ℓ T (X)|1 ∥X∥ ∞ >K n ] + E P X [|h(X)|1 ∥X∥ ∞ >K n ] = : τ 11 (T,h) + τ 12 (T,h). Now, from Lemma B.2, we know that there exist constants c 0 ,c 1 ,c 2 independent of n,T , such that we have |ℓ T (x)|≤ c 0 + c 1 ∥x∥ + c 2 ∥x∥ 2 for all x∈ R d . If the distribution P X is such that E P X [∥X∥ 2 ] <∞, then, this implies that τ 11 (T,h)≤ E P X [ c 0 + c 1 ∥X∥ + c 2 ∥X∥ 2 1 ∥X∥ ∞ >K n ] n→∞ −→ 0. Thus, the only condition needed to control τ 11 is that K n →∞, and that E P X [∥X∥ 2 ] <∞. Next, we look at the other term τ 12 . This is much easier to handle, since we know that ∥h∥ ∞ ≤ b n for all h∈H n . Hence, we have τ 12 ≤ b n P P X (∥X∥ ∞ > K n )≤ b n e −cK 2 n , for some constant c (depending on the Lipschitz constant of T ∗ ). Hence, a K n = Ω( √ logb n ) is sufficient to drive this to 0. 43 Tail term τ 2 . Finally, we consider the remaining tail term τ 2 , defined as τ 2 ≡ τ 2 (K n ,T ,H n ) = sup T∈T sup h∈H n E Q T h e h(X ′ ) − e ℓ T (X) 1 ∥X ′ ∥ ∞ >K n i For fixed T ∈T and h∈H n , observe that since h(x) = 0 for all x :∥x∥ ∞ > K n , we get E Q T h e h(X ′ ) − e ℓ T (X) 1 ∥X ′ ∥ ∞ >K n i = E Q T |1− e ℓ T |1 ∥X ′ ∥ ∞ >K n ≤ P T (∥X ′ ∥ ∞ > K n ) + E Q T e ℓ T 1 ∥X ′ ∥ ∞ >K n = P(∥X ′ ∥ ∞ > K n ) + P (∥X∥ ∞ > K n ). Thus, we have τ 2 ≤ sup T P ∥T −1 (Z)∥ ∞ > K n + P ∥(T ∗ ) −1 Z∥ > K n , which goes to 0 with n under the T -uniform Lipschitz bounds derived in Lemma B.1. B.6.4 Completing the Proof To complete the proof, we will show an explicit choice of the free parameters that drives the W 1 -metric between the true and the NF distributions to zero. In particular, we observe the following: • With m n =⌈e 4b n ⌉, we get e b n / √ m n ≤ e −b n . Plugging this in (19) then gives us r n ≲ b n √ n + b n e −b n ≲ b n √ n . • Now, set b n = C b γ d/4 K 2+d/2+ε n , and γ n = K 2+ε n for some ε > 0. This implies that b n ≍ K 2+d+ε ( 1+ d 4 ) n =⇒ r n ≲ K 2+d+ε ( 1+ d 4 ) n √ n . • The above choice of b n ensures that b n ≳ γ d/4 n K 2+d/2 n , and hence (25) applies to these parameter choices, which leads to δ n (K n ) ≲ K n √ γ n ≲ K −ε/2 n . • Due to the uniform Lipschitz condition on the elements of T along with the fact that Z ∼ N (0,I d ), we can show that τ 1 + τ 2 ≲ poly(K n )e −cK 2 n for some c > 0. With the choice of K n ≍ √ logn as stated in Theorem B.1, this reduces to τ 1 + τ 2 ≲ poly(logn)n −c . Thus, combining with the previous state, we get δ n + τ 1 + τ 2 ≲ (logn) −ε/4 + poly(logn)n −c . Thus, with the choice of the parameters in Theorem B.2, the estimation error converges to 0 with probability at least 1− δ. C Deferred Proofs from Section 3.2 C.1 Proof of Theorem 3.1 To simplify the notation, we will use P (A) to represent the conditional probability P (· | Y ∈ A). By the realizability assumption, we have P (A) Y,T ∗ z (X) = P (A) Y Z . The starting point of the proof is to use the realizability 44 and the uniform convergence assumptions to relate the population losses associated with the ERM model ˆ T n to the error terms u n and r n from Assumption 3.1. For any T ∈ T , let L(T ) = L y (T ) + λL z (T ), and b L n (T ) = b L y,n (T ) + λ b L z,n (T ), and observe the following chain: 0 = L(T ∗ )≥ b L n (T ∗ )− u n − λr n ≥ b L n ( ˆ T n )− u n − λr n ≥ L( ˆ T n )− 2(u n + λr n ). The equality follows from the realizability assumption (INN1), while the first and third inequalities use the uniform convergence assumption (INN2), and the second inequality follows form the definition of ˆ T n as the ERM model. This simple argument allows us to control the individual components of L( ˆ T n ) as follows: L y ( ˆ T n )≤ 2(u n + λr n ),and L z ( ˆ T n )≤ 2 r n + u n λ .(26) For the first inequality, we use the fact that L z (T )≥ 0 for all T as the critic class G n is assumed to contain the 0 element, while the second inequality uses the fact that L y (T )≥ 0 by definition of the squared loss. The next step of the proof is to use the Lipschitz property of the inverse map ˆ T −1 n to relate the W 1 metric between the posterior distributions P (A) ˆ T n (X) and P (A) Y Z . In particular, to further simplify the notation and use μ = P (A) ˆ T n (X) , ρ = P (A) Y Z , and let F denote the inverse map ˆ T −1 n : Y ×Z → X . Then, F # μ and F # ρ denote the true and INN posterior distributions on X by definition (recall that we use F # μ to represent the pushforward of measure μ on to the image space of F ; that is for any measurable E ⊂ X , we have (F # μ)(E) = μ(F −1 (E))). Let π ∈ Π(μ,ρ) denote any coupling on (Y×Z) 2 with marginals μ and ρ, and let π X ∈ Π(F # μ,F # ρ) denote the coupling by pushing π forward under F × F (recall that F = ˆ T −1 n ); that is, π X (E X × G X ) = π ((y,z) : F (y)∈ E X , F (z)∈ G X ). With these definitions, we immediately have W 1 (F # μ,F # ρ) =inf π ′ ∈Π(F # μ,F # ρ) Z X×X ∥x− x ′ ∥dπ ′ (x,x ′ ) =inf π X =(F×F ) # π π∈Π(μ,ρ) Z X×X ∥x− x ′ ∥dπ X (x,x ′ ) =inf π∈Π(μ,ρ) Z X×X ∥F (y,z)− F (y ′ ,z ′ )∥dπ((y,z), (y ′ ,z ′ )) ≤ Jinf π∈Π(μ,ρ) Z (Y×Z) 2 ∥(y,z)− (y ′ ,z ′ )∥dπ((y,z), (y ′ ,z ′ )) = JW 1 (μ,ρ). Putting back the values μ← P (A) ˆ T n (X) , ρ← P (A) Y Z , and F ← ˆ T −1 n , the inequality chain above implies W 1 P (A) X ,P (A) ˆ T −1 n (Y,Z) ≤ JW 1 P (A) ˆ T n (X) ,P (A) Y Z . Let us introduce the notation: ν = P (A) ˆ T n (X) , γ = P (A) Y, ˆ T n,z (X) ,and ω = P (A) Y,Z . Using the fact that W 1 is a metric, we can use triangle inequality to get W 1 P (A) ˆ T n (X) ,P (A) Y Z = W 1 (ν,ω)≤ W 1 (ν,γ) + W 1 (γ,ω).(27) To complete the proof, we need to obtain upper bounds on the two terms in the RHS above. To control the first term in (27), We begin by recalling the explicit definition of the first term W 1 (ν,γ) =inf π∈Π(ν,γ) Z ∥(y,z)− (y ′ ,z ′ )∥dπ((y,z), (y ′ ,z ′ )) Let us now construct the natural coupling π between ν and γ in the following steps: • Generate (X,Y )∼ P (A) XY . 45 • Let ν denote the conditional distribution of ˆ T n (X) = ( ˆ T n,y (X), ˆ T n,z (X)) • Let γ denote the conditional distribution of (Y, ˆ T n,z (X)). Thus ν and γ have a common second component. Using this particular coupling in the definition of W 1 (ν,γ), we observe that W 1 (ν,γ)≤ Z q ∥ ˆ T n,y (x)− y∥ 2 +∥ ˆ T n,z (x)− ˆ T n,z (x)∥ 2 dP (A) XY (x,y) = E (A) XY h ∥ ˆ T n,y (X)− Y∥ i ≤ r E (A) XY h ∥ ˆ T n,y (X)− Y∥ 2 i (28).1 ≤ q L y ( ˆ T n ) p P Y (A) ≤ p 2(u n + λr n ) p P Y (A) ,(28) where the last inequality follows from (26), and (28).1 uses the fact that L y ( ˆ T n ) = E XY [∥ ˆ T n,y (X)− Y∥ 2 ] = P Y (A)E (A) XY [∥ ˆ T n,y (X)− Y∥ 2 ] + P Y (A c )E (A c ) XY [∥ ˆ T n,y (X)− Y∥ 2 ]≥ P Y (A)E (A) XY [∥ ˆ T n,y (X)− Y∥ 2 ]. It remains for us to obtain a bound on the term W 1 (γ,ω). Introduce the conditional total variation (TV) distance between these two measures, ∆ A =∥P (A) Y Z − P (A) Y, ˆ T n,z (X) ∥ TV =∥ω− γ∥ TV , and observe that an application of the truncation argument stated in Lemma D.1 gives us W 1 (γ,ω)≤ C a 2R P Y (A) 1 1+a ∆ a 1+a A ,where C a = 2 a 1 1+a + a −a 1+a .(29) Next, we use the fact that the two distributions γ and ω share the marginal distribution of Y , which can be used to show that D f P (A) Y, ˆ T n,z (X) ∥ P (A) Y Z (30).1 ≤ 1 P Y (A) D f P Y, ˆ T n,z (X) ∥ P Y Z (30).2 ≤ 2(r n + u n /λ) + η n P Y (A) ,(30) where (30).1 is justified in Lemma D.2, and (30).2 follows from the bound obtained in (26), along with an application of the variational gap assumption (INN4) that relates L z ( ˆ T n ) to D f P Y, ˆ T n,z (X) ∥ P Y Z . Next, we use the Pinsker-type inequality to obtain ∆ A ≤ c f r 1 2 D f P (A) Y, ˆ T n,z (X) ∥ P (A) Y Z ≤ c f s (2r n + 2u n /λ + η n ) 2P Y (A) . Plugging this into (29), we get W 1 (γ,ω)≤ C a 2 2−a 2(1+a) (P Y (A)) − 2+a 2(1+a) c a 1+a f (2r n + 2u n /λ + η n ) a 2(1+a) (31) Combining (31) and (28) with (27) gives us the stated upper bound: W 1 (P (A) X ,P (A) ˆ T −1 n (Y,Z) )≤ J s 2(u n + λr n ) P Y (A) + K a (P Y (A)) − 2+a 2(1+a) (2r n + 2u n /λ + η n ) a 2(1+a) ! ,(32) where K a = 2 4+a 2+2a a 1 1+a + a −a 1+a c a 1+a f . For any finite a > 0, the dominant term is ≲ (r n + u n + η n ) a 2(1+a) , as claimed in the statement of Theorem 3.1. This completes the proof. C.2 Proof of Proposition 3.1 C.2.1 Proof of (7) The uniform convergence of the b L z,n (T ) follows from the boundedness of the function classes involved, along with some standard symmetrization and concentration techniques. In particular, we can show that with probability at least 1− δ/2, we have sup T∈T | b L z,n (T )− L z (T )|≤ 2R (1) n + 2R (2) n + (b n + A 2,n ) r 2 log(4/δ) n .(33) 46 We first introduce some notation to simplify the expressions. For any T and g, define α(g) = 1 n n X i=1 g(Y i ,Z i )− E[g(Y,Z)], ψ(T,g) = 1 n n X i=1 f ∗ (g(Y i ,T z (X i )))− E[f ∗ (g(Y,T z (X)))], and s(T,g) = E[g(Y,Z)]− E[f ∗ (g(Y,T z (X)))]. With these terms, we can write b L z,n (T ) = sup g∈G n s(T,g) + α(g)− ψ(T,g) ,and L z (T ) = sup g∈G n s(T,g). This leads to the following chain: sup T∈T b L z,n (T )− L z (T ) = sup T∈T sup g∈G n (s(T,g) + α(g)− ψ(T,g))− sup g∈G n s(T,g) ≤ sup T∈T sup g∈G n |α(g) + ψ(T,g)| ≤ sup g∈G n |α(g)| + sup T∈T sup g∈G n |ψ(T,g)|.(34) It remains to control the two terms in (34). For the first term, standard symmetrization arguments imply E[ sup g∈G n |α(g)|]≤ 2E Y n ,Z n ,ε n " sup g∈G n 1 n n X i=1 ε i g(Y i ,Z i ) # = 2R (1) n (G n ). Finally,theuniformboundednessassumptionon G n impliesthattherandomvariable sup g∈G n 1 n P n i=1 ε i g(Y i ,Z i ) satisfies the bounded difference property.Hence, an application of McDi- armid’s inequality leads to P sup g∈G n |α(α)|≥ 2R (1) n (G n ) + b n r 2 log(8/δ) n ! ≤ δ 4 .(35) Next, to bound the second term in (34), we use the shorthand m T,g (x,y) = f ∗ (g(y,T z (x)), and let M n = h T,g : g ∈ G n , T ∈ T denote the associated function, and obtain the following (dropping the subscript from m T,g ): E[sup T,g |ψ(T,g)|] = E " sup m∈M n 1 n n X i=1 m(X i ,Y i )− E[m(X,Y )] # ≤ 2E X n ,Y n ,ε n " sup m∈M n 1 n n X i=1 ε i m(X i ,Y i ) # , by another application of the symmetrization technique. Next, observe that by definition, u 7→ f ∗ (u) is Lipschitz with constant A 1,n , and thus using the vector contraction lemma (Fact A.1), we obtain E X n ,Y n ,ε n " sup m∈M n 1 n n X i=1 ε i m(X i ,Y i ) # ≤ 2A 1,n E X n ,Y n ,ε n " sup T,g 1 n n X i=1 ε i g(Y i ,T z (X i )) # = 2R (2) n (G n ,T ). Finally, observing that each |f ∗ (g(Y i ,T z (X i )))| ≤ A 2,n by assumption, another application of McDiarmid’s bounded difference inequality gives us the required concentration result: P sup g∈G n T∈T |ψ T,g |≥ 4A 1,n R (2) n (G n ,T ) + A 2,n r 2 log(8/δ) n ≤ δ 4 .(36) Combining (35) and (36) give us the required 1− δ/2 probability bound claimed in (7). 47 C.2.2 Proof of (8) The proof of (8) requires a careful truncation idea, as the squared loss function can be unbounded. We begin with the simple decomposition for some K n > 0: sup T | b L y,n (T )− L y (T )| = sup T | b L K n y,n (T )− L K n y (T )| | z : =Term1 + sup T | b L y,n (T )− b L K n y,n (T )| | z : =Term2 + sup T |L K n y (T )− L y (T )| | z : =Term3 . (37) We will bound the three terms in (37) separately. Bound on Term1 in (37). Let us introduce the function φ y (u) = ∥Π K n (u)− Π K n (y)∥ 2 , and its centered version ̄ φ y (u) = φ y (u)− φ y (0). Then, observe the following: E ε n sup T∈T 1 n n X i=1 ε i φ Y i (T y (X i )) = E ε n " sup T 1 n n X i=1 ε i ̄ φ Y i (T y (X i )) + 1 n n X i=1 ε i φ Y i (0) # = E ε n " sup T 1 n n X i=1 ε i ̄ φ Y i (T y (X i )) # + E ε n " 1 n n X i=1 ε i φ Y i (0) # = E ε n " sup T 1 n n X i=1 ε i ̄ φ Y i (T y (X i )) # ,(38) since each Rademacher random variable ε i is independent of Y i . Thus, we can conclude via the standard symmetrization argument that E [Term1]≤ 2E " sup T∈T 1 n n X i=1 ε i φ Y i (T y (X i )) # ≤ 2E " sup T∈T 1 n n X i=1 ε i ̄ φ Y i (T y (X i )) # . Now, we can verify that the functions φ y (u) are 4K n Lipschitz globally, which implies that we can apply the (scalar) contraction lemma to get E ε n " sup T∈T 1 n n X i=1 ε i φ Y i (T y (X i )) # ≤ 4K n E ε n sup T∈T ∥u i ∥≤1 1 n n X i=1 ε i ⟨u i ,T y (X i )⟩ = 4K n R n (F y ). Plugging this into (38), we get E[Term1]≤ 8K n R n (F y ). We can translate this expectation result into a high probability deviation bound via the standard bounded difference concentration inequality. In particular, by construction, the functions involved in Term1 satisfy a bounded differences property with parameter 2K n , which implies the following P Term1≥ 8K n R n (F y ) + 4K 2 n q 2 log 8 δ n ≤ δ 4 ,(39) by an application of Mcdiarmid’s inequality. Bound on Term2 in (37). From the definitions of b L K n y,n and b L y,n , we observe that b L y,n − b L K n y,n ≤ 1 n n X i=1 ∥T y (X i )− Y i ∥ 2 1 ∥T y (X i )∥∨∥Y i ∥>K n ≤ 1 n n X i=1 2∥T y (X i )∥ 2 1 ∥T y (X i )∥>K n + 2∥Y i ∥ 2 1 ∥Y i ∥>K n ≤ 2 n n X i=1 ∥T y (X i )∥ 2+β K β n + ∥Y i ∥ 2+β K β n , 48 where the last inequality uses the fact that 1 A>B ≤ (A/B) β . Thus, the assumption on the (2 + β) moment of T y (X i ) and Y i together imply that E[sup T | b L y,n (T )− b L K n y,n (T )|]≤ 4R y K β n .(40) For a given δ > 0, we then apply Markov’s inequality to conclude that P sup T | b L y,n (T )− b L K n y,n (T )| > 16R y K β n δ ≤ δ 4 .(41) Bound on Term3 in (37). The same argument that we used to obtain (40) is also applicable to this term, and we can conclude that sup T∈T |L y (T )− L K n y (T )|≤ sup T E[∥T y (X)− Y∥ 2 1 ∥T y (X)∥ 2 ∨∥Y∥ 2 >K n ] ≤ sup T 2 E[∥T y (X)∥ 2 1 ∥T y (X)∥ 2 >K n ] + E[∥Y∥ 2 1 ∥Y∥ 2 >K n ] ≤ 4R y K β n .(42) Combining (39), (41), and (42) with (37) and (33), we get the required result. C.3 Proof of Proposition 3.3 The general outline of the proof is similar to that of Proposition B.3. For any T ∈ T , let γ T denote D f (P ∥ Q T ), and observe that for any g ∈G n , we have ∆(T ) = D f (P ∥ Q T )− sup g∈G n E P [g]− E Q T [f ∗ (g)] = (E P [g ∗ T ]− E Q T [f ∗ (g ∗ T )])− sup g∈G n E P [g]− E Q T [f ∗ (g)] ≤ (E P [g ∗ T ]− E Q T [f ∗ (g ∗ T )])− (E P [g]− E Q T [f ∗ (g)]). Next, we subtract and add the clipped version of g ∗ T , denoted by ̄g T , to get ∆(T )≤ (E P [g ∗ T − ̄g T ] + E P [ ̄g T ])− (E Q T [f ∗ (g ∗ T )− f ∗ ( ̄g T )] + E Q T [f ∗ ( ̄g T )])− (E P [g]− E Q T [f ∗ (g)]) = (E P [g ∗ T − ̄g T ] + E Q T [f ∗ ( ̄g T )− f ∗ (g ∗ T )]) + (E P [ ̄g T − g] + E Q T [f ∗ (g)− f ∗ ( ̄g T )]) : = C 1 (T ) + C 2 (T,g). The first term C 1 (T ) is the “clipping error” and satisfies the inequality C 1 (T )≤ E P [(|g ∗ T |− b n ) + ] + A 1,n E Q T [ |g ∗ T |− b n ) + ]. In the above display, we have used the fact that |g ∗ T − ̄g T | = (|g ∗ T |− b n ) + and that f ∗ is A 1,n Lipschitz on [−b n ,b n ]. To analyze the term C 2 (T,g), we will consider its behavior inside and outside the ball B K n , which we denote by C 21 ≡ C 2,1 (T,g) and C 22 ≡ C 22 (T,g) respectively. C 21 = E P [( ̄g T − g)1 B K n ] + E Q T [ f ∗ (g)− f ∗ ( ̄g T ) 1 B K n ], and C 22 = E P [( ̄g T − g)1 B c K n ] + E Q T [ f ∗ (g)− f ∗ ( ̄g T ) 1 B c K n ]. To bound the term C 21 , we again appeal to the fact that f ∗ is A 1,n Lipschitz on [−b n ,b n ], and hence we have |C 21 |≤ E P [| ̄g T − g|1 B K n ] + A 1,n E Q T [| ̄g T − g|1 B K n ]. By definition, we have | ̄g T | ≤ b n and |f ∗ (·)| ≤ A 2,n , and also |g| ≤ b n due to the uniform boundedness of G n . Together these facts imply |C 22 |≤ b n P P B c K n + A 2,n P Q T B c K n . 49 Combining these, we get for any T ∈T , and g ∈G n : ∆(T )≤|C 1 (T )| + E P [| ̄g T − g|1 B K n ] + A 1,n E Q T [| ̄g T − g|1 B K n ] + b n P P B c K n + A 2,n P Q T B c K n . Taking the infimum over all g ∈G n , and supremum over all T ∈T , we get η n ≤ τ 2 (b n ) + (1 + A 1,n )δ n (b n ,K n ) + (b n + A 2,n )τ 1 (K n ). Hence, a sufficient condition for this to converge to zero is if τ 2 +A 1,n δ n + (b n +A 2,n )τ 1 converges to 0. This concludes the proof. C.4 Proof of Theorem 3.2 The proof of this result follows closely the proof of Theorem B.1 as given in Section B.3. Since we use the same model class T , we can follow the exact argument we used in Section B.6.2 to show that sup T∈T Lip(T ), Lip(T −1 ) ≤ J : = max (1 + s 2 ) M , (1− s 2 ) −M , andsup T∈T ∥T (0)∥ 2 ≤ J 0 : = s √ H + B (1 + s 2 ) M − 1 s 2 . Recall that the terms s,B,H, and M were introduced in Definition 3.1. Additionally, since our latent variable Z is a multivariate Gaussian, it follows from the realizability assumption that the conditions of Proposition 3.2 hold with all β > 0. To complete the proof, it remains to verify that Proposition 3.1 and Proposition 3.3 hold. C.4.1 Verification of Proposition 3.1 To verify this result, we need to control the supervised and unsupervised loss terms, and in particular, show the existence of u n ,r n → 0 such that sup T∈T | b L y,n (T )− L y (T )|≤ u n ,andsup T∈T | b L z,n (T )− L z (T )|≤ r n . Proposition 3.1 implies that it suffices to check the uniform control over 2 + β moment of T y (X) over all T ∈T , and that the Rademacher complexitiesR n (F y ),R (1) n (G n ), andR (2) n (G n ,T ) vanish with n. Finite 2 + β moment. From the global Lipschitz property of T , there exist finite constants J 0 ,J such that for all T ≡ (T y ,T z )∈T , we have ∥T y (x)∥≤ J 0 + J∥x∥,for all x∈X. Hence, for any β > 0, there exists a constant C β depending on J 0 ,J and β, such that ∥T y (X)∥ 2+β ≤ C β 1 +∥X∥ 2+β =⇒sup T∈T E∥T y (X)∥ 2+β ≤ C β 1 + E[∥X∥ 2+β ] <∞, using the finite 2 + β moment assumption on X. Rademacher complexityR n (F y ). For any T ∈T and u :∥u∥≤ 1, observe that E sup T∈T u:∥u∥≤1 1 n n X i=1 ε i ⟨u,T y (X i )⟩ ≤ 2JE " sup v:∥v∥≤1 1 n n X i=1 ε i ⟨v,X i ⟩ # ≤ 2J sup v:∥v∥=1 ∥v∥E " 1 n n X i=1 ε i X i # , 50 where the first inequality follows from the contraction lemma, and the fact that the map x 7→ ⟨u,T y (x)⟩ is J -Lipschitz, and the second inequality follows from an application of Cauchy-Schwarz. Next, by using Jensen’s inequality and the concavity of the map a7→ √ a, we have E " 1 n n X i=1 ε i X i # = E v u u t 1 n n X i=1 ε i X i 2 ≤ v u u u t E 1 n n X i=1 ε i X i 2 = v u u u t E 1 n 2 n X i=1 n X j=1 ε i ε j ⟨X i ,X j ⟩ . Since ε i ⊥ ε j for i̸= j, we get the bound R n (F y )≤ 2J v u u t 1 n 2 n X i=1 E[∥X∥ 2 ] =O 1 √ n , which converges to 0 with n. Rademacher ComplexitiesR (1) n (G n ) andR (2) n (G n ,T ). For the first term, observe that with U i = (Y i ,Z i ) and w i = 1 [−K n ,K n ] d x (U i ) and g(U i ) = h(U i )w i for some h∈H k : R (1) n (G n ) = E ε n ,U n " sup h∈H k :∥h∥ k ≤b n 1 n n X i=1 ε i ⟨h,k(U i ,·)⟩ k w i # ≤ b n n E U n v u u t n X i=1 w 2 i k(U i ,U i ) ≤ b n √ n , where the first inequality uses a standard argument for RKHSs presented in details in Section B.6.1. Next, for any T ∈ T and g ∈ G n , introduce the terms e U i = (Y i ,T z (X i )) and ew i = 1 [−K n ,K n ] d x ( e U i ), and observe that R (2) n (G n ,T )≤ E ε n , e U n " sup T,g 1 n n X i=1 ε i ◦ g( e U i ) # ≲ A 1,n b n √ n , where we used the (scalar) contraction lemma for Rademacher complexities. C.4.2 Verification of Proposition 3.3 We now obtain the terms τ 1 ,τ 2 and δ n from Proposition 3.3 for our specific example. Recall that we are using the notation P ≡ P Y,Z and Q T ≡ P Y,T z (X) First,weconsidertheterm τ 1 ,whichisequaltosup T∈T Q T (B c K n ) + P (B c K n )= sup T∈T Q T (u :∥u∥ > K n ) + P (u : ∥u∥ > K n ). Equivalently, with the random vector U T = (Y,T z (X)) for some T ∈ T , we can handle Q T (u :∥u∥ > K n ) = P (∥U T ∥ > K n ) simply by using the Markov’s inequality as P (∥U T ∥ > K n )≤ E[∥U T ∥ 2+β ] K 2+β n ≤ R K 2+β n , where the term R follows from the uniform moment condition that we verified using Proposition 3.2. An exact same argument also works for the term P (B c K n ). Hence, τ 1 converges to 0 as long as K n →∞ with n. Next, we look at the term τ 2 which is defined as τ 2 = sup T∈T E P [(|g ∗ T |− b n ) + ] + A 1,n E Q T (|g ∗ T |− b n ) + , where g ∗ T is the optimal (unconstrained) critic function (or witness function) in the definition of JS-divergence. A standard calculation shows that g ∗ T = log 2v T 1 + v T ,where v T = dP dQ T . 51 Next, observe that we g ∗ T is uniformly bounded from above by log 2, and we can use the simple inequality |g ∗ T | =|log 2 + ℓ T − log(1 + v T )|≤ log 2 +|ℓ T |,where ℓ T = logv T = log dP dQ T . As a result, for any b n > log 2, we have (|g ∗ T |− b n ) + ≤ (|ℓ T |− (b n − log 2)) + ≤|ℓ T |1 |ℓ T |>b n −log 2 . Now, following the exact argument for the NF case, we know from Lemma B.2, that there exist constants c 0 ,c 1 ,c 2 independent of n,T , such that we have |ℓ T (x)|≤ c 0 + c 1 ∥x∥ + c 2 ∥x∥ 2 for all x∈ R d . If the distribution P X is such that E P X [∥X∥ 2 ] <∞, then, this implies that E P X [|ℓ T |1 |ℓ T |>b n −log 2 ]≤ E P X [c 0 + c 1 ∥X∥ + c 2 ∥X∥ 2 :∥X∥ ∞ > b n − log 2] n→∞ −→ 0. Thus, the only condition needed to control the above term is that b n →∞, and that E P X [∥X∥ 2 ] <∞. The same argument also works for E Q T [|ℓ T |1 |ℓ T |>b n −log 2 ]→ 0 due to the bounded 2 + β moment condition. It remains to consider the approximation error term on the ball B K n = u ∈ R d x : ∥u∥ ≤ K n , defined as δ n ≡ δ n (b n ,K n ) = sup T inf g E P | ̄g T − g|1 B K n + E Q T | ̄g T − g|1 B K n ≤ 2 sup T inf g ∥ ̄g T − g∥ ∞,C K n , where recall that C K n = [−K n ,K n ] d x , ̄g T is the clipped version of g ∗ T at level b n , and∥ ̄g T −g∥ ∞,C K n denotes the sup-norm over the cube C K n . To complete the proof, it suffices to argue that for every T ∈ T , there exists a g T,n ∈G n , such that ∥ ̄g T − g T,n ∥ ∞,C K n =∥g ∗ T − g T,n ∥ ∞,C K n ≲ K n √ γ n . This is ensured by Lemma B.3. C.4.3 Completing the proof Recall that the statement of Theorem 3.2 makes the following choices of the various parameters: • K n = M (s √ H + B) + √ d x + √ 2 logn • γ n = K 2+ε n for some ε > 0 • b n = C b γ d x /4 n K 2+d x /2 n • A 1,n ≤ 1 and A 2,n ≤ b n Now, observe that from (8), these choices imply sup T∈T | ˆ L y (T )− L y (T )| ≲ r logn n + logn √ n + 1 (logn) 2/β n→∞ −→ 0 sup T∈T | ˆ L z (T )− L z (T )| ≲ b n √ n + (b n + A 2,n ) r log(1/δ) n ≲ (logn) d x ( 1 2 + ε 8 ) +1 √ n n→∞ −→ 0. This ensures that the conditions of Proposition 3.1 are satisfied by these parameters. To complete the proof, we will show that the terms τ 1 ,τ 2 , and δ n from Proposition 3.3 also converge to 0 for these parameters. As we showed in the previous section, 52 • τ 1 ≲ 1 K 2+β n for some β > 0, and hence τ 1 → 0 with n→∞ since K n ≍ √ logn. • τ 2 ≲ E P [∥X∥ 2 : ∥X∥ ∞ > b n − log 2] + E Q T [∥X∥ 2 : ∥X∥ ∞ > b n − log 2] which converges to 0 as lim n→∞ b n =∞. • Finally, we also showed that δ n ≲ K n √ γ n ≍ 1 K ε/2 n , which also vanishes as n→∞. This completes the verification of all the conditions for the INN example introduced in Section 3.2.2. D Details of Technical Lemmas D.1 A Truncation Argument Lemma D.1. Suppose U ∼ μ and V ∼ ν denote two random variables taking values in R j for some j ≥ 1, and suppose that max E[∥V∥ 1+a ], E[∥U∥ 1+a ] = R <∞, for some a > 0. Then, we have the following: W 1 (μ,ν) =sup f :R j →R Lip(f )≤1 Z fdμ− Z fdν ≲ ∆ a 2(1+a) , where ∆ =∥μ− ν∥ TV denotes the total variation distance between μ and ν. Proof. Let M ∈ R denote a finite positive real number to be specified later, and let us define the set E =x∈ R j :∥x∥≤ M. Then, we have the following with g = f − f (0): Z fdμ− Z fdν = Z fd(μ− ν)− f (0) Z d(μ− ν) = Z gd(μ− ν) = Z E gd(μ− ν) + Z E c gd(μ− ν) ≤ Z E ∥x∥d|μ− ν| + Z E c g(x)d(μ− ν) ,(43) where the last inequality uses the fact that g = f − f (0) is 1-Lipschitz. The definition of E leads to the following natural upper bound on the first term Z E ∥x∥d|μ− ν|≤ M Z E d|μ− ν|≤ M Z R j d|μ− ν|≤ 2M ∆.(44) The last inequality above uses the fact that 2∥μ− ν∥ TV = |μ− ν|(R d ). Next, we consider the second term in (43), and observe that Z E c gd(μ− ν) ≤ Z E c |g|(dμ + dν)≤ Z E c ∥x∥d(μ + ν) = E ∥U∥1 ∥U∥>M + E ∥V∥1 ∥V∥>M ≤ M −a E ∥U∥ 1+a 1 ∥U∥>M + M −a E ∥V∥ 1+a 1 ∥V∥>M ≤ M −a E ∥U∥ 1+a + M −a E ∥V∥ 1+a ≤ 2M −a R.(45) Combining the bounds in (44) and (45), we get W 1 (μ,ν) = sup g Z gd(μ− ν) ≤ 2M ∆ + 2M −a R.(46) 53 The function M 7→ 2M ∆ + 2M −a R is convex in M , and is optimized at M ∗ = (aR/∆) 1/(1+a) . Plugging this value back in (46) gives us the required W 1 (μ,ν)≤ 2 a 1 1+a + a −a 1+a R 1 1+a ∆ a 1+a : = C a R 1 1+a ∆ a 1+a . As a sanity check, note that if the supports of U and V are bounded and a→∞, we recover the bound that holds for bounded random variables; that is, W 1 (μ,ν) ≲ ∆. D.2 A Conditional f -Divergence Bound Lemma D.2. Let P,Q denote two probability measures on Y ×Z , and for any measurable set A ⊂ Y , let P Y (A) = P (A×Z) > 0. Define the measures P (A) (dy,dz) = 1 P Y (A) 1 y∈A P (dy,dz),and Q (A) (dy,dz) = 1 P Y (A) 1 y∈A Q(dy,dz). Then, we have the following: D f (P ∥ Q) ≤ 1 P Y (A) D f (P (A) ∥ Q (A) ). Note that in general, Q (A) may not be a probability measure, but the inequality still holds. Proof. Let μ denote a measure that dominates P,Q; for example, we may set μ = 1 2 (P + Q), and let p,q denote the densities of P,Q w.r.t. μ. It then follows that dP (A) dμ (y,z) = 1 y∈A P Y (A) p(y,z),and dQ (A) dμ (y,z) = 1 y∈A P Y (A) q(y,z). As a result, we obtain dP (A) dQ (A) (y,z) = p(y,z) q(y,z) = dP dQ (y,z) for all (y,z)∈ A×Z. Now, by the definition of f -divergences, we have D f (P (A) ∥ Q (A) ) = Z A×Z f dP (A) dQ (A) (y,z) Q (A) (dy,dz) = 1 P Y (A) Z A×Z f p(y,z) q(y,z) q(y,z)μ(dy,dz) (47).1 ≤ 1 P Y (A) Z Y×Z f p(y,z) q(y,z) q(y,z)μ(dy,dz) = 1 P Y (A) D f (P ∥ Q),(47) where (47).1 uses the nonnegativity of f and q. This completes the proof. D.3 Proof of Lemma B.1 Consider any j ∈ [M ], and let z j (x) = W j,1 x + b j,1 and D j (x) = diag(tanh ′ (z j (x))) = diag(sech 2 (z j (x))). Since sech 2 (z) ∈ [0, 1], it follows that 0 ≼ D j (x) ≼ I d j . Then, we can verify that the Jacobians of F j and G j , denoted by JF j and JG j respectively, are equal to JF j (x) = W j,2 D j (x)W j,1 ,and JG j (x) = I d + W j,2 D j (x)W j,1 . By the assumption that ∥W j,i ∥ op ≤ s∈ (0, 1) for all j ∈ [M ], we obtain ∥JF j (x)∥ op ≤∥W j,2 ∥ op ∥D j (x)∥ op ∥W j,1 ∥ op ≤ s 2 = ρ. Here we used the fact that ∥D j (x)∥ op ≤ ∥I d j ∥ op = 1. As a result, all the singular values of JG j (x) = (I d +JF j )(x) are in the range [1−ρ, 1 +ρ]. Now, by the chain rule of the flow, and with x 0 = x, x j = G j x j−1 for j ∈ [M ], we have JT (x) = 1 Y j=M JG j (x j−1 )=⇒log| det JT (x)| = M X j=1 log| det(I d + W j,2 D j (x j−1 )W j,1 )|. 54 Since for every j ∈ [M ], the singular values of JG j (x j−1 ) lie in [1− ρ, 1 + ρ], we obtain Md log(1− ρ)≤ M X j=1 log| det JG j (x j−1 )|≤ Md log(1 + ρ). Note that the above result is uniform in the input x, and uniform over all T ∈T . D.4 Proof of Lemma B.2 Let Z ∼ N (0,I d ), and denote its density by φ(z) = (2π) −d/2 e −∥z∥ 2 /2 . Then, for any diffeomorphism T : R d → R d , we have the following by the change of variables theorem: p T (x) = φ(T (x)) det JT (x). Since ℓ T = logp T ∗ /p T , this implies ℓ T (x) = logφ(T ∗ (x)) + log| det JT ∗ (x)|− logφ(T (x))− log| det JT (x)|. Since logφ(z) =−(1/2)∥z∥ 2 − d/2 log(2π), we obtain the following closed form expression for ℓ T : ℓ T (x) = 1 2 ∥T (x)∥ 2 −∥T ∗ (x)∥ 2 + log| det JT ∗ (x)|− log| det JT (x)|.(48) We now establish the uniform Lipschitz constant and the maximum value of ℓ T : T ∈ T when restricted to the cube C K = [−K,K] d for some K > 0. Lipschitz Constant. To derive the Lipschitz constant of ℓ T , we look at its gradient at any x∈ [−K,K] d : ∇ℓ T (x) = JT (x) T T (x)− JT ∗ (x) T T ∗ (x) +∇ log det JT ∗ (x)−∇ log det JT (x). For each coordinate i, we have ∂ x i log det JT (x) = tr (JT (x)) −1 ∂ x i JT (x) ,and ∂ x i log det JT ∗ (x) = tr (JT ∗ (x)) −1 ∂ x i JT ∗ (x) (49) We have already proved in Section B.6.2 that ∥T (x)∥≤ L √ dK + C 0 A,where L = (1 + s 2 ) M , C 0 = s √ H + B,and A = (1 + s 2 ) M − 1 s 2 . Hence, we get max ∥JT (x) T T (x)∥, ∥JT ∗ (x) T T ∗ (x)∥ ≤ L(L √ dK + C 0 A). We now need to analyze the norm of the gradient of the log-det terms in (49). We begin with the directional derivative along any direction v: D v (log det JT (x)) = tr (JT (x)) −1 D v JT (x) ≤ d∥JT (x) −1 ∥ op ∥D v JT (x)∥ op . For any residual block G j , we know that ∥JG j (u) −1 ∥ op ≤ (1− s 2 ) −1 ,and ∥D v JG j (u)∥ op ≤ 2s 3 ∥v∥. On chaining this for T = G M ◦·◦ G 1 , we get ∥D v log det JT (x)∥≤ M X j=1 d 1− s 2 2s 3 ∥J(G M ◦·◦ G 1 (x)v∥≤ 2ds 3 1− s 2 A∥v∥. Hence, we get the required bound max∥∇ log det JT (x)∥, ∥∇ log det JT ∗ (x)∥≤ 2ds 3 1− s 2 A.(50) Together, (49) and (50) imply the following for x,y ∈ [−K,K] d : |ℓ T (x)− ℓ T (y)|≤ 2L(L √ dK + C 0 A) + 4ds 3 1− s 2 A ∥x− y∥ = : L 1 ∥x− y∥.(51) Note that the Lipschitz constant L 1 is independent of T , and hence is valid uniformly over the family T . 55 Maximum value of ℓ T on [−K,K] d . From (48), we know that the value of ℓ T at 0 is ℓ T (0) = 1 2 ∥T (0)∥ 2 −∥T ∗ (0)∥ 2 + log det JT ∗ (0)− log det JT (0). From Lemma B.1, we know that Md log(1− s 2 ) ≤ log det JT (x) ≤ Md log(1 + s 2 ), and from Section B.6.2, we know that ∥T (0)∥≤ C 0 A. Combining these facts, we get the required bound sup T∈T |ℓ T (0)|≤ (C 0 A) 2 + 2Md log(1/(1− s 2 )) = : L 0 .(52) This fact, combined with the (uniform) Lipschitz constant derived in (51), leads to the following uniform bound on the maximum value of ℓ T : sup T∈T sup x∈[−K,K] d |ℓ T (x)|≤ L 0 + L 1 √ dK. This completes the proof. D.5 Proof of Lemma B.3 The high-level idea behind the proof is simple: given the function that we want to approximate (i.e., ℓ T ), we construct h σ by convolving this function with a Gaussian kernel with a scale parameter σ 2 (to be chosen later). Under certain conditions on σ, we can show that this smoothed function lies in H n . Hence, getting a bound on δ n reduces to that of evaluating the discrepancy between ℓ T and its smoothed version h σ , which can be obtained via standard arguments. A compactly supported Lipschitz extension. The first step in our proof is to construct a proxy for ℓ T that agrees with it on C K n = [−K n ,K n ] d , and is Lipschitz on the entire domain R d . To do this, we define ̄ ℓ T (x) = ( ℓ T (x),if x∈ C K n = [−K n ,K n ] d , inf y∈C K n ℓ T (y) + L∥x− y∥,if x̸∈ C K n . We also need to restrict this Lipschitz extension to a compactly supported domain. To do this introduce the function ν 1 : R→ [0, 1], and ν d : R d → [0, 1] as ν 1 (x) = 1,if |x|≤ K n , 1−|x|/2K n ,if |x|∈ (K n , 2K n ], 0,otherwise, and ν d (x) = d Y i=1 ν 1 (x i ), where x = (x 1 ,...,x n ). With this, we introduce f ν (x) = ̄ ℓ T (x)ν d (x). and observe that Lip(f ν )≤ Lip( ̄ ℓ T )∥ν d ∥ ∞ +∥ ̄ ℓ T ∥ ∞ Lip(ν d )≤ L 1 + d(L 0 + L 1 K n ) K n ≲ K n , and ∥f ν ∥ L 2 ≤∥f ν ∥ ∞ (4K n ) d/2 ≲ K 2+d/2 n .(53) The compact support property of f ν was important for controlling its L 2 norm that will be used later. Gaussian Convolution. Fix some σ > 0, and with φ σ (u) = (2πσ 2 ) −d/2 exp(−∥u∥ 2 /2σ 2 ), define the function h σ h σ (x) = (f ν ∗ φ σ )(x) = Z f ν (x− u)φ σ (u)du. 56 Hence, for any x∈ C K n = [−K n ,K n ] d , we have |f ν (x)− h σ (x)| = Z f ν (x)φ σ (u)du− Z f ν (x− u)φ σ (u)du ≤ Z |f ν (x)− f ν (x− u)|φ σ (u) (i) ≤ L 1 Z ∥u∥φ σ (u)du (i) ≤ L 1 Z ∥u∥ 2 φ σ (u)du 1/2 = L 1 √ dσ ≲ K n σ. The inequality (i) uses the fact that the function f ν is Lipschitz with constant L 1 ≲ K n over the entire domain R d (and not just over C K n like ℓ T ), and (i) uses Jensen’s inequality for the concave map x7→ √ x. RKHS norm of h σ . We know that the symmetric Fourier transform of the Gaussian kernel is ˆ k n (ω) = π γ n d/2 e −∥ω∥ 2 /4γ n . Hence, the RKHS norm of any h with ˆ h∈ L 2 is ∥h∥ 2 k n = 1 (2π) d/2 Z R d | ˆ h(ω)| 2 | ˆ k n (ω)| dω = γ n π d/2 1 (2π) d/2 Z | ˆ h(ω)| 2 e ∥ω∥ 2 /4γ n dω ≲ γ d/2 n Z | ˆ h(ω)| 2 e ∥ω∥ 2 /4γ n dω Now, returning to h σ , we know that ˆ h σ = ˆ f ν e −σ 2 ∥ω∥ 2 /2 , which implies that if σ 2 ≥ 1/4γ n , we have by Plancherel’s theorem, ∥h σ ∥ 2 k n ≲ γ d/2 n Z | ˆ f ν (ω)| 2 e −∥ω∥ 2 (σ 2 −1/4γ n ) dω ≤ γ d/2 n Z | ˆ f ν (ω)| 2 dω = γ d/2 n ∥f ν ∥ 2 L 2 ≲ γ d/2 n K 4+d n , where the last inequality uses the L 2 norm bound on f ν obtained in (53). Hence, h σ 1 [−K n ,K n ] d with σ = 1/2 √ γ n lies in our function class H n , and hence reduces our approximation bound to sup x∈[−K n ,K n ] d |f ν (x)− h σ (x)| ≲ K n √ γ n . This completes the proof. 57