figuret

# Marginalizable Density Models

Dar Gilboa
Harvard University
<EMAIL_ADDRESS>
&Ari Pakman
Columbia University
<EMAIL_ADDRESS>
&Thibault Vatter
Columbia University
<EMAIL_ADDRESS>

###### Abstract

Probability density models based on deep networks have achieved remarkable
success in modeling complex high-dimensional datasets. However, unlike kernel
density estimators, modern neural models do not yield marginals or
conditionals in closed form, as these quantities require the evaluation of
seldom tractable integrals. In this work, we present the _marginalizable
density model approximator_ (MDMA), a novel deep network architecture which
provides closed form expressions for the probabilities, marginals and
conditionals of any subset of the variables. The MDMA learns deep scalar
representations for each individual variable and combines them via learned
hierarchical tensor decompositions into a tractable yet expressive CDF, from
which marginals and conditional densities are easily obtained. We illustrate
the advantage of exact marginalizability in several tasks that are out of
reach of previous deep network-based density estimation models, such as
estimating mutual information between arbitrary subsets of variables,
inferring causality by testing for conditional independence, and inference
with missing data without the need for data imputation, outperforming state-
of-the-art models on these tasks. The model also allows for parallelized
sampling with only a logarithmic dependence of the time complexity on the
number of variables.

##  1 Introduction

Estimating the joint probability density of a set of random variables is a
fundamental task in statistics and machine learning that has witnessed much
progress in recent years. While traditional estimators such as histograms [1,
2] and kernel-based methods [3, 4] have appealing theoretical properties and
typically perform well in low dimensions, they become computationally
impractical above 5-10 dimensions. Conversely, recent density models based on
neural networks [5, 6, 7, 8, 9, 10] scale efficiently with the number of
random variables, but lack a crucial feature available to traditional methods:
the ability to compute probabilities, marginalize over subsets of variables,
and evaluate conditionals. These tasks require integrals of the estimated
density, which are intractable for modern neural density models. Thus, while
such operations are central to many applications (e.g., inference with missing
data, testing for conditional (in)dependence, or performing do-calculus [11]),
approaches based on neural networks require estimating separate models
whenever marginals or conditionals are needed.

Alternatively, one could model the cumulative distribution function (CDF),
making computing probabilities and marginalization straightforward. But
evaluating the density requires taking $d$ derivatives of the CDF, which
incurs an exponential cost in $d$ for a generic computational graph. This
observation has made direct CDF modeling traditionally challenging [12].

In this work, we present the _marginalizable density model approximator_
(MDMA), a novel deep network architecture preserving most of the expressive
power of neural models for density estimation, while providing closed form
expressions for the probabilities, marginals and conditionals of any subset of
the variables. In a nutshell, the MDMA learns many deep scalar representations
for each individual variable and combines them using hierarchical tensor
decompositions [13, 14] into a tractable multivariate CDF that can be fitted
using stochastic gradient descent. Additionally, sampling from MDMA can be
parallelized along the input dimension, resulting in a very low space
complexity and a time complexity that scales only logarithmically with the
number of variables in the problem (as opposed to linearly in naive
autoregressive sampling, see below in Section 2).

As could be expected, the architectural choices that allow for easy
marginalization take a minor toll in terms of performance. Indeed, while
competitive, our model admittedly does not beat state-of-the-art models in
out-of-sample log-likelihood of high-dimensional datasets. On the other hand,
it does beat those same models in a task for which the latter are ill-
prepared: learning densities from data containing missing values, a common
setting in some application areas such as genomics [15, 16]. While our model
is able to deal optimally with missing values by evaluating, for every data
point, the marginal likelihood over its non-missing values, other models must
resort to data imputation. Consequently, we significantly outperform state-of-
the-art neural density estimators trained using a number of common data
imputation strategies. We also show MDMA can be used to test for conditional
independence, which is useful for discovering the causal structure in graphs,
a task on which it outperforms existing methods, and that it enables
estimation of mutual information between arbitrary subsets of variables after
fitting a single model. Additionally, we prove that the model class is a
universal approximator over the space of absolutely continuous multivariate
distributions.

The structure of this paper is as follows. In Section 2 we review related
works. In Section 3 we present our new model and its theoretical properties.
We present our experimental results in Footnote 1, and conclude in Section 5.

##  2 Related work

Modern approaches to non-parametric density estimation, based on normalizing
flows [5, 6, 7, 8, 9, 10] (see [17, 18] for recent reviews), model expressive
yet invertible functions that transforms a simple (usually uniform or
Gaussian) density to the target density. Nonetheless, as previously mentioned,
such architectures lead to intractable derived quantities such as
probabilities, marginals and/or conditional distributions.

Moreover, many normalizing flow models rely on an autoregressive construction,
which makes the cost of generating samples scale linearly with the dimension.
This can be circumvented using inverse autoregressive flows [19], but in this
dual case a linear cost is incurred instead in density evaluations and hence
in training. Another solution is training a feed-forward network using the
outputs of a trained autoregressive model [20]. With MDMA, fast inference and
sampling is achieved without requiring this distillation procedure.

Tensor decompositions [13, 14], which are exploited in this work, have been
used in various applications of signal processing, machine learning, computer
vision, and more [21, 22, 23, 24, 25]. Recently, such decompositions have been
used to speed-up or reduce the number of parameters in existing deep
architectures [26, 27, 28, 29, 30, 31]. In addition to their practical appeal,
tensor methods have been widely studied to understand the success of deep
neural networks [32, 33, 34, 35, 36]

##  3 Marginalizable Density Models

###  3.1 Notations

In the following, we use a capital and lowercase Roman letter (e.g., $F$ and
$f$) or Greek letters along with dot above (e.g., $\varphi$ and
${\dot{\varphi}}$) to denote respectively absolutely continuous CDFs of
arbitrary dimensions and the corresponding densities. When dealing with
multivariate distributions, the marginal or conditional distribution over a
subset of variables will be indicated by the argument names (i.e.,
$F(x_{1}|x_{2},x_{3})=\mathbb{P}\left[X_{1}\leq
x_{1}|X_{2}=x_{2},X_{3}=x_{3}\right]$).

For a positive integer $p\in{\mathbb{N}}\setminus 0$, let
$[p]=\\{1,\dots,p\\}$. Denote the space of absolutely continuous univariate
and $d$-dimensional CDFs respectively by ${\mathcal{F}}_{1}$ and
${\mathcal{F}}_{d}$. For any $F\in{\mathcal{F}}_{1}$, the density
$f:{\mathbb{R}}\to{\mathbb{R}}_{+}$ is $f(x)=\partial F(x)/\partial x$.
Similarly, for any $F\in{\mathcal{F}}_{d}$, the density
$f(\bm{x}):{\mathbb{R}}^{d}\to{\mathbb{R}}_{+}$ is
$f(\bm{x})=\partial^{d}F(\bm{x})/\partial x_{1}\cdots\partial x_{d}$, and
$F_{j}(x)=\lim_{z\to\infty}F(z,\dots,x,\dots,z)\in{\mathcal{F}}_{1}$ for
$j\in[d]$ is the $j$th marginal distribution.

###  3.2 The bivariate case

For the task of modeling joint distributions of two variables supported on
$\mathbb{R}^{2}$, consider a family of univariate CDFs
$\\{\varphi_{i,j}\,\\}_{i\in[m],\,j\in[2]}$ with
$\varphi_{i,j}\in{\mathcal{F}}_{1}$, i.e., the functions
$\varphi_{i,j}:\mathbb{R}\rightarrow[0,1]$ satisfy

$\displaystyle\lim_{x\to-\infty}\varphi_{i,j}(x)=0\,,\qquad\lim_{x\to\infty}\varphi_{i,j}(x)=1\,,\qquad{\dot{\varphi}}_{i,j}(x)=\partial\varphi(x)/\partial
x\geq 0\,.$

These functions are our basic building block, and we model them using a simple
neural architecture proposed in [37] and described in Section 3.6. If $A$ is
an $m\times m$ matrix of nonnegative elements satisfying
$\sum_{i,j=1}^{m}A_{i,j}=1$, we can combine it with the univariate CDFs to
obtain

$\displaystyle
F(x_{1},x_{2})=\sum\limits_{i,j=1}^{m}A_{i,j}\varphi_{i,1}(x_{1})\varphi_{j,2}(x_{2}).$
(1)

The coefficients $A_{i,j}$ encode the dependencies between the two variables,
and the normalization ensures that $F$ is a valid CDF, that is
$F\in{\mathcal{F}}_{2}$. Even though in each summand the interaction is
modeled by a single scalar parameter, such a model can be used to approximate
well complex interactions between $x_{1}$ and $x_{2}$ if $m$ is sufficiently
large, as we show in Section 3.6. The advantage of this construction is that
$\\{{\dot{\varphi}}_{i,j}\\}_{i\in[m],\,j\in[2]}$, the family of densities
corresponding to the univariate CDFs, leads immediately to

$\displaystyle
f(x_{1},x_{2})=\sum\limits_{i,j=1}^{m}A_{i,j}{\dot{\varphi}}_{i,1}(x_{1}){\dot{\varphi}}_{j,2}(x_{2}).$

It is similarly straightforward to obtain marginal and conditional quantities,
e.g.:

$\displaystyle F(x_{1})$
$\displaystyle=\sum\limits_{i,j=1}^{m}A_{i,j}\varphi_{i,1}(x_{1}),\qquad
F(x_{1}|x_{2})=\frac{\sum\limits_{i,j=1}^{m}A_{i,j}\varphi_{i,1}(x_{1}){\dot{\varphi}}_{j,2}(x_{2})}{\sum\limits_{i,j=1}^{m}A_{i,j}{\dot{\varphi}}_{j,2}(x_{2})},$

and the corresponding densities result from replacing $\varphi_{i,1}$ by
${\dot{\varphi}}_{i,1}$. Deriving these simple expressions relies on the fact
that (1) combines the univariate CDFs linearly. Nonetheless, it is clear that,
with $m\to\infty$ and for judiciously chosen univariate CDFs, such a model is
a universal approximator of both CDFs and sufficiently smooth densities.

###  3.3 The multivariate case

To generalize the bivariate case, consider a collection of univariate CDFs
$\\{\varphi_{i,j}\\}_{i\in[m],\,j\in[d]}$ with
$\varphi_{i,j}\in{\mathcal{F}}_{1}$ for each $i$ and $j$, and define the
tensor-valued function $\Phi\colon{\mathbb{R}}^{d}\times[m]^{d}\to[0,1]$ by
$\Phi(\bm{x})_{i_{1},\dots,i_{d}}=\prod_{j=1}^{d}\varphi_{i_{j},j}(x_{j})$ for
$\bm{x}\in{\mathbb{R}}^{d}$. Furthermore, denote the class of normalized order
$d$ tensors with $m$ dimensions in each mode and nonnegative elements by

$\displaystyle{\mathcal{A}}_{d,m}=\\{A\in{\mathbb{R}}^{m\times\cdots\times
m}\colon A_{i_{1},\dots,i_{d}}\geq
0,\,\sum\limits_{i_{1},\dots,i_{d}=1}^{m}A_{i_{1},\dots,i_{d}}=1\\}.$ (2)

######  Definition 1 (Marginalizable Density Model Approximator).

For $\Phi\colon{\mathbb{R}}^{d}\times[m]^{d}\to[0,1]$ as above and
$A\in{\mathcal{A}}_{d,m}$, the marginalizable density model approximator
(MDMA) is

$\displaystyle F_{A,\Phi}(\bm{x})=\langle
A,\Phi(\bm{x})\rangle=\sum\limits_{i_{1},\dots,i_{d}=1}^{m}A_{i_{1},\dots,i_{d}}\prod\limits_{j=1}^{d}\varphi_{i_{j},j}(x_{j}).$
(3)

If is clear that the conditions on $\Phi$ and $A$ imply that
$F_{A,\Phi}\in{\mathcal{F}}_{d}$. As in the bivariate case, densities or
marginalization over $x_{j}$ are obtained by replacing each $\varphi_{i,j}$
respectively by ${\dot{\varphi}}_{i,j}$ or $1$. As for conditioning,
considering any disjoint subsets $R=\left\\{k_{1},\dots,k_{r}\right\\}$ and
$S=\left\\{j_{1},\dots,j_{s}\right\\}$ of $[d]$ such that $R\cap S=\emptyset$,
we have

$\displaystyle
F_{A,\Phi}(x_{k_{1}},\dots,x_{k_{r}}|x_{j_{1}},\dots,x_{j_{s}})$
$\displaystyle=\frac{\sum\limits_{i_{1},\dots,i_{d}=1}^{m}A_{i_{1},\dots,i_{d}}\underset{k\in
R}{\prod}\varphi_{i_{k},k}(x_{k})\underset{j\in
S}{\prod}{\dot{\varphi}}_{i_{j},j}(x_{j})}{\sum\limits_{i_{1},\dots,i_{d}=1}^{m}A_{i_{1},\dots,i_{d}}\underset{j\in
S}{\prod}{\dot{\varphi}}_{i_{j},j}(x_{j})}.$

For a completely general tensor $A$ with $m^{d}$ parameters, the expression
(3) is computationally impractical, hence some structure must be imposed on
$A$. For instance, one simple choice is
$A_{i_{1},\dots,i_{d}}=a_{i_{1}}\delta_{i_{1},\dots,i_{d}}$, which leads to
$F_{A,\Phi}(\bm{x})=\sum_{i=1}^{m}a_{i}\prod_{j=1}^{d}\varphi_{i,j}(x_{j})$,
with $a_{i}\geq 0$ and $\sum_{i=1}^{m}a_{i}=1$. Instead of this top-down
approach, $A$ can be tractably constructed bottom-up, as we explain next.

Assuming $d=2^{p}$ for integer $p$, define
$\varphi^{\ell}\colon{\mathbb{R}}^{d}\to{\mathbb{R}}^{m\times 2^{p-\ell+1}}$
for $\ell\in\\{1,\dots,p\\}$ recursively by

$\displaystyle\varphi_{i,j}^{\ell}(\bm{x})$
$\displaystyle=\begin{cases}\varphi_{i,j}(x_{j}),\qquad&\ell=1\\\
\sum\limits_{k=1}^{m}\lambda_{i,k,j}^{\ell-1}\varphi_{k,2j-1}^{\ell-1}(\bm{x})\varphi_{k,2j}^{\ell-1}(\bm{x}),\qquad&\ell=2,\dots,p\end{cases}$
(4)

for $i\in[m]$, $j\in[2^{p-\ell+1}]$, and where $\lambda^{\ell}$ is a non-
negative $m\times m\times 2^{p-\ell+1}$ tensor, normalized as
$\sum_{k=1}^{m}\lambda_{i,k,j}^{\ell}=1$. The joint CDF can then be written as

$\displaystyle F_{A^{\mathrm{HT}},\Phi}(\bm{x})$
$\displaystyle=\sum\limits_{k=1}^{m}\lambda_{k}^{p}\varphi_{k,1}^{p}(\bm{x})\varphi_{k,2}^{p}(\bm{x}),$
(5)

with $\lambda^{p}\in\mathbb{R}_{+}^{m}$ satisfying
$\sum_{k=1}^{m}\lambda_{k}^{p}=1$. It is easy to verify that the underlying
$A^{\mathrm{HT}}$ satisfies $A^{\mathrm{HT}}\in{\mathcal{A}}_{d,m}$ defined in
(2). A graphical representation of this tensor is provided in Figure 10 in the
supplementary materials.

For example, for $d=4$, we first combine $(x_{1},x_{2})$ and $(x_{3},x_{4})$
into

$\displaystyle\varphi^{2}_{i,1}(\bm{x})=\sum_{k=1}^{m}\lambda^{1}_{i,k,1}\varphi_{k,1}(x_{1})\varphi_{k,2}(x_{2})\,,\qquad\varphi^{2}_{i,2}(\bm{x})=\sum_{k=1}^{m}\lambda^{1}_{i,k,2}\varphi_{k,3}(x_{3})\varphi_{k,4}(x_{4})\,,$

and then merge them as

$\displaystyle
F_{A^{\mathrm{HT}},\Phi}(\bm{x})=\sum_{k=1}^{m}\lambda_{k}^{2}\varphi^{2}_{k,1}(\bm{x})\varphi^{2}_{k,2}(\bm{x}),$
(6)

from which we can read off that
$A^{\mathrm{HT}}_{i_{1},i_{2},i_{3},i_{4}}=\sum_{k=1}^{m}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\lambda_{k,i_{3},2}^{1}\delta_{i_{1},i_{2}}\delta_{i_{3},i_{4}}$.

Note that the number of parameters required to represent $A^{\mathrm{HT}}$ is
only $\mathrm{poly}(m,d)$. The construction is easily generalized to
dimensions $d$ not divisible by 2. Also, the number of $\varphi$ factors
combined at each iteration in (4) (called the _pool size_), can be any
positive integer. This construction is a variant of the hierarchical Tucker
decomposition of tensors [38], which has been used in the construction of
tensor networks for image classification in [32].

Given a set of training points $\\{\bm{x}\\}_{i=1}^{N}$, we fit MDMA models by
maximizing the log of the density with respect to both the parameters in
$\\{\varphi_{i,j}\\}_{i\in[m],j\in[d]}$ and the components of $A$. We present
additional details regarding the choice of architectures and initialization in
Appendix E.

###  3.4 A non-marginalizable MDMA

We can construct a more expressive variant of MDMA at the price of losing the
ability to marginalize and condition on arbitrary subsets of variables. We
find that the resulting model leads to state-of-the-art performance on a
density estimation benchmark. We define
$\mathbf{v}=\mathbf{x}+\mathbf{T}\sigma(\mathbf{x})$ where $\mathbf{T}$ is an
upper-triangular matrix with non-negative entries and $0$ on the main
diagonal. Note that
$\left|\frac{\partial\mathbf{v}}{\partial\mathbf{x}}\right|=1$. Given some
density
$f(v_{1},\dots,v_{d})=\prod\limits_{j=1}^{d}{\dot{\varphi}}_{j}(v_{j})$, we
have

$\displaystyle
f(x_{1},\dots,x_{d})=\left|\frac{\partial\mathbf{v}}{\partial\mathbf{x}}\right|f(v_{1},\dots,v_{d})=f(v_{1}(\mathbf{x}),\dots,v_{d}(\mathbf{x})).$
(7)

We refer to this model nMDMA, since it no longer enables efficient
marginalization and conditioning.

###  3.5 MDMA sampling

Given an MDMA as in 3, we can sample in the same manner as for autoregressive
models: from $u_{1},\dots,u_{d}$ independent $U(0,1)$ variables, we obtain a
sample from $F_{A,\Phi}$ by computing

$\displaystyle x_{1}=F_{A,\Phi}^{-1}(u_{1}),\quad
x_{2}=F_{A,\Phi}^{-1}(u_{2}|x_{1})\quad\cdots\quad
x_{d}=F_{A,\Phi}^{-1}(u_{d}|x_{1},\dots,x_{d-1}),$

where, unlike with autoregressive model, the order of taking the conditionals
does not need to be fixed. The main drawback of this method is that due to the
sequential nature of the sampling the computational cost is linear in $d$ and
cannot be parallelized.

However, the structure of $F_{A,\Phi}$ can be leveraged to sample far more
efficiently. Define by $R_{A}$ a vector-valued categorical random variable
taking values in $\left[m\right]\times\cdots\times\left[m\right]$, with
distribution

$\displaystyle\mathbb{P}\left[R_{A}=(i_{1},\dots,i_{d})\right]=A_{i_{1},\dots,i_{d}}.$

The fact that $A\in{\mathcal{A}}_{d,m}$ with ${\mathcal{A}}_{d,m}$ from (2)
ensure the validity of this definition. Consider a vector
$(\tilde{X}_{1},\dots,\tilde{X}_{d},R_{A})$ where $R_{A}$ is distributed as
above, and

$\displaystyle\mathbb{P}\left[\tilde{X}_{1}\leq x_{1},\dots,\tilde{X}_{d}\leq
x_{d}|R_{A}=\mathbf{r}\right]=\prod\limits_{i=1}^{d}\varphi_{r_{i},i}(x_{i}),$

for the collection $\\{\varphi_{i,j}\\}$ of univariate CDFs. Denoting the
distribution of this vector by $\tilde{F}_{A,\Phi}$, marginalizing over
$R_{A}$ gives

$\displaystyle\tilde{F}_{A,\Phi}(\bm{x})=\underset{\mathbf{r}\in[m]^{d}}{\sum}\mathbb{P}\left[R_{A}=\mathbf{r}\right]\prod\limits_{j=1}^{d}\varphi_{r_{j},j}(x_{j})=F_{A,\Phi}(\bm{x}).$

Instead of sampling directly from the distribution $F_{A,\Phi}$, we can thus
sample from $\tilde{F}_{A,\Phi}$ and discard the sample of $R_{A}$. To do
this, we first sample from the categorical variable $R_{A}$. Denoting this
sample by $\mathbf{r}$, we can sample from $\tilde{X}_{i}$ by inverting the
univariate CDF $\varphi_{r_{i},i}$. This can be parallelized over $i$.

The approach outlined above is impractical since $R_{A}$ can take $m^{d}$
possible values, yet if $A$ can be expressed by an efficient tensor
representation this exponential dependence can be avoided. Consider the HT
decomposition 5, which can be written as

$\displaystyle F_{A^{\mathrm{HT}},\Phi}(\bm{x})=\left\langle
A^{\mathrm{HT}},\Phi(\bm{x})\right\rangle=\prod\limits_{\ell=1}^{p}\prod\limits_{j_{\ell}=1}^{2^{p-\ell}}\sum\limits_{k_{\ell,j_{\ell}}=1}^{m}\lambda_{k_{\ell+1,\left\lceil
j_{\ell}/2\right\rceil},k_{\ell,j_{\ell}},j_{\ell}}^{\ell}\Phi(\bm{x})_{k_{1,1},\dots,k_{1,d/2}},$
(8)

that is a normalized sum of $O(md)$ univariate CDFs.

######  Proposition 1.

Sampling from (8) can be achieved in $O(\log d)$ time requiring the storage of
only $O(d)$ integers.

Note that the time complexity of this sampling procedure depends only
logarithmically on $d$. The reason is that a simple hierarchical structure of
Algorithm 1, where Mult denotes the multinomial distribution.

Result: $x_{j}$ for $j=1,\dots,d$

$k_{p,1}\sim\textrm{Mult}(\lambda_{*}^{p})$;

for  _$\ell\leftarrow p-1$ to $1$ by $-1$_ do

$k_{\ell,j}\sim\textrm{Mult}(\lambda^{\ell}_{k_{\ell+1,\left\lceil
j/2\right\rceil},*,j}),\qquad j=1,\dots,2^{p-\ell}$;

end for

$x_{j}\sim\varphi_{k_{1,\left\lceil j/2\right\rceil},j},\qquad
j=1,\dots,2^{p}$;

Algorithm 1 Sampling from the HT MDMA

The logarithmic dependence is only in the sampling from the categorical
variables, which is inexpensive to begin with. We thus avoid the linear
dependence of the time complexity on $d$ that is common in sampling from
autoregressive models. Furthermore, the additional memory required for
sampling scales like $\log m$ (since storing the categorical samples requires
representing integers up to size $m$), and aside from this each sample
requires evaluating a single product of univariate CDFs (which is independent
of $m$). In preliminary experiments, we have found that even for densities
with $d\leq 10$, this sampling scheme is faster by $1.5$ to $2$ orders of
magnitude than autoregressive sampling. The relative speedup should only
increase with $d$.

###  3.6 Universality of the MDMA

To model functions in ${\mathcal{F}}_{1}$, we use $\Phi_{l,r,\sigma}$, a class
of constrained feedforward neural networks proposed in [37] with $l$ hidden
layers, each with $r$ neurons, and $\sigma$ a nonaffine, increasing and
continuously differentiable elementwise activation function, defined as

$\displaystyle\Phi_{l,r,\sigma}=\\{\varphi\colon{\mathbb{R}}\to[0,1],\,\varphi(x)=\mbox{sigmoid}\circ
L_{l}\circ\sigma\circ L_{l-1}\circ\sigma\cdots\circ\sigma\circ
L_{1}\circ\sigma\circ L_{0}(x)\\},$

where $L_{i}\colon{\mathbb{R}}^{n_{i}}\to{\mathbb{R}}^{n_{i+l}}$ is the affine
map $L_{i}(x)=W_{i}x+b_{i}$ for an $n_{i+1}\times n_{i}$ weight matrix $W_{i}$
with _nonnegative_ elements and an $n_{i+1}\times 1$ bias vector $b_{i}$, with
$n_{l+1}=n_{0}=1$ and $n_{i}=r$ for $i\in[l]$. The constraints on the weights
and the final sigmoid guarantee that
$\Phi_{l,r,\sigma}\subseteq{\mathcal{F}}_{1}$, and for any
$\varphi\in\Phi_{l,r,\sigma}$, the corresponding density
${\dot{\varphi}}(x)=\partial\varphi(x)/\partial x$ can be obtained with the
chain rule. The universal approximation property of the class
$\Phi_{l,r,\sigma}$ is expressed in the following proposition.

######  Proposition 2.

$\cup_{l,r}\Phi_{l,r,\sigma}$ is dense in ${\mathcal{F}}_{1}$ with respect to
the uniform norm.

While the proof in the supplementary assumes that
$\lim_{x\to-\infty}\sigma(x)=0$ and $\lim_{x\to\infty}\sigma(x)=1$, it can be
easily modified to cover other activations. For instance, in our experiments,
we use $\sigma(x)=x+a\tanh(x)$ following [37], and refer to the supplementary
material for more details regarding this case. In the multivariate case,
consider the class of order $d$ tensored-valued functions with $m$ dimensions
per mode defined as

$\displaystyle\Phi_{m,d,l,r,\sigma}=\\{\Phi\colon{\mathbb{R}}^{d}\times[m]^{d}\to[0,1],\,\Phi(\bm{x})_{i_{1},\dots,i_{d}}=\textstyle\prod\limits_{j=1}^{d}\varphi_{i_{j},j}(x_{j}),\,\varphi_{i,j}\in\Phi_{l,r,\sigma}\\}.$

Combining $\Phi_{m,d,l,r,\sigma}$ with the ${\mathcal{A}}_{d,m}$, the
normalized tensors introduced in Section 3.3, the class of neural network-
based MDMAs can then be expressed as

$\displaystyle\textrm{MDMA}_{m,d,l,r,\sigma}=\\{F_{A,\Phi}\colon{\mathbb{R}}^{d}\to[0,1],\,F_{A,\Phi}(\bm{x})=\langle
A,\Phi(\bm{x})\rangle,\,A\in{\mathcal{A}}_{d,m},\,\Phi\in\Phi_{m,d,l,r,\sigma}\\}.$

######  Proposition 3.

The set $\cup_{m,l,r}\textrm{MDMA}_{m,d,l,r,\sigma}$ is dense in
${\mathcal{F}}_{d}$ with respect to the uniform norm.

The proof relies on the fact that setting $m=1$ yields a class that is dense
in the space of $d$-dimensional CDFs with independent components. All proofs
are provided in Appendix A.

##  4 Experiments

Additional experimental details for all experiments are provided in Appendix
C.111Code for reproducing all experiments is available at
https://github.com/dargilboa/mdma.

###  4.1 Toy density estimation

We start by considering 3D augmentations of three popular 2D toy probability
distributions introduced in [8]: two spirals, a ring of 8 Gaussians and a
checkerboard pattern. These distributions allow to explore the ability of
density models to capture challenging multimodalities and discontinuities [9,
6]. The results, presented in Figure 1, show that MDMA captures all marginal
densities with high accuracy, and samples from the learned model appear
indistinguishable from the training data.

|  |  |
---|---|---|---
|  |  |
|  |  |
Training data

Figure 1: Density estimation with closed-form marginals and conditionals.
First Row: Panels 1,2: Empirical histograms of training data. Panel 3: Samples
from the training data. Panel 4: Samples from the trained MDMA model. Second
Row: Panels 1,2: The marginal density learned by MDMA plotted on a grid.
Panels 3,4: Conditional densities learned by MDMA plotted on a grid. Third
Row: Results on additional datasets: Panels 1,2: Training data and learned
marginal density for a 3D checkerboard dataset. Panels 3,4: Similarly for a 3D
mixture of Gaussians.

###  4.2 Mutual information estimation

Given a multivariate probability distribution over some variables
$X=(X_{1}\dots,X_{d})$, estimating the mutual information

$I(Y;Z)=\int dp_{X}(x)\log\left(\frac{p_{X}(x)}{p_{Y}(y)p_{Z}(z)}\right),$ (9)

where $Y,Z$ are random vectors defined by disjoint subsets of the $X_{i}$,
requires evaluating $p_{Y},p_{Z}$ which are marginal densities of $X$.
Typically, $Y$ and $Z$ must be fixed in advance, yet in some cases it is
beneficial to be able to flexibly compute mutual information between any two
subsets of variables. Estimating both $I(Y,Z)$ and $I(Y^{\prime},Z^{\prime})$
may be highly inefficient, e.g. if $Y$ and $Y^{\prime}$ are highly overlapping
subsets of $X$. Using MDMA however, we can fit a single model for the joint
distribution and easily estimate the mutual information between any subset of
variables by simply marginalizing over the remaining variables to obtain the
required marginal densities. Thus a Monte Carlo estimate of (9) can be
obtained by evaluating the marginal densities at the points that make up the
training set. Figure 2 presents an example of this method, showing the
accuracy of the estimates.

Figure 2: Mutual information estimation between subsets of a random vector. We
fitted a single MDMA model to samples from a zero-mean $d=16$ Gaussian, with
covariance $\Sigma_{ij}=\delta_{ij}+(1-\delta_{ij})(i+j-2)/(5d)$. Monte Carlo
estimates of the mutual information (9) between $(X_{1},\dots,X_{k})$ and
$(X_{k+1},\dots,X_{d})$ for any $k=1,\dots,d-1$ are easily obtained and match
closely the exact values. For each $k$ we average over 5 repetitions of
drawing the dataset and fitting.

###  4.3 Density estimation with missing values

Dealing with missing values in multivariate data is a classical challenge in
statistics that has been studied for decades [39]. The standard solution is
the application of a data imputation procedure (i.e., “filling in the
blanks”), which requires making structural assumptions. In some cases, this is
natural, as for the matrix completion problem under a low-rank assumption [40,
41], where the imputed values are the main object of interest. But the
artifacts introduced by data imputation [42] are generally a price that one
must unwillingly pay in order to perform statistical inference in models that
require fully-observed data points. Two popular, generic techniques for
imputation are MICE [43] and $k$-NN imputation [44]. The former imputes
missing values by iteratively regressing each missing variable against the
remaining variables, while the latter uses averages over the non-missing
values at $k$-nearest datapoints.

More formally, let $X\in\mathbb{R}^{d}$ be distributed according to some
density $p$ with parameters $\theta$, let $X_{(0)},X_{(1)}$ be the non-missing
and missing entries of $X$ respectively, and $M\in\\{0,1\\}^{d}$ a vector
indicating the missing entries. In the missing-at-random setting (i.e. $M$ is
independent of $X_{(1)}$), likelihood-based inference using the full
likelihood of the model is equivalent to inference using the marginal
likelihood [39] $L(X_{(0)}|\theta)=\int p(X|\theta)dX_{(1)}$. Standard neural
network-based density estimators must resort to data imputation because of the
impossibility of computing this marginal likelihood. MDMA however can directly
maximize the marginal likelihood for any pattern of missing data at the same
(actually slightly cheaper) computational cost as maximizing the full
likelihood, without introducing any bias or variance due to imputation.

As a demonstration of this capability, we consider the UCI POWER and GAS
datasets, following the same pre-processing as [45]. We construct a dataset
with missing values by setting each entry in the dataset to be missing
independently with a fixed probability . We compare MDMA to BNAF [9], a neural
density model which achieves state-of-the-art results on a number of density
estimation benchmarks including GAS. We train MDMA directly on the log
marginal likelihood of the missing data, and BNAF by first performing data
imputation using MICE [43] and then training using the full log likelihood
with the imputed data. The validation loss is the log marginal likelihood for
MDMA and the log likelihood of the imputed validation set for BNAF. The test
set is left unchanged for both models and does not contain any missing values.
We train BNAF using the settings specified in [9] that led to the best
performance ($2$ layers and $40d$ hidden units where $d$ is the dimensionality
of the dataset). The results are shown in Figure 3.

Figure 3: Density estimation with missing data Test NLL on two density
estimation benchmarks, varying the proportion of entries in each data matrix
that are designated missing and not used for fitting. We compare MDMA which
can fit marginal densities directly with BNAF which achieves state-of-the-art
results on the POWER dataset, after performing data imputation using MICE. As
the proportion of missing data increases, MDMA outperforms BNAF.

We find that, as the probability of missingness increases, MDMA significantly
outperforms BNAF on both datasets. Note that, while the proportion of missing
values might seem extreme, it is not uncommon in some applications (e.g.,
proteomics data). We also trained BNAF using $k$-NN imputation [44], finding
that performance was worse than MICE imputation for all values of $\alpha$. A
comparison of the two methods is provided in Appendix B.

###  4.4 Conditional independence testing and causal discovery

Randomized control trials [46] remain the golden standard for causal
discovery. Nonetheless, experiments or interventions are seldom doable, e.g.
due to financial or ethical considerations. Alternatively, observational data
can help uncovering causal relationships [47, 48]. In this context, a class of
popular methods targeted at recovering the full causal graph, like PC or FCI
[47, 49], rely on conditional independence (CI) tests. Letting $X$, $Y$ and
$Z$ be random variables, the CI of $X$ and $Y$ given $Z$, denoted
$X\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}Y\mid Z$, means that given $Z$, no information
about $X$ (or $Y$) can be gained by knowing the value of $Y$ (or $X$). And
testing $H_{0}\colon X\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}Y\mid Z$ against $H_{1}\colon
X\not\mathchoice{\mathrel{\hbox to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}Y\mid Z$ is a problem tackled in econometrics
[50, 51], statistics [52, 53], and machine learning [54, 55].

Following [55], denote $U_{1}=F(X\mid Z)$ and $U_{2}=F(Y\mid Z)$. It is clear
that $H_{0}$ implies $U_{1}\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}U_{2}$, although the converse does not hold
[see e.g., 56]. Nonetheless, $U_{1}\not\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}U_{2}$ implies $H_{1}$, so a test based on the
independence between $U_{1}$ and $U_{2}$ can still have power. While the test
from [55] is based on estimating the conditional CDFs through quantile
regression, we proceed similarly, albeit using the MDMA as a plugin for the
conditional distributions. Our approach is especially appealing in the context
of causal discovery, where algorithms require computing many CI tests to
create the graph’s skeleton. Instead of having to regress for every test, MDMA
estimates the full joint distribution, and its lower dimensional conditionals
are then used for the CI tests.

In Table 1, we present results on inferring the structure of causal graphs
using the PC algorithm [47, 57, 58, 59, 60]. As a benchmark, we use the
vanilla (i.e., Gaussian) CI test, and compare it to the PC algorithm obtained
with the CI test from [55], albeit using the MDMA for the conditional
distributions. Synthetic random directed acyclic graphs (DAGs) along with
sigmoidal or polynomial mechanisms linking parents to children are sampled
using [61]. Each dataset is $d=10$ dimensional and contains 20,000
observations. We also compare the two algorithms on data from a protein
signaling network with $d=11$ [62] for which the ground truth causality graph
is known. Performance is assessed based on the structural Hamming distance
(SHD) [63], that is the $L^{1}$ norm of the difference between learned
adjacency matrices and the truth, as well as a variant of this metric for
directed graphs SHD(D) which also accounts for the direction of the edges.
Table 1 shows averages over 8 runs for each setting. In all cases, MDMA
outperforms the vanilla PC in terms of both metrics. For the synthetic data,
we note the large standard deviations, due in part to the fact that we sample
randomly from the space of DAG structures, which has cardinality super-
exponential in $d$. An example of the inferred graphs is presented in Appendix
B.

Table 1: MDMA for causal discovery. Conditional densities from a trained MDMA model can be used for causal discovery by allowing to test for conditional independence between variables. Both on synthetic DAG data and real data from a protein signaling network, MDMA infers the graph structure more accurately than a competing method based on quantile regression [55]. The metrics are the structural Hamming distance for the directed (SHD(D)) and undirected (SHD) graph.  Model | Sigmoidal DAG, d=10 | Polynomial DAG, d=10 | Sachs [62], d=11
---|---|---|---
| SHD(D) | SHD | SHD(D) | SHD | SHD(D) | SHD
Gaussian | $18.6\pm 3.0$ | $15.6\pm 2.7$ | $19.8\pm 4.1$ | $18.9\pm 4.2$ | 32 | 27
MDMA | $\mathbf{15.6}\pm 6.1$ | $\mathbf{12.8}\pm 5.2$ | $\mathbf{17.9}\pm 5.3$ | $\mathbf{15.0}\pm 4.5$ | $\mathbf{30.3}\pm 1.8$ | $\mathbf{25.8}\pm 0.7$

###  4.5 Density estimation on real data

We trained MDMA/nMDMA and the non-marginalizable variant described in Section
3.4 on a number of standard density estimation benchmarks from the UCI
repository,222http://archive.ics.uci.edu/ml/datasets.php following the pre-
processing described in [45]. Table 2 compares test log likelihoods of
MDMA/nMDMA with several other neural density models. We find the performance
of MDMA on the lower-dimensional datasets comparable to state-of-the-art
models, while for higher-dimensional datasets it appears to overfit. nMDMA
achieves state-of-the-art performance on the POWER ($d=6$) dataset, but at the
cost of losing the ability to marginalize or condition over subsets of the
variables. The width of MDMA was chosen based on a grid search over
$\\{500,1000,2000,3000,4000\\}$ for each dataset, and the marginal CDF
parameters by a search over
$\left\\{\left(l=2,w=3\right),\left(l=4,w=5\right)\right\\}$. All models were
trained using ADAM with learning rate $0.01$, and results for MDMA and nMDMA
are averaged over $3$ runs. Additional experimental details are provided in
Appendix C.

Table 2: General density estimation. Test log likelihood for density estimation on UCI datasets. The comparison results are reproduced from [10].  Model |  POWER [d=6] |  GAS [d=11] |  HEPMASS [d=21] |  MINIBOONE  [d=43]
---|---|---|---|---
Kingma et al. 2018 [5] | $0.17\pm.01$ | $8.15\pm.4$ | $-18.92\pm.08$ | $-11.35\pm.07$
Grathwohl et al. 2019 [8] | $0.46\pm.01$ | $8.59\pm.12$ | $-14.92\pm.08$ | $-10.43\pm.04$
Huang et al. 2018 [6] | $0.62\pm.01$ | $11.96\pm.33$ | $-15.08\pm.4$ | $-8.86\pm.15$
Oliva et al. 2018 [7] | $0.60\pm.01$ | $\mathbf{12.06}\pm.02$ | $-13.78\pm.02$ | $-11.01\pm.48$
De Cao et al. 2019 [9] | $0.61\pm.01$ | $\mathbf{12.06}\pm.09$ | $-14.71\pm.38$ | $-8.95\pm.07$
Bigdeli et al. 2020 [10] | $0.97\pm.01$ | $9.73\pm 1.14$ | $\mathbf{-11.3}\pm.16$ | $\mathbf{-6.94}\pm 1.81$
MDMA | $0.57\pm.01$ | $8.92\pm 0.11$ | $-20.8\pm.06$ | $-29.0\pm.06$
nMDMA | $\mathbf{1.78}\pm.12$ | $8.43\pm.04$ | $-18.0\pm 0.91$ | $-18.6\pm.47$

##  5 Discussion

MDMAs offer the ability to obtain, from a single model, closed form
probabilities, marginals and conditionals for any subset of the variables.
These properties enable one to straightforwardly use the model to solve a
diverse array of problems, of which we have demonstrated only a few: mutual
information estimation between arbitrary subsets of variables, inference with
missing values, and conditional independence testing targeted at multivariate
causal discovery. In addition to these, MDMA’s marginalization property can be
used for anomaly detection with missing values [64]. We have shown that MDMA
can fit data with missing values without requiring imputation, yet if one is
interested in data imputation for downstream tasks, the ability to sample from
arbitrary conditional distributions means that MDMA can be used for imputation
as well.

Additionally, in some application areas (e.g., financial risk management),
powerful models exist for the univariate distributions, and marginal
distributions are then glued together using copulas [65]. However, popular
copula estimators suffer from the same drawbacks as modern neural network
density estimators with regard to marginalization and conditioning. Using MDMA
for copula estimation (say by replacing the kernel density estimator by MDMA
in the formulation of [66]), one can then obtain copula estimators that do not
suffer from these deficiencies.

The main shortcoming of MDMA is the linearity in the combination of the
products of univariate CDFs which appears to limit the expressivity of the
model. The study of tensor decompositions is an active area of research, and
novel constructions, ideally adapted specifically for this task, could lead to
improvements in this regard despite the linear structure.

## Acknowledgements

The work of DG is supported by a Swartz fellowship. The work of AP is
supported by the Simons Foundation, the DARPA NESD program, NSF NeuroNex Award
DBI1707398 and The Gatsby Charitable Foundation.

## References

  * [1] David W Scott.  On optimal and data-based histograms.  Biometrika, 66(3):605–610, 1979.
  * [2] Gábor Lugosi, Andrew Nobel, et al.  Consistency of data-driven histogram methods for density estimation and classification.  Annals of Statistics, 24(2):687–706, 1996.
  * [3] Murray Rosenblatt.  Remarks on Some Nonparametric Estimates of a Density Function.  The Annals of Mathematical Statistics, 27(3):832 – 837, 1956.
  * [4] Emanuel Parzen.  On estimation of a probability density function and mode.  The annals of mathematical statistics, 33(3):1065–1076, 1962.
  * [5] Durk P Kingma and Prafulla Dhariwal.  Glow: Generative flow with invertible 1x1 convolutions.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
  * [6] Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville.  Neural autoregressive flows.  In International Conference on Machine Learning, pages 2078–2087. PMLR, 2018.
  * [7] Junier Oliva, Avinava Dubey, Manzil Zaheer, Barnabas Poczos, Ruslan Salakhutdinov, Eric Xing, and Jeff Schneider.  Transformation autoregressive networks.  In International Conference on Machine Learning, pages 3898–3907. PMLR, 2018.
  * [8] Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud.  Ffjord: Free-form continuous dynamics for scalable reversible generative models.  In International Conference on Learning Representations, 2019.
  * [9] Nicola De Cao, Wilker Aziz, and Ivan Titov.  Block neural autoregressive flow.  In Uncertainty in Artificial Intelligence, pages 1263–1273. PMLR, 2020.
  * [10] Siavash A Bigdeli, Geng Lin, Tiziano Portenier, L Andrea Dunbar, and Matthias Zwicker.  Learning generative models using denoising density estimators.  arXiv preprint arXiv:2001.02728, 2020.
  * [11] Judea Pearl.  Causality: Models, Reasoning and Inference.  Cambridge University Press, USA, 2nd edition, 2009.
  * [12] Pawel Chilinski and Ricardo Silva.  Neural likelihoods via cumulative distribution functions.  In Conference on Uncertainty in Artificial Intelligence, pages 420–429. PMLR, 2020.
  * [13] Wolfgang Hackbusch.  Tensor spaces and numerical tensor calculus, volume 42.  Springer, 2012.
  * [14] Andrzej Cichocki, Namgil Lee, Ivan Oseledets, Anh-Huy Phan, Qibin Zhao, and Danilo P Mandic.  Tensor networks for dimensionality reduction and large-scale optimization: Part 1 low-rank tensor decompositions.  Foundations and Trends® in Machine Learning, 9(4-5):249–429, 2016.
  * [15] Yun Li, Cristen Willer, Serena Sanna, and Gonçalo Abecasis.  Genotype imputation.  Annual review of genomics and human genetics, 10:387–406, 2009\.
  * [16] Jonathan Marchini and Bryan Howie.  Genotype imputation for genome-wide association studies.  Nature Reviews Genetics, 11(7):499–511, 2010.
  * [17] Ivan Kobyzev, Simon Prince, and Marcus Brubaker.  Normalizing flows: An introduction and review of current methods.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020\.
  * [18] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference.  Journal of Machine Learning Research, 22(57):1–64, 2021.
  * [19] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improved variational inference with inverse autoregressive flow.  In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
  * [20] Aaron Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George Driessche, Edward Lockhart, Luis Cobo, Florian Stimberg, et al.  Parallel wavenet: Fast high-fidelity speech synthesis.  In International conference on machine learning, pages 3918–3926. PMLR, 2018.
  * [21] M Alex O Vasilescu and Demetri Terzopoulos.  Multilinear analysis of image ensembles: Tensorfaces.  In European conference on computer vision, pages 447–460. Springer, 2002.
  * [22] Andrzej Cichocki, Rafal Zdunek, Anh Huy Phan, and Shun-ichi Amari.  Nonnegative matrix and tensor factorizations: applications to exploratory multi-way data analysis and blind source separation.  John Wiley & Sons, 2009.
  * [23] Animashree Anandkumar, Rong Ge, Daniel Hsu, Sham M Kakade, and Matus Telgarsky.  Tensor decompositions for learning latent variable models.  Journal of machine learning research, 15:2773–2832, 2014.
  * [24] Evangelos E Papalexakis, Christos Faloutsos, and Nicholas D Sidiropoulos.  Tensors for data mining and data fusion: Models, applications, and scalable algorithms.  ACM Transactions on Intelligent Systems and Technology (TIST), 8(2):1–44, 2016.
  * [25] Nicholas D Sidiropoulos, Lieven De Lathauwer, Xiao Fu, Kejun Huang, Evangelos E Papalexakis, and Christos Faloutsos.  Tensor decomposition for signal processing and machine learning.  IEEE Transactions on Signal Processing, 65(13):3551–3582, 2017\.
  * [26] Vadim Lebedev, Yaroslav Ganin, Maksim Rakhuba, Ivan Oseledets, and Victor Lempitsky.  Speeding-up convolutional neural networks using fine-tuned cp-decomposition.  In International Conference on Learning Representations. PMLR, 2015\.
  * [27] Cheng Tai, Tong Xiao, Yi Zhang, Xiaogang Wang, et al.  Convolutional neural networks with low-rank regularization.  In International Conference on Learning Representations. PMLR, 2016\.
  * [28] Alexander Novikov, Dmitry Podoprikhin, Anton Osokin, and Dmitry Vetrov.  Tensorizing neural networks.  In The 29-th Conference on Natural Information Processing Systems (NIPS), 2015.
  * [29] Yong-Deok Kim, Eunhyeok Park, Sungjoo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin.  Compression of deep convolutional neural networks for fast and low power mobile applications.  In International Conference on Learning Representations. PMLR, 2016\.
  * [30] Yongxin Yang and Timothy Hospedales.  Deep multi-task representation learning: A tensor factorisation approach.  2017\.
  * [31] Yunpeng Chen, Xiaojie Jin, Bingyi Kang, Jiashi Feng, and Shuicheng Yan.  Sharing residual units through collective tensor factorization to improve deep neural networks.  In IJCAI, pages 635–641, 2018.
  * [32] Nadav Cohen, Or Sharir, and Amnon Shashua.  On the expressive power of deep learning: A tensor analysis.  In Conference on learning theory, pages 698–728. PMLR, 2016.
  * [33] Benjamin D Haeffele and René Vidal.  Global optimality in tensor factorization, deep learning, and beyond.  CoRR, 2015.
  * [34] Majid Janzamin, Hanie Sedghi, and Anima Anandkumar.  Generalization bounds for neural networks through tensor factorization.  CoRR, abs/1506.08473, 1, 2015.
  * [35] Majid Janzamin, Hanie Sedghi, and Anima Anandkumar.  Beating the perils of non-convexity: Guaranteed training of neural networks using tensor methods.  arXiv preprint arXiv:1506.08473, 2015.
  * [36] Or Sharir and Amnon Shashua.  On the expressive power of overlapping architectures of deep learning.  In International Conference on Learning Representations. PMLR, 2018\.
  * [37] Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, and Nick Johnston.  Variational image compression with a scale hyperprior.  In International Conference on Learning Representations, 2018.
  * [38] Wolfgang Hackbusch and Stefan Kühn.  A new scheme for the tensor representation.  Journal of Fourier analysis and applications, 15(5):706–722, 2009\.
  * [39] Roderick JA Little and Donald B Rubin.  Statistical analysis with missing data, 3rd Ed, volume 793.  John Wiley & Sons, 2019.
  * [40] Emmanuel J Candès and Benjamin Recht.  Exact matrix completion via convex optimization.  Foundations of Computational mathematics, 9(6):717–772, 2009.
  * [41] Emmanuel J Candes and Yaniv Plan.  Matrix completion with noise.  Proceedings of the IEEE, 98(6):925–936, 2010.
  * [42] Lorenzo Beretta and Alessandro Santaniello.  Nearest neighbor imputation algorithms: a critical evaluation.  BMC medical informatics and decision making, 16(3):197–208, 2016\.
  * [43] S van Buuren and Karin Groothuis-Oudshoorn.  mice: Multivariate imputation by chained equations in r.  Journal of statistical software, pages 1–68, 2010.
  * [44] Olga Troyanskaya, Michael Cantor, Gavin Sherlock, Pat Brown, Trevor Hastie, Robert Tibshirani, David Botstein, and Russ B Altman.  Missing value estimation methods for dna microarrays.  Bioinformatics, 17(6):520–525, 2001.
  * [45] George Papamakarios, Theo Pavlakou, and Iain Murray.  Masked autoregressive flow for density estimation.  In Proceedings of the 31st International Conference on Neural Information Processing Systems, 2017.
  * [46] Ronald Aylmer Fisher.  Statistical methods for research workers.  Especially Section, 21, 1936.
  * [47] Peter Spirtes, Clark Glymour, and Richard Scheines.  Causation, Prediction, and Search.  MIT press, Causation2000, 2000.
  * [48] Marloes H. Maathuis and Preetam Nandy.  A review of some recent advances in causal inference.  In Handbook of Big Data. CRC Press, 2016.
  * [49] Eric V Strobl, Kun Zhang, and Shyam Visweswaran.  Approximate kernel-based conditional independence tests for fast non-parametric causal discovery.  Journal of Causal Inference, 7(1), 2019.
  * [50] Liangjun Su and Halbert White.  A consistent characteristic function-based test for conditional independence.  Journal of Econometrics, 141(2):807–834, 2007.
  * [51] Liangjun Su and Halbert White.  A nonparametric hellinger metric test for conditional independence.  Econometric Theory, pages 829–864, 2008.
  * [52] Tzee-Ming Huang et al.  Testing conditional independence using maximal nonlinear conditional correlation.  The Annals of Statistics, 38(4):2047–2091, 2010.
  * [53] Rajen D Shah, Jonas Peters, et al.  The hardness of conditional independence testing and the generalised covariance measure.  Annals of Statistics, 48(3):1514–1538, 2020.
  * [54] Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf.  Kernel-based conditional independence test and application in causal discovery.  UAI11, page 804–813. AUAI Press, 2011.
  * [55] Lasse Petersen and Niels Richard Hansen.  Testing conditional independence via quantile regression based partial copulas.  Journal of Machine Learning Research, 22(70):1–47, 2021.
  * [56] Fabian Spanhel and Malte S Kurz.  The partial copula: Properties and associated dependence measures.  Statistics & Probability Letters, 119:76–83, 2016.
  * [57] Markus Kalisch and Peter Bühlman.  Estimating high-dimensional directed acyclic graphs with the pc-algorithm.  Journal of Machine Learning Research, 8(3), 2007.
  * [58] Xiaohai Sun, Dominik Janzing, Bernhard Schölkopf, and Kenji Fukumizu.  A kernel-based causal learning algorithm.  In Proceedings of the 24th international conference on Machine learning, pages 855–862, 2007.
  * [59] Robert E Tillman, Arthur Gretton, and Peter Spirtes.  Nonlinear directed acyclic structure learning with weakly additive noise models.  In NIPS, pages 1847–1855, 2009.
  * [60] Naftali Harris and Mathias Drton.  Pc algorithm for nonparanormal graphical models.  Journal of Machine Learning Research, 14(11), 2013.
  * [61] Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta.  Causal discovery toolbox: Uncovering causal relationships in python.  Journal of Machine Learning Research, 21(37):1–5, 2020.
  * [62] Karen Sachs, Omar Perez, Dana Pe’er, Douglas A Lauffenburger, and Garry P Nolan.  Causal protein-signaling networks derived from multiparameter single-cell data.  Science, 308(5721):523–529, 2005.
  * [63] Ioannis Tsamardinos, Laura E. Brown, and Constantin F. Aliferis.  The max-min hill-climbing Bayesian network structure learning algorithm.  Machine Learning, 65(1):31–78, oct 2006.
  * [64] Thomas G Dietterich and Tadesse Zemicheal.  Anomaly detection in the presence of missing values.  ODD v5.0: Outlier Detection De-constructed Workshop, 2018.
  * [65] Alexander J McNeil, Rüdiger Frey, and Paul Embrechts.  Quantitative risk management: concepts, techniques and tools-revised edition.  Princeton university press, 2015.
  * [66] Gery Geenens, Arthur Charpentier, and Davy Paindaveine.  Probit transformation for nonparametric kernel estimation of the copula density.  Bernoulli, 23(3):1848–1873, August 2017.
  * [67] Charles Dugas, Yoshua Bengio, François Bélisle, Claude Nadeau, and René Garcia.  Incorporating functional knowledge in neural networks.  Journal of Machine Learning Research, 10(6), 2009.
  * [68] Hennie Daniels and Marina Velikova.  Monotone and partially monotone neural networks.  IEEE Transactions on Neural Networks, 21(6):906–917, 2010.
  * [69] Anirban DasGupta.  Asymptotic theory of statistics and probability.  Springer Science & Business Media, 2008.
  * [70] Elliott Ward Cheney and William Allan Light.  A course in approximation theory, volume 101.  American Mathematical Soc., 2009.
  * [71] Hien D Nguyen and Geoffrey McLachlan.  On approximations via convolution-defined mixture models.  Communications in Statistics-Theory and Methods, 48(16):3945–3955, 2019.
  * [72] David Maxwell Chickering.  Learning equivalence classes of bayesian-network structures.  The Journal of Machine Learning Research, 2:445–498, 2002.
  * [73] Robert W Robinson.  Counting unlabeled acyclic digraphs.  In Combinatorial mathematics V, pages 28–43. Springer, 1977.
  * [74] A Cichocki, A-H Phan, Q Zhao, N Lee, I V Oseledets, M Sugiyama, and D Mandic.  Tensor networks for dimensionality reduction and Large-Scale optimizations. part 2 applications and future perspectives.  August 2017.
  * [75] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington.  Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks.  In International Conference on Machine Learning, pages 5393–5402. PMLR, 2018.
  * [76] Yaniv Blumenfeld, Dar Gilboa, and Daniel Soudry.  Beyond signal propagation: Is feature diversity necessary in deep neural network initialization?  July 2020.
  * [77] Vittorio Giovannetti, Simone Montangero, and Rosario Fazio.  Quantum MERA channels.  April 2008.

Supplementary Material

##  Appendix A Proofs

###  A.1 Proof of Proposition 1

The proof follows directly from Algorithm 1. The distribution
$F_{A^{\mathrm{HT}},\Phi})=\left\langle A^{\mathrm{HT}},\Phi\right\rangle$ is
a mixture model, and thus in order to sample from it we can first draw a
single mixture component (which is a product of univariate CDFs) and then
sample from this single component. The mixture weights are the elements of the
tensor $A^{\mathrm{HT}}$ given by the diagonal HT decomposition 8. In the next
section, we add details on the sampling process for the sake of clarity.

####  A.1.1 Details on the sampling for the HT model

Define a collection of independent categorical variables
$R=\\{R^{\ell}_{i,j}\\}$ taking values in $[m]$, where $\ell\in[p],i\in[m]$
and for any $\ell$, $j\in[2^{p-\ell}]$. These variables are distributed
according to

$\displaystyle\forall\ell,i,j:\quad\mathbb{P}\left[R_{i,j}^{\ell}=k\right]=\lambda_{i,k,j}^{\ell},$

where $\\{\lambda^{\ell}\\}_{\ell=1}^{p}$ are the parameters of the HT
decomposition. The fact that the parameters are nonnegative and
$\sum_{k=1}^{m}\lambda_{k,i,j}^{\ell}=1$ ensures the validity of this
distribution.

With the convention $R_{k_{p+1,1},1}^{p}=R_{1,1}^{p}$, define the event

$\displaystyle\underset{\ell=1}{\overset{p}{\bigcap}}\left(\cap_{j_{\ell}=1}^{2^{p-\ell}}\left\\{R_{k_{\ell+1,\left\lceil
j_{\ell}/2\right\rceil},j_{\ell}}^{\ell}=k_{\ell,j_{\ell}}\right\\}\right)$
$\displaystyle=\left\\{R_{1,1}^{p}=k_{p,1}\right\\}$
$\displaystyle\phantom{=}\bigcap\left(\cap_{j=1}^{2}\left\\{R_{k_{p,1},j}^{p-1}=k_{p-1,j}\right\\}\right)$
$\displaystyle\phantom{=}\,\,\,\,\vdots$
$\displaystyle\phantom{=}\bigcap\left(\cap_{j=1}^{d/2}\left\\{R_{k_{2,\left\lceil
j/2\right\rceil},j}^{1}=k_{1,j}\right\\}\right).$

Let $(\tilde{X}_{1},\dots,\tilde{X}_{d},R)$ be a random vector such that

$\displaystyle\mathbb{P}\left[\tilde{X}_{1}\leq x_{1},\dots,\tilde{X}_{d}\leq
x_{d}\left|\underset{\ell=1}{\overset{p}{\bigcap}}\left(\cap_{j_{\ell}=1}^{2^{p-\ell}}\left\\{R_{k_{\ell+1,\left\lceil
j_{\ell}/2\right\rceil},j_{\ell}}^{\ell}=k_{\ell,j_{\ell}}\right\\}\right)\right.\right]=\prod_{i=1}^{d}\varphi_{k_{1,\left\lceil
i/2\right\rceil},i}(x_{i}),$ (10)

which implies that the distribution of $(\tilde{X}_{1},\dots,\tilde{X}_{d})$
obtained after conditioning on a subset of the $\\{R^{\ell}_{i,j}\\}$ in this
way is equal to a single mixture component in $F_{\mathrm{HT}}=\left\langle
A,\Phi\right\rangle$. Thus, based on a sample of $R$, one can sample
$\tilde{X}_{i}$ by inverting the univariate CDFs $\varphi_{k_{1,\left\lceil
i/2\right\rceil},i}$ numerically and parallelizing over $i$. Numerical
inversion is trivial since the functions are increasing and continuously
differentiable, and this can be done for instance using the bisection method.
It remains to sample a mixture component.

Assume that a sample $\\{R_{i,j}^{\ell}\\}$ for a sequence of variables as in
10 is obtained e.g. from Algorithm 1. With the convention
$\lambda_{k_{p+1,1},k,1}^{p}=\lambda_{k}^{p}$, since

$\displaystyle\mathbb{P}\left[\underset{\ell=1}{\overset{p}{\bigcap}}\left(\cap_{j_{\ell}=1}^{2^{p-\ell}}\left\\{R_{k_{\ell+1,\left\lceil
j_{\ell}/2\right\rceil},j_{\ell}}^{\ell}=k_{\ell,j_{\ell}}\right\\}\right)\right]=\prod\limits_{\ell=1}^{p}\prod\limits_{j_{\ell}=1}^{2^{p-\ell}}\lambda_{k_{\ell+1,\left\lceil
j_{\ell}/2\right\rceil},k_{\ell,j_{\ell}},j_{\ell}}^{\ell},$

sampling from the categorical variables in this fashion is equivalent to
sampling a mixture component. It follows that by first sampling a single
mixture component and then sampling from this component, one obtains a sample
from $F_{\mathrm{HT}}$.

The main loop in Algorithm 1 samples such a mixture component, and there are
$p=\log_{2}d$ layers in the decomposition, so the time complexity of the main
loop is $O(\log d)$, and aside from storing the decomposition itself this
sampling procedure requires storing only $O(d)$ integers. This logarithmic
dependence is only in sampling from the categorical variables which is
computationally cheap. This not only avoids the linear time complexity common
in sampling from autoregressive models (without using distillation), but the
space complexity is also essentially independent of $m$ since only a single
mixture component is evaluated per sample.

###  A.2 Proof of Proposition 2

Assume that the activation function $\sigma$ is increasing, continuously
differentiable, and such that $\lim_{x\to-\infty}\sigma(x)=0$ and
$\lim_{x\to\infty}\sigma(x)=1$. Proposition 2 then follows immediately from
Proposition 4 and the fact that
$\cup_{r}\Phi_{1,r,\sigma}\subseteq\cup_{l,r}\Phi_{l,r,\sigma}$.

######  Remark 1.

In practice, we use the activation $\sigma(x)=x+a\tanh(x)$ for some $a>-1$.
While it does not satisfy the assumptions, the arguments in the proof of
Proposition 5 can be modified in a straightforward manner to cover this
activation (see Remark 2).

######  Proposition 4.

$\cup_{r}\Phi_{1,r,\sigma}$ is dense in ${\mathcal{F}}_{1}$ with respect to
the uniform norm.

Letting

$\displaystyle\widetilde{{\mathcal{F}}}_{1}$
$\displaystyle=\\{\widetilde{F}\colon{\mathbb{R}}\to{\mathbb{R}},\,\widetilde{F}(x)=\log
F(x)/(1-F(x)),\,F\in{\mathcal{F}}_{1}\\},$
$\displaystyle\widetilde{\Phi}_{l,r,\sigma}$
$\displaystyle=\\{\widetilde{\varphi}\colon{\mathbb{R}}\to{\mathbb{R}},\,\widetilde{\varphi}(x)=\log\varphi(x)/(1-\varphi(x)),\,\varphi\in\Phi_{l,r,\sigma}\\},$

the proof of Proposition 4 relies on the following proposition.

######  Proposition 5.

$\cup_{r}\widetilde{\Phi}_{1,r,\sigma}$ is dense in
$\widetilde{{\mathcal{F}}}_{1}$ with respect to the uniform norm.

####  A.2.1 Proof of Proposition 4

This proof is similar to that of [67, Theorem 2], which deals with functions
with positive outputs. We want to show that, for any $F\in{\mathcal{F}}_{1}$,
compact $K\subset{\mathbb{R}}$, and $\epsilon>0$, there exists
$\varphi\in\cup_{r}\Phi_{1,r,\sigma}$ such that

$\displaystyle\|\varphi-F\|_{\infty,K}=\sup_{x\in
K}|\varphi(x)-F(x)|\leq\epsilon.$

Denote the sigmoid function by $\rho(x)=1/(1+e^{-x})$ and define the function
$\widetilde{F}\colon{\mathbb{R}}\to{\mathbb{R}}$ by $\widetilde{F}(x)=\log
F(x)/(1-F(x))$, so that $F=\rho\circ\widetilde{F}$. By Proposition 5, there
exists $\widetilde{\varphi}\in\cup_{r}\widetilde{\Phi}_{1,r,\sigma}$ such that

$\displaystyle\sup_{x\in K}|\widetilde{\varphi}(x)-\widetilde{F}(x)|\leq
4\epsilon.$

Thus, letting $\varphi=\rho\circ\widetilde{\varphi}$, we have

$\displaystyle|\varphi(x)-F(x)|=|\rho\circ\widetilde{\varphi}(x)-\rho\circ\widetilde{F}(x)|\leq\sup_{x\in
K}\rho(x)^{\prime}|\widetilde{\varphi}(x)-\widetilde{F}(x)|\leq\epsilon.$

Since
$\Phi_{1,r,\sigma}=\\{\rho\circ\widetilde{\varphi}\colon\widetilde{\varphi}\in\widetilde{\Phi}_{1,r,\sigma}\\}$,
the result follows.

####  A.2.2 Proof of Proposition 5

This proof is similar to that of [68, Theorem 3.1], which is incomplete and
only deals with the sigmoid activation. First, note that
$\widetilde{{\mathcal{F}}}_{1}$ is the space of strictly increasing and
continuously differentiable functions. Therefore, for any
$\widetilde{F}\in\widetilde{{\mathcal{F}}}_{1}$ and interval
$K=[K_{1},K_{2}]$, we can write, for any $x\in K$,

$\displaystyle\widetilde{F}(x)=\widetilde{F}(K_{1})+\int_{\widetilde{F}(K_{1})}^{\widetilde{F}(K_{2})}\mathds{1}_{\widetilde{F}(x)\geq
u}du=\widetilde{F}(K_{1})+\int_{\widetilde{F}(K_{1})}^{\widetilde{F}(K_{2})}\mathds{1}_{x\geq\widetilde{F}^{-1}(u)}du,$

where the existence of the inverse $\widetilde{F}^{-1}$ is guaranteed by the
fact that $\widetilde{F}$ is strictly increasing and continuous. Thus, for
$\widetilde{F}(K_{1})=u_{0}<u_{1}<\cdots<u_{k}=\widetilde{F}(K_{2})$ a
partition of $[\widetilde{F}(K_{1}),\widetilde{F}(K_{2})]$ with
$u_{j+1}-u_{j}\leq\epsilon/2(\widetilde{F}(K_{2})-\widetilde{F}(K_{1}))$,
$x_{j}=\widetilde{F}^{-1}(u_{j})$ and

$\displaystyle G(x)=\widetilde{F}(K_{1})+\sum_{j=1}^{k}\mathds{1}_{x\geq
x_{j}}(u_{j}-u_{j-1}),$

we have $|\widetilde{F}(x)-G(x)|\leq\epsilon/2$, namely the approximation
error of the Riemann sum for increasing functions. Let $a>0$ and
$\widetilde{\varphi}\in\widetilde{\Phi}_{1,k,\sigma}$ obtained by setting
$b_{1}=\widetilde{F}(K_{1})$, as well as $(W_{1})_{1,j}=(u_{j}-u_{j-1})/a>0$,
$(b_{0})_{j}=-ax_{j}$ and $(W_{0})_{j,1}=a$ for $1\leq j\leq k$, then

$\displaystyle|G(x)-\widetilde{\varphi}(x)|\leq\sum_{j=1}^{k}(u_{j+1}-u_{j})\left|\mathds{1}_{x\geq
x_{j}}-\sigma(a(x-x_{j}))\right|.$

By the assumptions on $\sigma$, it is clear that $|\mathds{1}_{x\geq
0}-\sigma(ax)|$ can be made arbitrarily small. Thus, taking $a$ large enough
so that $|G(x)-\widetilde{\varphi}(x)|\leq\epsilon/2$, we have

$\displaystyle|\widetilde{F}(x)-\widetilde{\varphi}(x)|\leq|\widetilde{F}(x)-G(x)|+|G(x)-\widetilde{\varphi}(x)|\leq\epsilon.$

######  Remark 2.

Let $\sigma(x)=x+a\tanh(x)$ for some $a>-1$ and
$\widetilde{\varphi}\in\widetilde{\Phi}_{1,k,\sigma}$ obtained by setting
$b_{1}=\widetilde{F}(K_{1})+1/2$, as well as
$(W_{1})_{1,j}=(u_{j}-u_{j-1})/2a>0$, $(b_{0})_{j}=-ax_{j}$ and
$(W_{0})_{j,1}=|a|$ for $1\leq j\leq k$, then

$\displaystyle|G(x)-\widetilde{\varphi}(x)|$
$\displaystyle\leq\sum_{j=1}^{k}(u_{j+1}-u_{j})\left|\mathds{1}_{x\geq
x_{j}}-\tanh(|a|(x-x_{j}))/2-1/2\right|$
$\displaystyle\phantom{=}+\sum_{j=1}^{k}(u_{j+1}-u_{j})|x-x_{j}|/2|a|.$

Because $a$ is arbitrary, one can take it large enough so that
$|G(x)-\widetilde{\varphi}(x)|\leq\epsilon/2$ as above.

###  A.3 Proof of Proposition 3

Consider the classes of order $d$ tensored-valued functions with $m$
dimensions per mode defined as

$\displaystyle\Phi_{m,d,l,r,\sigma}$
$\displaystyle=\\{\Phi\colon{\mathbb{R}}^{d}\times[m]^{d}\to[0,1],\,\Phi(\bm{x})_{i_{1},\dots,i_{d}}=\textstyle\prod\limits_{j=1}^{d}\varphi_{i_{j},j}(x_{j}),\,\varphi_{i,j}\in\Phi_{l,r,\sigma}\\},$
$\displaystyle{\mathcal{F}}_{m,d}$
$\displaystyle=\\{\Phi\colon{\mathbb{R}}^{d}\times[m]^{d}\to[0,1],\,\Phi(\bm{x})_{i_{1},\dots,i_{d}}=\textstyle\prod\limits_{j=1}^{d}F_{i_{j},j}(x_{j}),\,F_{i,j}\in{\mathcal{F}}_{1}\\}$

as well as the class of neural network-based and ${\mathcal{F}}_{1}$-based
MDMAs, that is

$\displaystyle\textrm{MDMA}_{m,d,l,r,\sigma}$
$\displaystyle=\\{F_{A,\Phi}\colon{\mathbb{R}}^{d}\to[0,1],\,F_{A,\Phi}(\bm{x})=\langle
A,\Phi(\bm{x})\rangle,\,A\in{\mathcal{A}}_{d,m},\,\Phi\in\Phi_{m,d,l,r,\sigma}\\},$
$\displaystyle\textrm{MDMA}_{m,d,{\mathcal{F}}_{1}}$
$\displaystyle=\\{F_{A,\Phi}\colon{\mathbb{R}}^{d}\to[0,1],\,F_{A,\Phi}(\bm{x})=\langle
A,\Phi(\bm{x})\rangle,\,A\in{\mathcal{A}}_{d,m},\,\Phi\in{\mathcal{F}}_{m,d}\\}.$

We can now state the following proposition.

######  Proposition 6.

$\cup_{l,r}\mathrm{MDMA}_{m,d,l,r,\sigma}$ is dense in
$\mathrm{MDMA}_{m,d,{\mathcal{F}}_{1}}$ with respect to the uniform norm

Proposition 3 then follows immediately from the fact that
$\cup_{m}\textrm{MDMA}_{m,d,{\mathcal{F}}_{1}}$ is the space of multivariate
mixture distributions admitting a density, which is dense in
${\mathcal{F}}_{d}$ with respect to the uniform norm (see e.g., [69, Theorem
33.2], [70, Theorem 5], or [71, Corollary 11]).

####  A.3.1 Proof of Proposition 6

With $A\in{\mathcal{A}}_{d,m}$, $\Phi_{1}\in\Phi_{m,d,l,r,\sigma}$, and
$\epsilon>0$ and a compact $K=K_{1}\times\cdots\times
K_{d}\subset{\mathbb{R}}^{d}$, we want to prove that there exists
$A_{2}\in{\mathcal{A}}_{d,m}$ and $\Phi_{2}\in{\mathcal{F}}_{m,d}$, such that
$\sup_{\bm{x}\in
K}|F_{A,\Phi_{1}}(\bm{x})-F_{A_{2},\Phi_{2}}(\bm{x})|\leq\epsilon$. Assuming
that we can show $\sup_{\bm{x}\in
K}|\Phi_{1}(\bm{x})_{i_{1},\dots,i_{d}}-\Phi_{2}(\bm{x})_{i_{1},\dots,i_{d}}|\leq\epsilon$,
the result would then follow from setting $A_{2}=A$ and the fact that
$F_{A,\Phi_{1}}(\bm{x})-F_{A,\Phi_{2}}(\bm{x})=\langle
A,\Phi_{1}(\bm{x})-\Phi_{2}(\bm{x})\rangle$ implies

$\displaystyle\sup_{\bm{x}\in
K}|F_{A,\Phi_{1}}(\bm{x})-F_{A,\Phi_{2}}(\bm{x})|=$
$\displaystyle\leq\sum_{i_{1},\dots,i_{d}}A_{i_{1},\dots,i_{d}}\sup_{\bm{x}\in
K}|\Phi_{1}(\bm{x})_{i_{1},\dots,i_{d}}-\Phi_{2}(\bm{x})_{i_{1},\dots,i_{d}}|=\epsilon.$

With $\delta=\epsilon^{1/d}$, by Proposition 2, there exists $l$, $w$, and
$\\{\varphi\\}_{i\in[m],j\in[d]}$ with $\varphi_{i,j}\in\Phi_{l,r,\sigma}$,
such that

$\displaystyle\max_{i\in[m],j\in[d]}\sup_{x_{j}\in
K_{j}}|\varphi_{i,j}(x_{j})-F_{i,j}(x_{j})|\leq\delta.$

Thus, we have that

$\displaystyle|\Phi_{1}(\bm{x})_{i_{1},\dots,i_{d}}-\Phi_{2}(\bm{x})_{i_{1},\dots,i_{d}}|=|\prod\limits_{j=1}^{d}\varphi_{i_{j},j}(x_{j})-\prod_{j=1}^{d}F_{i,j}(x_{j})|\leq\delta^{d}=\epsilon.$

##  Appendix B Additional experimental results

###  B.1 Toy density estimation

Figures 4 and 5 show more results on the popular checkerboard and 8 Gaussians
toy datasets studied in Figure 1.

|  |  |
---|---|---|---
|  |  |
|  |  |
Training data

Figure 4: Density estimation with closed-form marginals and conditionals. Top
Row: Samples from a 3D density, and 2D marginal histograms. Middle Row:
Samples from MDMA after fitting the density, and plots of the learned 2D
marginals. Bottom Row: Left: learned 1D marginals compared to 1D marginal
histograms of the training data. Right: Learned conditional densities.

|  |  |
---|---|---|---
|  |  |
|  |  |
Training data

Figure 5: Density estimation with closed-form marginals and conditionals. Top
Row: Samples from a 3D density, and 2D marginal histograms. Middle Row:
Samples from MDMA after fitting the density, and plots of the learned 2D
marginal. Bottom Row: Left: learned 1D marginals compared to 1D marginal
histograms of the training data. Right: Learned conditional densities.

###  B.2 Density estimation with missing data

We compare MICE imputation [43] to $k$-NN imputation (with $k=3$ neighbours)
[44] on the UCI POWER dataset in Figure 6, before performing density
estimation with BNAF [9]. Due to the size of the dataset, we were not able to
use $k$-NN imputation on the full dataset, but instead split it up into $100$
batches and performed the imputation per batch. Similar results were obtained
on the UCI GAS dataset, and for this reason we only compare MDMA to MICE
imputation in the main text.

Figure 6: A comparison of data imputation methods on the UCI POWER dataset
followed by density estimation with BNAF, showing that MICE imputation
outperforms $k$-NN. We subsequently use MICE in the comparison with MDMA in
the main text.

###  B.3 Causal discovery

In Figure 7, we present examples of completely partially directed acyclical
graphs (CPDAGs) learned using the PC algorithm, using either MDMA or the
vanilla (Gaussian) method for testing conditional independence used in [55].
See Section C.2 for additional details.

Figure 7: Recovered causal graphs: Top: Synthetic data from a random DAG with
sigmoidal causality mechanism. The graph inferred using MDMA PC had
directional SHD of 11, compared to 15 for the Gaussian PC. Bottom: Protein
signaling graph [62]. The graph inferred using MDMA PC had directional SHD of
27, compared to 32 for the Gaussian PC.

###  B.4 Density estimation on real data

To demonstrate how MDMA allows one to visualize marginal densities we show in
Figure 8 learned bivariate marginals from the UCI POWER and HEPMASS datasets.
The former is composed of power consumption measurements from different parts
of a house, with one of the variables ($X_{6}$) being the time of day.

Figure 8: Log marginal density on UCI datasets. The scatter plot is composed
of $500$ samples from the dataset. Left: The POWER dataset. One variable
corresponds to the time of day, and the other to power consuption from the
kitchen of a house. Note the small value of the density during night-time. The
data is normalized during training, yet the labels on the horizontal axis
reflect the value of the unnormalized variable for interpretability. Right:
The HEPMASS dataset. Despite MDMA not achieving state-of-the-art results on
test likelihood for this dataset, the model still captures accurately the non-
trivial dependencies between the variables.  Figure 9: Feature learning in
MDMA. We plot ten univariate PDFs $\dot{\varphi}_{ij}$ parameterized as in
section E.1 for $j=4$, both at initialization and after training on the UCI
POWER dataset. Overlaid is a histogram of computed from 500 datapoints. We
find that the features localize near the datapoints.

##  Appendix C Experimental details

All experiments were performed on Amazon Web Services using Tesla V-100 GPUs.
The total compute time was 3,623 hours, with the vast majority devoted to the
experiments on density estimation with missing values (Section 4.3), where
some of the runs of BNAF required over 72 hours to complete.

###  C.1 Mutual information estimation

$10^{6}$ samples from the true density are used for fitting MDMA and for
estimating the integral over the log marginals in order to compute the mutual
information. The MDMA model used had parameters $r=4,l=5,m=1000$ and was
trained with a batch size $500$ and learning rate $0.01$ for $2$ epochs.

###  C.2 Causal discovery

In many application areas, causal relationships between random variables can
be represented by a directed acyclical graph (DAG). The PC algorithm [47] is
an efficient algorithm for recovering sparse DAGs from observations. In
general, this recovery is complicated by the fact that two DAGs can induce the
same probability distribution, leading to them being called Markov equivalent.
Hence, observational data can only help infer the Markov equivalence class of
a given DAG. The equivalence class, known as a completely partially directed
acyclical graph (CPDAG, also called essential graph) [72], encodes all the
dependence information in the induced distribution. The object of the PC
algorithm is therefore the recovery of a CPDAG that is consistent with the
data. This is generally a hard problem, since the cardinality of the space of
DAGs is super-exponential in the number of variables [73].

The PC algorithm requires repeatedly testing for independence between pairs of
variables conditioned on subsets of the remaining variables. As mentioned in
the main text, testing for conditional independence can be reduced to an
independence test between variables that depend on conditional CDFs [55],
which can be obtained easily after fitting the joint density using MDMA. In
our experiments, the results of using MDMA as part of the PC algorithm for
testing conditional independence are compared to the results obtained by using
a Gaussian conditional independence test based on partial correlations.

The synthetic DAGs were generated using the the Causal Discovery
Toolbox.333https://fentechsolutions.github.io/CausalDiscoveryToolbox/html/index.html
When the sigmoidal causal mechanism is used, given a variable $Y$ and parents
$\\{X_{1},\dots,X_{s}\\}$, then
$Y=\underset{i=1}{\overset{r}{\sum}}w_{i}\sigma(X_{i}+w_{0})+\varepsilon$, and
if a polynomial mechanism is used then
$Y=\varepsilon\left(w_{0}+\underset{i=1}{\overset{s}{\sum}}w_{1i}X_{i}+\underset{i=1}{\overset{s}{\sum}}w_{2i}X_{i}^{2}\right)$,
where $w_{i},w_{ij},\varepsilon$ are random. MDMA was trained with
$m=1000,L=2,r=3$ for $50$ epochs and learning rate $0.1$ on all datasets. In
all experiments we find that the graphs recovered using MDMA are closer to the
truth than those recovered using Gaussian PC, as measured by the structural
Hamming distance. Example recovered graphs are shown in Figure 7.

###  C.3 Density estimation on real data

Table 3: Dimension and size of the UCI datasets, and the hyperparameters used for fitting MDMA on these datasets. $m$ is the width of the MDMA model, $l$ and $r$ are respectively the depth and width of the univariate CDF models described in Section E.1.  | POWER | GAS | HEPMASS | MINIBOONE
---|---|---|---|---
d | 6 | 8 | 21 | 43
Training set | 1659917 | 852174 | 315123 | 29556
Validation set | 184435 | 94685 | 35013 | 3284
Test set | 204928 | 105206 | 174987 | 3648
$m$ | 1000 | 4000 | 1000 | 1000
$l$ | 2 | 4 | 2 | 2
$r$ | 3 | 5 | 3 | 3

We train MDMA on four UCI datasets, details of the dataset sizes and
hyperparameter choices are presented in Table 3. In all experiments a batch
size of $500$ and learning rate of $0.01$ were used. We use the same pre-
processing as [45], which involves normalizing and adding noise. Details are
provided in the attached code.444The raw datasets are available for download
at https://zenodo.org/record/1161203#.YLUMImZKjuU The POWER dataset consists
of measurements of power consumption from different parts of a house as well
as the time of day. The GAS dataset contains measurements of chemical sensors
used to discriminate between different gases. The HEPMASS and MINIBOONE
datasets are both measurements from high-energy physics experiments, aiming
respectively for the discovery of novel particles and to distinguish between
different types of fundamental particles (electron and muon neutrinos).

##  Appendix D Tensor decompositions

In constructing the MDMA estimator, we are confronted with the problem of
combining products of univariate CDFs linearly, in a manner that is both
computationally efficient and expressive. The linearity constraint reduces
this to a tensor decomposition problem (with additional non-negativity and
normalization constraints). There is an extensive literature on such efficient
tensor decompositions (see [74] for a review).

The analogy with tensor decompositions becomes clear when we consider discrete
rather than continuous variables. Assume we wish to model the joint
distribution of $d$ discrete variables, each taking one of $S$ possible
values. The distribution is then a function
$\varphi_{S}:[S]\times\cdots\times[S]\rightarrow\mathbb{R}$, which can also be
viewed as an order $d$ tensor. A general tensor $\varphi_{S}$ will require
order of $S^{d}$ numbers to represent, and is thus impractical even for
moderate $d$. The continuous analog of such a tensor is a multivariate
function $F:\mathbb{R}^{d}\rightarrow\mathbb{R}$, with the value of $x_{j}$
corresponding to the discrete index $s_{j}$. We will thus use the same
notation for the continuous case. A graphical representation of tensors and of
the diagonal HT tensor used in MDMA is presented in Figure 10.

a)b)c)d) Figure 10: Tensor decompositions. a) Tensors of various orders (resp.
vectors, matrices, delta tensors). Each edge represents an index, and
connecting two edges represents contraction (summation over an index). b) The
set of univariate CDFs $\Phi$, which can be viewed as an order $2d$ tensor. c)
A general unstructured tensor of order 6. d) The hierarchical Tucker (HT)
decomposition 5. After suitable normalization, the tensor in d) can be
contracted with the tensor $\Phi$ shown in b) to give a multivariate CDF.

##  Appendix E Additional design details

###  E.1 Univariate marginal parameterization

We parameterize the univariate marginal CDF $\varphi(x)$ for some scalar
$x\in\mathbb{R}$ using a simple feed-forward network following [37]. Recall
from section Section 3.6 that we model the univariate CDFs as functions

$\displaystyle\varphi(x)=\mbox{sigmoid}\circ L_{l}\circ\sigma_{l-1}\circ
L_{l-1}\circ\sigma_{l-2}\cdots\circ\sigma_{1}\circ L_{1}\circ\sigma_{0}\circ
L_{0}(x),$

where $L_{i}\colon{\mathbb{R}}^{n_{i}}\to{\mathbb{R}}^{n_{i+l}}$ is the affine
map $L_{i}(x)=W_{i}x+b_{i}$ for an $n_{i+1}\times n_{i}$ weight matrix $W_{i}$
with _nonnegative_ elements and an $n_{i+1}\times 1$ bias vector $b_{i}$, with
$n_{l+1}=n_{0}=1$ and $n_{i}=r$ for $i\in[l]$. This is a slightly more general
form than the one in Section 3.6 since we allow the nonlinearities to depend
on the layer. For the nonlinearities, we use

$\sigma_{i}({x})={x}+{a_{i}}\odot\tanh({x})$

for some vector $a_{i}\in\mathbb{R}^{n_{i+1}}$ with elements constrained to
lie in $[-1,1]$ (the lower bound on $a_{i}$ is necessary to ensure that
$\sigma_{i}$ are invertible, but the upper bound is not strictly required).
This constraint, as well as the non-negativity constraint on the $W_{i}$, is
enforced by setting
${W}_{i}=\mathrm{softplus}(\tilde{{W}}_{i},10),{a_{i}}=\tanh(\tilde{{a}}_{i})$
in terms of some $\tilde{{W}}_{i}\in\mathbb{R}^{n_{i+1}\times
n_{i}},\tilde{{a}}_{i}\in\mathbb{R}^{n_{i+1}}$. The softplus function is
$\mathrm{softplus}(x,\beta)=\frac{1}{\beta}\log(1+\exp(\beta x))$ and is a
smooth, invertible approximation of the ReLU. We typically use small values
for $l,r$ in the experiments (see Appendix C).

###  E.2 Adaptive variable coupling

One degree of freedom in constructing a HT decomposition is the choice of
partitions of subsets of the variables at every layer over which the products
are taken. This imposes a form of weight-sharing, and it will be natural to
share weights between variables that are highly correlated. As a simple
example, let $d=4$ and consider two different HT decompositions

$\displaystyle
F(x_{1},x_{2},x_{3},x_{4})=\underset{i_{1},i_{2},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\lambda_{k,i_{2},2}^{1}\varphi_{i_{1},1}(x_{1})\varphi_{i_{1},2}(x_{2})\varphi_{i_{2},3}(x_{3})\varphi_{i_{2},4}(x_{4}),$
$\displaystyle\widetilde{F}(x_{1},x_{2},x_{3},x_{4})=\underset{i_{1},i_{2},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\lambda_{k,i_{2},2}^{1}\varphi_{i_{1},1}(x_{1})\varphi_{i_{1},3}(x_{3})\varphi_{i_{2},2}(x_{2})\varphi_{i_{2},4}(x_{4}),$

obtained by coupling $X_{1}$ in the first layer respectively with $X_{2}$ and
$X_{3}$. The univariate marginals for $X_{1}$, $X_{2}$ and $X_{3}$ can then be
written as

$\displaystyle
F(x_{1})=\underset{i_{1},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\varphi_{i_{1},1}(x_{1}),\
F(x_{2})=\underset{i_{1},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\varphi_{i_{1},2}(x_{2}),\
F(x_{3})=\underset{i_{2},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{2},2}^{1}\varphi_{i_{2},3}(x_{3}),$
$\displaystyle\widetilde{F}(x_{1})=\underset{i_{1},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\varphi_{i_{1},1}(x_{1}),\
\widetilde{F}(x_{2})=\underset{i_{2},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{2},2}^{1}\varphi_{i_{2},2}(x_{2}),\
\widetilde{F}(x_{3})=\underset{i_{1},k}{\sum}\lambda_{k}^{2}\lambda_{k,i_{1},1}^{1}\varphi_{i_{1},3}(x_{3}).$

Assume that the variables $X_{1}$ and $X_{2}$ are identical. In $F$, both of
their univariate marginals depend in an identical way on the tensor
parameters. In $\widetilde{F}$ however, additional parameters are required to
represent them. Hence $F$ is a more parsimonious representation of the join
distribution. If $X_{1}$ and $X_{3}$ are identical instead, then the converse
holds and $\widetilde{F}$ is the more parsimonious representation. This
property extends to any higher-dimensional (e.g., bivariate) marginals.

In data with spatial or temporal structure (e.g. if the variables are image
pixels) there is a natural way to couple variables based on locality. When
this is not present, we can adaptively construct the couplings based on the
correlations in the data using a simple greedy algorithm. After constructing
an empirical covariance matrix from a minibatch of data, we couple the two
variables that are most correlated and have not yet been paired. We repeat
this until we couple all the groups of variables. Then we "coarse-grain" by
averaging over blocks of the covariance matrix arranged according to the
generated coupling and repeat the process, this time coupling subsets of
variables. We find that this coupling scheme improves performance compared to
naive coupling that does not take correlations into account.

###  E.3 Initialization

As in the univariate case, the non-negativity constraint of the HT tensor
parameters $\lambda_{k,k^{\prime},j}^{i}$ is enforced by defining
$\lambda_{k,k^{\prime},j}^{i}=\mathrm{softplus}\left(\tilde{\lambda}_{k,k^{\prime},j}^{i},20\right)$
for some $\tilde{\lambda}_{k,k^{\prime},j}^{i}\in\mathbb{R}$.

As is standard, we initialize independently the elements of the univariate PDF
weights $\tilde{W}_{i}$ as zero-mean gaussians with variance
$1/n_{\mathrm{fan_{i}n}}$, the $\tilde{a}_{i}$ as standard gaussians and the
$b_{i}$ as 0. The initialization of the HT parameters is
$\tilde{\lambda}_{k,k^{\prime},j}^{i}=m*\delta_{k,k^{\prime}}$ for $1\leq
i\leq p-1$. This initialization is chosen so that after applying the softplus
the matrix $\lambda_{\cdot,\cdot,j}^{i}$ is close to an identity at
initialization, which we have found facilitates training compared to using a
random initialization. Benefits of such “orthogonal” initialization schemes
have also been shown for deep convolutional networks [75, 76]. The final layer
$\tilde{\lambda}_{k}^{p}$ are initialized as zero mean gaussians with variance
$0.3/m$.

###  E.4 From HT to MERA

The choice of the diagonal HT decomposition 5 is convenient, yet there is a
wealth of other tensor decompositions that can be explored. Here we highlight
one such decomposition that generalizes the diagonal HT and could potentially
lead to more expressive models. It is based on [77].

Let $F^{\mathrm{M},1}\colon{\mathbb{R}}^{d}\to{\mathbb{R}}^{m\times d}$ be a
matrix-valued function such that
$F^{\mathrm{M},1}_{i,j}(\bm{x})=F^{\mathrm{HT},1}_{i,j}(\bm{x})=\varphi_{i,j}(x_{j})$.
For $l\in\\{2,\dots,\log_{2}d\\}$, define the matrix-value functions
$F^{\mathrm{M},1}\colon{\mathbb{R}}^{d}\to{\mathbb{R}}^{m\times d/2^{l-1}}$
recursively by

$\displaystyle F^{\mathrm{M},l}_{i,j}(\bm{x})$
$\displaystyle=\sum\limits_{k}^{m}\lambda_{k,i,j}^{l-1}\\{\chi_{1,k,i,j}^{l-1}\varphi_{k,2j-1}^{\mathrm{M},l-1}(\bm{x})\varphi_{k,2j}^{\mathrm{M},l-1}(\bm{x})+\chi_{2,k,i,j}^{l-1}\varphi_{k,2j-1}^{\mathrm{M},l-1}(\bm{x})\varphi_{k+1,2j}^{\mathrm{M},l-1}(\bm{x})\\}$

with $k+1\equiv 1$ when $k=m$, $\lambda^{l}$ as in (4), and $\chi^{l}$ a
$2\times m\times m\times d/2^{l}$ tensor with nonegative elements satisfying
$\chi_{1,k,i,j}^{l}+\chi_{2,k,i,j}^{l}=1$. The MERA parametrization of a
distribution can then be written as

$\displaystyle F_{\mathrm{M}}(\bm{x})$
$\displaystyle=\sum\limits_{k=1}^{m}a_{k}\varphi_{k,1}^{\mathrm{M},\log_{2}d}(\bm{x})\varphi_{k,2}^{\mathrm{M},\log_{2}d}(\bm{x}),$

with $a\in\mathbb{R}^{m}$Since the conditions on $\lambda^{l}$ and $\chi^{l}$
imply
$\sum\limits_{k=1}^{m}\lambda_{k,i,j}^{l}(\chi_{1,k,i,j}^{l}+\chi_{2,k,i,j}^{l})=1$,
this parametrization clearly results in a valid CDF. Note that
$\chi_{1,k,i,j}^{l}+\chi_{2,k,i,j}^{l}=1$ leads to $\chi^{l}$ having only
$m\times m\times d/2^{l}$ free parameters. For $d=4$, we have

$\displaystyle F_{\mathrm{M}}(\bm{x})$
$\displaystyle=\sum\limits_{k_{1},k_{2},k_{2}^{\prime}=1}^{m}a_{k_{1}}\lambda_{k_{2},k_{1},1}^{1}\lambda_{k_{2}^{\prime},k_{1},2}^{1}\sum\limits_{i_{1},i_{2}=1}^{2}\chi_{i_{1},k_{2},k_{1},1}^{1}\chi_{i_{2},k_{2}^{\prime},k_{1},2}^{1}$
$\displaystyle\phantom{=\sum\limits_{k_{1},k_{2},k_{2}^{\prime},k_{3},k_{3}^{\prime},k_{3}^{\prime\prime},k_{3}^{\prime\prime\prime}=1}^{m}}\times\varphi_{k_{2},1}^{\mathrm{M},1}(\bm{x})\varphi_{k_{2}+i_{1}-1,2}^{\mathrm{M},1}(\bm{x})\varphi_{k_{2}^{\prime},3}^{\mathrm{M},1}(\bm{x})\varphi_{k_{2}^{\prime}+i_{2}-1,4}^{\mathrm{M},1}(\bm{x}).$