Analysis of Learning a Flow-based Generative Model from Limited Sample Complexity
Authors: Hugo Cui
Florent Krzakala
Eric Vanden-Eijnden
Lenka Zdeborova
Venue: ICLR 2024 poster (2024)
Keywords: statistical physics
flow-based generative model
stochastic interpolation
gaussian mixture
auto-encoder
Primary area: learning theory
TL;DR: A sharp analysis of a flow-based generative model, parametrized by a shallow auto-encoder, trained with limited data

Abstract: We study the problem of training a flow-based generative model, parametrized by a two-layer autoencoder, to sample from a high-dimensional Gaussian mixture. We provide a sharp end-to-end analysis of the problem. First, we provide a tight closed-form characterization of the learnt velocity field, when parametrized by a shallow denoising auto-encoder trained on a finite number $n$ of samples from the target distribution. Building on this analysis, we provide a sharp description of the corresponding generative flow, which pushes the base Gaussian density forward to an approximation of the target density. In particular, we provide closed-form formulae for the distance between the means of the generated mixture and the mean of the target mixture, which we show decays as $\Theta_n(\frac{1}{n})$. Finally, this rate is shown to be in fact Bayes-optimal.

Published as a conference paper at ICLR 2024
ANALYSIS OF LEARNING A FLOW-BASED GENERATIVE MODEL FROM LIMITED SAMPLE COMPLEXITY
Hugo Cui
Florent Krzakala
Statistical Physics of Computation Laboratory École Polytechnique Fédérale de Lausanne (EPFL) Lausanne, Switzerland
Information Learning and Physics Laboratory
École Polytechnique Fédérale de Lausanne (EPFL)
Lausanne, Switzerland
Eric Vanden-Fijnden
Lenka Zdehorová
Courant Institute of Mathematical Science New York University (NYU) New York, USA
Statistical Physics of Computation Laboratory École Polytechnique Fédérale de Lausanne (EPFL)
Lausanne, Switzerland
ABSTRACT
We study the problem of training a flow-based generative model, parametrized by two-layer autoencoder, to sample from high-dimensional Gaussian mixture. We provide sharp end-to-end analysis of the problem. First, we provide a tight flosed-form characterization of the learnt velocity field, when parametrized by a shallow denoising auto-encoder trained target distribution. Building on this analysis, we provide sharp description of the corresponding generative flow, which pushes the base Gaussian density forward to an approximation of the target density. In particular, we provide closed-for formulae for the distance between the means of the generated mixture and the mean of the target mixture, which we show decays as On.() (1/)). Finally, this rate is shown to be in fact Bayes-optimal.
finite number of samples from the
Flow and diffusion-based generative models have introduced shift in paradigm for density estimation and sampling problems, leading to state-of-the art algorithms e.g. in image generation (Rombach) al., 2022; Ramesh al., 20222; Saharia al., (022). Instrumental in these advances was the realization that the sampling problem could be recast as a transport process from a simple -typically Gaussian- base distribution to the target density. Furthermore, the velocity field governing the flow can be characterized as the minimizer of quadratic loss function, which can be estimated from data by (a) approximating the loss by its empirical estimate using available training data and (b) parametrizing the velocity field using denoiser neural network. These ideas have been fruitfully implemented as part of a number of frameworks, including score-based diffusion models (Song & Ermon 2019 Song al.) 2020 Karras al. 2022; Ho et al.) 2020), and stochastic interpolation (Albergo & Eijnden [002 Albergo et al. [023]] Lipman et al. 2022 Liu et al. 2022). A tight
analytical understanding of the learning of generative models from limited data, and the resulting generative process, is however still largely missing. This constitutes the research question addressed
the present manuscript.
A line of recent analytical works (Benton et al.) 2023 Chen et al. 2022: 2023a c Wibisono Yang 2022; Lee et al. 2022 2023 Li al. 2023 De Bortoli al.] 2021| De Bortoli 2022 Pidstrigach 2022] Block et al. 2020) have mainly ly focused on the study of the transport problem, and provide rigorous convergence guarantees, taking as starting point the assumption of an - estimate of the velocity or score. They hence bypass the investigation of the learning problem -and in particular the question of ascertaining the sample complexity needed to obtain such an accurate estimate. More importantly, the study of the effect of learning from limited sample complexity (and thus e.g. of possible network overfitting and memorization) on the generated density, is furthermore left unaddressed. On the other hand, very recent works (Cui & Zdeborová 2023 Shah et al. 2023) have characterized the learning of Denoising
(DAEs) (Vincent
al. [2010) Vincent 2011) in high dimensions Gaussian mixture densities. Neither work however studies the consequences on the generative process. Bridging that gap, recent works have offered
joint analysis of the learning and generative processes Oko et al. (2023) Chen al. (2023b): Yuan

---

Published as a conference paper at ICLR 2024
et al. (2023) derive rigorous bounds at finite sample complexity, under the assumption of data with a low- dimenssiiii structure. Closer to our manuscript, a concurrent work (Mei & Wu 2023) bounds the Kullback-Leibler distance between the generated and target densities, when parametrizing the flow using a ResNet, for high-dimensional graphical models. On the other hand, these bounds do not go
the sample complexity increases, and are a priori not tight.
The present manuscript aims at complementing and furthering this last body of works, by providing a tight end-to-end analysis of flow-based generative model starting from the study of the high- dimensional learning problem with finite number of samples, and subsequently elucidating the implications thereof on the generative process.
Main contributions- We study the problem of estimating and sampling a Gaussian mixture using a flow-based generative model, in the framework of stochastic interpolation (Albergo & Vanden Eijnden 2022; Albergo al.] 2023) Lipman al., 2022] Liu al. 2022). We consider the case where a non-linear two-layer DAE with one hidden unit is used to parametrize the velocity field of the associated flow, and is trained with a finite training set. In the high-dimensional limit,
We provide a sharp asymptotic closed-form characterization of the learnt velocity field, as a function of the target Gaussian mixture parameters. the stochastic interpolation schedule, and the number of training samples n.
We characterize the associated flow by providing a tight characterization of a small number of summary statistics, tracking the dynamics of a sample from the Gaussian base distribution as is transported by the learnt velocity field.
We show that even with a finite number of training samples, the learnt generative model allows to sample from mixture whose mean asymptotically approaches the mean of the target mixture as O, (2/)) in squared distance, with this rate being tight. Finally, we show that this rate is in fact Bayes-optimal.
The code used in the present manuscript is provided in this repository
RELATED WORKS
Diffusion and flow-based generative models Score-based diffusion models (Song & Ermon 2019
Song et al. 2020 Karras et al. 2022 Ho et al. 2020) build on the idea that any density can be mapped Gaussian density by degrading samples through an -hhieene-hhhleeckk process. Sampling from the original density can then be carried out by time-reversing the corresponding stochastic transport, provided the score is known or estimated. These ideas were subsequently refined in (Albergo &
Vanden-Eijnden 2022 Albergo et al., 2023 Lipman et al. 2022 Liu et al. 2022). which provide a flexible framework to bridge between two arbitrary densities in finite time.
Convergence bounds In the wake of the practical successes of flow and diffusion- based generative models, significant theoretical effort has been devoted to studying the convergence of such methods, by bounding appropriate distances between the generated and the target densities.
common
assumption of (Benton et al. 2023 |023|Chhen et al. 2022, 2023a c d Wibisono & Yang 2022 Lee et al.
2022, 2023 Li et al. 2023 De Bortoli et al. 2021 De Bortoli 2022 Pidstrigach 2022 Block et al. 2020) is the availability of good estimate for the score, Lee. estimate whose average (population)
squared distance with the true score is bounded by a small constant c. Under this assumption, Chen et al. (2022): ;; Lee et al.) (2022) obtain rigorous control the Wasserstein and total variation distances with very mild assumptions on the target density. Ghio et al. (2023) explore the connections between algorithmic hardness of the score/flow approximation and the hardness of sampling in number of graphical models.
Asymptotics for DAE learning The backbone of flow and diffusion-basse generative models the parametrization of the score or velocity by denoiser-type network, whose most standard realization is arguably the DAE (Vincent et al. 2010 Vinnent 2011). Very recent works have provided a detailed analysis of its learning on denoising tasks, for data sampled from Gaussian mixtures. Cui & Zdeborová (2023) sharply characterize how DAE can learn the mixture parameters with n = O2(d) training samples when the cluster separation is Oa(1). Closer to our work, for arbitrary cluster separation, Shah al. (2023) rigorously show that DAE trained with gradient descent on the denoising diffusion probabilistic model loss (Ho et al. 2020) can recover the cluster means with polynomial number of samples. While these works complement the aforediscussed
2

---

Published as a conference paper at ICLR 2024
convergence studies in that they analyze the effect of a finite number of samples, neither explores the flow associated to the learnt score.
Network-parametrinee models Tying together these two body of works, very recent line of research has addressed the problem of bounding, at finite sample complexity, appropriate distances between the generated and target densities, assuming a network-based parametrization. [Oko et al. (2023) provide such bounds when parametrizing the score using class of ReLUU networks. These bounds however suffer from the curse of dimensionality, Oko et al. (2023); Yuan et al. (2023); Chen al. (2033b) surmount this hurdle by assuming target density with low-dimnsiional structure. On a heuristic level. Biroll & Mézard (2023) estimate the order of magnitude of the sample complexity needed to sample from high-dimensinn
weirss model. Finally, work concurrent to ours (Mei & Wu 2023) derives rigorous bounds for a number of high-dimensional graphical models. On the other hand, these bounds are priori not tight, and do not go zero the sample complexity becomes large. The present manuscript aims at furthering this line of work, and provides a sharp analysis of a high-dimensional flow-based generative model.
SETTING
We start by giving a concise overview of the problem of sampling from target density P1 over R&d in the framework of stochastic interpolation (Albergo & Gunden-Eijndee 2022 Albergo et al. 2023)
Recasting sampling as an optimization problem Samples from #1 can be generated by drawing a sample from an casy-to-sample base density Po -henceforth taken to be a standard Gaussian density Pu (0, Iz),, and evolving it according to the flow described by the ordinary differential equation (ODE)
X, - b(X,,)))))
dt
for & E [0, 1]. Specifically, as shown in Albergo et al. (2023), if x;= ~ Po. then the final sample X has probability density Pl, the velocity field b b(x 1) governing the flow (1) is given by
b(x, () = B((x())=o + (())z) [a. = ær,
where we denoted x, (t) + (())x; and the conditional expectation bears over
Pl.
an ~ po, with 2o 1 21. The result holds for any fixed choice of schedule functions ax, 3 E ( 1))
satisfying (0)) (11 1, (x1) = 3(0) O, and a1)) + 3(() > 0 for all E [0, 1]. In addition
to the velocity field b(x; t). it convenient to consider the field f(x, t). related to b(x, f) by the
simple relation
b(x: t) ((()) (3()) a(t)
f(x,t) ((t) f a(t) a(t)
a(t) a(t)
Note that f(r can be alternatively expressed as E[ x], and thus admits a natural interpre- tation as denoising function, tasked with recovering the target value xı from the interpolated (noisy) sample x,. The denoiser f f(x, can furthermore characterized as the minimizer of the objective
R[f - / E fl - x 11| dt.
The loss is a simple sequence of quadractic denoising objectives.
Learning the velocity from data There are several technical hurdles in carrying out the minimiza- tion (4). First, since the analytical form of P1 is generically unknown, the population risk has to be approximated by its empirical version, provided dataset D {=f " of n training samples of (x=)) independently drawn from P1 (po) is available. Second, the minimization in (4) bears over a time-dependent vector field f. To make the optimization tractable, the latter can be parametrized at each time step ł by separate neural network fo, (-) with trainable parameters 0,. Under those approximations, the population risk (4) thus becomes
n
R R((OrtrE(o.))
|l fe. (xt)) xf1" dt.
(5)
3

---

Published as a conference paper at ICLR 2024
Remark that in practice, the time can enter imput of the neural network, and only one network then needs to be trained. In the present manuseript however, for technical reasons, we instead consider the case where separate network is trained for each time step t. Besides, note that since the base density Po is a priori easy to sample from, one could in theory augment the dataset with several samples from po for each available ax" For conciseness, technique in the present manuscript, and leave precise investigation thereof to future work. Denoting by {de [ne the minimizer of 5). the learnt velocity field b is related to the trained denoiser fa, by (4) as
do not examine such augmentation
b(æz t) (8(1)
a(t)
s(t) fa. (a)) + a(() a(t)
ait) x
(6)
The sampling can finally be carried out by using b as a proxy for the unknown b in 1
d X; = b(X,,t)
dt
Note that the solution X1 time 1 of the ODE 7 has law أr * P1 due to the two approximations in going from the population function-space objective (4) to the empirical parametric proxy (5). The present manuscript presents a sharp analysis of the learning problem (5) and the
resulting flow 77) for a solvable model, which we detail below.
Data model We consider the case of a target density P given by a binary isotropic and homoscedastic Gaussian mixture
1 2 2 o²
P1 2 ((
(8)
Each cluster is thus centered around its mean =u and has variance a². For definiteness, we consider here a balanced mixture, where the two clusters have equal relative probabilitiess, and defer the discussion of the imbalanced case to Appendix D Note that a sample ar" can then be decomposed as
2x = 8" f + 2", with 8 ~ [1-1,,111)))) +1}) and z² ~ (Ta)) Finally, note that the closed-form expression for the exact velocity field b (I) associated to the density p1 is actually known (see e.g.
Efron (201I); Albergo et al. (2023)) This manuscript explores the question whether a neural network can learn a good approximate b thereof without any knowledge of the density Pl. and only from finite number of samples drawn therefrom.
Network architecture We consider the case where the denoising function f is parametrized with a two-layer non-linear DAE with one hidden neuron, and -taking inspiration from modern practical architectures such as U-nets (Ronneberger et al. 2015 a trainable skip connection:
funees (x) = Gt x x + we x y x))
where yp assumed to tend to (resp. 1) as its argument tends to +oo (resp -ox). Sign, tanh and erf are simple examples of such activation function. The trainable parameters are therefore a E IR w, E Idd, Note that (9) is a special case of the architecture studied in Cui & Zdeborová (2023) It differs from the very similar network considered in Shah et al. (2023) in that it covers slightly broader range of activation functions (Shah et al.) (023) address the case ç = tanh), and in that the skip connection istrainable -rather than fixed- Since we consider the case where separate
network is trained at every time step, the empirical risk (5) decouples over the time index t. The parameters we Ct of the DAE (9) should therefore minimize
R, ct) = > 1ifer,nnr (xx"")
12 + l|w,l1?
(10)
where for generality we also allowed for the presence of a ly regularization of strength X. We remind
that a." = (t) + ((t) with { )p= (resp. {** in p n training samples independently drawn
from the target density P1 (8) (resp. the base density Pa (O, La)), collected in the training set D.
Asymptotic limit We consider in this manuscript the asymptotic limit d oo, with ||l| 2/1,/, (1( For definiteness, in the following, we set [|x]]] /a = 1. Note that Cui & Zdeborová (0023) consider the different limit || = @a(1). Shah et al. (202)) on the other hand address larger range of asymptotic limits, including the present one, but does not provide tight characterizations, nor an analysis of the generative process.
4

---

Published as a conference paper at ICLR 2024
2 LEARNING
In this section, we first provide sharp closed-form characterizations of the minimizers ćt, tü of the objective R, 10). The next section discusses how these formulae can be leveraged to access a tight characterization of the associated flow.
Result 2.1. (Sharp characterization of minimizers of (10) For any given activation yp satisfying ((x)) 2-+00 +1 and any t E [0, 1], in the limit d + oc, n, ll /dlli// /
Oa(1), the skip connection
strength é minimizing (70 is given by
è a (t) (())(((1 + 22) + (n - 1)02)
(11)
(())((((+++)))) (n-
Furthermore, the learnt weight vector w, is asymptotically contained in span([emm.. E) (in the sense that its projection on the orthogonal space span(pemp. E) has asymptotically vanishing norm). where
1 Σ
Ç =
Σ
8 xp.
Memp.
(12)
In other words. Hemp. is the empirical mean of the training samples. We remind that s +1 was defined below 88 and indicates the cluster the p-th sample æx belongs to. The components of w along each of these three vectors described by the summary statistics
m Pemp W
i
(13)
d(1++%
nd
which concentrate as oo the quantities characterized by the closed-form forrfffmflfformla
m X4n (tt)()tt a10)(2---1)) +
(14)
(10(1+00000)0)000000000000000001 + i Atn (())
E -((t)
The derivation of Result 2.1 is detailed in Ap- pendix and involves heuristic partition fune- tion computation, borrowing ideas from statis- tical physics. The theoretical predictions for the skip connection strength & and the component mi, of of the weight vector ıZ; are plotted as solid lines in Fig.[] and display good agreement with numerical simulations, corresponding to training the DAE (9) on the risk (TO) using the
1.-
0.6
of
Pytorch (Paszke al. 2019) implementation of the Adam optimizer Kingma & Ba 2014).
0.0
0.0
0.2
0.6
0.8
1.0
A notable consequence of (13) is that the weight vector vū is contained at all times in the two- dimensional subspace spanned by the empirical cluster mean Memp. and the vectors & (12) in other words, the learnt weights align some tent with the empirical mean, but still possess non-zero component along &. which is orthogo- nal thereto. & subsumes the aggregated effect of the base vectors (a() カ= used the train set. Rather remarkably, the training samples thus úZ; through (12. Since the vector
Figure 1: 1,0 0.9,
=.1,A=0.1,a()) =
t, (,) 1 tanh. Solid lines: theoretical predic- tions Result 2.T squared norm the DAE weight vec- tor ||o, (red), skip connection strength & (blue) co- sine similarity between the weight vector wh and the tar- get cluster mean μe, w; w; w, μl/ |/|||||||0,,, [|/|://:|| (green). components m,, of of w. along the vectors femp.. & (purple. pink, orange). Dots: numerical simulations in dimension d = 5 10 corresponding to training the DAF ( on the risk (1O) using the Pytorch imple- mentation of full-batch Adam, with learning rate 0.0001 over 4 x 10" epochs and weight decay X = 0.1. The experimental points correspond single instance of the model.
only enter in the characterization
the form of simple sums E is associated to the training samples, the fact that the learnt vector wù has non-zero compo- nents along hence signals a form of overfitting and memorization. Interestingly, Fig. shoow that the extent of this overfitting is non-monotonic
5

---

Published as a conference paper at ICLR 2024
in time, as l first increases then decreases. Finally, note that this effeet is expected mitigated
as the number of training samples n increases. From (14). for large n, my = On (1) while the components of is suppressed as e (2/x). These scalings are further elaborated upon in Remark B.3 in Appendix B. Finally, Result 2.1 and equation (6) can be straightforwardly combined to yield sharp characterization of the learnt estimate of the velocity field b (I. This characterization can be in turn leveraged to build tight description of the generative flow 7 This is the object of the following section.
GENERATIVE PROCESS
While Corollary 2.1 together with the definition (6). provides a concise characterization of the velocity field b, the sampling problem 7 remains formulated high-dimensional, and therefore hard to analyze, transport process. The following result shows that the dynamics of sample X, following the differential equation 7 can nevertheless be succinetly tracked using finite number of scalar summary statistics.
Result 3.1. (Summary statisties) Let X, be a solution of the ordinary differential equation 7 with initial condition Xo. For a given the projection of X₁ span( & is characterized by the summary statistics
X, pemp.
X nd
M. d(1 of)+))))
Qf
(15)
With probability asymptotically 1/2 the summary statistics Me, Qi 15 concentrate for all to the solution of the ordinary differential equations
at M₁ (ă(t)a + a(t)
a(I) (1 c,,((t)) M, + (ß(t) <(t)
a(3) B(t)
(16)
- = (B(t))= + a(0) (1 (,3(1)) Qf + (8(1) (en) 8(1) qi
= a(r)
with initial condition My = Q =0, 0, and with probability asymptotically 1/2 they concentrate to minus the solution of (16). Furthermore, the orthogonal component X, E span (emmp. E) obeys the simple linear differential equation
X (B(1)a+ + 0 -1- (--((1))))) X dl
d
(17)
Finally, the statistic Q = |x,,11 fa is given with high probability by
()(000 (100000)))000000000000000000001)
Q: = M, ((f))
(18)
A heuristic derivation of Result 3.1 is provided in Appendix B Taking a closer look at 16
might seem at first from equations (16) that there is singularity for t = 1 since a(1) 0 in the denominator. Remark however that both 1 - (t)) 11 and m, [4] are actually proportional to a(1)², and therefore (16) is in fact also well defined for = I. In practice, the numerical implementation of generative flow like 7 often involves a discretization thereof, given a discretization scheme [
1-V of [0, 1], where to and ty 1:
Xta+1 Xt, 6(X,, 4tt--11 tx).
(19)
The evolution of the summary statistics introduced in Result 3.1 Result3.1]aa be rephrased in more actionable form to track the discretized flow (19).
Remark 3.2. (Summary statistics for the discrete flow) Let {X sq J be a solution of the discretized
learnt flow 77 for an arbitrary discretization scheme VN (0, 1), where to 0 and in 1, with initial condition Xto Pn- The summary statistics introduced in Result 3.] are then equal to the solutions of the recursions
M = =Me, + 88x (B(1x)0r + o(te) (1 - Er B(t))) Mc. + 86x (3(1x) ditel ((te))((xx)
3(xx) mi (1 crg (1, 9(xx) Qi, + 8th (Ś(tz) alle) 3(tx)) gf.
a(te)
Qf Qfan -Qf, Stx (B(t))0ez
d(tu)
a(ta)
+ a(ta)
(20)

---

Published as a conference paper at ICLR 2024
23
0.3
X:AM
0.3
[- 1.0
R-
-01
n-
-02
0.0 0.4 0.6
0.8
1.0
-10
-0.5
0.0
0.5
1.0
D8 Hemp Millthrp lIVe)
Figure 2: In all three plots, == a.1, = 1,3(L 1, sign. (left) 1.5, 8. Temporal
evolution of the summary statistics Ms, Qi N,, x 15). Solid lines correspond to the theoretical prediction of (13) in Result [3.]] while dashed lines correspond to numerical simulations of the generative model, by discretizing the differential equation (7) with step size 0.01, and training separate DAE each time step using Adam with learning rate 0.01 for 2000 epochs. All experiments were conducted in dimension d = 5000, and single run is represented. (middle) = 2, 16. Projection of the distribution of X 7 (pamp E), transported by the velocity field b 6 learnt from data. The point clouds correspond to numerical simulations. The dashed line corresponds to the thcoretical prediction of the means of the cluster, as given by equation [T6) of Result 3.1 The target Gaussian mixture PI represented in red. The base zero-mean Gaussian density po (dark blue) is split by the flow (7) into clusters, which approach the target clusters (red) as time accrues. (right) a = PCA visualization
the generated density #r. by training the generative model on v. samples, for n E (4. 8, 16, 32, 61). Point clouds represent numerical simulations of the generative model. Crosses represent the theoretical predictions of Result 3.1 for the means of the clusters of Ộr. as given by equation (I6] of R. t 1. As the number of training samples increases, the generated clusters of øi approach the target clusters of p1. represented in red.
with probability 1/2, and the opposite theoreaf with probability 1/2. In equation (00 the initial conditions are understood as M. = O = O, and we have denoted öt, = +k++1 ts for clarity. Furthermore, the orthogonal component X E span( +) obeys the simple linear recursion
[1 + 8tx (B(t))C + a(tx) (1 cep B(ta)) X a(tx)
X X
(21)
Finally, the statistic Q+, |xx 11/a is given with high probability by
Q = M2 (1++2%) ((f.) + II
1+ ( 207 (1-2,, O(oe)) ((te)))0i 604
(22)
8=
Equations 20p.40 21 and 22) of Remark 3.2 are consistent discretizations of the continuous flows (7) and (18) of Result 3.1 respectively, and converge thereto the limit of small discretiza- tion steps max ot, + 0. A derivation of Remark 3.2 is detailed in Appendix B An important consequence of Result 3.T is that the transport of sample Xo ~ Po by 77) factorizes into the low-dimensional deterministic evolution of its projection on the low-rank subspace span span (Hemp.,) C)
as tracked by the two summary statistics Mr. Qf and the simple linear dynamics of its projection on the orthogonal space span(Nemp) - Result 3.1 thus reduces the high sigh-dimensional flow 7 into a set two scalar ordinary differential equations (16) and simple homogeneous línear differential equation (17). The theoretical predictions of Result (3.1 and Remark (3.2 for the summary statistics Mc. of, Qi are plotted in Fig 2. and display convincing agreement with numerical simulations, corre- sponding to discretizing the flow 7 in N 100 time steps, and training a separate network for each step described in Section [ A PCA visualization of the flow is further provided in Fig (middle).
one is now in a position to characterize the generated distribution Å1, which is the density effectively sampled by the generative model. In particular, Result 3.1 1 lestablishes that the distribution Pr is Gaussian over span (Memp.. E) -ince X is Gaussian and the flow is linear-, while the density in span(Hemp. . {) concentrates along the vector described by the components (16). The density أr is thus described by mixture of two
Leveraging the simple characterization of Result 3.I
7

---

Published as a conference paper at ICLR 2024
10-1
1-64u
10-2
10
10
1/
sign. Cosine asimilarity (left) and mean squared distance (right) between the of the generated mixture و and the mean μ of the target density Pl. function of the number of training samples n, for various variances Corollary Crosses represent numerical simulations the generative model, by discretizing equation with step size ót = 0.01. and training a separate DAE for each time step using the Pytorch implementation of the full-batch Adam optimizer, with learning rate 0.04 and weight decay X = 0.1 for 6000 epochs. All experiments were conducted dimension d = 5 10*, and single run represented. Dashed lines indicate the performance the Bayes-optimal estimator μ theoretically characterized in Remark 4.1 Dots indicate the performance the PCA estimator, which is found as in Cui & Zdeborová (2023) to yield performances nearly identical to the Bayes-optimal estimator.
Figure 3: (t) = (t) 6,
pr. Solid lines represent the theoretical characterization
differential
clusters, Gaussian along sharp characterization of the squared distance between the mean j of the generated density Å and the true mean μ of the target density P1-
directions, centered around + The following corollary provides
Corollary 3.3. (Mean squared error of the mean estimate) Let jı be the cluster mean of the density #i generated by the (continuous) learnt flow (7. In the asymptotic limit described by Result 2.1 the squared distance between and the true mean μs is given by
1 =lli= p1 M? + + 2M1,
(23)
with M1, Qi Qi being the solutions of the ordinary differential equations 16) evaluated time t = 1. Furthermore, the cosine similarity between ja and the true mean μ is given by
M1
(24)
VQ
Finally, both the Mean Squared Error (MSE) 1/alia i 11l2
23
and the cosine asimilarity 1 juL4
24 decay as O,, (2/v) for large number of samples n.
The heuristic derivation of Corollary
3.3 presented in Appendix A.1 The theoretical predictions
of the learning metrics 23 and (24) are plotted in Fig. 3 as a function of the number of samples, along with the corresponding numerical simulations, and display clear O,,)))) /x) decay, signalling the convergence of the generated density Ộ1 to the true target density P1 as the sample complexity accrues. PCA visualization of this convergence is further presented in Fig (right). Intuitively, this is because the DAE learns the empirical means up to a e,, component along &. and that the empirical means itself converges the true mean with rate O,,%))), While we focus on the MSE for conciseness, the rate of convergence in terms of a variant of the squared gaussian mixture Wasserstein distance (Delon & Desolneux 2020 Chen al. 2018) can similarly be derived to be O,, (2/n,, see Appendix
4 BAYES-OPTIMAL BASELINE
Corollary 3.3 completes the study of the performance of the DAE-parametrized generative model. It natural to wonder whether one can improve on the + (,,((a,, (1/x) rate that it achieves. useful baseline to compare with is the Bayes-optimal estimator fa* yielded by Bayesian inference when in addition to the dataset = {xx') 1==1, the form of the distribution (8) and the variance o are known, but not the mean μ -which for definiteness and without loss of generality will be assumed in this section to be have been drawn at random from (0, The following remark provides a tight characterization of the MSE achieved by this estimator.
8

---

Published as a conference paper at ICLR 2024
Remark 4.1. (Bayes-optimal estimator of the cluster mean) The Bayes-optimal estimator ji of p assuming knowledge of the functional form of the target density (8), the cluster variance o, and the training set D. is defined as the minimizer of the average squared error
arginfE = -=-N1010)
(D() -m1l2.
(25)
In the asymptotic limit of Result 2.7 the Bayes-optimal estimator ji (D) is parallel the empirical mean pemp. Its component m² Pemp, Ar (Di/at) + =2)/) concentrate asymptotically to
n
m²
(26)
o2
Finally, with high probability, the Bayes-optimal MSE reads
o²
allid - -0 n - 0²
(27)
In particular. 27 implies that the optimal MSE decays O,,()
Remark 4.1 whose derivation is detailed Appendix C thus establishes that the Bayes-optimal MSE decays as (2/v.) with the number of available training samples. Note that while the Bayes-optimal estimator is colinear to the empirical mean, it is differs therefrrom by a non-trivial multiplicative factor. On the other hand, the (%,) rate is intuitively due to the ('/))/) convergence of the empirical mean to the true mean. Contrasting to Corollary for the MSE associated to the mean fa of the density أi learnt by the generative model, it follows that the latter achieves the Bayes-optimal learning rate. The Bayes-optimal MSE (27) predicted by Remark 44 is plotted in dashed lines in Fig.3 alongside the MSE achieved by the generative model (see Corollary (.3). The common 1/n decay rate is also plotted in dashed black for comparison. Finally. we observe that the estimate of fe inferred by PCA, plotted as dots in Fig. leads to a cosine similarity which is very close to the Bayes-optimal one, echoing the findings Cui & Zdeborová (2023) in another asymptotic limit. We however stress an important distinction between the generative model analyzed in previous sections and the Bayes and PCA estimators dicussed in the present section. The generative model is tasked with estimating the full distribution P1 only from data, while being completely agnostic thereof. In contrast, PCA and Bayesian inference only offer estimate of the cluster mean, and require an exact oracle knowledge of its functional form (8) and the cluster variance σ. They do not, therefore, constitute generative models and are only discussed in the present section insightful baselines.
is a rather striking finding that the DAE succeeds in approximately sampling from p ( when trained on but n = (1) samples iinstead of simply generating back memorized training samples-, and further displays information-thhorrttiaalll lies in the fact that the architecture (9) is very close to the functional form of the exact velocity field b (I) as further detailed in Appendix B (see equation (67). and is therefore implicitly biased towards learning the latter while also not being expressive enough to too detrimentally overfit. A thorough exploration of this form of inductive bias for more complex architectures is an important and fascinating entreprise, which falls out of the scope of the present manuscript and left for future work.
optimal learning rates. The answer to this puzzle
CONCLUSION
We conduct tight end-to-end asympiotic analysis of estimating and sampling binary Gaussian mixture using a flow-based generative model, when the flow is parametrized by a shallow auto- encoder. We provide sharp closed-form characterizations for the trained weights of the network, the learnt velocity field, a number of summary statistics tracking the generative flow, and the distance between the mean of the generated mixture and the mean of the target mixture. The latter is found to display a e, (1/n) decay rate, where n is the number of samples, which is further shown to be the Bayes-optimal rate. In contrast to most studies of flow-based generative models in high dimensions, the learning and sampling processes are jointly and sharply analyzed in the present manuscript, which affords the possibility to explicitly investigate the effect of limited sample complexity the level of the generated density.
LEKNOWLEDGEMENT
We thank Michael Albergo, Nicholas Boffi, Joan Bruna, Arthur Jacot and Ahmed El Alaoui for insightful discussions. Part of this work was done during HC's visit in the Courant Institute in March

---

Published as a conference paper at ICLR 2024
2023. We acknowledge funding from the Swiss National Science Foundation grants OperaGOST (grant number 200390) and SMArtNet (grant number 212049). EVE supported by the National Science Foundation under awards DMR- 1420073, DMS-20125100 10, and Collaboration on Wave Turbulence, Grant No. 617006, and by a Vannevar Bush Faculty Fellowship.
by the Simons
REFERENCES
M. S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden, Stochastic interpolants: A unifying framework for flows and diffusions. ArXiv, 0bbs/2003.0899,, 2023.
Michael Albergo and Eric Bijden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571, 2022.
Joe Benton, George Deligiannidis, and Arnaud Doucet. Error bounds for flow matching methods. arXiv preprint arXiv:2305. 2023.
Giulio Biroli and Marc Mézard. Generative diffusion in very large dimensions. arXiv preprint arXiv:2306.035/1 2023.
Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto- encoders and langevin sampling. arXiv preprint arXiv:2002.00107, 2020.
Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. In International Conference on Machine Learning, pp. 4735-4763, PMLR, 2023a.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estima- tion and distribution recovery of diffusion models low-dimensional data. arXiv preprint arXiv:2302.07194, 2023b.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209.11215, 2022.
Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast, arXiv preprint arXiv:2305.11798 2023c.
Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degraddiin beyond linear diffusions: non-asymptotic analysis for ddim-type samplers. International Conference on Machine Learning, pp. 14622-448. PMLR, 2023d.
Yongxin Chen, Tryphon T Georgion, and Allen Tannenbaum. Optimal transport for gaussian mixture models. IEEE Access, 7:6269-6278, 2018.
Hugo Cui and Lenka Zdeborová. High digh--imensional asymptotics of denoising autoencoders. arXiv preprint arXiv:2305. 10te22305./1044 //041, 2023.
Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis.
arXiv preprint arXiv:2208.05314 2022.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34:17695-,7709, 2021.
Julie Delon and Agnes Desolneux. A wasserstein-type distance in the space of gaussian mixture models. SIAM Journal Imaging Sciences, 13(2)9936-970, 2020. 2020.
Bradley Efron. Tweedie's 's formula and selection bias. Journal of the American Statistical Association, 106:1609 16:2-00222 - 1614, 2011.
Davide Ghio, Yatin Dandi, Florent Krzakala, and Lenka Zdeborová. Sampling with flows, diffusion and autoregressive neural networks: A spin-glass perspective. arXiv preprint arXiv:2308. 14085, 2023.
10

---

Published as a conference paper at ICLR 2024
Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv, abs/2006. 1239, 2020.
Yukito lba. The nishimori line and bayesian statistics. Journal of Physics A: Mathematical and General, 32721)33755 1999.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. Advances in Neural Information Processing Systems, 555:26555-265777 2022.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 11141.96800 2014.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. Advances in Neural Information Processing Systems, 35:22870-2223 2022.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of based generative modeling for general data distributions. In International Conference Algorithmic Learning Theory, pp. 946-985. PMLR, 2023.
Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv preprint arXiv:2306.0925/ 2023.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Ham-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022.
Song Mei and Yuchen Wu. Deep networks denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309, 11420, 2023.
Hidetoshi Nishimori. Statistical physics spin glasses and information processing: introduction. Number 1 111. Clarendon Press, 2001.
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. arXiv preprint 2xxXi:2303331861 2023.
Victor M Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual review of statistics and its application, 6:405-431, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, perfomerform deep learning library. Advances in neural information processing systems, 32. 2019.
Jakiw Pidstrigach. Score-basee generative models detect manifolds. Advances in Neural Information Processing Systems, 35:35852-35865, 2022.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint arXiv:2204.06125. 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 11684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCA. 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 231-241. Springer, 2015.
11

---

Published as a conference paper at ICLR 2024
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494, 2022.
Kulin Shah, Sitan Chen, and Adam Klivans. Learning mixtures of gaussians using the ddpm objective. arXiv preprint arXiv:2307.01175 2023.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song. Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. 2020.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 2377716611-167, 2011.
Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, Pierre-Antoine Manzagol, and Léon Bottou. Stacked denoising autoencoders: Learning useful representations in deep network with local denoising criterion. Journal of machine learning research, 11(12), 2010.
Andre Wibisono and Kaylee Yingxi Yang. Convergence in kl divergence of the inexact langevin algorithm with application to score-based generative models. arXiv preprint 1.01512, 2022.
Hui Yuan, Kaixuan Huang, Chengzhuo Ni, Minshuo Chen, and Mengdi Wang. Reward-directed conditional diffusion: Provable distribution estimation and reward improvement. arXiv preprint arXiv:2307007055, 2023.
12

---

Published as a conference paper at ICLR 2024
A DERIVATION OF RESULT
In this appendix, we detail the heuristic derivation of Result (2.1). which provides sharp asymptotic
characterization of the parameters C,, Wy of the DAE minimizing the time / risk R. (10). In the
following, the time index E [0, 1] is considered fixed.
In order to characterize observables depending on the minimizers ét, iq of the risk R. 10 observe that for any test function ((c,, wt):
lim 1 o(C, ai,) 7-20 2 dwdck(c, w)e
( RR,(w,)))
(28)
where the normalization Z is
2 (D) ( de dwe dwe 11 (8())=) (a0))
(29)
We emphasized the dependence on the train set = { xt 2f1 moment generating function, and integrals of the form (28) deduced therefrom. In the following, therefore seek to establish an asymptotic characterization of ZDD).
J In ((D) can then be studied as
An important observation lies in the fact that the argument w x of the activation şo of the DAE expected in high dimensions d + oo to be very large. In particular, we shall self-consistently establish that it is more preciscly scaling like Oa(d). As result, only the asymptotic behaviour in too of yp matters, and by assumption fow" x) R sign(w) (w z) asymptotically. We shall therefore self-consistently take &p sign in the following.
A.1 COMPUTATION OF THE PARTITION FUNCTION
In the following, for clarity, we use the decomposition of s" + z", introduced below (8) in the main text, with s" E +1] and 2 ~ NN(0, o2 Under these notations, the partition function reads:
Z(D) / de dwe 3 E غلاي (u (1000)
7al E.r vix((> a=1
m " co()))))
12 u|l
an
1) p a(t)) +n a(0))
(30)
Note that we benignly introduced 1/u factor inside the sign function sign. One is now in position to introduce the overlaps
|
w t
20 d
W z d
w
s
(31)
Note that because of n (11) there is finite number these such overlaps. Besides, note that our starting assumption that the argument w of the activation yo is Oa(d) translates into the fact that all these order parameters should be (,(1), which we shall self-consistently show to be indeed the case. The partition function then reads
$49+ E (0 of +( nmm
Z(D) = =f dc dmdm dqdğ II #=1
ddg;; di" da
- -2 [||l|| -3W11
f| - mut En 18000" -x--)))"")))
dwe
-2d 2/ (anion cottoooottttttttccooeee - (1tttet(((ttmmmm=)))))///o oat
(32)
13

---

Published as a conference paper at ICLR 2024
Therefore
((D) = /
de dmdm dqdq II dd da day dap
4 diimmmd E (f +0 + श à
mut Σ (2 o6
e
31 I sixu ((C)( of())() sisu (01000 esco) ca(t)t]]
pill
(000 1m())+)))
(33)
The last term in the first exponent can be further simplified as
1
mp +
Σ (9t st
x² + 4th 24th
= m +
(GE) + (2
d
==1
#=1
1 26
+ (ity ste z4
+ 2
qe gu
d
d
#=
Σ
16 (oB) 4.7,00 (am) g"
To""
+0 4.
1,K=1
m +
(att) 3 (22) (LWa),
(34)
with the last line holding with high probability, using the fact that since z, zo are two independently drawn standard Gaussian vectors za/d = O2(1/va) with high probability. Finally,
B(D) / de dmdr dqdq II dāt datt dat da
Go drim (o of + マイイ m²+ E [
24 $ | Jisinn
( / u(()))²² 2 (scus ot racoof) pocosmid cattu]]]
e
(0402)0--000
In(*>+0)
(35)
Since all the terms in the exponent of the integrand scale like d. in the asymptotic limit d + oo the integral can be computed using a Laplace approximation.
A.2 SAMPLE-SYMMETEE ANSATZ
The partition function is given by taking the saddle point in 35 problem over 4n + 5 variables. Note that since n = Oa(1), this is a low dimensional thus priori tractable optimization problem, but which nevertheless remains cumbersome. However, the symmetries of the problem make it possible to determine the form of the maximizer, and thus drastically simplify the optimization problem. Note that indeed, asymptotically, the vectors 1,, { 2" }== involved in the definition of the overlaps m, {yE ) 1==1 (31) are all mutually asymptotically orthogonal - i.c. they have vanishing cosine similarity. Therefore, the parameters m, (g 2h 1n 1p= can be considered as independent variables. Since furthermore all the samples play interchangeable roles in high dimensions in that all data points asymptotically at the same angle with the cluster mean μr, which is the only relevant direction of the problem-, one can look for the saddle point assuming the symmetric ansatz
This involves maximization
Vp, o' = gg,
8 ge
(36)
Aμ, dh = Og
d/ = 4n
(37)
14

---

Published as a conference paper at ICLR 2024
This ansatz is further validated in numerical experiments, when training DAE with the Pytorch implementation of the Adam optimizer. Under this ansatz, the partition function reduces to
2(D) = dc dmddmn dqdıdũ dq, dage (nn---mmmmmddeeeeaaaa शहररकक (1t +n( +0
2 to) 2 (aiuulpttttttt)taaaaaaaa
1n())+00)
(38)
Note that the exponent is now independent of the dataset D. In other words, in the regime d + 00, 11 (1))), the log partition function concentrates with respect to the randomness associated with the sampling of the training set. The effective action (log partition function) therefore reads
In (D)
extr
qq + nim + n(9e + (g).)) + 2(72 + 4)
mm² 2 + + + o 2?)0²)
ch
n (1
2
y (11+00 (())) +a(()) - 1n(>> +0)
(39)
This expression has to be extremized with respect to c, Q, q, mi m, In In. in the y - co limit. Rescaling the conjugate variables
74 t 9.
râ. + ân.E.
ym m
(40)
the action becomes, in the 7 + oo limit (changing for readability the conjugates m, ân. +
ri, =)
1
In Z(D)
extr - mum n(āege + (p + c.q.q.m.m in.5 4n
[m n( + +
2(A
n
2
(41)
A.3 SADDLE-POINT EQUATIONS
The extremization of (DD) can be alternatively written as zero- pero-gradient equations on each of the parameters the extremization is carried over, yielding
v = (1) + (n) + a(t)ge
2 2 0 q- = (x+-))*
i = n
M sign(v) (1))(())))))))))))))))))))
X+1
(42)
95
Gn
9g X+Q
& -ca(() (iggvnvs)))))))))))))))))))){))
9n X a
C ( +1000"1
Note the identity
q= m2
n( + 22/0²)
(43)
which follows from the asymptotic orthogonality of the vectors and the Pythagorean theorem. This implics in particular that the square norm of w. measure by is only the sum of its projections along j (corresponding to m) & (corresponding to qe) and 17 (g,). Therefore, the norm of the orthogonal projection of á with respect to span(4, [) is asymptotically vanishing. In other words, i is asymptotically contained in span n7, E).
Remark that the symmetry m, m. 9n 97.5 unchanged, meaning that if m, m, 4, 9n is solution to the saddle- point equations, so is
m, - -n. Gv. leaves equations
42
m, m, 9n.5 In.. This is due to the symmetry between the clusters in the target density ( as
15

---

Published as a conference paper at ICLR 2024
1/ -
yields the same model. As convention, can thus suppose without loss of generality then simplifies to
V/ 2 in
nh + 04
i n
q
(x+-))
m = n - c3())
X
in ri
(44)
1 X+0
9g a(ttc
4n
(Ita (80)) Ho
(())==(()P((oo))
The skip connection strength thus satisfies the self-consistent equation
(1 (1+ -
(45)
a(t) + +(t)
which can be solved as
(())((( (n 1)02 a(t)2 - ()) + x²) + (n 1)02)
(46)
which recovers equation oag, 9n yields
of Result 2.1 Plugging this expression back to (44). and redefining
(2+x)(++=11)
m4 X+n (02X44400)4400((((40P00
9. o²
((()(()=-))
(47)
(00m ((n
q: A+n a(())(t)(A(( + (n
We have added subscripts emphasize the dependence on the time index t. Note that for > > O, which is self-consistent. For f = 0, 0 and the sign function in equation 422 becomes ill-
defined, signalling that the extremum of equation (4I) ceases be a critical point (i.e. differentiable).
However, one expects the extremum to still be given by the 0 limit of equation (42) as there is
a priori no singularity in the learning problem for = 0. This remark, together with recovers
equation (14) from Result 2.
A.4 METRICS
Result 2.1 provides a tight characterization of the skip conneciton strength è and of the vector dh.- The performance of the trained DAE fe,, (9) as a denoiser can be further quantified with number of metrics, for which we also provide sharp asymptotic characterizations below, for completeness.
Result A.1. (MSE) The test MSE of the learnt denoiser fer ee, is defined as the test error associated to the risk R₂ 10
MSC: Emp Ie (a(t)xo + B(t))1)) 21122
(48)
In the same asymptotic limit as Result 2.1 in the main text, this metric is sharply characterized by the closed-fi foomed-formfa formula
msc mi? + n=((g))2 (o²))0²)
2(1 - 2.3(t))m (1 -1- ((-223()))(11 + + 22a(t)² (49)
where é,, me; o qi were defined in Result 2.1 Furthermore, the MSE (48) lower-bounded by the oracle MSE
mso
Eaiepi Ilfi (o()))) + ((()m)))))))))))))))))())))))( x 211
(50)
where the oracle denoiser follows from an application of Tweedie's 's formula Efron (017)) Albergo al. (2023)
3(1)o2
a(1)
B(1)
f7 + oz a<t)2 3 x tanh
(51)
Finally, the oracle MSE msct admits the following asymptotic characterization:
mse;"- = aa a(t)))2a(t)² + +0)(1a
(52)
(043())+aa())))
16

---

Published as a conference paper at ICLR 2024
0.25
1.05
mM
2
1.00
0-55
n-
n 20
0.10
n=
n-3
0.05
42.
0.6 0.8 0.65
sin((t))))//////)))/ Solid lines: theoretical predictions for the MSE of Result (A.I and the cosine similarity of Result Result(A.2 Different colors correspond to different number of samples Dots: numerical simulations, corresponding to training the DAE ( on the risk (TO) using the Pytorch implementation
Figure 4: = 0.3, 0.1.
full-batch Adam, with learning rate 0.01 over 2000 epochs and weight decay = 0.1. The experimental points correspond to single instance of the model, and were collected in dimension 500. the left plot, the dashed line represent the oracle baseline 52
Derivation of Result
A.7
begin by detailing the characterizaiton of the DAE MSE (49):
mse; -BB 20 [||=01 [c+ x (3(1)z + a(l)ro)) + w, sign ((((()())) + aa())))))]
mi n n(()))
+
= m n ((g)) 20²) - 2(1 - (31() + (1 +1+)))
(53)
which recovers (49) Result A.I
51
follows directly from an application of Tweedie's formula
Efron (2011). The associated MSE
52
can be derived as
a11)a(1 + 02) o (1) a(1)))
a(1)1
A(e)²
msc*
sign
(o² + a(t)²)²
(o²
a(t)))
0²3(t)2 + a(t)²
2a(t)²
8(0)²
a(t)
023(1)² a1)3
sign
023(11² + a
0² -2 o
(()
(o² + (())) = al1)² + t (1())((11 - -a(())
a(t)a
(02B(1)2 + (1)2))
(o2 + a
(54)
which concludes the derivation of Result Result A.2. The cosine similarity wrLH = o7 ///a Illlell admits the asymptotic characterization
w 4p
(55)
V m + n(((f)) + 2g2
where mi, of oi are characterized in Result
2.1
Result Aollow directly from the definition of the summary statistics 1
These metrics are plotted in Fig. and contrasted to numerical simulations, corresponding to training the network (9) using the Pytorch implementation of full-batch Adam.
B DERIVATION OF RESULT
In this Appendix, we detail the heuristic derivation of Result a sample follows the transport 7
3.1 Given an initial condition X, Po-
d X, ((()) + a(l) (1 c, X, (((1)
a(!) a(t) (()) ii sign(w, Xt) (56)
dt
a(t)
17

---

Published as a conference paper at ICLR 2024
driven by the learnt velocity field b 66. This follows from Result 2.1 and 7 with M, &. 17,
Taking scalar products
x, at d d E
(B(1)) a(t)
at)) (1 &,((1))
x d + (B(1)
a(t)
a(3)
3(1)
(ign( X,) 15
d
(B())c= + a(t)
(1 a,((1))
x
(B())
a(Q
(())
sign (0 Xx)
(57)
di nd
ato
nd
a(t)
nd
d X n di adoa
(B(t)), all
a(t) (1 C,((t))
X, 7 ado2
+ (B(t)
a(E) ((t)
a())
sign 20 X1)
ndo2
reasonable assume the sign sign( X,) stays constant during the transport, and therefore takes value +1 with equal probability 1/2, according to the initial condition Xo. This is an heuristic assumption which is further confirmed numerically, Finally, plugging the definitions (16) and [3) in (57). one reaches
at Mr ((1) + <(t)
a(1) (1 e,3(1)) M.+ + (B(1) a())
B(1) m
a()
de O = ((t)) + 600)
a()) (1 a.((t))
Of + (Ó(1)
a(0)
3(t)
qi
(58)
at
Q7 (B(t)
400 (1 -.,(())) Q7 (B(t)
a(t) B(t))
al0)
a())
which recovers equation
(15) of Result 3.1 Noting that dù span(p, &, 1) (see Result 2.1 the differential equation (56) becomes, for the orthogonal component X, E span E,
d X, (B(t)
a(t) a(t)
(1- &i((t)) (1-))())))) x,
(59)
dt
3.1 This can be explicitly solved as which recovers 177 of Result
(á(0e+ 409 a (1-))
X =Xoet
(60)
Finally,
Q: |X
M² n( no mQ²) + |X, 12
2 ()(( (1 (1-))())))
M + (Qt))² (Q7) +
(61)
which concludes the derivation Result
B. DERIVATION REMARK
The derivation of Remark 3.2 follows identical steps, building the observation that the discretized flow 19 is explicitly expressed as
Xx, + 8la (8(l))e
a(tx) Gr2 B((x)) X
Xtat1
a(tx))
á(tx) B((x) vr, sign sign( Xex)
+ 8tu (8(tz)
(62)
al1x))
Taking overlaps with μ4, n yields
X = +5t ((t ). +
(1 Esg )(tx))
T
6lte) +5ta (b((z)
a(t2)) sign X mi5g)
d
o((k)
Xx 5 TT tōtz ((1x)) + [((g)]
altss
[1 -2 (((x))
Xi
+ōl+ ($(1))) -(() B((x)) sign(u.), X
nd Xt++1
nd " tee
ndo
_Tx tote (p(l))0r
a(t))
d(ta)
(1- A(tx)))
+ (B(tx) ac0e) p(tx) sign(w, X X
ndo
(63)
Like the continuous case, one makes the heuristic assumption that sign( [W]] X.,) stays constant along the flow, taking value + with equal probability, depending the initial condition. Doing so
18

---

Published as a conference paper at ICLR 2024
yields
Mani = 2 + +8x (g(tz)0 + 6(te)
a(t)) (1 En2 ((te))
Mr. + 8th (š(tz) d(ee)
alta)
3(tx))
mix
Phar =Qi, öt, (B(t))(:2 + d(ta) (1 - Er2 (tx))
2 + gt, (Š(tz)
d(le)
g(tx)
qi
a(ta)
asta)
Q 0t + 5tx ((tt))ee, + a(ta)
(1 ((xx))
2 + 8t ($(tx)
a(te) adlt))
B(ta))
It
a(t2)
(64)
which recovers equation (20 Equation (1 follows from equation (62) and the fact that úi span(μ, E, 7), see Result 2.1
This recursion can be explicitly solved as
X
X II (1+ (0)) (1-2, 8tz)
(65)
E=
Using the fact that || with high probability and the definition of the summary statistics Q finally yields equation (22)
B.2 DERIVATION OF COROLLARY
As implied by Result 3.1 the mean of the generated mixture is contained in span(ı, E, 7) and characterized by the summary statistics M1, Q7, Qf time Furthermore
1
_ 2 1 d +1
d 1 = =
= MM + n(Q))2 + no (Q)) +1.
(66)
This recovers
23 Equation
24 follows from the definition of the cosine similarity.
The derivation of the O, (2/n) decay of this distance require more work. The first step lies in the analysis of the exact flow (T)
Remark B.1. (exact velocity field) For the target density Pr b is given by Ejron 2017 et al. 2023)
Albergo
(B()) ((0))))
a(t)
a(t)o² ((()))++(1)000))
a(t) al1)2 3 x tanh(u' x)
a(t)
b(x, f)
a(l)
(67)
The formula 667 follows from an application of Tweedie's formula Efron (011 for the the density Note that with high probability for x ~ Po- or for any x such that fu >> I,
tanh( x) sign(u x) + oa(1).
(68)
One is now in position to characterize the exact flow TT.
Corollary B.2. (Summary statistics for the exact floctflow) Let X; be solution of the exact flow from initialization Xū Po- Consider the summary statistic
M 11
= Xi d
Xi
(69)
Asymptotically, Mi is equal with probability 1/2 the solution of the differential equation
di Mr = (B(t)c++ a(t) (0 M; + -0
d
all) ((0)
a(t)
a(t)²
(70)
a(1)2+(11)
and with probability 1/2 to the opposite thereof. We have introduced
=(t)
G
a(1))+
(71)
Corollary B.2 follows from equation 67 using a derivation identical to that Result 3.1 presented
in Appendix B provided the heuristic assumption is made that the tanh can always be approximated by a sign (68) along the flow. To show that the learnt flow 3.1 converges to the exact flow, observe the following scalings:
19

---

Published as a conference paper at ICLR 2024
Remark B.3. Let t > 0 and me. 2i 9: be defined by result 2.1 Then
a(t)² mr + O,(%), | - 10--al= - Ox(h)), of =x
== ='==x(''))
(72)
These observations immediately imply the following asymptotics, characterizing the difference between the learnt flow (7) and the exact flow (
Corollary B.4. (Convergence of the learnt flow) Let X (resp. X1) be solution of the exact exactflow
17) (resp. learnt flow (7)). from a common initialization Xo Po. Define the following sunmary statistics:
4 =" X7):
& nd (x xf);
a == ndo ( (X, -x))
(73)
Then with high probability these statistics obey the differential equations
d (š(t)c= + a(t)
a(t) (1 - c,B(t))
&: + On(+m)
är t
dt de ((())c + a(t)
(1 - c,3(t))
+x((k))
(74)
a(t)
dt 2 (3(1)c + a(t)
- c+3(())) C Om(1))
a(I)
from the initial condition Eg
m.C.. 0. Therefore at time = ef = Ox.%))
6 = On(n))
On(%))
(75)
Corollary follows from substracting the differential equations governing the learnt flow of Result 3.T and the true flow of Corollary [B.2] using the scaling derived in Remark B.3 Finally, noting that M by definition of the exact flow.
1 Jler +
l1a 1u12 = 2 + n(E)) + (2)) +O(1/a))
d
(76)
In the last line, we used Corollary B.4 This concludes the derivation of Corollary
3.3
Fig. (right)
gives PCA visualization of the convergence of the generated density År the target density Pr as the number available training samples n accrues.
C DERIVATION OF REMARK
In this appendix, we analyze the performance of the Bayes-optimal estimator of the cluster mean, defined as the estimator minimizing the average MSE knowing the train set {= x41) J==1 the clusters variance a,, but not the mean μ. This estimator yields the information theoretically minimal achievable MSE, and is known be given by the mean of the posterior distribution over the estimate w of the true mean /x:
PP(w(D, = -blll 11
w1l
1
2 1lx +w
+ ze
2
TT oz
|| 2
Σ
coch
(77)
where
2² n + o2
(78)
We remind the reader that the prior distribution over the cluster mean is supposed to be the standard Gaussian prior (O, ). In high dimensions, statistics associated to the posterior distribution (77
expected to concentrate. Again, it is useful to study the partition function (normalization) Z to
20

---

Published as a conference paper at ICLR 2024
some key summary statistics, which will in turn provide sharp characterization of the vector (D) extremizing the posterior (w)
The partition function reads
- 1e In cosh
=
dwe
= -f dqdgdmdm
2oo+a E wol
II
dap det
-4
Mun E si
-a Σ In nosh (/0 (m-
dwe
(79)
As in Appendix A, we have introduced the summary statistics
g W 2
w d
(80)
The integral of 79 be evaluated using Laplace approximation. Again, we assume the extremizer realized at the sample-symmetric point
V1 S p n, qt) 9u
A1 S p S n, qi qvy-
(81)
The partition function 779 then reduces to
de fo t +-mmi
-d E In cosh 1/d (.4 (m= (m+
/ dwe 411 --z ((mp-=((()+ we
dd Σ In cosh 1/d [, (m+
dgdādq, dG, dmdme tnn, dmmm 4mahh
62 (m +no +no?
I16-9
(82)
(1 + /²
Therefore dn. 9p m, must extremize the effective action
q4 2
6²
+-
+ ng mmn
In(1 + 0²0) 2(1 + 02
(m² n0²g²)
mm gyl,
(83)
o²
leading to
4n T vm o2
4n 020
2
nto²
(84)
n
m màd²
n
Refining g²qn < 4y so that
w T 1 4r ndo2
(8.5)
as in Remark 4.1 one finally reaches
1
71
4n 0²
(86)
72 + o2
Thus, remembering á (D) (w) (where the bracket notation denotes averages with respect to the posterior PD)
j (D)
w: 1
n
v) (T)
w TT
1
(87)
d
2
ndo²
ndo2
n +
2
21

---

Published as a conference paper at ICLR 2024
using the concentration of the bracketed quantities. Furthermore,
/ (D)l
1 [1))1
1 (w)
m.
(88)
d
d
We employed the Nishimori identity (Nishimori 2001 lba
Note further the identity:
A (D)Il
m²
no qi
mp 9,
(89)
which implies that the norm of ji (D) is equal to the norm of its projection in span( 77), which means that asymptotically the former contained in the latter. One is now in a position to derive the Bayes-optimal MSE of Remark 4.1 With high probability
1 2 (ll (D) - = 2m m
(90)
0²
This completes the derivation of Remark 4.
D FURTHER SETTINGS
D.I IMBALANCED CLUSTERS
0.8
0.2
0.0
0.2
0.6 0.8
1.0
= 1 (() = (. Imbalanced mixture with relative weights p = 0.24 and p 0.76. Solid lines: theoretical predictions of Result 2.1 squared norm of the DAE weight vector || (red), skip connection strength & (blue) cosine similarity between the weight vector eù, and the target cluster mean μ., rr,/p u |1 [| (green), components my of 4 of Wú along the vectors μt, &. 77 (purple, pink, orange). Dots: numerical simulations in d = 5 10a, corresponding to training the DAE 9 the risk (10) uisng the Pytorch implementation of full-batch Adam, with learning rate 0.001 over 20000 epochs and weight decay À 0.1. The experimental points correspond to single instance of the model.
Figure 5: n = 0 = 0.9, =
0.1,,(t)) 0.1,
In this appendix, address the case of binary homoscedastic but imbalanced mixture
P1 (N o2La) + (1 - p)
(91)
where P E (0, 1) controls the relative weights of the two clusters. The target density considered in main text thus corresponds to special case p 1/2.
It is immediate to verify that the derivations presented in Appendices A B carry through. In other
words, Result 2.1 Result 3.1 and Corollary 3.3 still exactly hold. Figure shows that the sharp characterization of Result 2 indeed still tightly captures the learning curves of DAE, trained on
imbalanced clusters, using the Pytorch Paazkee et al., 2019) implementation of the Adam (Kingmal
& Ba 2014) optimizer.
An important consequence of this observation that the generative model will generate a balanced density Ô1, failing capture the asymmetry of the target distribution P1. This echoes the findings of Biroli & Mézard (2023) in the related setting target ferromagnetic Curie- Weiss model, where they argue that the asymmetry of the ground state can only be learnt for >> d samples.
22

---

Published as a conference paper at ICLR 2024
D.2 DAE WITHOUT SKIP CONNECTION
In this appendix, we examine the importance of the skip connection in the DAE architecture ( More precisely, we consider the generative model parameterized by the DAE without skip connection
gv., (x) = wrsp(w, x))
(92)
where oo is an activation admitting horizontal asymptots at +1 (-1) in +oo (-oc). A tight charac- terization of the learnt weight úú can also straightlowardly be accessed, and is summarized in the following result, which is the equivalent of Result 2. for the DAE without skip connection (92)
Result D.I. (Sharp characterization of the trained weight of (92)) For any given activation satisfying (xx) = +1 and any E [0, 1], the limit d - JC, n. lelll /d, σ Oa(1), the learnt weight vector w; of the DAE without skip connection (92) trained on the loss 100 is asymptotically contained in span (u. n) (in the sense that its projection on the orthogonal space span(p, n) has asymptotically vanishing norm). The components of wh along each of these two vectors is given by the summary statistics
1 w.
B n
me
nide
(93)
which concentrate as d to the time-constant quantities characterized by the closed-form formulae
8"
(94)
+ m
X
0.8
0.4
0.8
Figure 6: n 4, 0.9, A = 0.1, (t) = 1 = t. Solid lines: theoretical predictions of Result 2.1 squared norm of the weight vector || 2 of the DAE without skip connection ( (red), skip connection strength & (blue) cosine similarity between the weight vector wū and the target cluster mean μe, r0:/U1 = u /4/ll |a|||| || (green), components m4 of qe of æh: along the
vectors μr, E. n (purple, pink, orange). Dots: numerical simulations in d 5 x 10", corresponding to training the DAE without skip connection (92) the risk (10) uisng the Pytorch implementation of full-batch Adam, with learning rate 0.001 over 20000 epochs and weight decay À 0.1. The experimental points correspond to single instance of the model.
Result D.I follows from a straightforward adapatation of the derivation of Result 2.1 presented in Appendix A In fact, it naturally corresponds to setting the skip connection strength to 0 in the expression of the log partition function equation (93) corresponds to the zero-gradient condition thereof.
striking consequence Result D.I is that asymptotically the trained vector à the DAE 92 does not depend on the time Fig. provides further support of this fact, the summary statisties measured in simulations training the DAE (92) using the Pytor orch implementation of full-batch Adam also observed to be constant in time, and furthermore to agree well with the theoretical prediction. As for the analysis presented in the main text, it is possible to track the generative flow with finite set of summary statisties. This is the object of the following result:
23

---

Published as a conference paper at ICLR 2024
Result D.2. (Summary statistics for the no-skip connection case) Le X₁ be solution of the ordinary differential equation (7) with initial condition Xo, when parametrized by the DAE without skip connection (92) For given t, the projection of Xt span(p, n) is characterized by the summary statistics
X d
X
M.
Q7
(95)
ndg2
With probability asymptotically 1/2 the summary statisties M,, Q7 solution of the ordinary differential equations
15) concentrate for all to the
d M, di
a(t)
M, (3(t)
a(1] a(0)
a(t)
a(t)
X+n 1
(96)
Q + (8()) a(t) (t)) d Q- ait) de
a(t)
att)
A+n
The derivation of Result D.2 can be made along the exact same lines as the one for Result 3.1 presented Appendix B An important observation is that the flows (96) are actually the exact flows corresponding to particular Gaussian mixture, as explicited in the following remark:
Remark D.3. (Generated density) The summary statistics evolution (96) are the same evolutions that would follow from the velocity field
(8(0)
a(t) alt) s(t) (f x tanh ())x)) + del))
b(x,t)=
x))
(97)
a<t)
where
11
x+n 7)
(98)
x+n
Comparing with equation (67 this is the exact velocity field associated to the singular Gaussian mixture
1 6(x - (1) 5(x + p))
(1
(99)
2
2
The generative model parameterized by the DAE without skip connection 92 thus learns singular density, which is sum of two Dirac atoms, centered at =p. It thus fails to generate a good approximation of the target (8) Note however that interestingly j remains a good approximation of the true mean μs, and actually converges thereto as - co. This is made more precise by the following result:
Remark D.4. (mse in the no-skip-connection case) Let j the cluster mean of the estimated density Pr. defined in Remark D.3 Then its squared distance the true mean μ is
à allá 1112= 2 n
(100)
n)
The minimum is achieved for o² and is equal to the Bayes MSE 4.1 In particular, this MSE decays ()))
Strikingly, the generative model parametrized by (2) manages to achieve the Bayes optimal O, (1/n) rate in terms of the estimation MSE over the cluster means, but completely fails to accurately estimate the true variance.
E A FULLY EXPRESSIVE MODEL MEMORIZES
In this appendix, show that the absence of memorization defined the ability of the generative model to generate new samples, and not just retrieve the training samples- is enabled by the network parametrization of the generative model. In fact. a fully expressive (flexible) model would in fact memorize the train set. Consider the network-parametrized minimization problem over the parameter space (0, rejo.1]
Σ Ezo Ilfo, (xc"' =xf11² dt.
(101)
24

---

Published as a conference paper at ICLR 2024
For ease of discussion, compared to equation the target Pl, we sample an infinity of noises 2y from the easy-to-sampa base Gaussian distribution Po. which corresponds to averaging over Zo in equation (TOI). Note that doing so, compared to the case where only one o is sampled for very z is expected to prevent the model from overfitting the noise and should only improve the performance. Now consider replacing the minimization equation (I01) by the minimization over the space of all denoising functions
consider the case where for each sample z of
R1
Exg [l - x"1l² dt
Ray~pi Eug [f t) - xi1 dt.
(102)
==1
In equation 1022 we denoted ýI the empirical distribution supported on the training samples
(11)) = Σ 8(æ1 - 24))
(103)
and remind that the distribution of the variable 22 follows from its definition 2t a0 =((t)=1- Finally, equation 102), instead of minimizing a function f. R + R for each E [0. 1]. we have without loss of generality rewritten fe(.) f(, (). The objective equation (102) be
as the limit of equation ( when the network infinitely flexible and can express any denoising function f. Comparing equation (102) to equation (4), it follows from Albergo al. (2023) that the minimizer of equation (T02) leads to flow mapping the base Gaussian distribution po to the density ýr equation (103), since equation (102) is the population objective for #1. Since Po and Ộ are both Gaussian mixtures (with the clusters of the latter being of vanishing variance). the corresponding velocity field is furthermore explicitly given Appendix of Alberoo al. (2023). Therefore, an infinitely expressive model minimizing the empirical risk equation (001) leads to generated density 71. In other words, only allows to generate samples x" from the training set, and the generated mixture has n clusters with 0 variance. In contrast, the DAE-parametrized model equation ( lears a bimodal mixture with non-zero variance.
F WASSERSETEIN DISTANCE
In this appendix, we derive a precise description of the generated distribution d and the target p1- We remind that the distribution of its projection in span(E, Memp.) follows the Gaussian distribution
(B()) 40 (1-2,()))
x+ ~ 0,e
1a-2
(104)
Observe that from Result 2.1
0²
& +(1)+ + 3(1222 (l/h))..
(105)
Thus.
a(t))(tt++ ((t) a(t) 3())202 + 0x(%/)) In ((t))))+)))))))1 0,(%) = Ino + +x(1/)).
2 = 2
(106)
Thus
0²=
(107)
We are now in a position to compute the Mixture Wasserstein distance between the generated density 9 and the target Pl. Because these are d- dimensional distributions, we normalize this distance by 1/a so as to have order metrics in the considered asymptotic limit - oo. Because of this, the precise distribution of the clusters of in the two dimensional space span(E, Memp.) does not matter, provided it does not involve moments diverging with d. We will show that this very reasonable assumption is indeed verified, after the computation of the distance.
25

---

Published as a conference paper at ICLR 2024
F.1 WASSERSTEIN DISTANCE
We aim at evaluating distance metric in the space of distributions to quantify the discrepancy between the true density P1 and the generated Å1. Many natural metrics (e.g. the KL divergence) are however intractable in our setting. We take inspiration from the Gaussian mixture Wasserstein distance 2V2, proposed in Delon Desolneux (2020) as variant of the W2 distance for Gaussian mixtures. Because P1 is mixture, but the clusters in span(E, Memp.) are only halves of Gaussian clusters, we define and employ a very similar metric for arbitrary mixtures:
Definition F.1. (mixture Wasserstein distance) Given two mixtures
Pijla and Σ Tays (with Jái M
not necessarily Gaussian densitics), the 22 distance is defined as
1 K
min MW2- JoERK (Pk p 1k (Ti
wky W (2)..
(108)
This is the same definition as Delon & Desolneux (2020), except that we allow for non-Gaussian Mi; Vi- Note that introduced without loss of generality normalization 1/4, since we comparing d- didimensional densities, and expect the distance to scale with d. With the normalization, the metric stays Oa(1) as d oo. In the present setting, this evaluates
MW (e,, 61) 1 (BT NN 0)), + WBBO
(109)
where introduced the densities أ 1/20T + 1/20T for the two clusters of Ộq centered at =p. We denote further decompose o p p into the product of the distribution pi in span(S, Hemp.) and the Gaussian d - 2 dimensional density p in span(E, Pemp.)' ). We can similarly decompose the target Gaussian density VV++,,, = (+ T² & N0,00 T 2). Using the the properties of Wasserstein distances between product measures Panaretos & Zemel (2019).
1 ( d
1 (N wo +I ( o.)) + W ON(
(110)
Note that since pf is Gaussian with variance o², the second term corresponds to the Wasserstein distance between two Gaussian distributions and read
1 d (et [ 66 Om(t/))
(111)
We now bound 4W3 +1 , ,²)). Note that the two densities are centered around å and μ. The discrepancy between these means will provide the dominant term in the distance. To see this, we introduce Vn (x) = 6(x A and v2(x)) = 5(x 11), two Diracs centered at the means, and upper-bound 4 +1l ,N(#tll,, o²)) using the triangular inequality
5 (t ob)) < 2 W +1l vn) + ( v2) + 2 N(tull (112)
The last term is asymptotically vanishing as (2(a/a) Under very mild assumption on p (which show are verified in the next subsection), the Wasserstein distance between two-dimensional distributions (2 m4) should be O2(1), so the first term also vanishes as Oa(1/a). The second term is equal to
2 W (v v2) = - -
On Oz//),
(113)
The derivation proceeds identically for the other pair of clusters f W ( 01,03) ))). Putting everything together, we reach
using result B.3
MM3(pr < (1/aa)
(114)
26

---

Published as a conference paper at ICLR 2024
Bayes optimal rate We briefly consider the mixture corresponding the bimodal mixture centered at the Bayes optimal mean estimator {D),,, and assuming perfect knowledge of the cluster covari- ances. Again, this is provided as an insightful baseline, and does not constitute generative model, since exact oracle knowledge of the form of #1 and a² is assumed. For the Bayes estimator:
2 MW Ai) d 1l ~||2= On (,))),,
(115)
using Result 4.1
We now briefly give two other examples for generative models differently parametrized, for which the generated density does not converge to the target p1 in MM2 distance.
Auto-encoder without skip connection As derived in D the generated density when the model is parametrized by a DAE without skip connection a degenerate mixture 1/20(- ()) + 1/20(- + p)).
which corresponds setting ē = 0 in the above derivation. Thus, it follows that
(pi, MW2 (n,, P1) (1),
(116)
i.e. without skip connection the generative model fails learn to generate the target mixture.
Fully expressive model discussed in Appendix
We now consider the case of a model which memorizes the train set, as In this case
(1
8(x
(117)
YL
which is a (degenerate) Gaussian mixture. It is straightforward to see that for any af. We[ 6(- =f))l 0² (1), and therefore
MW3 Pil 2 o² O1)
(118)
Thus أr is bounded away from the target f1.
We close the appendix by deriving the precise form of pi although the precise distribution in this two-dimensional space is asymptotically irrelevant for all the considered metrics, as we showed.
F.2 DISTRIBUTION IN span(E, Memp.)
We study in more detail the dynamics of x!! defined as the projection of X, in span(g, Pemp.). Since the initial Xo ~ Po is Gaussian, so is its projection x!! ~ (O Ig). Let us also call wi! the projection of wú in span(E, Hemp.). ). Projecting the dynamics equation 7 into span(E, Pemp.).
x (B(t) + a(1)
a(t) (1 2,8(t))
x!! + (tt)) a(1)
a(t)
3(t)
w!!
(119)
where the sign of the drift term is given by (X. (ig X₂) stays constant during the transport. This can be solved in closed form for 1 = as
I w!!)) Like in Appendix
B we assumed that
rttjtt + sign(X, X IT willee ((t)ds - 7())d (810)
a(0)3()
a(t) a(8) p(t)
x!
X, eo
wi dt (120)
where we used the shorthand
r(t) = [(t)a + a(t)
(1-
(121)
a(t)
The second term multiplied by the sign corresponds to á E span(E, Pemp.) ) as characterized by Result 3.1 The distribution of x!! follows from that of the Gaussian X, If x!! is in the half-space {o
E 22 x all 2 0} then x!! = xx!! + fu; If x is in the half-space {i E R² T will < 0} then x á. In other words, the distribution of X is mixture of two clusters, =... Each cluster corresponds to half a Gaussian cluster of variance d²1,, i.c. a Gaussian cluster cleft along
x1
hyperplane whose othogonal vector is wi as characterized by Result 2.1
27


Reviews:
[{"review_id": "5CqsOrwQLW", "reviewer": "Reviewer_v6CB", "summary": "The paper in detail analyses the case of training a stochastic interpolation model on few samples from a bimodal Gaussian mixture model. All theoretic predictions are accompanied by experiments.\n\nThe first result 2.1 characterizes the solution that is obtained when training on a bimodal Gaussian mixture. It shows that the weight vector is contained in the span of $\\mu$ (the displacement vector for the mixture components), $\\xi$ (the mean of all latent points), and $\\eta$ (the average offset of the data from the corresponding mean) and thus no other directions are relevant.\n\nThe second result 3.1 derives the resulting generative ODE and summarizes it in terms of the relevant space from result 2.1.\n\nThe third result 3.2 spells out the Euler integration of result 3.1.\n\nThe final result 3.3 shows that the distance and angle between the true $\\mu$ and the estimated $\\hat\\mu$ reduce by $\\Theta(1/n)$. This is the same as the Bayes optimal rate.", "questions": "1. How many dimensions are needed to span $\\hat w_t$ for all $t$? From a simple drawing of $\\eta, \\mu_+, \\mu_-$ I conjecture that two dimensions are enough (similar to first weakness).\n\n2. Why is $\\mu \\propto d$ a reasonable scaling? This seems like an unrealistic choice to me. In practice, data is often normalized say to a fixed range $[-1, 1]$ per dimension, so in order to obtain the scaling behavior the means have to be $\\mu = (\\pm 1, …, \\pm 1)$, i.e. both mixture components are at the corners of the hypercube. Alternative question: Does the sample complexity also transfer to $\\|\\mu\\|^2 < O(d)$, e.g. $O(1)$? I would guess that the other directions start playing a role then.\n\n3. What is the solution to this simple setup intuitively? Can you provide a simple drawing of $0, \\eta, \\mu_+, \\mu_-, \\mu, \\eta$ and a learnt trajectory? If the required dimension is indeed two, this should be easy.\n\n4. What is the shape of the learnt distribution within the two clusters, i.e. what is the local density in each cluster?\n\n5. Fig. 2: Why are the learnt means biased in one direction? Is more training data added sequentially?\n\n\nGiven the substantial weaknesses and the above questions, my *preliminary* vote for this paper is therefore not to accept it. I am happy to be corrected in any of the criticisms I raised and look forward to the authors' rebuttal.", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper considers an interesting question, how generative models generalize as a function of the number of training samples.\n\nThe solution for the bimodal Gaussian Mixture model appears sound and plausible.\n\nAll theoretical results are accompanied by experiments closely matching the prediction, increasing the credibility of the theoretical results.", "weaknesses": "I find the main technical result presented in a misleading way: Results 2.1 and 3.1 show that the weight vector and the ODE dynamic are not orthogonal to $\\mu$. I think a better way to describe the behavior of the system would be to span the relevant space via $\\eta$ and $\\hat\\mu_+$ and $\\hat\\mu_-$, corresponding to the empirical mean of the samples at $\\pm\\mu$. I assume that this is also sufficient to span the weight vector (please correct me if I’m wrong), and it would make clear that the model only has access to the empirical means. The remaining results could then be adapted to show how the sampling process is able to reproduce the empirical means of the two modes. Then, the Bayes-optimal baseline could be inserted to show that the empirical estimates go down with Theta(1/n), and that the flow is able to achieve the same rate.\n\nI also find that the research question formulated in the beginning is not really addressed in the main text, whether the network architecture determines whether a generative model memorizes training data. If my above understanding is correct, then the model actually does learn the training data by heart (i.e. it predicts the empirical means), and the rate to the true solution is essentially given how fast empirical means converge to the mean of the generative distribution. Also, the rate $\\Theta(1/n)$ is not affected by the regularization $\\lambda$, and the network architecture is not varied so as to judge wether this particular setup has a particular convergence rate.\n\nMinor points:\n- Formatting of contributions via itemize\n- Sentence? „Note that (9) is a special case of the architecture studied in Cui & Zdeborová (2023), and differs from the very similar network considered in Shah et al. (2023) in its slightly more general activation (Shah et al. (2023) address the case φ = tanh)“\n- Result 3.1 Le X_t -> Let X_t\n- finding that that the DAE on p. 9", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2Ph2Q0iIpn", "reviewer": "Reviewer_wFJz", "summary": "The authors analyze a generative flow-based model to sample from a mixture of Gaussians, where at each timestep the vector field is parameterized by a two-layer neural network.\nThey consider a high dimensional finite sample regime with the number of samples scaling linearly in the dimension. \nUsing tools from statistical physics they derive a precise characterization of the optimal performance.\nTheir experiments corroborate the theoretical findings.", "questions": "- Do you think this the same method could be used for more complex distributions, if yes what would be a concrete example?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The strongest point is the exact characterization of all important quantities. \n- Clear and concise presentation of the results\n- Solid experiments demonstrating the validity of the theoretical results", "weaknesses": "The approach seems to heavily rely on the gaussianity of the target distribution", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IdQsE5Crfr", "reviewer": "Reviewer_AsnX", "summary": "The paper studies a certain asymptotic limit of a flow-based generative model that uses weight-tied fully connected network with skip connection to approximate the velocity field. In such scenario, under the isotropic gaussian mixture assumption the authors provide a characterization of training dynamics in the finite sample complexity regime. Notably, the resulting asymptotic behaviour of the learned cluster means enjoys the Bayes optimal rate of $O(1/n)$.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- a complete characterization of training dynamics of shallow generative flow under the isotropic mixture of gaussians assumption\n- optimality of the resulting sample asymptotics\n- a neat symmetric ansatz", "weaknesses": "- lack of the correlation structure in the input data, which draws the conclusions to be less practical\n- minor: the approach is still an ansatz", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)