Paper: Analysis of Learning a Flow-based Generative Model from Limited Sample Complexity (ICLR-2024)

Reviewer comment (Experiments):
lack of the correlation structure in the input data, which draws the conclusions to be less practical

Author response:
We thank the reviewer for their comments and questions, which we address below.

>lack of the correlation structure in the input data, which draws the conclusions to be less practical

We believe that ideas similar to the ones used in the present work could be leveraged to address mixtures with structured clusters, see also the answer to reviewer wFJz. Such an extension would allow to analyze more realistic data distributions, and potentially quantitatively capture the learning curves of some simple real datasets, like in (Cui & Zdeborová, 23). This constitutes an interesting research direction which we will further pursue in future work.

Reviewer comment (Theory):
The approach seems to heavily rely on the gaussianity of the target distribution

Author response:
While the method indeed makes use of the fact that the target density is a Gaussian mixture, we believe that very similar ideas can be employed to address Gaussian mixture densities with structured covariances. This would provide a gateway to modeling more realistic setups, since as shown in (Cui & Zdeborová, 23), at the level of the learning process, the performance of the DAE on MNIST is quantitatively captured when modeling the latter by the matching Gaussian mixture density. Another natural extension which could be addressed with similar methods is the case of more clusters, and using an autoencoder with more hidden units. We leave the exploration of these exciting questions to future work.

Reviewer comment (Experiments):
Do you think this the same method could be used for more complex distributions, if yes what would be a concrete example?

Author response:
While the method indeed makes use of the fact that the target density is a Gaussian mixture, we believe that very similar ideas can be employed to address Gaussian mixture densities with structured covariances. This would provide a gateway to modeling more realistic setups, since as shown in (Cui & Zdeborová, 23), at the level of the learning process, the performance of the DAE on MNIST is quantitatively captured when modeling the latter by the matching Gaussian mixture density. Another natural extension which could be addressed with similar methods is the case of more clusters, and using an autoencoder with more hidden units. We leave the exploration of these exciting questions to future work.

Reviewer comment (Theory):
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.

Author response:
We thank the reviewer for this suggestion. The weight vectors can indeed be completely equivalently rewritten in terms of $\xi,\mu_{emp.} $ where $\mu_{emp.}$ is the empirical means
$$
\mu_{emp.}=\frac{1}{n}\sum\limits_{\mu=1}^n s^\mu x_1^\mu=\mu+\frac{1}{n}\eta.
$$
Note that this is just an equivalent way to present the equation, and that the underlying results remain totally unchanged. We chose the initial presentation in terms of $\xi, \eta, \mu$ to put further emphasis on the parameters of the target distribution.
We however agree with the reviewer that using $\{\eta, \mu_{emp.}\}$, rather than $\{\eta,\mu,\xi\}$, constitutes a more concise rephrasing of the results, and have rewritten the manuscript in this way. We have updated the pdf of our paper to reflect this.

Reviewer comment (Experiments):
I also find that the research question formulated in the beginning is not really addressed in the main text, namely 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 by 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 whether this particular setup has a particular convergence rate.

Author response:
The reviewer is correct that the model overfits, as the trained weights bear explicit dependence on the training samples, and the Gaussian noises $x_0$ employed during training. By "not memorizing", we rather mean that the generated density is not the discrete empirical distribution supported on the training data, but instead a Gaussian mixture. In other words, the generative model allows to generate novel samples, distinct from the training data. We emphasized this distinction in updated manuscript.

On the other hand, it is not entirely correct that the generated density has the empirical mean of the training data as cluster mean. It is already not the case for the Bayes-optimal estimate, which is the empirical means rescaled by a $1/(\sigma^2+n)$ factor. For the generative model, the cluster mean is a linear combination of $\xi$ (subsuming the effect of the noises used during training) and the empirical means, rescaled by multiplicative factors whose expression depend non-trivially on the model parameters - notably the schedule functions $\alpha(t),\beta(t)$, and the regularization $\lambda$. We believe that this precise characterization of the cluster mean of the generated density as a function of the parameters of the learning model is a strength of our analysis.

The main focus of the paper is not to study the dependence of the convergence rate as a function of the model architecture and parameters, but rather show, in a particular setting, how a learning model can manage to learn the target density with a reasonably fast rate only _from a few samples_. This is _allowed_ by the network architecture : indeed, had the minimization (5) been carried out over the space of denoising functions, instead of the parameter space of the considered learning model, the resulting generative model would memorize the training samples, and only allow to generate samples already present in the
train set, instead of sampling the Gaussian mixture. Furthermore, a discussion for another architecture can already be found in Appendix D of the supplementary material, where we detail the case of an auto-encoder without skip connection. We completely agree with the reviewer that the study of a greater number of architectures is an important research direction to be explored in future works.

This is indeed the correct intuition. We would however like to stress that the mean of the generated density is not equal to the empirical means, and that it is quite non-trivial that an auto-encoder can learn to generate a Gaussian mixture with such means just from a finite number of samples when trained on a denoising loss, and that the learning and transport processes can be sharply theoretically characterized.