Title: Learning a 1-layer conditional generative model in total variation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Generative models

distribution learning

maximum likelihood estimation
Abstract: A conditional generative model is a method for sampling from a conditional distribution $p(y \mid x)$.  For example, one may want to sample an image of a cat given the label ``cat''.  A feed-forward conditional generative model is a function $g(x, z)$ that takes the input $x$ and a random seed $z$, and outputs a sample $y$ from $p(y \mid x)$.  Ideally the distribution of outputs $(x, g(x, z))$ would be close in total variation to the ideal distribution $(x, y)$.

Generalization bounds for other learning models require assumptions on the distribution of $x$, even in simple settings like linear regression with Gaussian noise. We show these assumptions are unnecessary in our model, for both linear regression and single-layer ReLU networks.  Given samples $(x, y)$, we show how to learn a 1-layer ReLU conditional generative model in total variation. As our result has no assumption on the distribution of inputs $x$, if we are given access to the internal activations of a deep generative model, we can compose our 1-layer guarantee to progressively learn the deep model using a near-linear number of samples.

Learning a 1-layer conditional generative model in total variation

Abstract

A conditional generative model is a method for sampling from a conditional distribution p ( y | x ) . For example, one may want to sample an image of a cat given the label 'cat'. A feed-forward conditional generative model is a function g ( x, z ) that takes the input x and a random seed z , and outputs a sample y from p ( y | x ) . Ideally the distribution of outputs ( x, g ( x, z )) would be close in total variation to the ideal distribution ( x, y ) .

Generalization bounds for other learning models require assumptions on the distribution of x , even in simple settings like linear regression with Gaussian noise. We show these assumptions are unnecessary in our model, for both linear regression and single-layer ReLU networks. Given samples ( x, y ) , we show how to learn a 1-layer ReLU conditional generative model in total variation. As our result has no assumption on the distribution of inputs x , if we are given access to the internal activations of a deep generative model, we can compose our 1-layer guarantee to progressively learn the deep model using a near-linear number of samples.

1 Introduction

Generative models are in the midst of an explosion in accessibility, as models like DALL-E [31] or Stable Diffusion [32] capture the attention of millions. In many cases, these generative models can be succinctly represented by a fundamental mathematical object-the conditional distribution p ( y | x ) . In the example of text-to-image generative models, x can represent a text prompt or its Word2Vec embedding [26], and the model can be seen as sampling an image y from its conditional distribution p ( y | x ) . With large numbers of people accessing these models, a massive amount of sample pairs ( y i , x i ) , are becoming available online. A natural question to ask is: How many samples ( y i , x i ) does it take to learn the conditional generative model p ( y | x ) ?

Recent empirical studies such as Stanford Alpaca [34], which attempt to learn GPT-3.5 from limited samples, indicate that the number of samples needed may be within a practical range. In this paper, we attempt to address this problem from a fundamental perspective grounded in a concept from classical theoretical statistics: the Maximum Likelihood Estimator (MLE). Specifically, we focus on feed-forward generative models from a relatively simple family and ask: with no assumptions on x , how many samples are required to efficiently learn the conditional generative model p ( y | x ) ?

Linear Regression. Consider ordinary linear regression with Gaussian noise: you observe independent samples ( x i , y i ) ∈ R k × R of the form

Ananya Uppal

"Beautiful

sunset in

New Orleans"

How many samples does it take to get a 'good" solution? For standard metrics-such as parameter distance in w ∗ or prediction error on y -the sample complexity of linear regression depends on properties of the distribution (such as the conditioning of x , or the variance of y ) that could be unbounded and cannot be tested. For example, in the classic analysis (see [7], Chapter 3), sample complexity depends on the design matrix, which in turn results in a dependence on the expectation of ∥ x ∥ 2 . The typical approach to deal with this would require bounded moments for x . For learning the conditional distribution, the more natural metric is total variation (TV) distance: the parameters w induce a distribution p w ( x, y ) = p ( x ) p w ( y | x ) where p ( x ) is the true distribution of x , and we would like to find w such that

This ensures that when the input x come from the user in the true unknown distribution p ( x ) , the model generates a conditional sample that is close in TV to the true model. It turns out that this goal, unlike parameter distance, can be solved with no assumption on the distribution.

Theorem 1.1 (Informal version of Theorem 4.1) . The MLE (i.e., least squares regression) achieves (1) with O ( 1 ε 2 k log 1 ε ) samples, regardless of the distribution of x .

To our knowledge, all previous guarantees for linear regression either require some assumptions on, or give guarantees in terms of, the distribution of x or y . We avoid this dependence on the x distribution by adopting a similar analysis to Theorem 11.2 in [20].

Main Result: 1-Layer Networks. Modern feed-forward conditional generative models (e.g., StyleGAN2) are more complicated than linear regression. They are stochastic neural networks with layers of the form:

for x ∈ R k , y ∈ R d , ϕ ( x ) = max( x, 0) is the ReLU activation, and some weights W ∗ ∈ R d × k , Σ ∗ ∈ R d × d .

We show that 1-layer generative models of the form (2) can also be estimated using the MLE (which is concave), with small TV error and without any assumption on x .

Theorem 1.2 (Informal version of Theorem 4.5) . Suppose y is drawn according to (2) where Σ ∗ has condition number at most κ . Using O ( kd + d 2 ε 2 log kdκ ε ) samples, the distribution generated by the MLE has TV error ε , regardless of the distribution of x .

Like Theorem 1.1, Theorem 1.2 shows that a conditional generative model can be learned in TV regardless of the distribution of the input x . It generalizes Theorem 1.1 in three ways: (a) the output y is d -dimensional rather than 1 -dimensional; (b) the covariance of the noise Σ ∗ is learned rather than specified; and (c) the ReLU nonlinearity ϕ imparts additional complex structure.

Indeed, the point (c) means that parameter distances are poor metrics for this problem. For example, when an element of W ∗ x has a large negative bias, the corresponding coordinate in y will almost

̂

always be 0 . This means sample pairs ( x i , y i ) provide little information about W ∗ , but still provide useful information about the conditional distribution. As we will later see, exploiting this valuable information in the truncated samples allows us to significantly outperform prior works such as [35], which do not exploit these 0 -valued samples.

Multilayer Networks Given Activations. Given our distribution-free results on 1-layer networks, it is possible to extend our results to deep multi-layer networks. Given access to the internal activations of a neural network (but not the weights), our results can be applied layer-wise. Intermediate layers may have poorly conditioned input distributions, but since our result does not depend on the input distribution, we achieve strong guarantees for layer-wise learning in Theorem 4.6.

1.1 Proof Approach

In this outline we focus on the case of 1 -dimensional y , and and standard Gaussian noise η ∼ N (0 , 1) . Our proof approach is inspired by learning bounds that exploit finite VC dimension. We would like to show (1), or equivalently,

when we see n samples x i of x , and one sample y i for each x i . We do this in two stages. First, we show that the empirical distance between w and ̂ w is small i.e.,

where ˜ E x [ f ( x )] = 1 n ∑ n i =1 f ( x i ) denotes the empirical expectation over x .

Second, we show that the empirical distance is a good proxy for the true distance, i.e.,

which gives (3).

Linear Case. In the linear case, both stages are straightforward. The linear regression solution has an explicit form, and it is well known and easy to show that

Since d TV ( p w ∗ ( y | x ) , p ̂ w ( y | x )) = Θ(min(1 , | x · ( w ∗ -̂ w ) | )) , Jensen's inequality implies (4) for n > k/ε 2 .

Secondly, f w ( x ) := d TV ( p w ∗ ( y | x ) , p w ( y | x )) is bounded and unimodal in w . Thus, it suffices to bound the deviation of the empirical average from the true f w ( x ) with Chernoff's inequality.

ReLU Case. In the ReLU case, we have y = ϕ ( w ∗ · x + η ) , and both stages of the previous analysis are more difficult.

The most interesting part of our proof is showing the first stage for the ReLU case, which states that the ̂ w maximizing

satisfies (4). Now, for any w not satisfying (4),

where the first inequality follows from Pinsker's inequality. Unfortunately, L ( w ) -L ( w ∗ ) does not concentrate well, by virtue of the KL-divergence being unbounded. However, we can upper bound it via the Bernstein inequality, such that for a fixed w not satisfying (4), and given n = 1 ε 2 log( 1 δ ) samples, we have

with probability 1 -δ . Using a careful covering argument and n = ( k/ε 2 ) log(1 /δ ) samples, we can uniformly extend this to all w not satisfying Eq (4). By definition, the MLE has L ( ̂ w ) ≥ L ( w ∗ ) , and by our uniform bound, it must satisfy (4).

The second stage changes because f w ( x ) depends on x · w and x · w ∗ in a more complicated way than through x · ( w -w ∗ ) . This makes showing a bounded VC dimension more difficult; however, unpacking the proof that VC implies generalization, we can still show that the net (normally given by Sauer's lemma) is bounded. This generalization holds as long as f w ( x ) is unimodal in x · w .

2 Contributions

We show that MLE can perform distribution learning in the setting of linear regression and multi-layer ReLU networks. Our bounds do not make assumptions on the distribution of x or the condition number of W ∗ , and achieve a sample complexity polynomial in the system parameters.

We improve the sample complexity bound in [35], which estimates the parameters of a one-layer ReLU network but suffers an exponential dependence on the W ∗ term. In contrast, as we seek to estimate the distribution of ( x, y ) , rather than the parameter W ∗ , we are able to avoid this. See Section 4.2 for more details.

Our algorithm for learning multi-layer ReLU networks is considerably simpler than [1], who learn discriminators that are engineered to perform moment-matching on the output of each layer of the network. Furthermore, [1] impose a strong requirement on the sparsity and independence of the activations at each layer, which essentially allows standard techniques in sparse coding to recover these activations.

3 Related Work

Generative adversarial networks (GANs) [19, 2, 30] are a popular family of generative models that train a generator and discriminator in an adversarial training framework. The seminal result by [22] proposed progressive growing of GANs (PGGANs) as a way to stabilize and accelerate the training phase of these models. Future results, such as StyleGAN [23] introduce more complicated architectures and 'style' variables. Additionally, these models add noise at each layer of the generator in order to introduce greater stochasticity in the generated images, which is important for textures such as hair and skin.

Distributional Learning Most theoretical results have focused on the min-max optimality of GANs [15, 28, 29], characterizing their stationary points [21, 17], or characterizing their generalization once they have reached a global minimum [3, 5]. The closest result to ours is [35]. Setting x to a deterministic scalar in our problem statement reduces it to [35], who consider y = ϕ ( b + η ) s.t. b, η ∈ R d , and they seek to learn the covariance of η along with the bias vector b . However, their sample complexity bound suffers an exponential dependence on ∥ b ∥ 2 ∞ .

Single layer networks have attracted recent attention, as they provide a tractable formulation for studying the dynamics and generalization of adversarial learning [24, 18, 25, 11]. The recent results of [1] show that multi-layered models that satisfy a property known as forward super-resolution (such as PGGANs) can be learned in polynomial time and sample complexity using stochastic gradient descent-ascent. In this case, the discriminator is designed to detect differences between higher order moments of the generated and training distribution. Deep models have also been considered in [12, 10].

4 Main Results

In this section we first show that the MLE of the parameters learns the input-output joint distribution for linear regression. Then, we extend this guarantee to the case where the ReLU activation function ϕ is applied to the multi-dimensional output y . Finally, we show that we can compose the 1-layer ReLU guarantee to learn the distribution generated by a multi-layer model.

4.1 Linear Regression

We begin with the classic linear regression problem of learning the parameter w ∗ from a linear model

where η ∼ N (0 , σ 2 ) , w ∈ R k , σ is known and x ∈ R k with some distribution. This problem has been studied for centuries. Our novelty is that we view (6) as a conditional generative process, and instead of studying error in Euclidean distance in parameter space, i.e., minimizing ∥ ̂ w · x -w ∗ · x ∥ 2 , we focus on error in d ( w ∗ , ̂ w ) as defined in (3), which only captures the error in ̂ w insofar as it impacts our distribution estimate. Given data { ( x i , y i ) } n i =1 generated by (6), the MLE is:

The following theorem establishes that the MLE is close in TV distance. The proofs for all results in this section are in Appendix A.

Theorem 4.1. Let { ( x i , y i ) } n i =1 be i.i.d. random variables generated from the linear model (6) , and assume that σ is known. Then, for a sufficiently large constant C > 0 ,

samples suffice to ensure that with probability 1 -e -Ω( ε 2 n ) over the data,

Note that one cannot hope to get such a guarantee in the classical setting where error is measured in ∥ ̂ w · x -w ∗ · x ∥ 2 without additional assumptions on the distribution of x because if x is badly conditioned, the error may be dominated by very rare directions of x that we never sample. For example, the bounds in [7], Chapter 3, require the second moment of x to be bounded.

Since we only wish to learn the distribution of y in total variation, no single x can contribute much to our loss and we get a distribution-free result. This is possible as the total variation distance is bounded, and we can invoke Theorem 11.2 in Györfi et al [20].

We now state two lemmas needed to prove Thereom 4.1, assuming without loss of generality that σ 2 = 1 . We split the proof into two stages. In the first stage, we bound ˜ d ( ̂ w,w ∗ ) , which denotes the empirical TV distance (4) over the training set { x i } n i =1 .

Lemma 4.2. Let { ( x i , y i ) } n i =1 be i.i.d. random variables such that y i = x i · w ∗ + N (0 , 1) . Then, for n ≥ k 2 , with probability 1 -e -Ω( n ) , the MLE ̂ w satisfies

The proof relies on the fact that p w ∗ ( y | x i ) and p ̂ w ( y | x i ) are Gaussian distributions, so

Using the explicit form of the MLE, we can show that, with high probability,

and Lemma 4.2 follows from Jensen's inequality.

The second stage shows that the empirical average of the TV distance ˜ d ( ̂ w,w ∗ ) is close to the population average d ( ̂ w,w ∗ ) .

Lemma 4.3. Let { x i } n i =1 be i.i.d. random variables such that x i ∼ D x . For a sufficiently large constant C > 0 , and for n = C k ε 2 log 1 ε with n ≥ k 2 , we have:

Note the probability in the above statement is with respect to the distribution of x , and does not depend on y . The proof follows Theorem 11.2 in Györfi et al [20]: it relies on the fact that the TV distance is bounded and a unimodal function of w · x . This implies that for each x i , we are able to partition the space of w with O (1 /ε ) hyperplanes such that within each cell the TV distance varies by at most ε . As we have n samples and w ∈ R k , the number of cells induced in R k is ∝ ( n/ε ) k , and it is sufficient to provide concentration bounds for one representative in each cell. This approach is similar to bounding the VC dimension of a set of binary functions. Setting n = Θ( k ε 2 log 1 ε ) and combining Lemma 4.2 with Lemma 4.3 gives d ( ̂ w,w ∗ ) ≤ ε.

4.2 ReLU Case

Now consider the single-layer ReLU. We observe ( x, y ) ∈ R k × R d such that:

where η ∈ R d , W ∗ and ϕ ( · ) = max( · , 0) is applied coordinate-wise. The matrices W ∗ ∈ R d × k and Σ ∗ ∈ R d × d are unknown , and we do not observe η . The variable x is drawn from an arbitrary probability distribution D x , and we make no additional assumptions on D x : this is important, as we will progressively cascade layers, and one should think of D x being the distribution of activations at each layer.

Given a sample ( x, y ) ∈ R k × R d , let S denote the co-ordinates of y that are zero-valued, and let S c denote the compliment of S . Then, the log-likelihood of W, Σ , on this sample is given by

where c is a normalization constant which does not depend on W or Σ . This function is a mixed density: in the coordinates of y that are 0, i.e., in the set S , we integrate the Gaussian density over the negative orthant, as W ∗ x + η could have been any negative value in those coordinates.In Lemma F.1 in the Appendix, we show that Eqn (8) is concave after an invertible reparameterization of W, Σ .

In this setting, proving an analogue of Theorem 4.1 poses multiple challenges:

The output y is d -dimensional rather than a scalar, and η in Eqn (7) introduces correlations between the coordinates of W ∗ x , such that we cannot decompose the log-likelihood in Eqn (8) per coordinate.

We do not know the covariance matrix Σ ∗ , and it must be estimated.

Lemma 4.2 requires the explicit form of the MLE in linear regression. In the absence of such a closed-form solution, we need to directly analyze the log-likelihood, which is a mixed density and involves integrating the Gaussian likelihood over the zero-valued coordinates of y . In order to handle this, we use the Györfi approach again on the log-likelihood. This is challenging because the variables we concentrate are KL-divergences, which are unbounded and require Bernstein type inequalities.

Recovering the true parameters W ∗ , Σ ∗ is difficult: if we see a zero in y , we do not know its magnitude in W ∗ x + η before the ReLU. This manifests in the results in [35], where it is assumed that each entry in W ∗ is positive - otherwise, their sample complexity scales as e ∥ W ∥ 2 ∞ .

Nonetheless, we can handle most of these difficulties, and the only assumption we make is that the condition number of Σ ∗ is bounded and known to our estimator.

Assumption 4.4 (Condition number bound) . Let λ ∗ max , λ ∗ min denote the largest and smallest singular values of Σ ∗ . We assume there exists κ < ∞ such that λ ∗ max λ ∗ min ≤ κ . We further assume that the value of κ is known to our estimator.

Note that the condition number only allows us to control the correlation between the coordinates of W ∗ x + η . The other challenges introduce by the ReLU, such as the lack of a closed form MLE, the need to estimate Σ ∗ and a mixed density log-likelihood remain.

Under Assumption 4.4, the following theorem shows that the MLE ̂ W, ̂ Σ achieves a small total variation distance. The proof of this theorem is in Appendix C.

Theorem 4.5. Let R d × d κ denote the set of positive definite matrices with condition number κ . Given n samples { ( x i , y i ) } n i =1 satisfying Assumption 4.4, where x i ∼ D x i.i.d., and y i is generated according to (7) , let ̂ W, ̂ Σ := arg max W ∈ R d × k , Σ ∈ R d × d κ 1 n ∑ i log p W, Σ ( y i | x i ) . Then, for a sufficiently large constant C > 0 ,

samples suffice to ensure that with probability 1 -δ , we have

Comparison to [35]. Our result is closely related to [35]. Our ReLU model reduces to their model by setting W ∗ ∈ R d as a vector, and x = 1 . In order to learn the distribution of y , they first estimate the parameters W ∗ , Σ ∗ , in ℓ 2 norm and then convert the ℓ 2 error to a TV error. The parameter recovery is done per coordinate of W ∗ , Σ ∗ , by performing MAP estimation on the positive samples in y . This crucially assumes that each coordinate has enough positive samples, and to that end, they assume that each entry in W ∗ is positive-otherwise, their sample complexity scales as e ∥ W ∗ ∥ 2 ∞ . Our results do not make any assumptions on W ∗ and handles a wider class of matrix valued W ∗ . Additionally, the objective function (8) does not discard the zero valued samples in y , making it more sample efficient.

We assumed that the covariance matrix Σ ∗ has condition number κ , and our sample complexity scales as log κ . Hence, even if κ = e poly( d,k ) , we only pay a poly( d, k ) penalty. This improves on the result in [35], where the sample complexity scales as κ 2 . While our statistical guarantees are strictly better, [35] gives poly-time and poly-space algorithmic guarantees for their estimator. We discuss the empirical limitations of our algorithm in Section 6.

Lower Bounds Ignoring log factors, the complexity factor of kd is obviously required. Furthermore, learning a Gaussian with unknown covariance matrix in total variation takes ˜ Ω( d 2 ) samples; see [6]. Our Theorem 4.5 would solve their lower bound instance, the same lower bound applies to our problem.

Extension to Multi-Layer Generative Models. Consider the following ( L +1) -layered generative model.

We can compose the guarantees provided by Theorem 4.1 and 4.5 to show that we can learn this model.

Theorem 4.6. Given n i.i.d. samples of ( x 0 , . . . , x L +1 ) , such that each matrix Σ ∗ ℓ satisfies Assumption 4.4, let ̂ W ℓ , ̂ Σ ℓ , be the MLE estimates of W ∗ ℓ , Σ ∗ ℓ learned from samples of ( x ℓ , x ℓ +1 ) . Define m := max ℓ d 2 ℓ . Then,

samples suffice to ensure that with probability 1 -δ ,

Comparison to [1]. The modelling assumptions in [1] are similar to ours - the authors learn a generative model per layer, using images produced per layer. The key differences of their model are: (i) each layer is deterministic (there is no η ℓ ), (ii) their learning algorithm does not require access to the activations of each layer, (iii) their algorithm performs moment-matching by crafting the discriminator strategically.

In order to avoid requiring activations at each layer, [1] imposes a sparsity assumption on the activations: this allows them to leverage tools from existing results in sparse coding [4], such that

10 2

10 3

10 4

10 5

10 -3

10 -2

10 -1

(a) Total Variation vs.

n

10 1

10 -2

10 -1

(b) Total Variation vs.

k

sparse activations at each layer can be recovered using images produced by the layer. This assumption can be somewhat strong, as it implies that the activations are roughly independent of one another, and the sparsity remains constant over layers, despite the layers themselves expanding by a factor of 4, i.e., d ℓ ≥ 4 d ℓ -1 ∀ ℓ ≤ L -1 .

5 Simulations

We now numerically verify our theoretical claims and compare against other approaches. A detailed description of simulation methods are included in the appendix. Our code is available at https: //github.com/basics-lab/learningGenerativeModels.git .

5.1 Scaling in n and k

Figure 2 numerically investigates how TV distance of the MLE scales with the number of samples n and the input dimension k . We consider a model with 1 -dimensional output and a k -dimensional input: y = ϕ ( x · w ∗ + η ) , for w ∗ ∈ R k and η ∼ N (0 , σ 2 ) . We set σ 2 = 1 and w ∗ = 1 k × 1 , both unknown to the optimizer, which has samples ( y i , x i ) n i =1 . Figure 2a, which plots the error in TV distance against n on a log-log plot, has a slope of roughly -1 / 2 as predicted by our theory. Similarly, Figure 2b has a slope of 1 / 2 , which is in line with our theory on scaling with respect to input dimension k . We defer simulations involving scaling in d to the appendix because computing TV distance becomes increasingly difficult as d becomes large, and we must resort to using upper bounds.

5.2 Distribution Independence

The fact that our guarantee does not depend on the distribution of x suggests that the expected TV error of the distribution learned from the MLE may be similar for all distributions over x . To test this we consider both x ∼ N (0 , I k ) and x ∼ ⊗ k Lap(0 , 1) , i.e., each element of x is drawn independently standard Laplace. Figure 2a verifies our hypothesis, showing only very slight differences in our observed empirical average TV error between the two distributions over a wide range of n , and for k = 5 and k = 20 .

-4

-3

-2

-1

0

1

2

3

4

5

10 -4

10 -3

10 -2

10 -1

10 -3

10 -2

10 -1

(a) Total Variation vs. Bias

50

100

150

200

250

300

350

400

450

500

550

0.012

0.014

0.016

0.018

0.02

0.022

0.024

0.026

0.028

0.03

0.032

(b) Total Variation vs. Condition Number

5.3 Scaling with Bias and Condition Number of Covariance Matrix

A key feature of the MLE is that it makes use of truncated samples. This is in contrast to [35], which leverages results on learning truncated normal distributions [14] where truncated samples are not observed. This leads to a stark difference in performance as the number of truncated samples becomes large. To show this, we consider a model with a d -dimensional output and 1 -dimensional input. We let x = 1 almost surely, and then take w ∗ = b 1 d × 1 for some bias b ∈ R and η = N (0 , Σ) with Σ = I d , thus y = ϕ ( η + b 1 d × 1 ) . As b becomes more negative, the number of truncated samples increases. Figure 3a shows the differing behavior of MLE and that of [35] as b becomes negative. For ease of computing the TV distance we set d = 3 , and restrict optimization over diagonal Σ . The solid blue line depicts the performance of the MLE. We observe that the TV error is constant for b > 0 and begins to decrease rapidly for b < 0 . This happens because as b becomes more negative, the truncation places more probability mass at y = 0 . Indeed, the dashed blue lines indicate that even as the TV error is decreasing rapidly, the mean square estimation error of the covariance and mean increase, however, since most of the probability mass is at zero, this does not significantly impact the TV. In contrast, the method of [35] rapidly deteriorates as b < 1 as the number of untruncated samples decreases. We also point out that even when the bias is large, [35] is still significantly worse. We attribute this to the fact that even when there is no truncated samples, [35] is still minimizing a different MAP objective. More discussion of this is provided in the appendix.

Robustness to Condition Number of Σ . Another concern is how TV error scales as a function of the condition number of Σ ∗ . Poorly conditioned Σ ∗ can put significant probability masses on small sets, and potentially cause large error. We consider a similar environment to the one described above, but fix b = 1 and alter diagonal entries of Σ ∗ such that one entry is √ κ , another is √ κ -1 and the rest are 1 , making the condition number κ . Figure 3b shows that the MLE is not measurably impacted by the changing condition number over the range plotted. This is not true of [35], where we observe that TV does grow with condition number.

6 Limitations

This work is only a first step to understanding fundamental limits of learning generative models. We showed that our theorems for single-layer networks can be composed to get sample complexity bounds on deep networks, with a critical caveat: we require access to not just the input and output pairs, but also the intermediate activations. This is not practical in many scenarios, and removing this restriction will be an important direction for future research. Beyond this, we assume that the learner has an understanding of the model architecture. In many cases, however, a learner may not be aware

of the number of layers a network has or a vast number of other architectural details. Additionally, we inherit known problems with the MLE, such as exacerbation of biases that exist in the training data (Chapter 24.1.3 in [33]).

Our results place emphasis on sample complexity over computational complexity . Though we show that the MLE problem is concave, this work does not provide a thorough analysis of the optimization problem. It is possible that similar results to [14, 35] can be derived for the MLE problem. Indeed, empirically, we find that a similar projected stochastic gradient ascent performs well in our problem. A careful analysis must consider factors like the distribution over x , the condition number of Σ ∗ and the truncation probability, all of which are likely to impact the optimization.

7 Conclusions

We have studied the problem of learning conditional generative models from a limited number of samples. We have shown that it is possible to learn a 1-layer ReLU conditional generative model in total variation, with no assumption on the distribution of the conditioning variable x using the MLE. We have also shown that this result can be extended to multi-layer ReLU networks, given access to the internal activations. Our results suggest that MLE is a promising approach for learning feed-forward generative models from limited samples.

References

Zeyuan Allen-Zhu and Yuanzhi Li. Forward super-resolution: How can gans learn hierarchical generative models for real-world distributions. arXiv preprint arXiv:2106.02619 , 2021.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan. arXiv preprint arXiv:1701.07875 , 2017.

Sanjeev Arora, Rong Ge, Yingyu Liang, Tengyu Ma, and Yi Zhang. Generalization and equilibrium in generative adversarial nets (gans). In International Conference on Machine Learning , pages 224-232. PMLR, 2017.

Sanjeev Arora, Rong Ge, Tengyu Ma, and Ankur Moitra. Simple, efficient, and neural algorithms for sparse coding. In Conference on learning theory , pages 113-149. PMLR, 2015.

Sanjeev Arora, Andrej Risteski, and Yi Zhang. Do gans learn the distribution? some theory and empirics. In International Conference on Learning Representations , 2018.

Hassan Ashtiani, Shai Ben-David, Nicholas JA Harvey, Christopher Liaw, Abbas Mehrabian, and Yaniv Plan. Near-optimal sample complexity bounds for robust learning of gaussian mixtures via compression schemes. Journal of the ACM (JACM) , 67(6):1-42, 2020.

Francis Bach. Learning theory from first principles. Draft of a book, version of Sept , 6:2021, 2021.

Stéphane Boucheron, Gábor Lugosi, and Pascal Massart. Concentration inequalities: A nonasymptotic theory of independence . Oxford university press, 2013.

Stephen P Boyd and Lieven Vandenberghe. Convex optimization . Cambridge university press, 2004.

Sitan Chen, Jerry Li, and Yuanzhi Li. Learning (very) simple generative models is hard. arXiv preprint arXiv:2205.16003 , 2022.

Sitan Chen, Jerry Li, Yuanzhi Li, and Raghu Meka. Minimax optimality (probably) doesn't imply distribution learning for gans. arXiv preprint arXiv:2201.07206 , 2022.

Sitan Chen, Jerry Li, Yuanzhi Li, and Anru R Zhang. Learning polynomial transformations. arXiv preprint arXiv:2204.04209 , 2022.

Annie AM Cuyt, Vigdis Petersen, Brigitte Verdonk, Haakon Waadeland, and William B Jones. Handbook of continued fractions for special functions . Springer Science & Business Media, 2008.

Constantinos Daskalakis, Themis Gouleakis, Chistos Tzamos, and Manolis Zampetakis. Efficient statistics, in high dimensions, from truncated samples. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science (FOCS) , pages 639-649. IEEE, 2018.

Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, and Haoyang Zeng. Training gans with optimism. arXiv preprint arXiv:1711.00141 , 2017.

Luc Devroye, Abbas Mehrabian, and Tommy Reddad. The total variation distance between high-dimensional gaussians. arXiv preprint arXiv:1810.08693 , 2018.

Farzan Farnia and Asuman Ozdaglar. Do gans always have nash equilibria? In International Conference on Machine Learning , pages 3029-3039. PMLR, 2020.

Soheil Feizi, Farzan Farnia, Tony Ginart, and David Tse. Understanding gans in the lqg setting: Formulation, generalization and stability. IEEE Journal on Selected Areas in Information Theory , 1(1):304-311, 2020.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural information processing systems , pages 2672-2680, 2014.

László Györfi, Michael Köhler, Adam Krzy˙ zak, and Harro Walk. A distribution-free theory of nonparametric regression , volume 1. Springer, 2002.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. In International Conference on Learning Representations , 2018.

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 4401-4410, 2019.

Qi Lei, Jason Lee, Alex Dimakis, and Constantinos Daskalakis. Sgd learns one-layer networks in wgans. In International Conference on Machine Learning , pages 5799-5808. PMLR, 2020.

Yuanzhi Li and Zehao Dou. Making method of moments great again?-how can gans learn distributions. arXiv preprint arXiv:2003.04033 , 2020.

Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. arXiv preprint arXiv:1301.3781 , 2013.

John P Mills. Table of the ratio: area to bounding ordinate, for any portion of normal curve. Biometrika , pages 395-400, 1926.

Aryan Mokhtari, Asuman E Ozdaglar, and Sarath Pattathil. Convergence rate of o(1/k) for optimistic gradient and extragradient methods in smooth convex-concave saddle point problems. SIAM Journal on Optimization , 30(4):3230-3251, 2020.

Vaishnavh Nagarajan and J Zico Kolter. Gradient descent gan optimization is locally stable. Advances in neural information processing systems , 30, 2017.

Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434 , 2015.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning , pages 8821-8831. PMLR, 2021.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms . Cambridge university press, 2014.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Alpaca: A strong, replicable instruction-following model. 2023.

Shanshan Wu, Alexandros G Dimakis, and Sujay Sanghavi. Learning distributions generated by one-layer relu networks. Advances in neural information processing systems , 32, 2019.

A Proofs of Linear Case

Throughout the appendix, for ease of notation, we overload the definition of the function d TV ( · , · ) . When inputs are random variables, it represent the TV distance between the distributions of those random variables.

Lemma 4.2. Let { ( x i , y i ) } n i =1 be i.i.d. random variables such that y i = x i · w ∗ + N (0 , 1) . Then, for n ≥ k 2 , with probability 1 -e -Ω( n ) , the MLE ̂ w satisfies

The proof of this lemma requires Lemma A.1, which characterizes the distribution of the residual error of the MLE.

Lemma A.1. Given y ∈ R n , X ∈ R n × k satisfying y = Xw ∗ + η , where η ∼ N (0 , σ 2 I n ) , the least square solution ̂ w satisfies

Proof. The least squares solution is given by

Multiplying on the left by X , we have

Since η is i.i.d. Gaussian with variance σ 2 , we have,

This implies

Proof of Lemma 4.2. The KL divergence between two Gaussians P = N ( µ 1 , Σ) and Q = N ( µ 2 , Σ) is:

By Pinsker's inequality, this implies

Hence, the empirical TV on the dataset can be bounded by

where the second line follows from Jensen's inequality.

which implies that with probability 1 -e -Ω( n ) , we have

Substituting in the earlier inequality, we get

Lemma 4.3. Let { x i } n i =1 be i.i.d. random variables such that x i ∼ D x . For a sufficiently large constant C > 0 , and for n = C k ε 2 log 1 ε with n ≥ k 2 , we have:

Proof. The proof is inspired by Theorem 11.2 in [20], with modifications to our setting. Let Since f w ( x ) is bounded, for any fixed w , the Chernoff bound gives

for any α > 0 . The challenge lies in constructing a 'net' to be able to union bound over R k without assuming any bound on w or the covariate x . A net is a partitioning of an space, where within each part, points are close together in some way. In this case, we construct a net using what we will refer to as 'ghost' samples.

Ghost samples. First, we construct a 'ghost' dataset D ′ x consisting of n new samples, drawn i.i.d. { x ′ i } i ∈ [ n ] of D x . This gives another metric ˜ d ′ ( · , · ) . Instead of directly considering the distance between ˜ d ( w,w ∗ ) and d ( w,w ∗ ) , it is sufficient to consider the difference between ˜ d ( w,w ∗ ) and ˜ d ′ ( w,w ∗ ) i.e.,

To see this, let ¯ w maximize ˜ d ( w,w ∗ ) -˜ d ′ ( w,w ∗ ) . Since ¯ w and { x ′ i } i ∈ [ n ] are independent, by the Chernoff bound,

for any ( D , ¯ w ) and large enough n . Thus,

which implies (12).

Symmetrization. Since D x and D ′ x each have n independent samples, we could instead draw the datasets by first sampling 2 n elements x 1 , . . . , x 2 n from D x , then randomly partition this sample into two equal datasets. Let s i ∈ {± 1 } so s i = 1 if z i lies in D ′ x and -1 if it lies in D x . Then

For a fixed w and x 1 , . . . , x 2 n , the random variables ( s 1 , . . . , s 2 n ) are a permutation distribution, so negatively associated. Then the variables s i · d TV ( p w ( y | x i ) , p w ∗ ( y | x i )) are monotone functions of s i , so also negatively associated. They are also bounded in [ -1 , 1] . Hence we can apply a Chernoff bound:

for any fixed w .

Constructing a net. We partition R k the space of w s.t. if w,w ′ are in the same partition then,

for each x in the dataset x 1 , . . . , x 2 n . Then take the intersection of all 2 n partitions to construct a net over R k .

As the total variation distance is a unimodal function of x i · w -x i · w ∗ , we partition w the sets

where j goes from 0 to 1 /α -1 . So the space of w , R k is partitioned by 2 n sets of 1 /α parallel hyper-planes. Then the total number of cells is at most

We define a net N by choosing one representative of each cell in the partition, so | N | ≤ e 2 k log n αk . By (13),

Finally, for any w ∈ R let ¯ w ∈ N be the representative of its cell. By definition of the cells,

for all i ∈ [2 n ] . Thus

∣

′

∗

(

∣

˜

d

(

w,w

)

d

(

w,w

-

∣

and so

Setting α = ε/ 4 , we have that

suffices for

˜

∗

)

)

-

(

′

d

˜

( ¯

w,w

∗

)

-

d

( ¯

˜

w,w

∗

)∣

∣

∣

∣

∣ ˜

∣

d

(

w,w

∗

)

-

d

( ¯

˜

w,w

∗

)

∣

∣

∣

+

∣

∣ ˜

∣

′

d

(

w,w

∗

)

-

d

′

˜

( ¯

w,w

∗

)

∣

∣

∣

2

α

B ReLU Activation with Scalar y

In this section, we consider the model of

where w ∗ , x ∈ R k , y, η ∈ R . We are given samples ( x, y ) ∈ R k × R , and want to estimate a ̂ w that estimates the distribution of y in TV.

The most challenging aspect of the ReLU setting is that we do not have an expression for the TV suffered by the MLE, such as Lemma 4.2 in the linear case. This forces us to directly analyze the log-likelihood.

For a fixed x, w, , the expectation of the log-likelihood ratio over y is

where the last inequality is via Pinsker's inequality. This equation implies that if w is ε -far from w ∗ , then the expected log-likelihood ratio(LLR) is < -2 ε 2 . By definition, the MLE has a non-negative LLR. Hence, if the empirical LLR is close to the expectation, this would imply that the MLE has small TV.

However, we only receive a single sample of y per x . For a fixed w , we can prove a Bernstein inequality, showing that given 1 /ε 2 log(1 /δ ) samples, the empirical LLR is < -ε 2 for w that are ε -far.

Lemma B.1. Let p 1 , . . . , p n and q 1 , . . . , q n be distributions with E i [ d TV ( p i , q i )] ≥ ε , where we use the uniform measure on i ∈ [ n ] . Let x i ∼ p i for i ∈ [ n ] . Then w.p. 1 -δ , E i [log q i ( x i ) p i ( x i ) ] ≤ -ε 2 4 for n ≥ O ( 1 ε 2 log 1 δ ) .

The proof of this Lemma, as well as other Lemmas in this section, can be found in Appendix B.1.

In order to extend this to all w ∈ R k that are ε -far, we will construct a cover over R k depending on the values the log-likelihood ratio can take, and then apply the Bernstein inequality to each element in the cover.

In order to construct the cover, we first show that the log-likelihood ratio is bounded above by the magnitude of noise in y . For ease of notation, for a fixed x ∈ R k , and each w ∈ R k , define

and let θ ∗ = ⟨ x, w ∗ ⟩ .Similar to the notation for w , for each θ ∈ R , define p θ as the distribution of ϕ ( θ + η ) for η ∼ N (0 , 1) . Define the log likelihood ratio

The following Lemma states that for a fixed datapoint ( x, y ) , the log-likelihood ratio is bounded by the noise in y :

Lemma B.2. For any y = ϕ ( θ + η ) ,

Now, as γ is bounded above by | η | 2 2 , and it is concave wrt θ , the following Lemma shows that we can partition θ into O ( A ε ) intervals, such that in each interval, γ changes by atmost ε , or is very negative, i.e., γ < -A .

Lemma B.3 (One-dimensional net) . Let A > B 2 > 1 . There exists a partition of R into O ( A/ε ) intervals such that, for each interval I in the partition and every y = ϕ ( θ ∗ + η ) with | η | ≤ B , one of the following holds:

and therefore, for all y ,

For all θ ∈ I , γ θ ( y ) ≤ -A

For all θ, θ ′ ∈ I , | γ θ ( y ) -γ θ ′ ( y ) | ≤ ε

Using Lemma B.2 and Lemma B.3, we can form a uniform bound, such that all w that are ε -far from w ∗ in distribution will have log-likelihood ratio smaller than -ε 2 4 on the training set. With some additional arguments, we can now show that as the MLE has positive log-likelihood ratio, it has small empirical TV.

Lemma B.4. Let x 1 , . . . , x n be fixed, and y i ∼ ϕ ( x T i w ∗ + η i ) for η i ∼ N (0 , 1) . For n ≥ 1 ε 2 k log 1 ε , the MLE ̂ w satisfies

This sample complexity guarantees that the MLE is good for the set of empirical x i ∼ D x , and we need to extend this to the expectation over x ∼ D x , for which we use a Lemma similar to Lemma 4.3 in the linear case, and this completes the proof.

A straight forward combination of Lemma 4.3 and Lemma B.4 gives the following Theorem.

Theorem B.5. Let y = ϕ ( x T w ∗ + η ) , for w ∗ ∈ R k , x ∼ D x , and η ∼ N (0 , 1) . Then for a sufficiently large constant C > 0 ,

samples of { ( y i , x i ) } n i =1 suffices to guarantee that the MLE ̂ w satisfies

B.1 Proofs

Lemma B.1. Let p 1 , . . . , p n and q 1 , . . . , q n be distributions with E i [ d TV ( p i , q i )] ≥ ε , where we use the uniform measure on i ∈ [ n ] . Let x i ∼ p i for i ∈ [ n ] . Then w.p. 1 -δ , E i [log q i ( x i ) p i ( x i ) ] ≤ -ε 2 4 for n ≥ O ( 1 ε 2 log 1 δ ) .

Proof. Define γ i ( x ) = log q i ( x ) p i ( x ) and a i ( x ) := max( γ i ( x ) , -2) . We have that

and want to show that E i [ γ i ( x )] ≤ -ε 2 / 4 with high probability. Note that a i ( x ) ≥ γ i ( x ) , so it suffices to show E i [ a i ( x )] ≤ -ε 2 / 4 . We will do this with Bernstein's inequality, for which we need bounds on the moments of a i ( x ) .

To simplify notation, fix a particular i and consider p = p i , q = q i , a = a i , and x ∼ p .

For a random variable v , define v + , v -to be the positive/negative parts of v , respectively, so v = v -+ v + . Define ∆( x ) = q ( x ) p ( x ) -1 . We have that E x ∼ p [∆( x )] = 0 , and

Now, consider the function b ( z ) := max(log(1 + z ) , -2) -z . This function is nonpositive over z ≥ -1 , and b ( z ) ≤ -z 2 / 2 for z ≤ 0 . Since

or by (14),

Bounding the positive higher moments. We have that p ( x ) e a ( x ) = max( q ( x ) , e -2 p ( x )) so

In the following, we use that e t ≥ 1 + t for all t , as well as e t = 1 + t + ∑ ∞ k =2 1 k ! t k . Therefore

so

We now show that the Pr[ a ( x ) = -2] is smaller than the E [ -a ( x )] term, by relating to -b . When a ( x ) = -2 , ∆( x ) ≤ -1+1 /e 2 , and b (∆( x )) = -2 -∆( x ) ≤ -1 . Since -b (∆( x )) is non-negative, and at least 1 whenever a ( x ) = -2 ,

In particular, E [ a k + ( x )] ≤ 2 k ! E [ -a ( x )] for all k ≥ 2 .

Bounding the second moment of a . We have that

and E [ a 2 + ( x )] ≤ 4 E [ -a ( x )] by (16). We now bound E [ a 2 -( x )] . Note that | a -( x ) | ≤ 2 1 -1 /e 2 | ∆ -( x ) | by the construction of a . Therefore and so by (15),

Thus

Now we can apply Bernstein's inequality (Theorem 2.10 of [8]).

We apply the theorem to X i := a i ( x i ) , which are independent. The theorem uses that

and hence by (17), and since

so we can set c = 1 . Applying the theorem, we have that S = ∑ a i ( x i ) -E [ a i ( x i )] satisfies

with probability 1 -δ . Plugging in v and rescaling by n , with probability 1 -δ we have:

By our assumption on n for a sufficiently large constant in the big O , this implies

Since by (15), ε ≤ √ E i [ d TV ( p i , q i ) 2 ] ≤ √ E i,x [ -2 a i ( x )] , this means √

as desired.

Lemma B.2. For any y = ϕ ( θ + η ) , and therefore, for all y ,

Proof. Let Φ( x ) be the cdf of a standard Gaussian. For y > 0 ,

Thus:

Now suppose | η | ≤ B . We can upper bound γ θ ( y ) for all θ :

If y = 0 , then -θ ∗ ≥ -B , so

If y > 0 , then

as desired.

Lemma B.3 (One-dimensional net) . Let A > B 2 > 1 . There exists a partition of R into O ( A/ε ) intervals such that, for each interval I in the partition and every y = ϕ ( θ ∗ + η ) with | η | ≤ B , one of the following holds:

For all θ ∈ I , γ θ ( y ) ≤ -A

For all θ, θ ′ ∈ I , | γ θ ( y ) -γ θ ′ ( y ) | ≤ ε

Proof. To define our partition, we actually define two partitions, depending on whether y = 0 , then intersect them for our final partition.

First, consider y = 0 . By Lemma B.2, γ θ (0) is monotonically decreasing in θ , from its maximum of at most B 2 / 2 . We can thus define a partition P 1 consisting of intervals of the form I i := { θ | γ θ (0) ∈ ( B 2 / 2 -( i +1) ε, B 2 / 2 -iε ) } , for i ∈ { 0 , 1 , . . . , ( A + B 2 / 2) /ε } , plus a special interval I ′ of { θ | γ θ (0) < -A } . When y = 0 , this partition satisfies the desired conclusion to the lemma: | γ θ (0) -γ θ ′ (0) | ≤ ε for all θ, θ ′ ∈ I i , while γ θ (0) < -A for θ ∈ I ′ . Call this partition P 0 , which has size O ( A/ε ) .

Second, consider y > 0 . Define R = √ 2 A + B . Note that R 2 ≲ A and ( R -B ) 2 ≥ 2 A . Therefore for | θ -θ ∗ | ≥ R ,

Consider any θ, θ ′ ∈ [ θ ∗ -R,θ ∗ + R ] with α := | θ -θ ′ | . We have

Thus, for α = ε 2 R , this is at most ε . If we partition [ θ ∗ -R,θ ∗ + R ] into lengthα intervals, we get a size O ( R 2 /ε ) = O ( A/ε ) partition P 1 of R that has the desired property for all y > 0 .

Our final partition is defined by all endpoints in either P 0 and P 1 . This has size O ( A/ε ) , and within each interval the conclusion holds for both y = 0 and y > 0 , as needed.

Proof. For any w ∈ R k , and a sample ( x i , y i ) , let p w ( y | x i ) be the conditional distribution of y = ϕ ( ⟨ x i , w ⟩ + η ) , and let γ i,w be the log-likelihood ratio between w and w ∗ on this sample:

Then

Define

Concentration. From Lemma B.1, we see that if ˜ d ( w ∗ , w ) ≥ ε , then for n ≥ O ( 1 ε 2 log 1 δ ) ,

with probability 1 -δ .

Of course, whenever γ w < 0 , the likelihood under w ∗ is larger than the likelihood under w . Thus, for each fixed w with ˜ d ( w ∗ , w ) ≥ ε , maximizing likelihood would prefer w ∗ to w with probability 1 -δ if n ≥ O ( 1 ε 2 log 1 δ ) .

Nothing above is specific to our ReLU-based distribution. But to extend to the MLE over all w , we need to build a net using properties of our distribution.

Building a net. First, with high probability, | η i | ≤ B = O ( √ log n ) for all i . Suppose this happens. For each i , by an abuse of notation, let γ i,w ( y ) = γ ⟨ x i ,w ⟩ ( y ) where the value of θ ∗ when considering i is ⟨ x i , w ∗ ⟩ . By Lemma B.2,

for all i . Let A = O ( n log n ) > nB 2 . By Lemma B.3, for each i ∈ [ n ] , there exists a partition P i of R into O ( A/ε 2 ) intervals, such that for interval I ∈ P i , and any w,w ′ with x T i w,x T i w ′ ∈ I , either

These individual partitions P i on ⟨ x i , w ⟩ induce a partition P on R k , where w,w ′ lie in the same cell of P if ⟨ x i , w ⟩ and ⟨ x i , w ∗ ⟩ are in the same cell of P i for all i ∈ [ n ] . Since P is defined by n sets of O ( A ε 2 ) parallel hyperplanes in R k , the number of cells in P is:

We choose a net N to contain, for each cell in P , the w in the cell maximizing ˜ d ( w ∗ , w ) . This has size

By (18), for our n ≥ O ( 1 ε 2 k log k ε ) , we have with high probability that γ w ≤ -ε 2 4 , for all w ∈ N with ˜ d ( w ∗ , w ) ≥ ε . Suppose that both this happens, and | η i | ≤ B for all i . We claim that the MLE ̂ w must have ˜ d ( w ∗ , ̂ w ) < ε .

Consider any w ∈ R d with ˜ d TV ( w ∗ , w ) ≥ ε . Let w ′ ∈ N lie in the same cell of P . By our choice of N , we know ˜ d TV ( w ∗ , w ′ ) ≥ ˜ d TV ( w ∗ , w ) ≥ ε , so γ w ′ ≤ -ε 2 . Now we consider two cases. In the first case, there exists i with γ i,w ( y i ) < -A . Then

Otherwise, by (19),

In either case, γ w < 0 and the likelihood under w ∗ exceeds that under w . Hence the MLE ̂ w must have ˜ d ( w ∗ , ̂ w ) ≤ ε .

Theorem B.5. Let y = ϕ ( x T w ∗ + η ) , for w ∗ ∈ R k , x ∼ D x , and η ∼ N (0 , 1) . Then for a sufficiently large constant C > 0 ,

samples of { ( y i , x i ) } n i =1 suffices to guarantee that the MLE ̂ w satisfies

Proof. Let D x denote the dataset { x i } i ∈ [ n ] that is used to find the MLE. Notice that the MLE is found using this finite subset, but we would like to make a claim about D x without making any parametric or simplifying assumptions on the distribution D x .

An application of Lemma 4.3 tells us that with probability 1 -e -Ω( nε 2 ) , the expectation over the distribution D x and the dataset D x are within ε/ 2 of one another:

Now, all we need to show is that the MLE has a small TV distance on the finite dataset, and Lemma B.4 tells us that with probability 1 -e -Ω( nε 2 ) ,

Substituting in the above inequality, we get d ( ̂ w,w ∗ ) ≤ ε .

C ReLU Activations with d > 1 , Unknown Covariance

We recommend the reader review Appendix B, which contains the proof recipe for the case of scalar y . The proofs in this section generalize those of Appendix B.

Consider a sample ( x, y ) ∈ R k × d , with

where W ∗ ∈ R d × k , and noise η ∼ N (0 , Σ ∗ ) . The matrices W ∗ and Σ ∗ are unknown. For each matrix W ∈ R d × k , let θ = Wx ∈ R d , denote a reparametrization of W , and let θ ∗ denote θ ∗ = W ∗ x . Let S denote the co-ordinates of y that are zero-valued. Then the log-likelihood for each θ, Σ is given by

where c is a normalization constant which does not depend on θ or Σ . Let

and let P ∗ be the precision matrix of the noise η , and P S , P SS c , P S c S , P S c be the block matrices of P corresponding to the index sets S and its complement S c .

By some arithmetic involving completion of squares, we can decompose the integral in f into the sum of two functions g, h , such that

The first term g corresponds to the quadratic term involving the observed positive-valued coordinates y S c :

As the matrix P S c -P S c S ( P S ) -1 P SS c = (( P -1 ) S c ) -1 = Σ -1 S c is the precision matrix of η S , if Σ were the covariance of η , we can simplify the above equation as

The second term corresponds to the probability under θ, P of observing zero-valued coordinates corresponding to the index set S , given the positive coordinates y S c :

The log-likelihood ratio is the difference between f θ, Σ and f θ ∗ , Σ ∗ , which we denote by

Over a dataset { ( x i , y i ) } i ∈ [ n ] , the average log-likelihood ratio is given by

Remark C.1. For ease of analysis, we will interchange between the precision matrix P in γ θ,P and the covariance matrix Σ in γ θ, Σ , and it should be understood that P = Σ -1 . The same applies to the functions g θ, Σ and h θ, Σ . Finally, the matrix P ∗ refers to the ground truth precision matrix (= Σ ∗-1 ) .

Analogous to Appendix B, we start by showing that the log-likelihood ratio is bounded by the noise in the sample. The proofs of results in this Section are in Subsection C.1.

Lemma C.2. Assume P ∗ := Σ ∗-1 satisfies Assumption 4.4.

For all y = ϕ ( θ ∗ + η ) such that S denotes the zero-coordinates of y , and η such that ∥ P ∗ 1 2 S η S ∥ , ∥ P ∗ 1 2 S c η S c ∥ ≤ B , if the max eigenvalue λ max ( P ) satisfies

For the ease of stating the next Lemma, we assume that across the samples of y in the training data, at least one coordinate has sufficiently many positive samples. The proof of our theorem separately handles cases violating this assumption.

Assumption C.3. Let δ ∈ (0 , 1) be a parameter corresponding to the failure probability of our algorithm. Then, there exists a coordinate j ∈ [ d ] , such that for at least n ′ = O ( log 1 δ ) samples y i 1 , . . . , y i n ′ in the dataset, the j -th coordinate is positive.

This is a very weak assumption: if it is violated, then W = 0 d × k , Σ = 0 will achieve a TV distance smaller than 2 ε 2 d .

Appendix B assumed that the variance in y was 1. Since Section 4.2 considers an unknown Σ ∗ , we need the following Lemma to show that the MLE will select a precision matrix P , whose eigenvalues are reasonably bounded wrt Σ ∗-1 .

Lemma C.4. Under Assumption 4.4, C.3, consider P ∈ R d × d + such that λ max ( P ) λ min ( P ) ≤ κ and

Then, for all W ∈ R d × k , and for all y i = ϕ ( W ∗ x i + η i ) with ∥ P ∗ 1 2 S c η S c ∥ , ∥ P ∗ 1 2 S η S ∥ ≤ B , we have

Lemma C.2 and Lemma C.4 show that the MLE will only select precision matrices P that have max eigenvalues in a certain range of the true precision matric P ∗ .

Now, for matrices in the above eigenvalue range, we first construct a geometric net over the max eigenvalue ρ of the precision matrix, and then cover the matrices whose max eigenvalue is smaller than ρ .

Lemma C.5 ( Σ cover) . For B > 1 , and 0 < L < U , let A > max {√ log 1 ε , B 2 Uκ, d 2 log ( κU L ) , 1 } .

Let P ∗ := Σ ∗-1 be the precision matrix of η . Let Ω ⊂ R d × d + denote the set of positive definite matrices P ∈ R d × d + with condition number κ and whose maximum eigenvalue lies in [ Lλ min ( P ∗ ) , U · λ max ( P ∗ )] .

Then, there exists a partition of Ω of size

such that for all θ ∈ R d and all y = ϕ ( θ ∗ + η ) ∈ R d with ∥ P ∗ 1 2 S η S ∥ , ∥ P ∗ 1 2 S c η S c ∥ ≤ B , and each cell I in the partition, one of the following holds:

for all P ∈ I , γ θ,P ( y ) < -A , or

for all P, P ′ ∈ I , we have | γ θ,P ( y ) -γ θ,P ′ ( y ) | ≤ ϵ .

then for all θ ∈ R d , we have

Analogous to Appendix B, we now construct a partition over W for a fixed precision matrix P , such that each cell in the partition has very small log-likelihood (in which case the MLE will not choose it) or the log-likelihood changes slowly.

Lemma C.6 ( W -net) . Let η S c , η S be such that

for B 1 , B 2 ≥ 0 .

Let A > max { B 2 1 , B 2 2 , poly( C, κ ) } . Let P ∗ = Σ ∗-1 be the precision matrix of η . For a fixed matrix P ∈ R d × d whose condition number satisfies Assumption 4.4 and whose eigenvalues satisfy λ max ( P ) ∈ [ e -2 A d λ min ( P ∗ ) , Cλ max ( P ∗ )] , there exists a partition I of R d with size

such that for each interval I ∈ I , we have one of the following:

for all θ ∈ I , γ θ,P ( y ) < -A , or

for all θ, θ ′ ∈ I , | γ θ,P ( y ) -γ θ ′ ,P ( y ) | ≤ ϵ .

Using the above lemmas, we can show that the MLE will only pick out ̂ W, ̂ P such that they have small TV on the dataset of { x i } .

Lemma C.7. Let x 1 , . . . , x n be fixed, and y i = ϕ ( W ∗ x i + η i ) for η i ∼ N (0 , Σ ∗ ) , and W ∗ ∈ R d × k with Σ ∗ ∈ R d × d satisfying Assumption 4.4 and Assumption C.3. For a sufficiently large constant C > 0 ,

samples suffice to guarantee that with high probability, the MLE ̂ W, ̂ Σ satisfies

Lemma C.8. Let { x i } n i =1 be i.i.d. random variables such that x i ∼ D x .

Let P ∗ := Σ ∗-1 . Let λ ∗ min , λ ∗ max be the minimum and maximum eigenvalues of P ∗ . For 0 < L < U , let Ω denote the following set of precision matrices

Then, for a sufficiently large constant C > 0 , and for

we have:

Theorem 4.5. Let R d × d κ denote the set of positive definite matrices with condition number κ . Given n samples { ( x i , y i ) } n i =1 satisfying Assumption 4.4, where x i ∼ D x i.i.d., and y i is generated according to (7) , let ̂ W, ̂ Σ := arg max W ∈ R d × k , Σ ∈ R d × d κ 1 n ∑ i log p W, Σ ( y i | x i ) . Then, for a sufficiently large constant C > 0 ,

samples suffice to ensure that with probability 1 -δ , we have

Proof of Theorem 4.5. First, we consider the cases violating Assumption C.3.

As n ∝ d 2 ε 2 log 1 δ , if assumption C.3 is violated, then it implies that each coordinate is non-zero in atmost a ε 2 /d 2 fraction of the samples, and a union bound implies that the probability of seeing a non-zero vector is atmost ε 2 /d . Hence, with high probability over the draws of the data, returning the all-zeros vector always will achieve a TV distance smaller than 2 ε 2 d .

Let ̂ P,P ∗ = ̂ Σ -1 , Σ ∗-1 . Now, if Assumption C.3 holds, Lemma C.7 guarantees that the MLE has small TV on the x i observed in the dataset:

The above result is over the finite x i observed in our dataset. To generalize it over x ∼ D x , we use Lemma C.8, which gives

Rescaling ε gives the conclusion of the Theorem.

C.1 Proofs of Appendix C.

Lemma C.2. Assume P ∗ := Σ ∗-1 satisfies Assumption 4.4.

For all y = ϕ ( θ ∗ + η ) such that S denotes the zero-coordinates of y , and η such that ∥ P ∗ 1 2 S η S ∥ , ∥ P ∗ 1 2 S c η S c ∥ ≤ B , if the max eigenvalue λ max ( P ) satisfies

then for all θ ∈ R d , we have

Proof. We have

Substituting in Eqn (23), we get

As ( P ∗ S ) -1 = Σ ∗ S -Σ ∗ SS c Σ ∗-1 S c Σ ∗ S c S , by the matrix determinant rule, we have log | Σ ∗ | +log | P ∗ S | = log | Σ ∗ S c | .

This gives

This gives

As the matrix Σ -1 S c is positive definite, we trivially get

Substituting in Eqn (24), we get

As the Lemma assumes we get

which gives

Lemma C.9. Consider the function g defined in Eq (21) . For the ground truth parameters θ ∗ , Σ ∗ , the function g θ ∗ , Σ ∗ satisfies

which is, with probability 1 -e -Ω( d ) ,

Proof. As y S c are the positive valued coordinates in y , we have

As η S c is Gaussian with covariance Σ ∗ S c , the expected norm is | S c | 2 , which implies that with probability 1 -e -Ω( | S c | ) , we have

Lemma C.10. Consider y generated according to Eqn (20) by

For all θ ∈ R d , Σ ∈ R d × d + , and the function h θ, Σ defined in Eqn (22) , the difference h θ, Σ ( y ) -h θ ∗ , Σ ∗ ( y ) satisfies

where P ∗ = Σ ∗-1 is the precision matrix of η .

Proof. For θ ∈ R d , Σ ∈ R d × d + , and P = Σ -1 , we have

where the last step follows from the integral of a Gaussian pdf. This gives a sufficient upper bound on h θ, Σ ( y ) , and now we will focus on lower bounding h θ ∗ , Σ ∗ ( y ) .

For the coordinates of y in S c , we have y S c -θ ∗ S c = η S c . Substituting in Eqn (22), we get

Using ∥ a + b ∥ 2 ≤ 2 a 2 +2 b 2 , we get

Set u := P ∗ 1 2 S ( t -θ ∗ S ) , and by the change of variables formula, we get:

For i ∈ S , we have θ ∗ i + η i ≤ 0 . This gives

P ∗-1 2 S u ≤ η S ⇒ P ∗-1 2 S u + θ ∗ S ≤ 0 ⇒ log ∫ P ∗-1 2 S u + θ ∗ S ≤ 0 exp ( -∥ u ∥ 2 ) ≥ log ∫ P ∗-1 2 S u ≤ η S exp ( -∥ u ∥ 2 ) , using which we get

which gives

From Eqn (26) -Eqn (27), we get

Lemma C.4. Under Assumption 4.4, C.3, consider P ∈ R d × d + such that λ max ( P ) λ min ( P ) ≤ κ and

Then, for all W ∈ R d × k , and for all y i = ϕ ( W ∗ x i + η i ) with ∥ P ∗ 1 2 S c η S c ∥ , ∥ P ∗ 1 2 S η S ∥ ≤ B , we have ¯ γ W,P := 1 n ∑ i ∈ [ n ] γ Wx i ,P ( y i ) < 0 .

Proof of Lemma C.4. For each W ∈ R d × k , let

From Eqn (24) in Lemma C.2, for each i ∈ [ n ] , we have,

Now consider

where the second inequality comes from the eigenvalue interlacing Theorem, and the last line follows from the condition number assumption on Σ , P .

By Assumption C.3, there exist at least ε 2 n samples for a coordinate j such that ( y i ) j > 0 . Averaging g θ i , Σ , by Lemma A.1, we get that with high probability,

which gives

This gives

Completing the squares, we get

For

the above inequality satisfies

Lemma C.11. Assume P ∗ := Σ ∗-1 satisfies Assumption 4.4 with condition number κ .

For all y = ϕ ( θ ∗ + η ) such that S denotes the zero-coordinates of y , and η such that ∥ P ∗ 1 2 S η S ∥ , ∥ P ∗ 1 2 S c η S c ∥ ≤ B , consider precision matrices P whose max eigenvalue λ max ( P ) satisfies

Let A ≥ 4 max { d 2 log C, 3 B 2 } . Then for V := ( P -1 ) S c and R P defined as

Proof of Lemma C.11. Consider Eqn (24) in Lemma C.2. We have

where the last inequality follows from λ max ( P ) λ min ( P ∗ ) ≤ C in the statement of the Lemma. By the definition of g θ,P , we have

We can rewrite the matrix ( P S c -P S c S P -1 S P SS c ) as

we have

By setting we can rewrite g θ,P as

Now, as y S c = η S c + θ ∗ S c , we have

Ignoring the ∥ η S c ∥ 2 V term, we get

By the Cauchy-Schwartz inequality and the eigenvalue interlacing theorem, we have

By the statement of the Lemma, we have ∥ P ∗ 1 2 S c η S c ∥ ≤ B = ⇒ ∥ η S c ∥ 2 ≤ B λ 1 2 min ( P ∗ S c ) . Substituting in the above inequality, we get

Substituting in the function g θ,P , we get

Hence, for θ satisfying we get

In order to cover our precision matrices, we will consider a subset of matrices whose entries are quantized by an interval size β :

Definition C.12 (Quantized Precision Matrices) . For κ > 0 , define Ω ⊂ R d × d as the set of positive definite matrices with condition number κ .

For ρ > 0 , define the set Ω ρ ⊂ Ω as

For a quantization size β > 0 , define ˜ Ω ρ,β ⊂ Ω ρ as:

Lemma C.5 ( Σ cover) . For B > 1 , and 0 < L < U , let A > max {√ log 1 ε , B 2 Uκ, d 2 log ( κU L ) , 1 } .

Let P ∗ := Σ ∗-1 be the precision matrix of η . Let Ω ⊂ R d × d + denote the set of positive definite matrices P ∈ R d × d + with condition number κ and whose maximum eigenvalue lies in [ Lλ min ( P ∗ ) , U · λ max ( P ∗ )] .

Then, there exists a partition of Ω of size

such that for all θ ∈ R d and all y = ϕ ( θ ∗ + η ) ∈ R d with ∥ P ∗ 1 2 S η S ∥ , ∥ P ∗ 1 2 S c η S c ∥ ≤ B , and each cell I in the partition, one of the following holds:

for all P ∈ I , γ θ,P ( y ) < -A , or

for all P, P ′ ∈ I , we have | γ θ,P ( y ) -γ θ,P ′ ( y ) | ≤ ϵ .

Proof. In order to construct the net over the precision matrices, we will consider geometrically spaced values of ρ ∈ [ L · λ min ( P ∗ ) , U · λ max ( P ∗ )] , and for each ρ , we will construct a net over matrices that have max eigenvalue ≤ ρ .

Now consider ρ > 0 that lies in the following discrete set:

This set is a geometric partition over the possible max eigenvalues that the MLE can return.

For the current ρ , let Ω ρ follow Definition C.12. Now consider P ∈ Ω ρ .

Constructing the interval for which γ θ,P < -A . By Lemma C.11, for V = ( P -1 ) S c , and R P = O ( B √ ρ λ min ( P ∗ ) + √ A ) = O ( √ A ) , we have

For any θ , notice that the set of matrices P satisfying ∥ θ S c -θ ∗ S c ∥ V ≥ R P is connected (as its complement is compact). This forms the set I for which γ θ,P < -A .

Constructing intervals for which | γ θ,P -γ θ,P ′ | ≤ ε . We will now construct a partition over those P which satisfy ∥ θ S c -θ ∗ S c ∥ V < R P , and show that the log-likelihood changes by atmost ε for each cell in this partition.

If P ∈ Ω ρ , then each of its elements P ij ∈ [ -ρ, ρ ] . For a parameter β > 0 that we will specify later, consider the partition ˜ Ω ρ,β of Ω ρ , following Definition C.12. Clearly, the size of ˜ Ω ρ,β can be upper bounded by

We will now analyze the effect of rounding down P ∈ Ω ρ to its nearest element in ˜ Ω ρ,β .

By Claim C.13, for γ = 2 κβd 2 , we have

Consider the log-likelihood at θ, P ′ :

We will use the LHS of Eqn (30) to show that

and deal with the log | P ′ | term later. The lower bound for the log-likelihood at P ′ can be obtained via analogous proof using the RHS of Eqn (30).

By the LHS of Eqn (30), we get

Rearranging the terms, we get

The non-integral term corresponds to g θ,P in Eqn (21), while the integral term corresponds to h θ,P in Eqn (22).

Handling the non-integral term. As we are only considering θ such that ∥ y S c -θ S c ∥ Σ S c ≤ R P , we have that for

the non-integral term corresponds to g θ,P + ε , which gives

Handling the integral. Now we consider the integral term. Define the integral

for γ = 2 d 2 κβ and µ = θ S -P -1 S P SS c ( y S c -θ S c ) .

Define the analogous integral that does not have the (1 -γ ) term in the exponential:

We only need to consider µ such that ∥ µ ∥ ∞ ≤ O ( √ Aρ ) : otherwise the likelihood will be smaller than -A .

Clearly, I 1 ≥ I 2 .

By Lemma C.14, for γ = O ( ε poly( A ) ) , we have

Handling the log-determinant term. Now consider the log | P | term. As we are decreasing each element by atmost βρ , none of the eigenvalues can increase. Moreover, as

we can conclude that each eigenvalue decreases by at most -dβρ . Also, as ρ ≤ κλ j ( P ) ∀ j ∈ [ d ] , we can conclude that each eigenvalue satisfies

Hence, the log-determinant satisfies

This finally gives

Bounding the size of the net As β = O ( ε poly( A ) ) , and the max radius is also O (poly( A )) , we have a cover of size ( poly( A ) ε ) per entry of the precision matrix (for a fixed Ω ρ ).

Intersecting the d 2 nets means that for each Ω ρ , we have a net of size

As we are considering poly( A ) many Ω ρ s, the size of the net remains the same as the above.

Claim C.13. In the setting of Lemma C.5, if P ∈ Ω ρ and P ′ ∈ ˜ Ω ρ,β is its nearest neighbor, then for γ = 2 κβd 2 , we have

where Σ := P -1 , Σ ′ := P ′ -1 .

Proof. Consider P ∈ Ω ρ and P ′ ∈ ˜ Ω ρ,β such that P = P ′ +∆ . Since P ′ is the rounding down of P , we have ∆ ij ∈ [0 , βρ ] .

As Tr (∆) ∈ [0 , ρβd ] , and ∥ ∆ ∥ F ≤ ρβd , we have

This implies that when considering untruncated Gaussians with precision matrices P, P ′ , we have that for all t, θ ∈ R d ,

Since λ min ( P ) ≥ ρ 2 κ , we have

Substituting in the previous inequality, we get

For the sake of symmetry, we will use the weaker bound of

Setting γ = 2 κβd 2 completes the proof.

Lemma C.14. Consider a bounded mean vector µ with ∥ µ ∥ ∞ ≤ α and precision matrix P with max eigenvalue ρ and condition number κ .

Proof of Lemma C.14. Wlog, consider µ ≥ 0 . The case where the entries are possibly negative follow a similar proof.

Define the integral on the LHS and RHS of the Lemma statement by I 1 and I 2 respectively.

By a change of variables, we set t ′ = √ 1 -γ ( t -µ ) + µ in I 1 , to get

Since γ < 1 , we have (1 - √ 1 -γ ) µ < γµ . Substituting in I 1 , and for γ = O ( ε ) , we get

The integrating set in the above inequality can be split into two parts: one over the negative orthant (which is exactly to e I 2 ) and another over the shell

This gives

In the above inequality, let e I 3 denote the integral over the shell C . We will now show that I 3 satisfies e I 3 ≤ εe I 2 .

Let f ( x ) denote the Gaussian density with mean µ and precision matrix P .

̸

For a subset of co-ordinates S ⊆ [ d ] , S = ∅ , and t ∈ R d , let x + , x -∈ R d be such that

By the monotonicity of the Gaussian density, the integral over the shell C can be upper bounded by breaking up into a sum of integrals over lower-dimensional strips, where for a fixed subset S ∈ [ d ] , the variables t S c are integrated over ( -∞ , γµ S c ] , while the variables in S are fixed to γµ S .

This gives

By Claim C.15, for any S , and γ = O ( min { ε α √ ρd , ε α 2 d 2 ρ }) each summand satisfies f ( x + ,S ) ≤ 2 f ( x -,S )

Furthermore, for any S ⊆ [ d ] , we have

Rescaling ε ← ε 3 d completes the proof.

Claim C.15. Let f be the Gaussian density with mean µ ∈ [0 , α ] d and precision matrix P ∈ R d × d with max eigenvalue ρ and condition number κ .

̸

Let γ = O ( min { ε α √ ρd , ε α 2 d 2 ρ }) . For any subset of co-ordinates S ⊆ [ d ] , S = ∅ , and t ∈ R d , let x + , x -∈ R d be such that

we have

Proof. WLOG, let S be a contiguous set such that we can separate the coordinates of x + and x -into disjoint sets. For the coordinates belonging to S , let µ S denote the coordinates of µ belonging to µ , and µ S c the coordinates not belonging to S (similarly for t S and t S c ).

Taking the logarithm on both sides of the claimed inequality, we want to show that

Let a and b denote the vectors whose norms correspond to the log-densities in the claimed inequality, and let δ = a -b .

This gives

This gives

We want to show that

For γ = O ( ε α √ ρd ) , we get where

for B 1 , B 2 ≥ 0 .

Let A > max { B 2 1 , B 2 2 , poly( C, κ ) } . Let P ∗ = Σ ∗-1 be the precision matrix of η . For a fixed matrix P ∈ R d × d whose condition number satisfies Assumption 4.4 and whose eigenvalues satisfy λ max ( P ) ∈ [ e -2 A d λ min ( P ∗ ) , Cλ max ( P ∗ )] , there exists a partition I of R d with size

such that for each interval I ∈ I , we have one of the following:

for all θ ∈ I , γ θ,P ( y ) < -A , or

for all θ, θ ′ ∈ I , | γ θ,P ( y ) -γ θ ′ ,P ( y ) | ≤ ϵ .

As ∥ P ∥ ≤ ρ and ∥ µ ∥ ∞ ≤ α , we have

Similarly, consider the inner product ⟨ δ, b ⟩ in Eqn (33). By the trace trick, we get

Notice that the diagonal elements of ∆ are all non-negative. This implies that all singular values are non-negative. The trace of ∆ is

Hence, by Von Neumann's trace inequality, we get

For γ = O ( ε α 2 | S | ρd ) , this gives

Substituting Eqn (34) and Eqn (35) in Eqn (33) completes the proof.

Lemma C.6 ( W -net) . Let η S c , η S be such that

Proof of Lemma C.6. Recall that the log-likelihood ratio γ θ can be decomposed into the difference of two terms that depend on θ :

Without loss of generality, consider the net for the first coordinate θ 1 . The final net will be the intersection of the per-coordinate nets.

We will construct three partitions: the first is I h, 0 for h when y 1 = 0 , the second is is I h, 1 for h when y 1 > 0 , and the last is I g for g when y 1 > 0 . The final partition will be the intersection of these partitions.

Case 1: Net over h , y 1 = 0 . As y 1 = 0 , we have 1 ∈ S . For θ ∈ R d , we have

By Claim C.16, if θ 1 ≥ Θ( √ CκA p 1 ) , then the log-likelihood is smaller than -A .

Now, consider θ 1 < O ( √ CκA p 1 ) . Let θ ′ = θ + αe 1 for α > 0 . As h is monotonically decreasing per coordinate, θ 1 < θ ′ 1 = ⇒ h θ,P ≥ h θ ′ ,P . We would like to now upper bound h θ,P in terms of h θ ′ ,P . Let

In the function h θ , break the integrating set into two domains: one where t -µ is small,

for some r > 0 that we will specify later.

Let I 1 and I 2 denote the integrals over Ω 1 and Ω 2 respectively.

I 2 corresponds to the tail of an unnormalized Gaussian distribution, and hence we have

We can simplify I 2 be comparing | P | to | P ∗ | :

By Lemma C.10, we have and another where it is large:

As we are only consider θ ′ such that h θ ∗ ,P ∗ ( y ) -A < h θ ′ ,P ( y ) , for

we have

Subsituting in h θ,P ( y ) , we get

Now consider the integral I 1 = ∫ t ≤ 0 ,t ∈ Ω 1 exp ( -∥ P 1 2 S ( t -µ ) ∥ 2 / 2 ) .

′

By Claim C.17, as Ω 1 is defined for t bounded by r , and µ -µ = αe 1 , we have

Substituting in the expression for h θ,P , we get

As log( ε + e x ) ≤ ε + x for x ≥ 0 , we have

This shows that h θ,P changes by at most ε for the considered net. We need to defined the other end point for the net. By a similar argument to the positive end point, if θ 1 = -O ( √ Cκ log ( 1 ε ) p 1 ) , the log-likelihood ratio changes by at most ε until θ 1 = -∞ .

As we are only trying to cover θ such that | θ 1 | ≤ O ( √ CκA p 1 ) , this net has size

Case 2: Net over h , y > 0 . A similar argument to Case 1 works here as well.

Case 3: Net over g , y 1 > 0 . By Lemma C.11, if | θ -θ ∗ | P S c > R 1 for

Setting α = O ( ε p 1 r ) , we get then

Now consider θ such that and θ, θ ′ such that θ -θ ′ = αe 1 .

The difference in g θ -g θ ′ is

The second and third terms in the RHS can bounded by observing that

and hence we get

Now, for the first term in the RHS, we have

This further implies that

Setting we get

As we are covering a set of size R 1 using a grid size of α , the size of this partition is

Claim C.16. In the setting of Lemma C.6, we have λ max ( P ) ≤ Cλ max ( P ∗ ) . Let p 1 denote the first diagonal element of P .

If θ 1 ≥ Θ( √ CκA p 1 ) , then the function h θ,P is such that

Proof of Claim C.16. Recall that the function h θ,P is defined as:

Consider the term ∥ P 1 2 S ( t -θ S ) + ( P S ) -1 / 2 P SS c ( y S c -θ S c ) ∥ . By the triangle inequality, we have

where the last inequality follows as

Similarly, the function g θ,P only considers θ such that ∥ P 1 2 S c ( θ S c -θ ∗ S c ) ∥ ≤ √ CκA ( otherwise, the log-likelihood ratio is smaller than -A by virtue of g θ,P , irrespective of h θ,P ). For these θ , we have

which gives

Hence, if θ 1 ≥ O ( CκA p 1 ) , then we have

and hence the Gaussian integral is at most (2 π ) | S | / 2 ∣ ∣ ∣ P -1 2 S ∣ ∣ ∣ e -Ω( CκA )

.

By Lemma C.10, we have h θ ∗ ,P ∗ ≥ -1 2 log | P ∗ S | -O ( A ) , which gives

This gives a contiguous interval over θ 1 for which γ θ,P < -A .

Claim C.17. In the setting of Lemma C.6, let µ, µ ′ be such that µ -µ ′ = αe 1 Then, for all t such that

and p 1 := P 11 , we have

Proof of Claim C.17. Consider the term ∥ P 1 2 S ( t -µ ) ∥ 2 .

Adding and subtracting

As µ -µ ′ = αe 1 , we have

By the Cauchy-Schwartz inequality, and since µ -µ ′ = αe 1 , the inner product can be lower bounded as

Substituting, we get

This completes the proof.

Lemma C.18. Following Definition C.12 let Ω ρ be the set of precision matrices with condition number κ satisfying λ max ( P ) ∈ [ ρ 2 , ρ ] , and let ˜ Ω ρ,β be the quantized net with quantization level β .

For any P ∈ Ω ρ , let ˜ P ∈ ˜ Ω ρ,β be its element-wise rounding down. Then, for any µ ∈ R d , we have

Proof of Lemma C.18. Let Σ = P -1 and ˜ Σ = ˜ P -1 .

By Theorem 1.1 in [16], the TV between two Gaussians with the same mean is

where ξ i are the eigenvalues of ˜ Σ -1 Σ -I d .

We can convert the bound on the eigenvalues to the Frobenius norm of ˜ Σ -1 Σ -I d :

Recall that ˜ P is the rounding down per entry of P . Hence,

Taking the Frobenius norm, we get

where the first inequality follows as each element of

P

has condition number

κ

This completes the proof.

Lemma C.7. Let x 1 , . . . , x n be fixed, and y i = ϕ ( W ∗ x i + η i ) for η i ∼ N (0 , Σ ∗ ) , and W ∗ ∈ R d × k with Σ ∗ ∈ R d × d satisfying Assumption 4.4 and Assumption C.3. For a sufficiently large constant C > 0 ,

samples suffice to guarantee that with high probability, the MLE ̂ W, ̂ Σ satisfies

Proof of Lemma C.7. For any W ∈ R d × k , Σ ∈ R d × d and a sample ( x i , y i ) , let p i,W, Σ ( y | x i ) be the conditional distribution of y = ϕ ( Wx + η ) , and let γ i,W, Σ be the log-likelihood ratio between ( W, Σ) and ( W ∗ , Σ ∗ ) on this sample:

, and the third follows as

ν

P

is at most

∈

Ω

ρ

βρ

, the second inequality follows as

ρ

(

P

)

ρ

2

.

=

Then

Concentration. From Lemma B.1, we see that if d TV (( W ∗ , Σ ∗ ) , ( W, Σ)) ≥ ε , then for n ≥ O ( 1 ε 2 log 1 δ ) ,

with probability 1 -δ .

Of course, whenever γ W, Σ < 0 , the likelihood under W ∗ , Σ ∗ is larger than the likelihood under W, Σ . Thus, for each fixed W, Σ with d TV (( W ∗ , Σ ∗ ) , ( W, Σ)) ≥ ε , maximizing likelihood would prefer W ∗ , Σ ∗ to W, Σ with probability 1 -δ if n ≥ O ( 1 ε 2 log 1 δ ) .

Nothing above is specific to our ReLU-based distribution. But to extend to the MLE over all W, Σ , we need to build a net using properties of our distribution.

Building a net. First, for a given sample y i , let S i be the set of coordinates of y i that are zero, and S c i be its complement. Then, with high probability,

where P ∗ = Σ ∗-1 , and P ∗ S , P ∗ S c are the block matrices in P ∗ .

Supposing the above event happens, we will construct a net over the precision matrices P = Σ -1 . Note that as we are only considering matrices with bounded condition number, this is a bijective mapping.

Net over Σ -1 . By Lemma C.4, any precision matrix P = Σ -1 satisfying Assumption 4.4 and whose max eigenvalue satisfies

Similarly, by Lemma C.2, any precision matrix satisfying Assumption 4.4 and whose max eigenvalue satisfies

for L = e -O ( κ log n ) has γ i,W,P -1 ≤ 0 for all i ∈ [ n ] , and hence its average ¯ γ W,P -1 is also < 0 .

This shows that for all precision matrices P whose max eigenvalue is extremely small / large when compared to the min / max eigenvalues of P ∗ , has

irrespective of W , and the MLE, which has non-negative ¯ γ , will never pick these P .

Let A = poly( n, d, κ, 1 ε ) be large enough such that A > n ( d log( Uκ ) + B 2 ) , and such that it meets the requirements of Lemma C.5 and Lemma C.6.

Then, by Lemma C.5 with U = poly( d, κ, n ) and log( 1 L ) = poly( κ, n ) , there exists a partition P of precision matrices whose max-eigenvalue lies in [ L · λ max ( P ∗ ) , U · λ max ( P ∗ )] into ( poly ( d, κ, n, 1 ε )) d 2 cells, such that for each cell I ∈ P , and P, P ′ ∈ I , the following holds for all i ∈ [ n ] and W ∈ R d × k :

or γ i,W,P ( y i ) < -A .

Using Lemma C.18, we also have that for all W ,

We can choose a net N consisting of precision matrices from each cell in P . This net has size

This gives a sufficient net over the precision matrices.

Now we will construct a net over W for each precision matrix in the net.

W -net. Now, for each ˜ P ∈ N P , by Lemma C.6, for each i ∈ [ n ] , there exists a partition P ˜ P,i of R d into ( poly ( d, k, κ, n, 1 ε )) d cells such that for each cell I ∈ P ˜ P,i , and W,W ′ ∈ I , one of the following holds:

or γ i,W, ˜ P ( y i ) < -A .

Let W j be the j -th row of W . The individual partitions P ˜ P,i on ⟨ x i , W j ⟩ induce a partition P ˜ P,i,j on R k , where W j , W ′ j lie in the same cell of P ˜ P,i,j if ⟨ x i , W j ⟩ and ⟨ x i , W ′ j ⟩ are in the same cell of P ˜ P,i for all i ∈ [ n ] . Since P ˜ P,i,j is defined by n sets of ( poly ( d, k, κ, n, 1 ε )) parallel hyperplanes in R k , the number of cells in P ˜ P,i,j is

As there are d rows in W , we can intersect P ˜ P,i,j over j ∈ [ d ] , which induces poly( d, k, κ, n, 1 ε ) kd cells in R k . We choose a net N ˜ P to contain, for each cell in ⋂ j ∈ [ d ] P ˜ P,i,j , the W in the cell maximizing d TV (( W ∗ , P ∗ ) , ( W, ˜ P )) . This has size

Proving MLE works. By (37), for our n ≥ O ( ( kd + d 2 ) ε 2 log kdκ ε ) , we have with high probability that

for all P ∈ N and for all W ∈ N ˜ P with d TV (( W ∗ , P ∗ ) , ( W,P )) ≥ ε . Suppose that both this happens, and

We claim that the MLE ̂ W, ̂ Σ must have d TV (( W ∗ , Σ ∗ ) , ( ̂ W, ̂ Σ)) < 17 16 ε .

Consider any W ∈ R d × k and P ∈ R d × d with d TV (( W ∗ , P ∗ ) , ( W,P )) ≥ 17 16 ε . Using our net on precision matrices, we can find ˜ P ∈ N such that

Recall that we are only currently considering W,P such that d TV (( W ∗ , P ∗ ) , ( W,P )) ≥ 17 16 ε . By Eqn (39), we have d TV (( W,P ) , ( W, ˜ P )) ≤ ε 2 16 , which gives

Now, for this ˜ P , we can find a ˜ W ∈ N ˜ P , and by our choice of N ˜ P , we know that

and by (37), we have ¯ γ ˜ W, ˜ P ≤ -ε 2 2 .

Now we consider two cases. In the first case, there exists i with γ i,W,P ( y i ) < -A . Then

Otherwise, by Eqn (38) and Eqn (40), we have

In either case, γ W,P < 0 and the likelihood under w ∗ exceeds that under w . Hence the MLE ̂ w must have d TV ( w ∗ , w ) ≤ 17 16 ε . Rescaling ε gives the conclusion of the Lemma.

Lemma C.8. Let { x i } n i =1 be i.i.d. random variables such that x i ∼ D x .

Let P ∗ := Σ ∗-1 . Let λ ∗ min , λ ∗ max be the minimum and maximum eigenvalues of P ∗ . For 0 < L < U , let Ω denote the following set of precision matrices

Then, for a sufficiently large constant C > 0 , and for

we have:

Proof of Lemma C.8. For P = Σ -1 and P ∗ = Σ ∗-1 , let

and

Since the function is bounded, for any fixed W,P , the Chernoff bound gives

for any α > 0 . The challenge lies in constructing a net to be able to union bound over R k without assuming any bound on W or the covariate x . As before, we do so by constructing a 'ghost' sample, symmetrizing, and constructing a net based on these samples.

Ghost sample. First, we construct a 'ghost' dataset D ′ x consisting of n fresh samples IID samples { x ′ i } i ∈ [ n ] of D x . This gives another metric

Similar to the proof in Lemma 4.3, it is sufficient to consider the difference between f n ( W,P ) and f ′ n ( W,P ) i.e.,

Symmetrization. Since D x and D ′ x each have n independent samples, we could instead draw the datasets by first sampling 2 n elements x 1 , . . . , x 2 n from D x , then randomly partition this sample into two equal datasets. Let s i ∈ {± 1 } so s i = 1 if z i lies in D ′ x and -1 if it lies in D x . Then

For a fixed W,P and x 1 , . . . , x 2 n , the random variables ( s 1 , . . . , s 2 n ) are a permutation distribution, so negatively associated. Then the variables s i · d TV ( p W,P ( y | x i ) , p W ∗ ,P ∗ ( y | x i )) are monotone functions of s i , so also negatively associated. They are also bounded in [ -1 , 1] . Hence we can apply a Chernoff bound:

for any fixed W,P .

Constructing a net. We will first construct a net over the precision matrices P (independent of W ), and then for each element in the P -net, we will construct a net over W .

Net over Σ -1 . In the following, λ max ( P ) denotes max eigenvalue of a matrix P , λ min ( P ) denotes the min eigenvalue, and λ i ( P ) denotes the i -th eigenvalue, in decreasing order.

In order to construct the net over the precision matrices, we will consider geometrically spaced values of λ ∈ [ L · λ min ( P ∗ ) , U · λ max ( P ∗ )] , and for each λ , we will construct a net over matrices that have max eigenvalue ≤ λ .

Now consider λ > 0 that lies in the following discrete set:

This set is a geometric partition over the possible max eigenvalues that the MLE can return.

Following definition C.12, let Ω λ denote the subset of positive definite matrices in R d × d that have condition number κ and max-eigenvalue in [ λ 2 , λ ] . Similarly, following Definition C.12, let ˜ Ω λ,β denote the gridded version of Ω λ , where entries in the matrix are multiples of λβ .

For any P ∈ Ω λ , let ˜ P ∈ ˜ Ω λ,β be the matrix obtained by rounding down every element in P .

By the Data Processing Inequality, for any W ∈ R d × k , we have

By Lemma C.18, we can upper bound the RHS of the above inequality by

we have a partition of size O ( ( d 2 κ/ε ) d 2 ) per λ such that:

We will now construct a net over W , so as to show Eqn (43) for all W,P .

W -net. By repeated triangle inequalities, we have

Using the cover over P , the first and last term on the RHS are O ( ε ) . This gives

Setting

For a fixed ˜ W, ˜ P , we will have (43) using a Chernoff bound. Since ˜ P is already finite, we will now construct a net over W for each ˜ P .

It is sufficient to bound d TV ( p W, ˜ P ( y | x i ) , p W ′ , ˜ P ( y | x i )) as the triangle inequality implies that this is larger than the RHS above.

We want, for any W,W ′ in a cell,

for all i ∈ [2 n ] . It is also sufficient to bound d TV ( p W, ˜ P ( y | x i ) , p W ′ , ˜ P ( y | x i )) as the triangle inequality implies that this is larger than the left hand side above.

Lemma C.19 implies that we can find O ( d ϵ 3 / 2 √ log( 2 d ϵ ) ) per row of W such that for any W,W ′ in a cell, either

which implies or

Therefore, for each i regardless of the value of W ∗ j z i there are at most O ( d ϵ 3 / 2 √ log( 2 d ϵ ) ) partitioning hyperplanes.

We then take the intersection of all 2 n partitions (for each data point z i ). The cells of this partition are defined by 2 n sets of O ( d ϵ 3 / 2 √ log( 2 d ϵ ) ) parallel hyperplanes. Since z ∈ R k , the number of cells is at most O ( ( nd ϵ 3 / 2 √ log ( 2 d ϵ ) ) k ) .

Hence the total number of cells for d rows is at most O ( ( nd ϵ 3 / 2 √ log ( 2 d ϵ ) ) dk ) .

Putting everything together. Finally, for any W ∈ R d let ˜ W ∈ N ˜ P be the representative of its cell. Recall that each representative ˜ P of P induces a different cover N ˜ P over W . Let N be the net over the precision matrices P .

By definition of the cells,

As there are log κ U L partitions over P (corresponding to the maximum possible eigenvalue of P ), each with ( O ( d 2 κ ε )) d 2 elements, we have

and each cover N P over W has size

This implies that suffices for

Lemma C.19. Let y = ϕ ( µ ∗ + η σ ∗ ) where µ ∗ , σ ∗ are fixed, and y µ,σ = ϕ ( µ + η σ ) . We partition the space R of µ s.t. for µ, µ ′ in a cell, either

or

Then the number of cells is at most O ( d ϵ 3 / 2 √ log( 2 d ϵ )) .

Proof. In one dimension

where c is a constant. So, we can assume WLOG that σ ∗ = 1 . The number of cells in the grid only depends on σ/σ ∗ .

Now, we show that, regardless of the value of µ ∗ we only need to make a grid on a segment of length at most 3 max( σ, 1) √ log(2 d/ϵ ) . This is because for any µ outside the ranges specified below the d TV ( p µ,σ ( y ) , p µ ∗ ,σ ∗ ( y )) ≥ 1 -ϵ .

If µ ∗ ≤ -√ log(2 d/ϵ ) and for any µ such that µ ≥ σ √ log(2 d/ϵ ) , the d TV ( p µ ∗ ,σ ∗ ( y ) , p µ,σ ( y )) ≥ the difference in the probabilities at 0 which is bigger than 1 -ϵ .

If 0 ≥ µ ∗ ≥ -√ log(2 d/ϵ ) and for any µ s.t. µ ≥ max( σ, 1) √ log(2 d/ϵ ) , the d TV ( p µ ∗ ,σ ∗ ( y ) , p µ,σ ( y )) is the same as in the linear case and since, µ -µ ∗ ≥ max( σ, 1) √ log(2 d/ϵ ) , the d TV ( p µ ∗ ,σ ∗ ( y ) , p µ,σ ( y )) ≥ 1 -ϵ .

If 0 ≤ µ ∗ ≤ √ log(2 d/ϵ ) , for any µ s.t. µ -µ ∗ ≥ max( σ, 1) √ log(2 d/ϵ ) , the d TV ( p µ ∗ ,σ ∗ ( y ) , p µ,σ ( y )) ≥ 1 -ϵ .

If µ ∗ ≥ √ log(2 d/ϵ ) then for any µ s.t. µ -µ ∗ ≥ max( σ, 1) √ log(2 d/ϵ ) using the same argument as above, we have, the d TV ( p µ ∗ ,σ ∗ ( y ) , p µ,σ ( y )) ≥ 1 -ϵ . Moreover, this is also true for µ s.t. -σ √ log(2 d/ϵ ) ≤ µ ≤ µ ∗ -max( σ, 1) √ log(2 d/ϵ ) . Therefore, in this case, we have an additional cell.

In addition to the above, for any µ, µ ′ ≤ -σ √ log(2 d/ϵ ) , the d TV ( p µ,σ ( y ) , p µ ′ ,σ ( y )) ≤ ϵ/ 2 d since both y µ,σ , y µ ′ ,σ are only non-zero with probability at most ϵ/ 2 d . Therefore, for all the above cases, we only need to partition a segment of length at most 3 max( σ, 1) √ log(2 d/ϵ )

Moreover, for σ sufficiently small we can do better. We only need to partition a space of σ √ log(2 d/ϵ ) . This is primarily because when σ sufficiently small, for any µ in the linear case we have that d TV ( p µ,σ ( y ) , p µ ∗ ,σ ∗ ( y )) ≥ 1 -ϵ .

It is easy to see that d TV ( p µ,σ ( y ) , p 0 , 1 ( y )) ≥ d TV ( p 0 ,σ ( y ) , p 0 , 1 ( y )) . The PDFs of N (0 , σ ) and N (0 , 1) intersect at x = ± σ √ log(1 /σ 2 ) 1 -σ 2 . To show that d TV ( p 0 ,σ ( y ) , p 0 , 1 ( y )) ≥ 1 -ϵ , it is now sufficient to show that

where Φ( x ) is the CDF of the standard normal distribution. By using classical bounds on Φ( x ) , we have that

which is ≤ ϵ/ 2 if σ 2 ≤ ϵ/ 2 . And, which is ≥ (1 -ϵ ) / 2 if

When µ ∗ ≤ -√ log(2 d/ϵ ) the same argument as above shows that if µ > σ √ log(2 d/ϵ ) then the d TV ( p µ ∗ ,σ ∗ ( y ) , p µ,σ ( y )) ≥ the difference in the probabilities at 0 which is bigger than 1 -ϵ .We consider the case where µ ∗ ≥ -√ log(2 d/ϵ .

Since, when σ is small, for any µ in the linear case the TV distance is large it is sufficient to have µ large enough so that the intersection of the PDFs are positive and we are in the linear case.

The point of intersection assuming σ ∗ = 1 is given by

which is positive whenever or

d TV ( p µ,σ ( y ) , p µ ∗ ,σ ∗ ( y )) ≥ 1 -ϵ and d TV ( p µ ′ ,σ ( y ) , p µ ∗ ,σ ∗ ( y )) ≥ 1 -ϵ.

Then the number of cells is the max of O ( d √ log(2 d/ϵ ) /ϵ ) (when σ small) or O ( d √ log(2 d/ϵ ) /σϵ ) ≤ d √ log(2 d/ϵ ) /ϵ 3 / 2 (when σ large).

D Proof of composition of layers

Proof. We can use the triangle inequality to compose our single layer guarantees. Suppose, for layer j and j +1 we have

then,

where we use the fact that d TV ( f ( X ) , f ( Y )) ≤ d TV ( X,Y ) .

For the rest of the space, we partition µ ∈ R k s.t. for any µ, µ ′ in a cell,

This implies that for any µ, µ ′ in a cell, either,

6

8

10

12

14

16

18

20

1.5

2

2.5

3

3.5

4

10 -3

(a) Total Variation Upper Bound vs

d

10 2

10 3

10 4

10 -2

10 -1

(b) Total Variation vs.

n

for Different

D

x

E Simulations

E.1 Additional Simulations

In this section, we provide additional simulations to supplement some of the discussion in Section 5.

E.2 Simulation Details

E.2.1 Figure 2

In these experiment, we set d = 1 , and plot the results for various values of the number of samples n in Figure 2a and various values of the input dimension k in Figure 2b. For each plot, we fix the true σ ∗ = 1 and the w ∗ = 1 k × 1 . In each case the MLE is solved via gradient descent with backtracking line search, and we check a first order condition ∥∇ w,σ log p w,σ (( y | x )) ∥ 2 < δ = 10 -3 as the exit condition. We verify that increasing or decreasing δ by one order of magnitude makes no difference to the Figure.

The expected total variation distance for the two distributions is calculated as follows. We sample x according to the true distribution (in this case either Laplace or Normal). Then we compute d TV ( p ˆ w, ˆ σ ( y | x ) , p w,σ ( y | x )) via the MATLAB integral function which uses vectorized adaptive quadrature. We repeat this a total of 100 times and take the average to compute our expected total variation. We then repeat the entire process 2000 times, each time optimizing to find an MLE, and then compute its average total variation distance. Lines indicate the average of these experiments, and the error bars, (not easily visible due to their size) indicates one standard error.

E.2.2 Figure 3

In these experiments we fix d = 3 to retain reasonable complexity for computing the TV distance, and take input dimension k = 1 with deterministic x in order to compare with [35]. In Figure 3a we fix Σ ∗ = I d and take let W ∗ = b 1 d × 1 , where b will vary across our experiments. We set n = 10000 . In Figure 3b we set n = 5000 and adjust Σ such that one diagonal entry is κ 1 / 2 , and the other is κ -1 / 2 , making the total condition number κ .

In both of these experiments, we restrict the MLE computation to be over diagonal Σ only. This is not because computation of the MLE is too difficult, but rather because computing the TV distance is greatly simplified in this case. The algorithm of Wu et al. is hence modified to use the knowledge that the output must be diagonal. This is simply done, because the procedure of Wu et al. essentially first

estimates the diagonal entries of the matrix as if it were diagonal and then computes the correlations. Removing this second phase allows us to achieve our goal.

Since x is deterministic, we do not need to consider randomness in computing the expected TV distance, though other challenges remain. Since our distribution is degenerate, we must be very careful in computing the TV distance in higher dimensions. Specifically, in the diagonal case, the TV may be written as:

where y i is the i th element of y . Though at first glace it seems that this is a single high-dimensional integral, the reality is that due to the truncation, the probability mass on the boundary of the nonnegative orthant cone R 3 ≥ 0 has a complex structure that cannot be ignored. Instead we perform a series of integrals of continuous bounded functions, which are much more amenable to Monte-Carlo integration techniques:

Essentially, for each possible support of y , S ′ , we integrate over the absolute deviation in those coordinates.

E.2.3 Figure 4

In Figure 4a, we plot an upper bounds for the TV distance of the MLE as the output dimension d grows. We set the input dimension k = 1 with deterministic x and fix the number of samples n = 5000 . To estimate the KL divergence, we repeatedly sample y according to the true distribution, and compute the empirical average log-likelihood ratio.

In Figure 4b we fix the output dimension d = 1 and input dimension k = 5 , and compute the TV over a range of values of n . In addition to x sampled i.i.d. from the Normal and Laplace distributions, we also plot a performance with a Normal mixture, where with probability 0 . 01 , the normal distribution has mean shifted by 100. We observe, as our theory suggests, that in all cases, there is only very minor differences in the expected TV distance.

Note that in the case where x is distributed according to a Normal mixture, we observe that the optimization may become very challenging, and in the plot above, we have omitted some of the instances where optimization failed due to lack of smoothness in the objective and numerical imprecision. Omitting these point may lead to a small systematic error in the figure, which may explain why it is lower than the other plots. In practice, for a fixed optimization budget, we may observe meaningful differences in TV for different distributions of x , since computing the MLE becomes more challenging for more complex heavy-tailed distributions.

F The Likelihood Function

In this section, we discuss the likelihood function, proving log-concavity, as well as discussing computational challenges.

F.1 One Dimensional Case

In this section, we consider the case where the output dimension d = 1 , with some σ ∗ and some W ∗ ∈ R 1 ,k and describe how to compute the likelihood function. We defer the proof of the

log-concavity to the follwing section, which covers the more general case with d ≥ 1 When we re-parameterize as u = -W/σ , v = 1 /σ , the likelihood function is written as:

where in this case we let S = { i | y i = 0 } , where y i is the i th sample in the set { y i , x i } n i =1 . Note this is distinct from how we define S and S ′ in the multidimensional case, where it corresponds to the zero and non-zero coordinates of a single sample y i . The case of d > 0 with uncorrelated η follows a similar approach.

Numerical concerns. In (46), the term log Φ( -u · x i ) presents some numerical concerns when u · x i ≫ 0 if we naively compute Φ( -u · x i ) and then take the logarithm. Instead we compute it from the mills ratio m ( x ) [27], defined to be the ratio of the standard normal pdf and the complementary cdf. The mills ratio is easily computed, with many well-known expansions, see for example, [13]. Then we can write:

Since m ( x ) changes relatively slowly in x compared to Φ( -x ) , this greatly improves numerical stability.

F.2 Multidimensional Case

In the multi-dimensional case, we will generally use the more standard natural parameters:

Note that in the one-dimensional case, we could have also use the natural parameters, but due to the truncation structure, the parameters we used make the computation simpler, in a way that does not apply to the multidimensional case. Also note that here we are considering a fixed x and writing v as a vector. In full generality, we should take V = -Σ -1 W , however, this is a simple extension which we omit here for readability. It turns out that density is log-concave in these natural parameters:

Lemma F.1. The log-likelihood function in Eqn (8) is concave in the natural parameter space.

Proof. First, let's write the un-truncated density in terms of these parameters:

Thus, the untrucated conditional density can be written as:

where A ( U, v ) is the cumulant function (note this is distinct from the related cumulant generating function). A well known result is that A is jointly convex in its arguments, U and v . Taking logs and using this fact, shows us that f U,v ( y ) is log-concave in U, v .

Our truncated density is simply:

(50)

For any log-concave density f ( x ) , integration over a convex subset of the coordinates preserves log-concavity ([9], Example 3.42-3.44). Thus the objective is log-concave.

Then the likelihood function at U, v can be rewritten as

Separating the terms corresponding to S and S ′ , we get

The last term resembles the log Φ term that appears in the univariate case. This resemblance can be made more clear as follows. Let r S = U SS ′ y ′ S + 1 2 v S and M T M = U S .

Putting this together, f U,v can be written as:

Thus, it appears that evaluating the likelihood for even a single sample involves the high-dimensional integral that is the rectangular cdf in equation (51).

F.3 Computing Gradients

F.3.1 One Dimensional Case

In the one-dimensional case, the gradient with respect to u is easily computed as:

where we have previously defined m ( x ) as the mills ratio. Furthermore, we have:

F.3.2 Multidimensional Case

First we consider the non-integral terms in the likelihood. Differentiating each term wrt U , we get

Differentiating each term wrt v , we get

Now consider the integral term. Differentiating wrt U , we get

Let M be a matrix such that

Then via completion of squares in the exponential term, we get

Notice that this density is Gaussian, with mean and covariance:

And hence, the gradient can be estimated as

where t S is the truncation of

to the negative quadrant. A similar calculation gives the gradient for v as

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper investigates the sample complexity of learning conditional generative models without assumptions on the input distribution. It applies the Maximum Likelihood Estimator (MLE) to linear regression and 1-layer networks with ReLU activation. The results show that the MLE achieves small total variation error with sample complexities of $O(k/\\epsilon^2 log(1/\\epsilon))$ for linear regression and O((kd + d^2) / \\epsilon^2 log(kd\\kappa/\\epsilon)) for 1-layer networks. The paper also discusses the extension to multilayer networks,  given access to the internal activations. The results suggest that MLE is a promising approach for learning feed-forward generative models from limited samples, though the authors did mention that the computational aspects of the optimization problem are not thoroughly analyzed in the paper.\n\nSTRENGTHS:\nThis paper provides a solid theoretical foundation for understanding the sample complexity of learning multi-layer ReLU networks using MLE method. The derived bounds do not make assumptions on the distribution of X or the condition number of W, and achieve a sample complexity polynomial in the system parameters. The developed algorithm and theories show considerable improvement over those in the existing literature. Though there are limitations that are also noted by the authors, I think the paper is innovative and provides interesting insights into efficient learning of conditional generative models. \nMoreover, the paper is well-written and presents its concepts and results in a clear and concise manner.\n\nWEAKNESSES:\nThe weakness and limitations of the paper are well noted and discussed by the authors. \n\n1. To extend the theory to multilayer neural networks, it requires access to intermediate activations, which is impractical\n2. It assumes that the learner has an understanding of the model architecture, which might not be the case in practice\n3. It might be challenging to perform MLE on some of the models being considered, such as neural networks, and the computational aspects of the optimization problem are not thoroughly examined in this paper.\n\nQUESTIONS:\nWhat are the practical implications of the sample complexity bounds derived in the paper? How do these bounds impact the feasibility and scalability of learning generative models in real-world applications?\n\nI understand that the paper is limited by space restrictions, but can you provide some insights into how to go about addressing the requirement for accessing intermediate activations in order to apply the proposed algorithm to solving deep neural nets? Is it going to be a deal breaker for the practical applicability of the proposed method?\n\nLIMITATIONS:\nNo potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies conditional distribution learning: given iid samples (x,y) where x ~ D and y ~ p(y|w*,x), the goal is to find some estimate w such that the distributions p(y|w*,x) and p(y|w,x) are close in expectation over x ~ D, or equivalently the learned distribution of (x,y) (where x ~ D) is close to the true distribution.\n\nSpecifically, this work studies the 1-layer conditional generative model y = max(W* x + eta, 0) where W* is a matrix and eta is a multivariate mean-zero Gaussian with some unknown covariance Sigma*. The main result shows that for an arbitrary covariate distribution D, and arbitrary W*, so long as Sigma* has covariance at most kappa, the sample complexity of MLE (needed to learn up to total variation distance epsilon) is polynomial in the dimension parameters, log(kappa), and 1/epsilon.\n\nAn straightforward extension to multi-layer networks is also given (under the assumption that intermediate activations are known).\n\nSTRENGTHS:\n- This work introduces a (to my knowledge) novel perspective to the problem of learning generative models: distribution learning rather than parameter estimation. This more accurately addresses the problem that actually matters in practice for generative models, and avoids needing to worry about identifiability issues.\n- Due to the new goal, no distributional assumptions are needed (aside from the tame bound on condition number of the noise).\n- The paper is well-written, with a toy example of linear regression given for intuition. I did not have time to check the proofs, but the approach and proof sketch seem reasonable.\n\nWEAKNESSES:\n- As the authors acknowledge, the paper only addresses the statistical question. It's claimed that the MLE is concave; however, I could not find a proof of this fact in the paper. Moreover, looking at the log-likelihood function (8) I do not see why it should be concave.\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
In this paper, the authors consider the problem of linear regression and ReLU applied to linear regression. The goal is to recover the weight vector such that the resulting distributions are close as opposed to recovering the weight vector under certain norms such as $\ell_2$ which has been thoroughly studied. For the linear regression, the authors show that the MLE estimator learns the distribution in total variation distance. For the ReLU regression, the authors show again the MLE estimator works when the covariance matrix of the noise is well-conditioned. The resulting algorithm is sample-efficient but not time-efficient.

STRENGTHS:
The problem considered is fundamental and it has important connections to learning practically important problems.

WEAKNESSES:
I think the scientific novelty and contribution of this paper is very limited. The proposed results were well-known in the literature in my opinion. As an example, the distributional learning guarantee of linear regression in TV distance was analyzed in [arXiv:2107.10450]. Similarly, one-layer ReLU networks have been analyzed in the works of [Diakonikolas et al, Klivans et al, and Arora et al]. Moreover, the condition number assumption and the high running time of the proposed algorithms are quite restrictive in my opinion.

QUESTIONS:
None.

LIMITATIONS:
None,
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe article provides complexity bounds for learning the conditional distribution y|x. One of the main novelties claimed by the authors is that the control of the TD distance between the estimated distribution and the ground truth is more meaningful. Thus, they are able to provide bounds independently of the distribution of label x. The article is well written and the flow is quite pleasant. \nSee my comments below for more critical details\n\nSTRENGTHS:\nThe document is clearly written and the main arguments are transparent and easy to follow (even if they remain fairly technical). The authors did a great job explaining the intuition before getting into formal details. \n\nThe distribution free result (wrt to label x) is quite remarkable.\n\nWEAKNESSES:\nMy main concern is that the results presented might be difficult to compare to classical results on the subject. Below I describe some of my missunderstanding.\n\nQUESTIONS:\n- Usually x is the features and y the labels. The notation might be misleading. \nAs a matter of fact the product $x \\cdot w^*$ is even weird when the label does not belong to a vectorial space.\n- On the same flavor, it is difficult to compare with classical results, since the assumptions are usually in the feature space ie distribution assumption on the objet $x$ and the label is generated conditional on $x$ and some noise. It is unclear how the results fit on the same ground.\n- The lines 32-35 need explicit references and maybe explicit classical bounds to discuss.\n\n- The complexity presented is not explicitly compared with (the criticized) classical ones. Hence it makes the contributions harder to appreciates\n\n- Theorem 4.1 (and probably) others must states the full assumptions eg $n \\geq k/2$ which usually does not hold in high dimensional regime where the number of features is way larger than the sample size.\n\n- As a main proof technique, the authors rely on the relation between TV distance and prediction error in the equation below line 159.\nMainly $d(\\hat w, w^*) \\approx |x\\cdot \\hat w - x \\cdot w^*|$. And then we directly fall into the OLS analysis as in https://www.di.ens.fr/%7Efbach/ltfp_book.pdf Chapter 3. Can the authors comments on the main novelties after this. Usually the prediction error does not require much restriction on the design matrix (compared to estimation error control).\n\n- In Lemma 4.3, can the authors specify what the probability Pr is?\nAlso, since the (empirical) distance is Lipschitz, can't we deduce this lemma directly from Talagrand inequality?\n\n- In which practical settings do we know the ground-truth conditioning number?\n- Also in Theorem 4.5, $\\hat W$ and $\\hat \\Sigma$ might be non unique. Can the authors comment on that?\n\nSome failure of MLE are quite known, for example in Chapter 24.1.3, https://www.cs.huji.ac.il/~shais/UnderstandingMachineLearning/understanding-machine-learning-theory-algorithms.pdf an explicit overfitting example is shown. Can the authors discuss how their proposition escape such a situation?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n1. This paper shows how MLE can perform distribution learning in the setting of linear regression and in multi-layer ReLU networks.\n2. This does not take a distribution on labels to be any specific form but rather unknows and tries to derive the sample complexity for a small total variational distance between the model’s conditional distribution and actual conditional distribution.\n3. Improves the sample complexity of the previous work which suffers an exponential dependence on the ||W|| term.\n4. Generalized how one layer Relu layer sample complexity can be extended to multilayer Relu networks.\n\nSTRENGTHS:\n1. This work is a nice extension of the previous work mentioned in the reference[27]. It relaxes the assumption of previous work which assumed a fixed distribution of the labels.\n2. The quality and clarity of the paper is good. Gives good proof techniques using the ideas from the learning theory.\n3. Nice result for the linear regression in Theorem-4.1 where the sample complexity is linear in k: dimensionality of the labels which is some finite value in most of the cases.\n4. This paper can be significant where the label distribution is unknown but the data distribution is known.\n\nWEAKNESSES:\n1. The proof techniques assumed the distribution of the data is Gaussian which is rarely the case because the actual data distribution for generative models thought of some complex unknown distribution.\n2. Even though theorem 4.2 proved the result nicely there is a dependence of the square of the dimensionality of data. So the curse of dimensionality remains and these bounds can be very loose.\n\nQUESTIONS:\n1. Theoretically the distribution of labels x being unknown is correct. But in most cases, labels follow some multinomial distribution with some parameters. Does this assumption make the bound better? It would be better if some analysis is being done with the distribution of x being multinomial.\n\n2. The proofs rely on the distribution of y being Gaussian. What would happen if you take the distribution of y as unknown and rather x being some multinomial as it happens in many practical settings of time series and computer vision?\n\n3. There is some confusion in the notation. In line 36 it is mentioned that x denotes the labels and y is the data generated using that label. In linear regression, authors tried to predict y  using x which is correct. In theorem 4.1. is that notation of y being the data and x being the labels being followed or is it reversed?\n\nLIMITATIONS:\n1. There is some limitation where we can have the conditional distribution of data to be Gaussian in those cases the proof techniques do not work.\n2. Mostly in the literature of generative models the data distribution is taken to be some unknown complex distribution and labels are multinomial distribution with some appropriate parameters. These bounds don’t address that. So that is why there will be very limited use cases.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-26"}