VISCOS FLOWS: VARIATIONAL SCHUR CONDI-
Year: 2021

Abstract:
We present a method for conditional sampling for (possibly pre-trained) normalizing flows when only part of an observation is available. We derive a lower bound to the conditioning variable log-probability using Schur complement properties in the spirit of Gaussian conditional sampling. Our derivation relies on partitioning flow’s domain in such a way that the flow restrictions to subdomains remain bijective, which is crucial for the Schur complement application. Simulation from the variational conditional flow then amends to solving an equality constraint. Our contribution is three-fold: a) we provide detailed insights on the choice of variational distributions; b) we discuss how to partition the input space of the flow to preserve bijectivity property; c) we propose a set of methods to optimise the variational distribution. Our numerical results indicate that our sampling method can be successfully applied to invertible residual networks for inference and classification.

heading: 1 INTRODUCTION
text: Conditional data generation is a ubiquitous and challenging problem, even more so when the data is high dimensional. If the partitioning of the conditioned data and conditioning itself can be established in advance, a wide variety of inference tools can be used ranging from Bayesian inference (Gelman, 2013) and its approximations (e.g. variational inference (Klys et al., 2018)) to Gaussian (Williams & Rasmussen, 2006) or Neural Processes (Garnelo et al., 2018). The task is considerably more challenging when the partitioning cannot be anticipated. This is met in a wide range of applications, from few shots learning (Wang et al., 2020) to super resolution image generation (Bashir et al., 2021), high-dimensional Bayesian optimisation (Moriconi et al., 2020; Gomez-Bombarelli et al., 2018), learning with incomplete datasets (Yoon et al., 2018; Li et al., 2019; Richardson et al., 2020), image inpainting (Elharrouss et al., 2020) and many more.
If we can assume that the data is normally distributed, the conditional distribution can be computed in closed form using Gaussian elimination and the Schur complement. This observation paved the way for Gaussian Processes (GPs) solutions (Williams & Rasmussen, 2006) and alike. Although they constitute a solid component of the modern machine learning toolbox, the Gaussian assumption is quite restrictive in nature. Indeed: GPs in their vanilla form are used when the observations are uni-dimensional and when inference has to be performed on a subset of these variables. As a consequence, alternative tools may need to be used when dealing with even moderately dimensional data, or when part of the input data can be missing.
It is well-established that normalizing flows (NFs) (Rezende & Mohamed, 2015) can model arbitrarily complex joint distributions. Furthermore, they naturally embed together the joint probability of any partition of data that could occur during training or at test time. Hence, it is tempting to make use of NFs to perform conditional sampling, but only a few works have provided solutions each with some limitations. (Cannella et al., 2020) proposed Projected Latent Monte Carlo Markov Chains (PL-MCMC) to generate samples whose probability is guaranteed to converge to the desired conditional distribution as the chain gets longer. While such convergence is certainly an advantage, one can identify several important drawbacks: first, as for every other Monte Carlo sampling method, the mixing time cannot be known in advance, and hence one may wonder whether the samples that are gathered truly belong to the conditional distribution. Second, although the observed part of the imputed data converges towards the true values, a difference between the two may persist. Third,
training a model comprising of missing data with PL-MCMC is achieved through a Monte Carlo Expectation Maximisation (MCEM) scheme (Dempster et al., 1977; Neath, 2013). Under appropriate assumptions, MCEM is guaranteed to converge to a local maximiser of the log-likelihood function, but this is conditioned on the quality of the data generated by the Monte Carlo algorithm. Consequently, as the optimisation progresses towards the optimum, longer chains may be required to ensure convergence or even to obtain convincing results (Neath, 2013). MCFlow (Richardson et al., 2020) relies on an auxiliary feedforward neural network whose role is to produce latent embeddings with maximum a posteriori likelihood values. Those values are constrained to lie on the manifold of latent vectors whose mapping to the observed space match the observations. MCFlow produces state-of-the-art results in terms of image quality or classification accuracy when compared to GAN-based methods such as (Yoon et al., 2018; Li et al., 2019). However, this method requires a set of adjustments to the model and needs retraining for any additional incomplete data. On the other hand, ACFlow Li et al. (2020) learns all conditional distributions for all possible masking operations, which can be quite computationally expensive. As such, both MCFlow and ACFlow cannot be applied to post-training data completion as PL-MCMC.
Several conditioning techniques have been used with NFs in contexts where the joint probability distribution of the conditioning and conditioned random variable is of no interest. For instance, (Rezende & Mohamed, 2015; van den Berg et al., 2018) extended the amortisation technique used in Variational Auto-Encoders (VAEs) (Kingma & Welling, 2013) to the parameters of the flow. (Winkler et al., 2019) extended this use to the scenario of conditional likelihood estimation, while (Trippe & Turner, 2018) provided a Variational Bayes view on this problem by using a Bayesian Neural Network for conditioning feature embedding. On a different line of work, (Kingma & Dhariwal, 2018) use a heuristic form of a posteriori conditioning. The generative flow is trained with a classifier over the latent space, which forces the latent representation location to be indicative of the class they belong to. At test time, some parametric distribution is fitted on the latent representations to account for the attributes of the images that are to be generated. (Nguyen et al., 2019) used a hybrid approach, whereby the conditioning and conditioned variables joint distribution is modelled by a normalizing flow in such a way that conditional sampling of one given the other is straightforward to apply. Still, either with this approach or the others above, the knowledge of what part of the data constitutes the conditioning random variable is required beforehand.
Our Contribution: We propose a conditional NF sampling method in the following setting: a) conditional data generation is performed after a model has been trained without taking conditional sampling into account; b) training data may be itself incomplete, in the sense that some training features might be missing from examples in the training dataset; c) the subset of missing data may not be known in advance, and it could also be randomly distributed across input data. Importantly, we are interested in deriving a method whereby the distribution of the generated data faithfully reflects the Bayesian perspective. Our derivations heavily rely on the Schur complement properties in the spirit of the conditional Gaussian distributions. To highlight this feature we call our approach VISCOS Flows: VarIational Schur COnditional Sampling with normalizing Flows. The use of a variational posterior brings some advantages: for example, with a single fitted posterior, multiple samples can quickly be recovered. We also show how to amortise the cost of inference across multiple partially observed items by using inference networks (Kingma & Welling, 2013).

heading: 2 VISCOS FLOWS: VARIATIONAL SCHUR CONDITIONAL SAMPLING WITH NORMALIZING FLOWS

heading: 2.1 PRELIMINARIES AND PROBLEM FORMULATION
text: Preliminaries: Consider a C1-diffeomorphism f(X) : X → Y where X ⊆ Rd, Y ⊆ Rd are open sets and define the inverse of f to be g(y) ≡ f−1 : Y → X . Consider the case where the random variable Y = f(X) while X is distributed according to the base distribution P0(X) with density p0(x). According to the change of variable rule Y has a log-density given by the following formula:
log p(y) = log p0(g(y)) + log |det ∇yg(y)| = log p0(x) − log |det∇xf(x)|.
(1)
We will use two set of complementary indexes covering [d]: observable O ⊂ [d] with cardinality #O = dO, and hidden H ⊂ [d] with #H = dH and O ∪H = [d]. In what follows, without loss of
generality we assume O = {0, . . . , dO − 1}) and H = {dO, . . . , d}, and use the following notation:
f(x) = ( fO(xO;xH) fH(xH ;xO) ) , ∇xf(x) = J(x) = ( JOO(x) JOH(x) JH O(x) JHH(x) ) ,
and we will use a similar partition for g(y) and its Jacobian G(y). We will use the notation mOO(J) = JOO to denote sub-matrix masking. The central to our approach is the Schur complements and its properties and, in particular, the following well-known identities:
GHH(f(x)) = ( JHH(x)− JH O(x) ( JOO(x) )−1 JOH(x) )−1 , (2)
det(J) = det ( JOO ) det ( JHH − JH O ( JOO )−1 JOH ) = det ( JOO ) det (GHH) . (3)
We denote by A the “detached” version of a matrix-valued function A(x), i.e. a matrix whose values match those ofA but have a null gradient.
Problem formulation: Consider a set Y containing some samples of interest, for example, images. Assume also a pre-trained normalizing flow, i.e., a map f : X → Y with X ,Y ⊆ Rd. In what follows we relax the assumption on the pre-trained flow, but for streamlining the presentation we assume that the normalizing flow is given at this point. Let us now assume that the observation sample y is partially missing or masked (for instance, photographer’s finger is covering part of the image). In particular, the observation is split into observed YO = {yO | y ∈ Y} and hidden YH = {yH | y ∈ Y} regions, whereO,H are observed and hidden indexes, respectively. Our goal is to sample yH from the conditional distribution over a partial observation Y O ∈ YO. While this distribution has a density given by p(yH | yO) = p(y)
p(yO) , computing p(yO) =
∫ p(yO,yH)dyH
is in general intractable. Since it is often easier to sample from and optimise in the latent space we will aim to express all distributions in the space X . First, we introduce a partition into observed and hidden indexes in the latent space X .
Assumption A0. Both spaces X and Y are partitioned as O = {0, . . . , dO − 1} and H = {dO, . . . , d}. We stress that Assumption A0 is made only to simplify the exposition. In general, we can have arbitrary partitions of the observed space, which will not affect our algorithm. While this partition seems somewhat artificial, it enables the use of the Schur complement in our derivations. Choosing the same partition in X and Y is natural for the invertible residual networks (iResNet) (Behrmann et al., 2019; Chen et al., 2019) due the
structure of the flow. This observation also holds for any flows sharing a residual structure x+h(x) with an activation h (e.g. ReLU), such as Planar and Radial flows (Rezende & Mohamed, 2015), Sylvester flows (van den Berg et al., 2018), continuous flows (Chen et al., 2018; Onken et al., 2021; Mathieu & Nickel, 2020) or – to some extent – Implicit NFs (Lu et al., 2021), where at each layer t of the flow, the derivative of the output xit+1 with respect to the input layer xt will be dominated by the ith component of this input. We note that the partition in latent space X can be made different from the partition in the observed Y for numerical purposes. For example, picking a latent space partitioning with a maximal value of log |det(JOO)| will ensure that the matrix JOO is not ill-conditioned. Indeed, for any matrix A we have 2 log |det(A)| = log(det(AAT )) = ∑ i log(σi(A)), where σi are the singular values of A. Thus promoting larger singular values will improve conditioning. We empirically found that across several architectures (iResNet, Planar, Radial, Sylvester and Implicit flows), this approach leads to a well-behaving algorithm. Maximising log |det(JOO)| to choose partitioning can also potentially be used for coupling flow architectures (e.g. Glow (Kingma & Dhariwal, 2018) or RealNVPs (Dinh et al., 2016)), however, this requires additional research. While Assumption A0 is made for convenience of notation, the following assumption on f is necessary.
Assumption A1. The map f : X → Y and its restriction fO(XO;XH) : XO → YO for all xH ∈ XH are C1-diffeomorphisms.
Let us now formally define our reparametrisation from Y space to X as follows:
xH(yH ,yO) := gH(yH ;yO), (4)
xO(xH ,yO) such that fO(xO;xH)− yO = 0. (5) Our reparametrisation is well-posed as stated in the following lemma with the proof in Appendix A.
Lemma 1 If Assumption A1 holds then the reparametrisation in Equations 4 and 5 is well-posed, i.e, gH(yH ,yO) is a diffeomorphism, while the map xO(xH , yO) exists and it is differentiable.
Now we can rewrite Equation 1 using Equation 3 in simpler terms:
log p(y) = log p0(x H) + log ∣∣detGHH ∣∣+ log p0(xO)− log ∣∣det JOO∣∣ . (6) Here we used the fact that the base distribution p0(x) can be factorised as a product of independent univariate distributions, and hence p0(x) = p0(xO)p0(xH). Now we can define the variational distribution in terms of the latent samples as follows: q(yH ;yO) = q(xH)
∣∣detGHH ∣∣. Combining Equation 6 with our variational distributions we obtain the evidence lower bound:
LELBO = Eq(yH ;yO) [ log p(yH ,yO) ] +H [ q(yH ;yO) ] = E
q(xH)
[ log p0(xH)
q(xH) + log p0(xO)− log ∣∣∣det JOO(xO,xH)∣∣∣ ], (7) where H[·] is the entropy and we have taken advantage of the fact that the log-absolute determinant (LAD) of the Jacobian belonging to the approximate posterior q(yH ;yO) cancels with the one from the joint log-probability given by the model (i.e., the terms withGHH cancel each other out). Furthermore, we have eliminated the dependence on the variable yH completely instead we treat xH as a variational variable. We note here that our, at the first sight, artificial partition of the latent space into observable xO and hidden xH parts was the key allowing for the reparametrisation from Y space to X space. In particular, we used the invertibility of sub-jacobiansGHH and JOO allowing to map between observable and hidden samples in the X and Y spaces. Our approach is in the spirit of Gaussian conditional sampling and Schur complement, which justifies the method’s name. Note that this ELBO loss can be added to other losses, e.g., to a classifier loss as we do in what follows. Note that this ELBO can be used in combination with other losses. Thus the normalizing flow can be trained simultaneously while performing conditional sampling of missing data. This way we can treat masked data during training as well as after training.
There are still two technical issues that we cover in the following subsections before presenting our algorithm. First, we need to get xO by solving Equation 5. Second, we need to discuss the computation of the ELBO gradient with respect to xH . Specifically, we need to reparametrise xO and compute the LAD gradient. As the following two sections are quite technical the reader interested in the algorithm itself can skip to Subsection 2.4.

heading: 2.2 SOLVING THE EQUALITY CONSTRAINT
text: We consider two root-finding approaches: a new fixed-point iterative algorithm and Newton-Krylov methods. In practice, we used the latter as a fallback option when the former did not converge below the desired convergence threshold.
Fixed-point iterative algorithm As our method heavily relies on solving equality constraints, one may wish to reduce the computational burden of this operation. We therefore seek for an efficient gradient-free method that can achieve this in some settings. In particular, we solve simultaneously for xO and yH the equations xO = gO(yO,yH) and yH = fH(xO,xH), instead of solving yO = fO(xO,xH) for xO. This allows to consider a fixed-point iterative approach in Algorithm 1, whose convergence properties are analysed in Theorem 1 proved in Appendix B:
Theorem 1 Consider Algorithm 1, assume that there exists a unique solution x̃O = gO(yO, ỹH), ỹH = fH(x̃O,xH) around which gO and fH are continuously-differentiable, and the global Lipschitz constants of fH and gO are equal to La and Lb, respectively. Then the algorithm converges to the unique solution for any α ∈ (0, 1], β ∈ (0, 1], if LaLb < 1.
Algorithm 1: Iterative algorithm solving for xO
1 Input: Invertible flow f = g−1, partial (masked) observation yO = mO(y), conditional latent hidden sub-vector xH , mixing coefficients {α, β} ∈ (0, 1] 2 Output: yH and xO satisfying gO(yO,yH) = xO, fH(xO,xH) = yH for given xH and yO. Initialisation: χO = 0; ỹH = fH(χO,xH);
x̃O = xO = gO(yO,yH); while has not converged do
yH = fH(x̃O,xH);
ỹH = αyH + (1− α)ỹH ; xO = gO(yO, ỹH);
x̃O = βxO + (1− β)x̃O; end
Newton-Krylov methods (Knoll & Keyes, 2004): In this approach, the value xOk is updated in an iterative manner for a given step size s
xOk+1 = x O k − s
( JOO )−1 (fO(xOk ;x
H)− yO), which is derived using the first order expansion of the flow restriction. The reliance of this method on the inverse Jacobian-vector product (JOO)−1u requires us to invert a sub-matrix (defined by the input and output masks) that, in most cases, is only accessible via left or right vector-matrix product, due to the constraints imposed by the use of backward propagation algorithms. Moreover, explicit computation of this sub-Jacobian is also usually prohibitively expensive to retrieve in closed form. Therefore, to compute this product involving an inverse sub-Jacobian, we rely on the generalized minimal residual method (GMRES) (Saad & Schultz, 1986), a Krylov subspace method (Simoncini & Szyld, 2007). At their core, Newton-Krylov methods (Knoll & Keyes, 2004) rely on the computational amortisation of the GMRES-based inversion step by caching intermediate values between successive iterations of the solver (Baker et al., 2005) accelerating the algorithm convergence.
Still, computing a single Jacobian vector product can be quite computationally expensive, for instance when using invertible residual networks (Chen et al., 2019; Behrmann et al., 2019). To this end we employ an identity reciprocal to Equation 2, i.e., (JOO)−1 = GOO−GOH(GHH)−1GHO, giving an alternative representation of (JOO)−1. First, if dO dH , then inverting the sub-Jacobian GHH is more computationally efficient, e.g., using GMRES or a similar technique. In this case, the inverse of the sub-Jacobian can be obtained or approximated only by querying Jacobian-vector products involving G, and not J . Otherwise, when computing the plain inverse of JOO using GMRES, we usedGOO as a preconditioner. We found this to speed up the computation by a factor of 2 to 4, depending on the size of the problem and the architecture used.

heading: 2.3 ELBO GRADIENT
text: Taking the gradient of ELBO relies on solving two technical issues: computing the gradient of xO and the gradient of LAD with respect to variational posterior parameters θ.
Reparametrising xO. According to the implicit function theorem and Equation 5, we can reparameterise xO as a function of xH , yO and θ to:
∇xHxO = −(JOO)−1JOH . The inverse of the sub-Jacobian can be obtained via GMRES algorithm with similar considerations regarding methods to speed up this computation to the Newton-Krylov method in Section 2.2. Reparameterising as a function of the variational posterior parameters θ is then trivially achieved via pathwise or implicit gradient calculation.
Log-Absolute Determinant gradient Estimators. First recall that the gradient of the logdeterminant can be represented using trace as follows (Petersen & Pedersen, 2008):
∇x log |detA(x)| = ∇xTr [ A−1A(x) ] ,
Algorithm 2: VISCOS Flows: Variational Schur Conditional Sampling with Normalizing Flows 1 Input: Base distribution p0(x), invertible flow f = g−1 : X → Y , partial (masked) observation yO = mO(y) 2 Output: Conditional sample yH ∼ p(· | yO)
Initialization: set qθ(xH), select partitioning function xO −mO(x); while has not converged do
sample xH ∼ qθ(xH); compute xO as in Section 2.2; reparameterise xO(xH ,yO,θ); compute stochastic gradient estimate∇θL(q) and update variational posterior;
end
which allows us to use built-in functions if the inverse of A is readily available. Hence, although the value of log |detA(x)| cannot be estimated in closed form, its gradient only requires us to differentiate the trace of the product ofA−1 andA(x), where only the latter is differentiated through.
In many cases, one learns the map g(·) and hence J is only implicitly defined by its inverseG (e.g. J is obtained through Neumann series or Krylov methods). In such cases, higher order derivatives (i.e., ∇xJ(x)) cannot be accessed readily by backpropagating through the graph. Therefore, it is necessary to express J(x) throughG(y). In particular, if dO dH we can obtain the Natural LAD gradient Estimator (NLADE) of the sub-Jacobian log |det JOO(xH)|:
∇x log |det JOO(x)| = ∇xTr [ (JOO)−1JOO(x) ] =
∇yTr [ (JOO)−1mOO ( G−1G(y)G−1 )] ∣∣∣ y=f(x) J(x), (8)
where we used ∇x(J(x)) = ∇x(G(f(x)))−1 = ∇x(G−1G(f(x))G−1) = ∇y(G−1G(y)G−1)∇xf(x) (Petersen & Pedersen, 2008). Using NLADE can be inefficient if dH dO. In this case, one could express (JOO)−1 = GOO − GOH(GHH)−1GH O, but this approach can be executed more efficiently using the identity log |det JOO(x)| = log |detGHH(f(x))| − log |detG(f(x))| in the first place. This leads to the Corrected LAD gradient Estimator (CLADE):
∇x log |det JOO(x)| = ∇y ( Tr [ (GHH)−1GHH(y) ] − Tr [ G−1G(y) ]) ∣∣∣ y=f(x) J(x). (9)
Both expressions have their advantages and possible drawbacks: if dO dH , then NLADE would be more efficient. On the other hand, if dH dO, then CLADE appears to be a better option. Furthermore, as CLADE does not use implicit matrix value (JOO)−1 (defined throughG), its variance is potentially lower than NLADE’s. In other cases, the choice would depend on the efficiency of matrix inversions of J ,G, JOO,GHH .

heading: 2.4 ALGORITHM
text: We suggest a two-pass iterative algorithm aimed at sampling from a conditional normalizing flow using a parametric posterior qθ(yH | yO): first, a set of conditional samples is drawn by solving the equations gO(yO,yH) − xO = 0, fH(xO,xH) − yH = 0 or fO(xO,xH) − yO = 0. Next, this sample is reparameterised given the variational posterior parameters, and lastly an optimisation step is made given the stochastic estimate of the ELBO gradient. The summary is given in Algorithm 2.
We used two distinct designs of approximate posterior depending on the task at hand. When possible, each data to be completed was treated separately, with a single variational posterior configuration being fitted to each item in the dataset. When the amount of incomplete items was too large to be completed on a one-to-one basis, we amortised the cost of inference by acquiring the variational posterior parameters through a inference network that was trained to minimize the KL divergence in Equation 7. More details can be found in Appendix C.

heading: 3 EXPERIMENTS
text: Experimental setup We tested our variational conditional sampling algorithm on a classification and data completion task when the training and test data were partially observed, and on a posttraining data completion task. Due to the restriction imposed by Assumption A1, we only considered the iResNet (Behrmann et al., 2019; Chen et al., 2019) architecture to test our approach on in the main text. Appendix D provides results with Implicit NFs, showing that our approach can be extended to other types of architectures with possibly higher Lipschitz constants. Several steps of our gradient computation rely on computing the inverse of the Jacobian of the transform. This quantity was queried using Neumann series: at each layer i ∈ [L] of a L-deep network, we locally computed the product of a vector with the inverse Jacobian as
uJ−1i = u ∞∑ j=0 − (∇xh(x))j .
In practice, the sum was truncated when the difference between two successive iterations was below a predefined threshold (10−5 for single and 10−10 for double precision). Jacobian-vector products were computed using the identity J(x)u = u∇v [vJ(x)] for some v. As this formula requires a graph to be built on top of another one, and to avoid memory overflow, we computed this quantity locally at each layer of the network. The convergence threshold for Algorithms 1 and 2 was 10−3.
At early stages of training, Algorithm 1 quickly converged with α = β = 1. However, lower values were needed later on, or when using a fully trained model (in the post-training data imputation setting). Therefore, we adopted the following schedule throughout the training process: an initial mixing rate was set to α = β = 0.5, and at each iteration of Algorithm 1, the mixing rate was decayed by a factor of 0.95. In less than 1% of the iterations, our fixed-point algorithm did not reach the convergence threshold, and then we used a Newton-Krylov solver. In the case of MNIST dataset, each optimisation step took approximately 20 seconds, roughly equally split between solving the equality constraint retrieving the gradient estimate.
Quantitative experiment: Training on incomplete MNIST data To test the capability of Variational Schur Conditional Sampling (VISCOS Flows) to learn from incomplete data, we used the MNIST dataset (LeCun & Cortes, 2010) where a percentage of the pixels of each digit image was missing. We considered only the challenging setting where the missing pixels were randomly spread across the image. The pre-trained NF component was similar to (Chen et al., 2019) and consisted of an iResNet model with 73 convolutional layers with 128 channels each, followed by four fully connected residual layers. The trained inference network consisted of a 5-layer convolutional network with a 3× 3 kernel, SAME-padding strategy and with [4, 8, 16, 32, 8] channels at each layer, respectively. For this network, the ReLU activation function (Nair & Hinton, 2010) was used. Finally, we used a simple classifier consisting of the sequence of a squeezing layer, a 1-dimensional batch-norm layer (Ioffe & Szegedy, 2015) and a bias-free linear layer. We used the Adam optimiser (Kingma & Ba, 2015) with a learning rate of 10−2, decaying by a factor of 0.5 at every epoch. The batch size was set to 64 digits, and the components were trained simultaneously for a total of 5 epochs. The NLADE gradient estimator was used for these experiments with the Hutchinson stochastic trace
estimator (Hutchinson, 1990), taking a single sample for each gradient computation. We compared our results to the ones presented in PL-MCMC (Cannella et al., 2020), MCFlow (Richardson et al., 2020) and MisGAN (Li et al., 2019). Due to the limits of our computational resources, we resort to the performances displayed in (Richardson et al., 2020; Cannella et al., 2020). All the presented results are computed on MNIST official test dataset, which was not used for training. Figure 2 gives a visual account of the performance of the VISCOS Flows data completion capability. In all cases, VISCOS Flows was either first or second ranked in terms of Top-1 classification accuracy (Table 1), Fréchet Inception Score (FID (Heusel et al., 2017), Table 2) or Root-Mean-Squared-Error (RMSE, Table 3). It is important to note that, although MCFlow outperformed our approach on some occasions, this algorithm relies on the assumption that one can train a separate classifier on complete data, and then use this classifier on the completed dataset. Our classifier was, instead, trained on-the-fly on the incomplete dataset only.
Qualitative experiment: post-training data imputation We now turn to the problem of completing a partially observed sample using the VISCOS Flow technique. We trained a standard iResNet-164 architecture on the CIFAR-10 dataset (Krizhevsky, 2009) for 250 epochs. VISCOS Flow completion capability was compared to PL-MCMC, as MCFlow cannot perform post-training data imputation. The task was similar to the one presented in (Cannella et al., 2020): a masked centered square of size 8 × 8 had to be completed. PL-MCMC completion was achieved using the same hyperparameters provided by the authors. For VISCOS Flows, we used a Gaussian approximate posterior with a sparse approximation to the covariance matrix encoded by 50 Householder layers. The parameters were trained using the Adam optimiser with a learning rate of 10−2, and at each step, a batch of 8 completed images were generated to obtain a gradient estimate. A total of 500 steps showed to be sufficient to reach convergence. Figure 3 shows how our algorithm compared with PL-MCMC for a set of test images. We also included the first solution obtained by the initial standard Gaussian distribution to show that, although training did make the images look more convincing, the initial guess was already somewhat well fitted to the rest of the image. As a qualitative measure of fitness, we measured the RMSE between generated and true images for both algorithms, and obtained a similar value for both (VISCOS Flow (NLADE): 0.2291 ± 0.1570, PL-MCMC: 0.2287 ± 0.1961, where the confidence interval is 0.95% of standard deviation). By comparison, at the first iteration of the VISCOS Flow algorithm, an RMSE of 0.2802± 0.1188 was measured. As the gradient of the variational posterior parameters can be retrieved with two different methods (CLADE or NLADE), we compared one against the other, but observed little qualitative or quantitative difference between the two (RMSE for VISCOS Flow (CLADE): 0.2320±0.1793). A noticeable advantage of VISCOS Flow is that with a single fit of the variational posterior can generate a wide variety of candidate
Figure 3: Upper row: Original images (left panel) and PL-MCMC reconstruction (right panel). Lower row: VISCOS Flow after the first iteration (left panel), VISCOS Flow after the final iteration
Figure 4: Random samples obtained from a VISCOS Flow trained variational posterior.
images with little effort, as only a few seconds are needed to solve the equality constraint depicted above. Figure 4 gives some flavour of this capability.

heading: 4 CONCLUSION
text: This article presents VISCOS Flows, a new method for data completion with NFs. This technique can be successfully applied to architectures when the units of the flow are made of instances of residual layers. Unlike other approaches, such as MCFlow (Richardson et al., 2020) and ACFlow (Li et al., 2020), we do not need to include the conditioning in the training process. This feature improves algorithm’s modularity by allowing to decouple NF training and conditional sampling. While PL-MCMC (Cannella et al., 2020) also enables such modularity, it relies on MCMC algorithms, which are arguably harder to use than our gradient descent based algorithm. For some important classes of NFs (e.g. composed of coupling flows, such as Glow (Kingma & Dhariwal, 2018) or RealNVP (Dinh et al., 2016)) the partitioning choice and solving for xO may be more involved than for NFs with residual architectures. Future work should focus on finding alternative ways to fit low-dimensional variational posteriors for data completion in those cases.

heading: A DIFFERENTIABILITY OF THE REPARAMETRISATION
text: For convenience we state the lemma
Lemma 1 If Assumption A1 holds then the reparametrisation in Equations 4 and 5 is well-posed, i.e, gH is a diffeomorphism, while the map xO(xH , yO) exists and it is differentiable.
First, let us show that gH is aC1 diffeormorphism. Recall that we use the notation JOO = ∇xOfO, JOH = ∇xHfO, JH O = ∇xHfO, JHH = ∇xHfH . Note that the matrices JOO, J are invertible since f(·) and fO(·;xH) for all xH are C1-diffeomorphisms. Now, according to the Guttman rank additivity formula (Guttman, 1946) if JOO is an invertible submatrix of an invertible matrix J , then
rank(J) = rank(JOO) + rank ( JHH − JHO ( JOO )−1 JOH ) .
It is easy to see that in this case, the block JHH − JHO ( JOO )−1 JOH is also full-rank, and its
inverse is given by GHH which is the Jacobian matrix ∇yHgH(yH ;yO). As these identities are valid for all x, the Jacobian matrix of gH(yH ;yO) for yO is always invertible and hence gH is a C1-diffeomorphism, as well.
Now let us show the second part of the statement. Since fO(xO;xH) is a diffeomorphism, due to implicit function theorem there exists a locally differentiable map xO(xH ,yO) for all xH , yO.

heading: B THEOREM 1: A CONVERGENCE RESULT
text: For convenience we restate the theorem.
Theorem 1 Consider Algorithm 1, assume that there exists a unique solution x̃O = gO(yO, ỹH), ỹH = fH(x̃O,xH) around which gO and fH are continuously-differentiable, and the global Lipschitz constants of fH and gO are equal to La and Lb, respectively. Then the algorithm converges to the unique solution for any α ∈ (0, 1], β ∈ (0, 1], if LaLb < 1.
We argue convergence of the iterative procedure using dynamical systems theory and the notion of stability (cf., (Sontag, 2013)). In general, stability is a stronger property than convergence, as stability requires convergence under perturbations in the initial conditions of an iterative procedure. We will prove the theorem in Proposition 1, in what follows, but first we need to develop theoretical notions of stability.
Recall that the updates in the algorithm are as follows:
ỹH := αfH(x̃O,xH) + (1− α)ỹH
x̃O := βgO(yO, ỹH) + (1− β)x̃O
By letting η, ξ stand for ỹH , x̃O, respectively, and a(·) = fH(·,xH) and b(·) = gO(yO, ·), our iterative procedure can be written without loss of generality as the following dynamical system:
ηk+1 = (1− α)ηk + αa(ξk), ξk+1 = (1− β)ξk + βb(ηk+1).
(10)
We now will review stability results, which are presented for the systems ζk+1 = C(ζk) with the initial state ζ0 = ν, we will denote systems’ trajectories as ζk(ν). In our bounds, we will make use of the class of KL functions. The function γ ∈ KL if γ : R≥0 × R≥0 → R≥0, γ is continuous and strictly increasing in the first argument, continuous and strictly decreasing in the second argument, with limt→∞ γ(ζ, k) = 0 for every ζ, γ(0, k) = 0 for all k ≥ 0.
Definition 1 A system is called globally asymptotically stable at ζ∗ if there exists a function γ ∈ KL such that for all ν and all k:
‖ζk(ν)− ζ∗‖ ≤ γ(‖ν − ζ∗‖, k),
Stability theory offers many additional tools, for instance, the ability of studying how trajectories behave in comparison to each other, which can be done using contraction theory (Lohmiller & Slotine, 1998) or using a similar concept of incremental stability(Tran et al., 2016).
Definition 2 A system is called globally asymptotically incrementally stable if there exists a function γ such that for all ν1, ν2 and all k:
‖ζk(ν2)− ζk(ν1)‖ ≤ γ(‖ν2 − ν1‖, k),
In our derivation, we will also employ local convergence concepts and results
Definition 3 A system is called locally asymptotically stable at ζ∗ if there exist a function γ and ε > 0 such that for all ‖ν − ζ∗‖ < ε and all k:
‖ζk(ν)− ζ∗‖ ≤ γ(‖ν − ζ∗‖, k),
Checking local stability is rather straightforward given some regularity conditions. In particular, we need to check the magnitude of the spectral radius of ∂C(ζ). Recall that spectral radius ρ(D) is defined as ρ = maxi |λi(D)|, where λi are the eigenvalues ofD.
Lemma 2 Let A be continuously-differentiable around ζ∗. If the spectral radius ρ (∂C(ζ∗)) is strictly smaller than one, then the system is locally asymptotically stable around ζ∗. If C(ζ) is a linear function, then the system is also globally stable.
We also have this lemma linking incremental stability and stability properties, which can be shown in a straightforward manner.
Lemma 3 If the system ζk+1 = C(ζk) is locally asymptotically stable around ζ∗ and globally asymptotically incrementally stable, then it is also globally asymptotically stable.
Finally, we need the following lemma, which is a combination of two results discussed in (Karow et al., 2006): properties of nonnegative matrices (the property ρ2) and Lemma 4.1., which follows the introduction of these properties.
Lemma 4 Let Cjk ∈ Rlj×lk , ‖ · ‖ is the singular value matrix norm, and ‖Cjk‖ ≤ cjk then
ρ  C11 . . . C1m... ... Cm1 . . . Cmm   ≤ ρ   ‖C11‖ . . . ‖C1m‖... ... ‖Cm1‖ . . . ‖Cmm‖   ≤ ρ   c11 . . . c1m... ... cm1 . . . cmm   .
Now we are ready to proceed with the proof.
Proposition 1 The system equation 10 is globally asymptotically incrementally stable for any α, β ∈ (0, 1) if LaLb < 1. Furthermore, if there exists ξ∗, η∗ such that ξ∗ = a(η∗), η∗ = b(ξ∗), then the system is globally asymptotically stable.
Proof: First let us show incremental stability. Let δηk = ηk1 − ηk2 , δηk+1 = ηk+11 − η k+1 2 , δξk = ξk1 − ξk2 , and δξk+1 = ξk+11 − ξ k+1 2 , then we have
‖δηk+1‖ = ‖ηk+11 − η k+1 2 ‖ = ‖(1− α)(ηk1 − ηk2 ) + α(a(ξk1 )− a(ξk2 ))‖ ≤
(1− α)‖ηk1 − ηk2‖+ α‖a(ξk1 )− a(ξk2 ))‖ ≤ (1− α)‖δηk‖+ αLa‖δξk‖.
Similarly
‖δξk+1‖ = ‖ξk+11 − ξ k+1 2 ‖ = ‖(1− β)(ξk1 − ξk2 ) + β(b(η k+1 1 )− b(η k+1 2 ))‖ ≤
(1− β)‖ξk1 − ξk2‖+ β‖b(ηk+11 )− b(η k+1 2 )‖ ≤ (1− β)‖δξk‖+ βLb‖δηk+1‖ ≤
((1− β) + αβLaLb)‖δξk‖+ (1− α)βLb‖δηk‖.
Summarising we have the following bounds
‖δηk+1‖ ≤ (1− α)‖δηk‖+ αLa‖δξk‖, ‖δξk+1‖ ≤ (1− α)βLb‖δηk‖+ ((1− β) + αβLbLa)‖δξk‖,
(11)
which we can equivalently write in the matrix form:( ‖δηk+1‖ ‖δξk+1‖ ) ≤ ( 1− α αLa (1− α)βLb (1− β) + αβLbLa )( ‖δηk‖ ‖δξk‖ ) . (12)
As these bounds are valid for every step k, we can study convergence of ‖δηk‖, ‖δξk‖ using the system: (
zk+11 zk+12
) = ( 1− α αLa
(1− α)βLb 1− β + αβLbLa ) ︸ ︷︷ ︸
C
( zk1 zk2 ) .
Indeed, if |zk1 |2 + |zk2 |2 converges to zero as k →∞, then ‖δηk‖2 + ‖δξk‖2 also converges to zero as k → ∞. According to Lemma 2 as long as the matrix C has the largest absolute value of the eigenvalues smaller than one, the sequence {|zk1 |2 + |zk2 |2} converges to zero as k → ∞ for any initialization z01 = ‖δη0‖, z02 = ‖δξk‖. This would imply that ‖δηk‖2 + ‖δξk‖2 converges to zero as k →∞ and hence the system equation 10 is global asymptotic incremental stability. Now all we need to show is that the spectral radius is strictly smaller than one. First, we make
a simple transformation obtained by a change of variables z2 = √
(1−α)βLb αLa
z̃2 resulting in the following dynamical system(
zk+11 z̃k+12
) = ( 1− α √ α(1− α)βLaLb√
α(1− α)βLaLb 1− β + αβLbLa ) ︸ ︷︷ ︸
C̃
( zk1 z̃k2 ) .
Since the system matrix C̃ is symmetric, all we have to do is derive the conditions when( 1− α √ α(1− α)βLaLb√
α(1− α)βLaLb 1− β + αβLbLa
) ≺ ( 1 0 0 1 ) ,
which is equivalent to( α √ α(1− α)βLaLb√
α(1− α)βLaLb β − αβLaLb
) 0⇐⇒
αβ − α2βLaLb > α(1− α)βLaLb ⇐⇒ αβ > αβLaLb ⇐⇒ 1 > LaLb.
This proves that {|zk1 |2 + |z̃k2 |2} and consequently ‖δηk‖2 + ‖δξk‖2 converge to zero as k →∞. Now we need to show that the system is locally asymptotically stable around any solution to ξ∗ = a(η∗), η∗ = b(ξ∗). Let us compute the Jacobian of the dynamical system equation 10. Denoting A = ∂a(η∗)B = ∂b(ξ∗) we have the following Jacobian
C =
( 1− α αA
(1− α)βB (1− β) + αβBA ) Due to the Lipschitz constraint on the function a and b, we have ‖A‖ ≤ La and ‖B‖ ≤ Lb. According to Lemma 4, this leads to:
ρ(C) ≤ ρ ((
1− α α‖A‖ ‖(1− α)βB‖ ‖(1− β) + αβBA‖
)) ≤
ρ
(( 1− α αLa
(1− α)βLb (1− β) + αβLbLa
)) ,
and finally to ρ(C) < 1, local asymptotic stability around the fixed point, as well as global asymptotic stability according to Lemmas 2 and 3.

heading: C ON THE FORM OF THE VARIATIONAL POSTERIOR
text: Post-training conditional inference A broad range of variational algorithms impose a fully factorisable form for the approximate posterior, which is usually referred to as the mean-field assumption (Hoffman et al., 2013; Blei et al., 2017; Kingma & Welling, 2013; Jaakkola & Jordan, 1996; Wainwright & Jordan, 2008; Knowles & Minka, 2011). In this context, one would usually opt in favour of a diagonal Gaussian distribution for qθ base distribution. To relax this assumption and bring flexibility to the variational posterior, we relied on Householder flows (Tomczak & Welling, 2016) as a sparse approximation to a full covariance matrix:
x = µ+ n∏ i=1 HiD(σ) (13)
where the function D maps the vector σ to a diagonal matrix. Householder flows consist of a series of n orthogonal transformationHi of a base vector ∼ N (0, I). In principle, such maps can model any orthogonal matrix when n is sufficiently large. The main advantage of this approximate posterior formulation is its time and memory computational cost, as it only requires inner vector-vector and vector-scalar products.
Learning from incomplete data A common usage in VAEs is to amortise the construction of parametric variational posteriors by building a so-called inference network that maps each datapoint to its correspondent parameter configuration. In the case of incomplete data, one usually has to deal with missing values that are randomly spread in the data space and whose dimensionality can be hard to foreseen. Inputting such sparse data in a neural network can be challenging. To solve this, we filled every missing value with the median value of the data tensor, thereby creating a workable input for the inference network. In turn, this inference network produced a fully-factorised variational posterior which was indexed according to the latent space partition.

heading: D MNIST DATA COMPLETION WITH IMPLICIT NORMALIZING FLOWS
text: We tested the VISCOS Flow algorithm on a data completion task with Implicit Normalizing Flows (INF) (Lu et al., 2021) to show that our method can be extended beyond iResNet architectures. Unlike iResNet, the Lipschitz constant of a single block of INF is unbounded, making it potentially considerably more expressive. To efficiently train INF, we note that the original problem that is to be solved for one block, which reads
F (x, z) = x− z + f(x)− g(z) = 0
is equivalent to solving z = h−1(x+ f(x))
where h(z) = z+ g(z) is a iResNet block. We know how to invert h (Behrmann et al., 2019) using fixed-point iterations, and we have already presented how to propagate gradients through this transformation using Neumann series (see Section 3). These two solutions showed to be considerably faster and more memory efficient than the joint use of Brodyen method and linear system solving algorithms presented in the original work.
We trained a convolutional INF network that was defined in a similar manner as the iResNet model designed for the CIFAR task in the main text, with the following notable amendments. First, to match the INF structure, the direction of one every two iResNet blocks was flipped (fi → f−1i ) such that plain and inverted residual blocks alternated. ActNorm layers (Kingma & Dhariwal, 2018) were interleaved in between iResNet blocks. To compute the higher order gradients required in Equation 8, and since the gradient of inverted iResNet block was computed using Neumann series, we relied once more on the following identity:
∇yG(y) = ∇y (−GJ(x(y))G)
where G = J−1 = (I +∇xf(x))−1 is the inverse Jacobian obtained through the truncated series∑n j=0− (∇xf(x))
j , and the dependence of G (or x) on y is made explicit where needed. We trained this INF network on the complete MNIST training dataset using the Adam optimiser for 20 epochs with a learning rate of 10−3. We used the VISCOS Flow algorithm to complete digits
where 90% of the pixels were missing. The CLADE gradient estimator was used to approximate the gradient of the LAD of the partial Jacobian. Results of this task are displayed in Figure 5. We noticed that the fixed-point algorithm was slightly less effective with INF networks than with iResNet: for a small proportion of the iterations during the optimisation (< 30%), our algorithm did not converge to the solution and required a few Newton-Krylov iterations (< 5) to complete the process. It is worth emphasising that, in all cases, the fixed-point algorithm always played a major role in finding the solution of the equality constraint, reducing the distance between the imposed values yO, xH and their estimate fO(x̃), gH(ỹ) by more than 90% with respect to the original values. Also, using the same partition of the observed and latent space led us to find a workable solution in every single completion case.

Review summaries:
This paper presents a variant of flow, VISCOS flow, utilizing the Schur complement. This flow model is applied to the in-painting task, which estimates the unobserved feature variables given the observed features. Given this in-painting task, the inference requires the modeling on the latent variable, which results in the ELBO derivation on such latent variables as variational variables.

This paper proposes a technique for performing data imputation using pre-trained normalizing flows. They do so by fitting a variational distribution over a subset of the 'base' variables of the normalizing flow. Along with the observations, a sample from this variational distribution is sufficient to fully specify a sample from the normalizing flow model. The authors present techniques for constructing samples in this way, and computing the derivatives required to optimise the variational distribution.