VISCOS F LOWS : V ARIATIONAL S CHUR C ONDITIONAL S AMPLING WITH N ORMALIZING F LOWS
Anonymous authors
Paper under double-blind review
A BSTRACT
We present a method for conditional sampling for (possibly pre-trained) normalizing ﬂows 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 ﬂow’s domain in such a way that the ﬂow restrictions to subdomains remain bijective, which is crucial for the Schur complement application.
Simulation from the variational conditional ﬂow 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 ﬂow 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 classiﬁcation.
I NTRODUCTION
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 ﬂows (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: ﬁrst, 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 classiﬁcation 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 ﬂow.
(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 ﬂow is trained with a classiﬁer 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 ﬁtted 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 ﬂow 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 reﬂects 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 ﬁtted 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).
VISCOS F LOWS : V ARIATIONAL S CHUR C ONDITIONAL S AMPLING WITH N ORMALIZING F LOWS 2. P RELIMINARIES AND P ROBLEM F ORMULATION
Preliminaries : Consider a C 1 -diffeomorphism f
( X ) :
X → Y where X ⊆ R d , Y ⊆
R d are open sets and deﬁne 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 P 0 ( X ) with density p 0 ( x ) .
According to the change of variable rule Y has a log-density given by the following formula:
[Equation]
We will use two set of complementary indexes covering [ d ] : observable O ⊂ [ d ] with cardinality # O =
d O , and hidden H ⊂
[ d ] with # H = d H and O ∪ H =
[ d ] .
In what follows, without loss of generality we assume O = { 0 , . . .
, d O − 1 } ) and H = { d O , . . .
, d } , and use the following notation:
[Equation]
and we will use a similar partition for g ( y ) and its Jacobian G ( y ) .
We will use the notation m O O ( J ) =
J O O 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:
[Equation]
We denote by A the “detached” version of a matrix-valued function A ( x ) , i.e. a matrix whose values match those of A 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 ﬂow, i.e., a map f : X → Y with X , Y ⊆ R d .
In what follows we relax the assumption on the pre-trained ﬂow, but for streamlining the presentation we assume that the normalizing ﬂow is given at this point.
Let us now assume that the observation sample y is partially missing or masked (for instance, photographer’s ﬁnger is covering part of the image).
In particular, the observation is split into observed Y O =
{ y O | y ∈ Y} and hidden Y H =
{ y H | y ∈ Y} regions, where O , H are observed and hidden indexes, respectively.
Our goal is to sample y H from the conditional distribution over a partial observation Y O ∈ Y O .
While this distribution has a density given by p ( y H | y O )
= p ( y ) p ( y O ) , computing p ( y O )
= (cid:82) p ( y O , y H )
d y H 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 , . . .
, d O − 1 } and H = { d O , . . .
, 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 artiﬁcial, 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 ﬂow.
This observation also holds for any ﬂows sharing a residual structure x
+ h ( x ) with an activation h (e.g. ReLU), such as Planar and Radial ﬂows (Rezende & Mohamed, 2015),
Sylvester ﬂows (van den Berg et al., 2018), continuous ﬂows (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 ﬂow, the derivative of the output x it +1 with respect to the input layer x t will be dominated by the i th 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( J O O ) | will ensure that the matrix J O O is not ill-conditioned.
Indeed, for any matrix A we have 2 log | det( A ) | =
log(det( AA T ))
= (cid:80) 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 ﬂows), this approach leads to a well-behaving algorithm.
Maximising log | det( J O O ) | to choose partitioning can also potentially be used for coupling ﬂow 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
f O ( X O ; X H ) :
X O → Y O for all x H ∈
X H are C 1 -diffeomorphisms.
Let us now formally deﬁne our reparametrisation from Y space to X as follows:
[Equation]
Our reparametrisation is well-posed as stated in the following lemma with the proof in Appendix A.
Lemma
If Assumption A1 holds then the reparametrisation in Equations 4 and 5 is well-posed, i.e, g H ( y H , y O ) is a diffeomorphism, while the map x O ( x H , y O ) exists and it is differentiable.
Now we can rewrite Equation 1 using Equation 3 in simpler terms:
[Equation]
Here we used the fact that the base distribution p 0 ( x ) can be factorised as a product of independent univariate distributions, and hence p 0 ( x ) = p 0 ( x O ) p 0 ( x H ) .
Now we can deﬁne the variational distribution in terms of the latent samples as follows:
q ( y H ; y O ) =
q ( x H ) (cid:12)(cid:12) det G H H (cid:12)(cid:12) .
Combining
Equation 6 with our variational distributions we obtain the evidence lower bound:
[Equation]
[Table] , (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 ( y H ; y O ) cancels with the one from the joint log-probability given by the model (i.e., the terms with G H H cancel each other out).
Furthermore, we have eliminated the dependence on the variable y H completely instead we treat x H as a variational variable.
We note here that our, at the ﬁrst sight, artiﬁcial partition of the latent space into observable x O and hidden x H parts was the key allowing for the reparametrisation from Y space to X space.
In particular, we used the invertibility of sub-jacobians G H H and
J O O 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 justiﬁes the method’s name.
Note that this ELBO loss can be added to other losses, e.g., to a classiﬁer loss as we do in what follows.
Note that this ELBO can be used in combination with other losses.
Thus the normalizing ﬂow 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 x
O by solving Equation 5.
Second, we need to discuss the computation of the ELBO gradient with respect to x H .
Speciﬁcally, we need to reparametrise x
O 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.
2. S OLVING THE EQUALITY CONSTRAINT
We consider two root-ﬁnding approaches: a new ﬁxed-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 efﬁcient gradient-free method that can achieve this in some settings.
In particular, we solve simultaneously for x O and y H the equations x
O =
g O
( y O , y H ) and y H =
f H ( x O , x H ) , instead of solving y O =
f O ( x O , x H ) for x O .
This allows to consider a ﬁxed-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 (cid:102) x
O =
g O
( y O , (cid:102) y H ) , (cid:102)
y H
=
f H ( (cid:102) x
O , x H ) around which g O and f H are continuously-differentiable, and the global Lipschitz constants of f H and g O are equal to L a and L b , respectively.
Then the algorithm converges to the unique solution for any α ∈ (0 , 1] , β ∈ (0 , 1] , if L a L b < 1 . while has not converged do [Equation] x O ; end
Newton-Krylov methods (Knoll & Keyes, 2004) :
In this approach, the value x Ok is updated in an iterative manner for a given step size s
[Equation]
which is derived using the ﬁrst order expansion of the ﬂow restriction.
The reliance of this method on the inverse Jacobian-vector product ( J OO ) − 1 u requires us to invert a sub-matrix (deﬁned 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., ( J OO ) − 1 =
G OO − G OH ( G HH ) − 1 G HO , giving an alternative representation of ( J OO ) − 1 .
First, if d O (cid:29) d H , then inverting the sub-Jacobian
G HH is more computationally efﬁcient, 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 J OO using GMRES, we used G OO as a preconditioner.
We found this to speed up the computation by a factor ofto 4, depending on the size of the problem and the architecture used.
2. ELBO G RADIENT
Taking the gradient of ELBO relies on solving two technical issues: computing the gradient of x O and the gradient of LAD with respect to variational posterior parameters θ .
Reparametrising x O .
According to the implicit function theorem and Equation 5, we can reparameterise x O as a function of x H , y O and θ to:
[Equation]
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):
[Equation]
1 Input: Base distribution p 0 ( x ) , invertible ﬂow f = g − 1 : X → Y , partial (masked) observation y O =
m O ( y )
2 Output: Conditional sample y H ∼ p ( · | y O )
Initialization: set q θ ( x H ) , select partitioning function x O −
m O ( x ) ; while has not converged do sample x H ∼ q θ ( x H ) ; compute x O as in Section 2.2; reparameterise x O ( x H , y O , θ ) ; 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 | det A ( x ) | cannot be estimated in closed form, its gradient only requires us to differentiate the trace of the product of A − 1 and A ( x ) , where only the latter is differentiated through.
In many cases, one learns the map g ( · ) and hence J is only implicitly deﬁned by its inverse G (e.g.
J is obtained through Neumann series or Krylov methods).
In such cases, higher order derivatives (i.e., ∇ x J ( x ) ) cannot be accessed readily by backpropagating through the graph.
Therefore, it is necessary to express J ( x ) through G ( y ) .
In particular, if d O (cid:28) d H we can obtain the Natural LAD gradient Estimator (NLADE) of the sub-Jacobian log | det J OO ( x H ) | :
[Equation]
, (8) where we used ∇ x ( J ( x )) =
∇ x ( G ( f ( x ))) −= ∇ x ( G − 1 G ( f ( x ))
G − 1 ) =
∇ y
( G − 1 G ( y ) G − 1 ) ∇ x f ( x ) (Petersen & Pedersen, 2008).
Using NLADE can be inefﬁcient if d H (cid:28) d O .
In this case, one could express ( J O O ) − 1 =
G O O − G O H ( G H H ) − 1 G H O , but this approach can be executed more efﬁciently using the identity log | det J OO ( x ) | = log | det G HH ( f ( x )) | − log | det G ( f ( x )) | in the ﬁrst place.
This leads to the Corrected LAD gradient Estimator (CLADE):
[Equation]
Both expressions have their advantages and possible drawbacks: if d O (cid:28) d H , then NLADE would be more efﬁcient.
On the other hand, if d H (cid:28) d O , then CLADE appears to be a better option.
Furthermore, as CLADE does not use implicit matrix value ( J O O ) − 1 (deﬁned through G ), its variance is potentially lower than NLADE’s.
In other cases, the choice would depend on the efﬁciency of matrix inversions of J , G , J O O , G H H .
2. A LGORITHM
We suggest a two-pass iterative algorithm aimed at sampling from a conditional normalizing ﬂow using a parametric posterior q θ ( y H | y O ) : ﬁrst, a set of conditional samples is drawn by solving the equations
g O ( y O , y H ) − x O = 0 , f H ( x O , x H ) − y H = 0 or f O ( x O , x H ) − y O = 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 conﬁguration being ﬁtted 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.
E XPERIMENTS
Experimental setup We tested our variational conditional sampling algorithm on a classiﬁcation 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
[Equation]
In practice, the sum was truncated when the difference between two successive iterations was below a predeﬁned threshold (cid:15) ( 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 overﬂow, 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 .
In less than 1% of the iterations, our ﬁxed-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.
[Figure]
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 classiﬁer 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
[Table]
[Table] [Table]
performance, while the symbol ”-” represents a missing experiment.
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 ofﬁcial test dataset, which was not used for training.
Figuregives a visual account of the performance of the VISCOS Flows data completion capability.
In all cases, VISCOS Flows was either ﬁrst or second ranked in terms of Top-1 classiﬁcation accuracy (Table 1),
Fr ´ echet 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 classiﬁer on complete data, and then use this classiﬁer on the completed dataset.
Our classiﬁer was, instead, trained on-the-ﬂy 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 iResNetarchitecture 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 ofcompleted images were generated to obtain a gradient estimate.
A total of 500 steps showed to be sufﬁcient to reach convergence.
Figure 3 shows how our algorithm compared with PL-MCMC for a set of test images.
We also included the ﬁrst 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 ﬁtted to the rest of the image.
As a qualitative measure of ﬁtness, 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 conﬁdence interval is 0 .
95% of standard deviation).
By comparison, at the ﬁrst 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 ﬁt of the variational posterior can generate a wide variety of candidate images with little effort, as only a few seconds are needed to solve the equality constraint depicted above.
Figure 4 gives some ﬂavour of this capability.
C ONCLUSION
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 ﬂow 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 ﬂows, such as Glow (Kingma & Dhariwal, 2018) or
RealNVP (Dinh et al., 2016))
the partitioning choice and solving for x
O may be more involved than for NFs with residual architectures.
Future work should focus on ﬁnding alternative ways to ﬁt low-dimensional variational posteriors for data completion in those cases.
A D IFFERENTIABILITY OF THE REPARAMETRISATION
For convenience we state the lemma
Lemma
If Assumption A1 holds then the reparametrisation in Equations 4 and 5 is well-posed, i.e, g H is a diffeomorphism, while the map
x
O
( x H , y O ) exists and it is differentiable.
First, let us show that g H is a C 1 diffeormorphism.
Recall that we use the notation J O O = ∇ x O f O ,
J O H = ∇ x H f O , J H
O = ∇ x H f O , J H H = ∇ x H f H .
Note that the matrices J O O , J are invertible since f ( · ) and f O ( · ; x H ) for all x H are C 1 -diffeomorphisms.
Now, according to the
Guttman rank additivity formula (Guttman, 1946) if J OO is an invertible submatrix of an invertible matrix J , then
[Equation]
It is easy to see that in this case, the block J HH − J HO (cid:0) J OO (cid:1) − 1 J OH is also full-rank, and its inverse is given by G HH which is the Jacobian matrix ∇
y H
g H ( y H ; y O ) .
As these identities are valid for all x , the Jacobian matrix of g H ( y H ; y O ) for y O is always invertible and hence g H is a
C 1 -diffeomorphism, as well.
Now let us show the second part of the statement.
Since f O ( x O ; x H ) is a diffeomorphism, due to implicit function theorem there exists a locally differentiable map
x
O
( x H , y O ) for all x H , y O .
B T HEOREM 1: A CONVERGENCE RESULT
For convenience we restate the theorem.
Theorem 1 Consider Algorithm 1, assume that there exists a unique solution (cid:102) x
O =
g O
( y O , (cid:102) y H ) , (cid:102)
y H
=
f H ( (cid:102) x
O , x H ) around which g O and f H are continuously-differentiable, and the global Lipschitz constants of f H and g O are equal to L a and L b , respectively.
Then the algorithm converges to the unique solution for any α ∈ (0 , 1] , β ∈ (0 , 1] , if L a L b < 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 ﬁrst we need to develop theoretical notions of stability.
Recall that the updates in the algorithm are as follows:
[Equation]
By letting η , ξ stand for (cid:102) y H , (cid:102) x O , respectively, and a ( · ) =
f H ( · , x H ) and b ( · ) =
g O ( y O , · ) , our iterative procedure can be written without loss of generality as the following dynamical system:
[Equation]
We now will review stability results, which are presented for the systems ζ k +
=
C ( ζ k ) with the initial state ζ
=
ν , 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 ≥
× R ≥
→ R ≥
[Equation]
, γ is continuous and strictly increasing in the ﬁrst argument, continuous and strictly decreasing in the second argument, with lim t →∞ γ ( ζ , k )
=
for every ζ , γ (0 , k ) =
for all k ≥ 0 .
Deﬁnition 1 A system is called globally asymptotically stable at ζ ∗ if there exists a function γ ∈ KL such that for all ν and all k :
[Equation]
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).
Deﬁnition
A system is called globally asymptotically incrementally stable if there exists a function γ such that for all ν 1 , ν 2 and all k :
[Equation]
In our derivation, we will also employ local convergence concepts and results
Deﬁnition
A system is called locally asymptotically stable at ζ ∗ if there exist a function γ and ε > 0 such that for all (cid:107) ν − ζ ∗ (cid:107) <
ε and all k :
[Equation]
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 deﬁned as ρ = max i | λ i ( D ) | , where λ i are the eigenvalues of D .
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
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 C jk ∈ R l j × l k , (cid:107) · (cid:107) is the singular value matrix norm, and (cid:107) C jk (cid:107)
≤ c jk then ρ    
C. . .
C 1 m .
C m. . .
C mm     ≤
ρ     (cid:107) C 11 (cid:107) . . .
(cid:107) C 1 m (cid:107) .
(cid:107) C m 1 (cid:107) . . .
(cid:107) C mm (cid:107) 
   ≤ ρ     c. . .
c 1 m .
c m. . .
c mm     .
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 L a L b < 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 = η k 1 − η k 2 , δ η k +
=
η k + − η k
+ , δ ξ k = ξ k 1 − ξ k 2 , and δ ξ k +1 =
ξ k + − ξ k + , then we have (cid:107) δ
η k +1 (cid:107) = (cid:107) η k
+ − η k
+ (cid:107) =
(cid:107) (1 − α )
( η k 1 − η k 2 )
+ α ( a ( ξ k 1 ) − a ( ξ k 2 ))
(cid:107) ≤ (1 − α ) (cid:107) η k 1 − η k 2 (cid:107) + α (cid:107) a ( ξ k 1 ) − a ( ξ k 2 ))
(cid:107) ≤ (1 − α ) (cid:107) δ η k (cid:107) + αL a (cid:107) δ ξ k (cid:107) .
Similarly (cid:107) δ ξ k +1 (cid:107) =
(cid:107) ξ k + − ξ k
+
(cid:107) = (cid:107) (1 − β )( ξ k 1 − ξ k 2 )
+ β ( b ( η k + ) − b ( η k + )) (cid:107) ≤ (1 − β ) (cid:107)
ξ k 1 − ξ k 2 (cid:107) + β (cid:107) b ( η k + )
− b ( η k + ) (cid:107) ≤ (1 − β ) (cid:107) δ ξ k (cid:107) + βL b (cid:107) δ η k +1 (cid:107) ≤ ((1 − β )
+ αβL a L b ) (cid:107) δ ξ k (cid:107) + (1 − α ) βL b (cid:107) δ η k (cid:107) .
Summarising we have the following bounds
[Equation]
which we can equivalently write in the matrix form:
[Equation]
As these bounds are valid for every step k
[Equation]
we can study convergence of (cid:107) δ η k (cid:107) , (cid:107) δ ξ k (cid:107) using the system:
[Equation]
Indeed, if | z k 1 | 2 + | z k 2 | 2 converges to zero as k → ∞ , then (cid:107) δ η k (cid:107)
+ (cid:107) δ ξ k (cid:107) 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 {| z k 1 | 2 + | z k 2 | 2 } converges to zero as k → ∞ for any initialization z
= (cid:107) δ η 0 (cid:107) , z 02 = (cid:107) δ ξ k (cid:107) .
This would imply that (cid:107) δ η k (cid:107)
+ (cid:107) δ ξ k (cid:107) 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 z 2 = (cid:113) (1 − α ) βL b αL a ˜
z
2 resulting in the following dynamical system
[Equation]
Since the system matrix (cid:101) C is symmetric, all we have to do is derive the conditions when
[Equation]
which is equivalent to
[Equation]
a L b (cid:19) (cid:31)
⇐⇒ αβ − α 2 βL a L b > α (1 − α ) βL a L b ⇐⇒ αβ > αβL a L b ⇐⇒ 1 >
L a L b .
This proves that {| z k 1 | 2 + | ˜
z
k 2 | 2 } and consequently (cid:107) δ η k (cid:107) 2 + (cid:107) δ ξ k (cid:107) 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
[Equation]
Due to the Lipschitz constraint on the function a and b , we have (cid:107) A (cid:107)
≤ L a and (cid:107) B (cid:107) ≤ L b .
According to Lemma 4, this leads to:
[Equation]
and ﬁnally to ρ ( C ) < 1 , local asymptotic stability around the ﬁxed point, as well as global asymptotic stability according to Lemmas 2 and 3.
[Equation]
C O N THE FORM OF THE VARIATIONAL POSTERIOR
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-ﬁeld 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 ﬂexibility to the variational posterior, we relied on Householder ﬂows (Tomczak & Welling, 2016) as a sparse approximation to a full covariance matrix:
[Equation]
where the function D maps the vector σ to a diagonal matrix.
Householder ﬂows consist of a series of n orthogonal transformation H i of a base vector (cid:15) ∼ N (0 , I ) .
In principle, such maps can model any orthogonal matrix when n is sufﬁciently 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 conﬁguration.
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 ﬁlled 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.
D MNIST DATA COMPLETION WITH I MPLICIT N ORMALIZING F LOWS
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 efﬁciently train INF, we note that the original problem that is to be solved for one block, which reads
[Equation]
is equivalent to solving
[Equation]
where h ( z ) =
z + g ( z ) is a iResNet block.
We know how to invert h (Behrmann et al., 2019) using ﬁxed-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 efﬁcient 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 deﬁned 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 ﬂipped ( f i → f − 1 i ) 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:
[Equation]
where G = J − 1 = ( I + ∇ x f ( x )) − 1 is the inverse Jacobian obtained through the truncated series (cid:80) nj =0 − ( ∇ x f ( 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 ﬁxed-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 ﬁxed-point algorithm always played a major role in ﬁnding the solution of the equality constraint, reducing the distance between the imposed values y O , x H and their estimate f O ( (cid:101) x ) , g H ( (cid:101) 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 ﬁnd a workable solution in every single completion case.