[Paper Review] A Universal Approximation Theorem of Deep Neural Networks for Expressing Probability Distributions

Yulong Lu, Jianfeng Lu|arXiv (Cornell University)|Apr 19, 2020

Generative Adversarial Networks and Image Synthesis56 references77 citations

TL;DR

The paper proves that a ReLU deep neural network can push-forward a source distribution to approximate a target distribution arbitrarily well under 1-Wasserstein, MMD, or Kernelized Stein Discrepancy, with explicit network-size bounds, depending on the chosen metric.

ABSTRACT

This paper studies the universal approximation property of deep neural networks for representing probability distributions. Given a target distribution $\\pi$ and a source distribution $p_z$ both defined on $\\mathbb{R}^d$, we prove under some assumptions that there exists a deep neural network $g:\\mathbb{R}^d\ ightarrow \\mathbb{R}$ with ReLU activation such that the push-forward measure $(\ abla g)_\\# p_z$ of $p_z$ under the map $\ abla g$ is arbitrarily close to the target measure $\\pi$. The closeness are measured by three classes of integral probability metrics between probability distributions: $1$-Wasserstein distance, maximum mean distance (MMD) and kernelized Stein discrepancy (KSD). We prove upper bounds for the size (width and depth) of the deep neural network in terms of the dimension $d$ and the approximation error $\\varepsilon$ with respect to the three discrepancies. In particular, the size of neural network can grow exponentially in $d$ when $1$-Wasserstein distance is used as the discrepancy, whereas for both MMD and KSD the size of neural network only depends on $d$ at most polynomially. Our proof relies on convergence estimates of empirical measures under aforementioned discrepancies and semi-discrete optimal transport.

Motivation & Objective

• Motivate the study of neural networks as generators to represent probability distributions beyond function approximation.
• Show that a ReLU DNN can map a source distribution to approximate a target distribution via push-forward by the gradient of the network output.
• Provide quantitative complexity bounds (depth/width) for achieving a given approximation accuracy under three IPMs.
• Connect empirical measure convergence and semi-discrete optimal transport to construct explicit neural-network-based transport maps.

Proposed method

• Construct a neural-network-based potential u whose gradient pushes p_z to approximate π under IPMs.
• Use empirical measures P_n to approximate π and bound d_FD(P_n, π) for W1, MMD, and KSD.
• Apply semi-discrete optimal transport to show the optimal map T = ∇φ that pushes a continuous μ to a discrete ν, with φ being a maximum of affine functions.
• Express φ as a neural network via a result that max_j{x·y_j + m_j} can be represented by a DNN.
• Provide explicit depth L = ⌈log2 n⌉ and width N = 2^L for the DNN, with dependency of the target metric on sample size n.

Experimental results

Research questions

• RQ1 Can deep ReLU networks express a given target distribution π as a push-forward of a base distribution p_z via the gradient of a neural-network-defined potential?
• RQ2 How does the choice of IPM (Wasserstein, MMD, KSD) affect the required network size to achieve a given approximation error ε?
• RQ3 What are the quantitative bounds on network depth/width (complexity) to approximate π under each IPM?
• RQ4 How can semi-discrete optimal transport be leveraged to construct neural-network-based transport maps?
• RQ5 What are the convergence rates of empirical measures under the three IPMs used for evaluation?

Key findings

• There exists a deep ReLU neural network with d inputs and one output such that the push-forward of p_z under the gradient map is ε-close to π in the chosen IPM.
• For 1-Wasserstein distance, the required network size n scales as C/ε^2 in d=1, C log^2(ε)/ε^2 in d=2, and C^d/ε^d for d≥3 (under finite third moment).
• For MMD, with a kernel satisfying Assumption K2, n ≤ C/ε^2.
• For KSD, with a kernel satisfying Assumption K3 and π satisfying Assumptions 1 and 2, n ≤ C d/ε^2.
• The transport map achieving the semi-discrete optimum is the gradient of a piecewise affine function, exactly representable by a DNN.
• The neural-network-based potential φ(x) = max_j{x·y_j + m_j} can be realized by a DNN with depth ⌈log n⌉ and width 2^⌈log n⌉.