Universal Approximation of Residual Flows in Maximum Mean Discrepancy
Abstract
Normalizing ﬂows are a class of ﬂexible deep generative models that offer easy likelihood computation.
Despite their empirical success, there is little theoretical understanding of their expressiveness.
In this work, we study residual ﬂows, a class of normalizing ﬂows composed of Lipschitz residual blocks.
We prove residual ﬂows are universal approximators in maximum mean discrepancy.
We provide upper bounds on the number of residual blocks to achieve approximation under different assumptions.
1. Introduction
Normalizing ﬂows are a class of generative models that learn an invertible function to transform a predeﬁned source distribution into a complex target distribution (Tabak et al., 2010; Tabak and Turner, 2013; Rezende and Mohamed, 2015).
One category of normalizing ﬂows called residual ﬂows use residual networks (He et al., 2016) to construct the transformation (Rezende and Mohamed, 2015; Van Den Berg et al., 2018; Behrmann et al., 2019; Chen et al., 2019).
These models have shown great success in complicated real-world tasks.
However, to ensure invertibility, these models apply additional Lipschitz constraints to each residual block.
Under these strong constraints, how expressive these models are remains an open question.
Formally, can they approximate certain target distributions to within any small error?
In this paper, we carry out a theoretical analysis on the expressive power of residual ﬂows.
We prove there exists a residual ﬂow F that achieves universal approximation in the mean maximum discrepancy (MMD, (Gretton et al., 2012)) metric.
Formally, given a target distribution, we provide upper bounds on the number of residual blocks in F such that applying F over the source distribution can approximate the target distribution in squared MMD (see (4)).
Although residual networks are universal approximators (Lin and Jegelka, 2018), the proof of approximation uses a non-invertible construction and therefore does not apply to residual ﬂows.
This reﬂects the main difﬁculty in analyzing residual ﬂows: under strong Lipschitz and invertibility constraints, they become a very restricted function class.
As an illustration, take the set of piecewise constant functions.
Classical real analysis shows that piecewise constant functions can approximate any Lebesgue-integrable function and therefore any probability density.
However, the invertible subset of all piecewise constant functions is the empty set!
Consequently, this universal approximation result does not apply to normalizing ﬂows.
This difﬁculty leads to many negative results for normalizing ﬂows: they are either unable to express or ﬁnd it hard to approximate certain functions (Zhang et al., 2019; Koehler et al., 2020; Kong and Chaudhuri, 2020).
To tackle this problem, we adopt a new construction that satisﬁes the strong Lipschitz constraints in Behrmann et al.
Speciﬁcally, we construct the residual blocks by multiplying a small (cid:15) to a pre-speciﬁed Lipschitz function.
Therefore, as long as (cid:15) is small enough, the strong Lipschitz constraints are satisﬁed.
We then analyze the following quantity: how much can the MMD be reduced if a new residual block is appended?
Since this quantity is a function of (cid:15) , we can analyze its Taylor expansion.
With a ﬁrst-order analysis and under mild conditions, we show there is an
F with Θ (cid:16) 1 δ (cid:0) log 1 δ (cid:1) 2 (cid:17) residual blocks that achieves (4) (see Theorem 1), where δ is the ratio between the ﬁnal squared MMD and the initial squared MMD.
With a secondorder analysis and under more conditions, we show there is a shallower F with only Θ (cid:0) log 1 δ (cid:1) residual blocks that achieves (4) (see Theorem 2).
To sum up, we show residual ﬂows are universal approximators in MMD under certain assumptions and provide explicit bounds on the number of residual blocks.
2. Related Work
The classic universal approximation theory for fully connected or residual neural networks in the function space are widely studied (Cybenko, 1989;
Hornik et al., 1989; Hornik, 1991; Montufar et al., 2014; Telgarsky, 2015; Lu et al., 2017; Hanin, 2017; Raghu et al., 2017; Lin and Jegelka, 2018).
However, these results do not generalize to residual ﬂows (Rezende and Mohamed, 2015; Van Den Berg et al., 2018; Behrmann et al., 2019; Chen et al., 2019) for two reasons.
First, the approximation theory for normalizing ﬂows analyzes how well they can transform between distributions , rather than their ability to approximate a target function in the function space .
Despite that L p universality in the function space may lead to distributional universality for triangular ﬂows (Teshima et al., 2020), there is no similar results for non-triangular ﬂows including residual ﬂows.
Second, the classic results do not consider the invertibility or the Lipschitz constraints of the neural networks, which greatly restrict the expressive power.
There are also universal approximation results for Lipschitz networks (Anil et al., 2019; Cohen et al., 2019; Tanielian et al., 2020).
These results are related because in this work, we assume the expressive power of each Lipschitz residual block is large.
However, these results only apply to functions deﬁned on compact sets.
Because compact sets are bounded, it is “easier” to satisfy the Lipschitz constraints.
It is not trivial to extend their results to functions deﬁned on R d .
Concerning the expressive power of generative networks, there are prior works showing feed-forward generator networks can approximate certain distributions (Lee et al., 2017; Bailey and Telgarsky, 2018; Lu and Lu, 2020;
Perekrestenko et al., 2020).
However, the results are again based on non-invertible constructions, so they do not apply to normalizing ﬂows.
In the literature of normalizing ﬂows, there are universal approximation results for several models including autoregressive ﬂows (Germain et al., 2015; Kingma et al., 2016;
Papamakarios et al., 2017; Huang et al., 2018; Jaini et al., 2019), coupling ﬂows (Teshima et al., 2020; Koehler et al., 2020), and augmented normalizing ﬂows (Zhang et al., 2019; Huang et al., 2020) 1 .
There is also a continuous-time generalization of normalizing ﬂows called neural ODEs (Chen et al., 2018; Dupont et al., 2019) with a universal approximation result (Zhang et al., 2019).
We do not consider these ﬂows in this paper.
In addition, (M¨uller, 2020) suggests residual networks can approximate neural ODEs, but the invertibility is again not considered in this case.
On the expressive power of residual ﬂows, all existing theoretical analysis present negative results for these models (Zhang et al., 2019; Koehler et al., 2020; Kong and Chaudhuri, 2020).
These results indicate residual ﬂows are either unable to express certain functions, or unable to approx- imate certain distributions even with large depths.
Compared to these results, our paper presents positive results for standard residual ﬂows: given a source distribution q , they can approximate a target distribution p in the MMD metric (Gretton et al., 2012) under certain conditions.
We provide explicit upper bounds on the number of residual blocks (see
Theorem 1 and Theorem 2).
3. Preliminaries
We ﬁrst deﬁne the maximum mean discrepancy (MMD) metric between distributions below.
Deﬁnition 1 (MMD, (Gretton et al., 2012)) .
Let q, p be two distributions on R d .
Then,
[Equation]
for some kernel function K ( · , · ) .
Let φ : R d → R d φ be the feature map associated with K : K ( x, z ) =
φ ( x ) (cid:62) φ ( z ) , where we assume d φ < ∞ .
Then, the squared MMD can be simpliﬁed as
[Equation]
Next, we deﬁne a residual ﬂow as a composition of invertible layers parameterized as Id + f , where Id is the identity map and f is 12 -Lipschitz 2 .
The class of residual ﬂows include planar ﬂows (Rezende and Mohamed, 2015), Sylvester ﬂows (Van Den Berg et al., 2018), and the more general invertible residual networks (Behrmann et al., 2019; Chen et al., 2019).
In these models every f i is parameterized as a certain kind of fully-connected neural network.
Since the expressive power of ( 12 -)Lipschitz neural networks on R d remains an open problem, in this paper we assume every f i can be selected as any 12 -Lipschitz function.
Formally, we make the following deﬁnition.
Deﬁnition 2 (Residual ﬂows) .
The set of N -block residual ﬂows is deﬁned as
[Equation]
Now we state the main problem.
Let q source and p target be two distributions on R d , where q source is the source distribution and p target is the target distribution.
We aim to answer the following problem in this paper.
Problem Statement.
Let δ > 0 be a small number.
For any pair of distributions q source and p target on R d satisfying
For algebraic convenience, we restrict the Lipschitz constant to be at most 12 .
MMD( q source , p target ) < ∞ , does there exist an N and
F ∈ F N such that
[Equation]
where F # q refers to the distribution obtained by applying
F over q ?
In this paper, we prove existence of such F with a loose bound on N using ﬁrst-order analysis under mild assumptions (see Section 4), and provide a tighter bound on N using second-order analysis under more assumptions (see
Section 5).
4. A Bound with First-Order Analysis
In this section, we show under mild conditions, there exists a residual ﬂow F with N = Θ (cid:16) 1 δ (cid:0) log 1 δ (cid:1) 2 (cid:17) residual blocks that achieves (4).
The idea is to show that a single residual block can reduce the squared MMD by a certain fraction, so F is obtained by stacking an enough number of these residual blocks.
To begin with, we make the follow deﬁnition.
Deﬁnition 3.
For distributions q , p , and a 12 -Lipschitz function f , we deﬁne the improvement of the squared MMD by
Id + f as
[Equation]
Then, if ∆( q, p ; f ) > 0 , the residual block Id + f is helpful for reducing the squared MMD.
It is straightforward to see that sup { ∆( q, p ; f ) : f is 12 -Lipschitz } ≥ 0 .
In order to construct an f that has a large ∆( q, p ; f ) , we choose f = ˆ f (cid:15) deﬁned below.
Deﬁnition 4.
Deﬁne ψ ( p, q ) =
( E x ∼ p − E x ∼ q ) φ ( x ) , g ( z ) =
ψ ( p, q ) (cid:62) φ ( z ) , and ˆ f (cid:15) =
(cid:15) · ∇ g , where (cid:15) > 0 .
Then, MMD( q, p ) = (cid:107) ψ ( p, q ) (cid:107) .
In addition, ˆ f (cid:15) ( z ) =
(cid:15)J φ ( z ) ψ ( p, q ) , where J φ is the Jacobian matrix of φ .
Then, according to (2) and (5),
[Equation]
Note that ∆( q, p ; ˆ f (cid:15) ) is a function of (cid:15) .
We then analyze the ﬁrst-order Taylor expansion of ∆( q, p ; ˆ f (cid:15) ) at (cid:15) =
+ , denoted as ∆ 1 ( q, p ; ˆ f (cid:15) ) .
Then, ∆( q, p ; ˆ f (cid:15) ) = ∆ 1 ( q, p ; ˆ f (cid:15) )
+
O (cid:0) (cid:15) 2 (cid:1) .
With some arithmetic, we have
[Equation]
p
[Equation]
We have the following bound on ∆ 1 ( q, p ; ˆ f (cid:15) ) .
Lemma 1.
If d φ < ∞ , and the minimum singular value σ min ( J φ ( z )) ≥ √ b > 0 holds for any z ∈ R d , then
[Equation]
Since ∆( q, p ; ˆ f (cid:15) ) =
∆ 1 ( q, p ; ˆ f (cid:15) )
+ O (cid:0) (cid:15) 2 (cid:1) , when (cid:15) is small, the residual block Id + ˆ f (cid:15) can indeed reduce the squared
MMD by a certain fraction ( ≥ 2 (cid:15)b ).
Next, as we require f = ˆ f (cid:15) to be 12 -Lipschitz, we show under certain conditions the Lipschitz constant of ˆ f (cid:15) is O ( (cid:15) ) in the following lemma.
Lemma 2.
If for any z ∈ R d , the Lipschitz constant of each element in J φ ( z ) is no more than a universal constant L Jac , then
[Equation]
With these tools, we can construct an F ∈ F N that achieves (4) in the following theorem.
Theorem 1.
Under the conditions of Lemma 1 and Lemma 2, there exists an F ∈ F N with N = Θ (cid:16) 1 δ (cid:0) log 1 δ (cid:1) 2 (cid:17) that achieves (4) .
The proof is deferred to Section A.3.
The main idea in the proof is to construct each f i iteratively based on fthrough f i − 1 , so that adding this residual block can reduce the squared MMD by a certain fraction as indicated in
Lemma 1.
The bound is obtained by carefully balancing (cid:15) , δ , and N .
5. A Tighter Bound with Second-Order Analysis
In this section, we show under a few additional assumptions, there exists a much smaller N = O (cid:0) log 1 δ (cid:1) and F ∈ F N such that F achieves (4).
The idea is to bound the secondorder remainder of the Taylor expansion of ∆( q, p ; ˆ f (cid:15) ) :
[Equation]
Once ∆ 2 ( q, p ; ˆ f (cid:15) ) is explicitly bounded we can pick a small constant (cid:15) for every residual block 3 so ∆( q, p ; ˆ f (cid:15) ) is lower bounded by a universal constant times MMD( q, p ) 2 .
This then yields the O (cid:0) log 1 δ (cid:1) bound for N .
Now, we provide an explicit bound on ∆ 2 ( q, p ; ˆ f (cid:15) ) in the following lemma.
Lemma 3.
Let B, C, L feat be positive constants.
If for any z ∈ R d , the maximum singular value σ max ( J φ ( z )) ≤ √ B , the absolute value of any eigenvalue | λ ( ∇ 2 φ i ( z ))
| ≤ C for any 1 ≤ i ≤ d φ , and φ is L feat -Lipschitz, then
[Equation]
[Equation] Given this explicit bound on ∆ 2 ( q, p ; ˆ f (cid:15) ) , we can pick a small (cid:15) such that | ∆ 2 ( q, p ; ˆ f (cid:15) ) | ≤ 12 ∆ 1 ( q, p ; ˆ f (cid:15) ) so that ∆( q, p ; ˆ f (cid:15) ) ≥ 12 ∆ 1 ( q, p ; ˆ f (cid:15) ) .
Once this lower bound on ∆( q, p ; ˆ f (cid:15) ) is achieved, the squared MMD is multiplied by at most a universal constant less than 1 when the new residual block Id + ˆ f (cid:15) is added.
We formalize the result in the following theorem.
Theorem 2.
Under the conditions of Lemma 1, Lemma 2, and Lemma 3, there exists an F ∈ F N with N = Θ (cid:0) log 1 δ (cid:1) that achieves (4) .
The proof is deferred to Section A.5.
The main idea in the proof is to construct each f i in a similar way as in Theorem 1, but (cid:15) is selected as a universal constant according to
Lemma 3.
6. Conclusions
Normalizing ﬂows are a class of ﬂexible deep generative models that offers easy likelihood computation.
Despite their empirical success, there is little theoretical understanding on whether they are universal approximators in transforming between probability distributions.
In this work, we prove residual ﬂows are indeed universal approximators in maximum mean discrepancy.
Upper bounds on the number of residual blocks to achieve approximation are provided.
Under mild conditions, we show Θ (cid:16) 1 δ (cid:0) log 1 δ (cid:1) 2 (cid:17) residual blocks can achieve (4) (see Theorem 1).
Under more conditions, we show as few as Θ (cid:0) log 1 δ (cid:1) residual blocks can achieve (4) (see Theorem 2).
There are a large number of open problems.
One extension is to build universal approximation theory for residual ﬂows in more general probability metrics such as the integral probability metrics (M ¨ uller, 1997) and the f -divergences (Csisz´ar and Shields, 2004).
Another direction is to extend the proposed universal approximation theory to other classes of normalizing ﬂows such as autoregressive ﬂows.
A ﬁnal open problem is to look at normalizing ﬂows in real-world applications, and analyze their expressive power under practical assumptions. Acknowledgement
We thank Songbai Yan for helpful feedback.
A. Omitted Proofs A.1. Proof of Lemma
Proof.
According to (7) and the chain rule,
[Equation]
A.2. Proof of Lemma
Proof.
For any x, y ∈ R d , (cid:107)
[Equation]
y − x
[Equation] = [Equation] Therefore, by taking the supreme over the left-hand-side, we have the Lipschitz constant of ˆ f (cid:15) is upper bounded by the right-hand-side.
A.3. Proof of Theorem
Proof.
Let r > 0 and (cid:15)
= r/N .
Deﬁne
D n ( r ) =
MMD(( Id + f n ) ◦ · · · ◦ ( Id + f 1 )# q source , p target )where each f i ( z )
=
(cid:15)J φ ( z ) ψ ( p target , ( Id + f
i − 1 ) ◦ · · · ◦ ( Id + f 1 )# q source ) .
Note that each f i is exactly the ˆ f (cid:15) in Deﬁnition 4 for q
= ( Id + f
i − 1 ) ◦ · · · ◦ ( Id + f 1 )# q source and p = p target .
By Lemma 1,
[Equation]
D
[Equation]
Therefore,
D N ( r ) ≤
N (cid:89) n = (cid:18)
[Equation]
For small δ >
[Equation]
we choose
[Equation]
Then, we have
[Equation]
Note that
[Equation]
Therefore [Equation] by Lemma
[Equation]
when δ is small enough, the Lipschitz constant of each f n is less than 12 .
A.4. Proof of Lemma
Proof.
According to (1) and (5),
[Equation]
There is a closed-form expression for ∆ 2 ( q, p ; ˆ f (cid:15) ) .
According to the remainder of multivariate Taylor polynomials
[Equation]
there exist two maps ξ 1 , ξ 2 : R d → (0 , 1) such that
[Equation]
First, we bound |
[Equation]
Deﬁne
[Equation]
Since φ is L feat -Lipschitz and 0 < ξ 1 ( x ) < 1 , we have
[Equation]
Therefore,
[Equation]
For any z (cid:48) , v ∈ R d ,
[Equation]
By letting z (cid:48)
=
[Equation] z ) [Equation] ( z ) and v = ˆ f (cid:15) ( z ) , we have
[Equation]
Next, we bound ∆ 32 .
Observe that
[Equation]
Therefore,
[Equation]
Combining these bounds, we have
[Equation]
A.5. Proof of Theorem
Proof.
Let q 0 =
q source and q m =
( Id + f m ) ◦ · · · ◦ ( Id + f 1 )# q source , where each
[Equation]
Deﬁne ψ 0 = ψ ( p target , q 0 ) and assume (cid:107) ψ ( p target , q m ) (cid:107) ≤ (cid:107) ψ 0 (cid:107) (which we will prove by induction).
Note that
[Equation]
According to Lemma 1 and Lemma 3, we have ∆( q m , p target ; f m )
[Equation]
φ
[Equation]
32 CL feat
[Equation] [Equation]
32 CL feat [Equation]
When
[Equation]
CL
[Equation]
we have
[Equation]
Next, by Lemma 2, in order to satisfy the Lipschitz condition, we require
[Equation]
This is satisﬁed when we assign
[Equation]
Now, we set
[Equation]
Then, we have
[Equation]
which also implies (cid:107) ψ ( p target , q m +1 ) (cid:107) ≤ √ 1 − b ˆ (cid:15) (cid:107) ψ 0 (cid:107) ≤ (cid:107) ψ 0 (cid:107) .
Finally, in order to satisfy (4), we only need to take the number of residual blocks as
[Equation]