SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ajil Jalal Justin Kang
{ajiljalal, justin_kang}@berkeley.edu
Ananya Uppal
ananya.uppal09@gmail.com
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.
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 \mid 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 \mid 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 \mid 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 \mid x)$ ?
Linear Regression. Consider ordinary linear regression with Gaussian noise: you observe independent samples  $(x_{i},y_{i})\in \mathbb{R}^{k}\times \mathbb{R}$  of the form
$$
y = x \cdot w ^ {*} + \eta \quad \text {f o r} \eta \sim N (0, 1).
$$
Figure 1: A conditional distribution defined by a conditional generative model. To sample from the conditional distribution  $p(y \mid x)$ , we perform inference. Due to the stochastic nature of the model, each output is different.
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 \mid x)$  where  $p(x)$  is the true distribution of  $x$ , and we would like to find  $\widehat{w}$  such that
$$
d _ {T V} \left(p _ {w ^ {*}} (x, y), p _ {\hat {w}} (x, y)\right) \leq \varepsilon . \tag {1}
$$
This ensures that when the input  $x$  comes 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\left(\frac{1}{\varepsilon^2} k \log \frac{1}{\varepsilon}\right)$  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:
$$
y = \phi \left(W ^ {*} x + \eta\right) \quad \text {f o r} \eta \sim \mathcal {N} \left(0, \Sigma^ {*}\right) \tag {2}
$$
for  $x \in \mathbb{R}^k$ ,  $y \in \mathbb{R}^d$ ,  $\phi(x) = \max(x, 0)$  is the ReLU activation, and some weights  $W^* \in \mathbb{R}^{d \times k}$ ,  $\Sigma^* \in \mathbb{R}^{d \times 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  $\Sigma^{*}$  has condition number at most  $\kappa$ . Using  $O\left(\frac{kd + d^2}{\varepsilon^2} \log \frac{k d \kappa}{\varepsilon}\right)$  samples, the distribution generated by the MLE has TV error  $\varepsilon$ , 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  $\Sigma^{*}$  is learned rather than specified; and (c) the ReLU nonlinearity  $\phi$  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 standard Gaussian noise  $\eta \sim \mathcal{N}(0,1)$ . Our proof approach is inspired by learning bounds that exploit finite VC dimension. We would like to show (1), or equivalently,
$$
d \left(w ^ {*}, \widehat {w}\right) := \underset {x} {\mathbb {E}} \left[ d _ {T V} \left(p _ {w ^ {*}} (y \mid x), p _ {\widehat {w}} (y \mid x)\right) \right] \leq \varepsilon , \tag {3}
$$
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  $\widehat{w}$  is small i.e.,
$$
\widetilde {d} \left(w ^ {*}, \widehat {w}\right) := \underset {x} {\widetilde {\mathbb {E}}} \left[ d _ {T V} \left(p _ {w ^ {*}} (y \mid x), p _ {\widehat {w}} (y \mid x)\right) \right] \leq 0. 5 \varepsilon , \tag {4}
$$
where  $\widetilde{\mathbb{E}}_x[f(x)] = \frac{1}{n}\sum_{i = 1}^{n}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.,
$$
d \left(w ^ {*}, \widehat {w}\right) \leq \widetilde {d} \left(w ^ {*}, \widehat {w}\right) + 0. 5 \varepsilon \leq \varepsilon , \tag {5}
$$
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
$$
\widetilde {\mathbb {E}} \left(x ^ {T} \left(w ^ {*} - \widehat {w}\right)\right) ^ {2} \propto k / n.
$$
Since  $d_{TV}(p_{w^*}(y \mid x), p_{\widehat{w}}(y \mid x)) = \Theta(\min(1, |x \cdot (w^* - \widehat{w})|))$ , Jensen's inequality implies (4) for  $n > k / \varepsilon^2$ .
Secondly,  $f_{w}(x) \coloneqq d_{TV}(p_{w^{*}}(y \mid x), p_{w}(y \mid 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 = \phi (w^{*}\cdot x + \eta)$ , 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  $\widehat{w}$  maximizing
$$
L (w) := \frac {1}{n} \sum_ {i = 1} ^ {n} \log p _ {w} \left(y _ {i} \mid x _ {i}\right)
$$
satisfies (4). Now, for any  $w$  not satisfying (4),
$$
\begin{array}{l} \underset {y} {\mathbb {E}} [ L (w) - L (w ^ {*}) ] = - \underset {x} {\widetilde {\mathbb {E}}} d _ {K L} (p _ {w ^ {*}} (y \mid x) \| p _ {w} (y \mid x)) \\ \leq - 2 \widetilde {\mathbb {E}} _ {x} [ d _ {T V} (p _ {w ^ {*}} (y \mid x), p _ {w} (y \mid x)) ^ {2} ] \leq - 2 \varepsilon^ {2}, \\ \end{array}
$$
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 = \frac{1}{\varepsilon^2} \log \left( \frac{1}{\delta} \right)$  samples, we have
$$
L (w) - L \left(w ^ {*}\right) \leq - \varepsilon^ {2},
$$
with probability  $1 - \delta$ . Using a careful covering argument and  $n = (k / \varepsilon^2)\log (1 / \delta)$  samples, we can uniformly extend this to all  $w$  not satisfying Eq (4). By definition, the MLE has  $L(\widehat{w})\geq L(w^{*})$ , and by our uniform bound, it must satisfy (4).
The second stage changes because  $f_{w}(x)$  depends on  $x \cdot w$  and  $x \cdot w^{*}$  in a more complicated way than through  $x \cdot (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 \cdot w$ .
1. 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.
2. 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.
3. 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.
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 = \phi(b + \eta)$  s.t.  $b, \eta \in \mathbb{R}^d$ , and they seek to learn the covariance of  $\eta$  along with the bias vector  $b$ . However, their sample complexity bound suffers an exponential dependence on  $\| b \|_{\infty}^{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].
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  $\phi$  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
$$
y = x \cdot w ^ {*} + \eta , \tag {6}
$$
where  $\eta \sim \mathcal{N}(0,\sigma^2)$ ,  $w\in \mathbb{R}^k$ ,  $\sigma$  is known and  $x\in \mathbb{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  $\| \widehat{w}\cdot x - w^{*}\cdot x\|_{2}$ , we focus on error in  $d(w^{*},\widehat{w})$  as defined in (3), which only captures the error in  $\widehat{w}$  insofar as it impacts our distribution estimate. Given data  $\{(x_i,y_i)\}_{i = 1}^n$  generated by (6), the MLE is:
$$
\widehat {w} := \underset {w \in \mathbb {R} ^ {k}} {\arg \max} \sum_ {i \in [ n ]} \log p _ {w} (y _ {i} | x _ {i}) \equiv \underset {w \in \mathbb {R} ^ {k}} {\arg \min} \sum_ {i \in [ n ]} \frac {(y _ {i} - w \cdot x _ {i}) ^ {2}}{\sigma^ {2}}.
$$
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)\}_{i=1}^n$  be i.i.d. random variables generated from the linear model (6), and assume that  $\sigma$  is known. Then, for a sufficiently large constant  $C > 0$ ,
$$
n = C \frac {k}{\varepsilon^ {2}} \log \frac {1}{\varepsilon}
$$
samples suffice to ensure that with probability  $1 - e^{-\Omega (\varepsilon^2 n)}$  over the data,
$$
d (\widehat {w}, w ^ {*}) \leq \varepsilon .
$$
Note that one cannot hope to get such a guarantee in the classical setting where error is measured in  $\| \widehat{w} \cdot x - w^{*} \cdot 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  $\sigma^2 = 1$ . We split the proof into two stages. In the first stage, we bound  $\widetilde{d}(\widehat{w}, w^*)$ , which denotes the empirical TV distance (4) over the training set  $\{x_i\}_{i=1}^n$ .
Lemma 4.2. Let  $\{(x_i, y_i)\}_{i=1}^n$  be i.i.d. random variables such that  $y_i = x_i \cdot w^* + \mathcal{N}(0,1)$ . Then, for  $n \geq \frac{k}{2}$ , with probability  $1 - e^{-\Omega(n)}$ , the MLE  $\widehat{w}$  satisfies
$$
\widetilde {d} (\widehat {w}, w ^ {*}) \leq \sqrt {\frac {k}{2 n}}.
$$
The proof relies on the fact that  $p_{w^*}(y|x_i)$  and  $p_{\tilde{w}}(y|x_i)$  are Gaussian distributions, so
$$
d \left(p _ {\hat {w}} \left(y | x _ {i}\right), p _ {w ^ {*}} \left(y | x _ {i}\right)\right) = \Theta \left(\min  \left\{1, \left| x _ {i} ^ {T} \left(\hat {w} - w ^ {*}\right) \right| \right\}\right).
$$
Using the explicit form of the MLE, we can show that, with high probability,
$$
\frac {1}{n} \sum_ {i} \big (x _ {i} ^ {T} (\widehat {w} - w ^ {*}) \big) ^ {2} \leq \frac {k}{2 n},
$$
and Lemma 4.2 follows from Jensen's inequality.
The second stage shows that the empirical average of the TV distance  $\widetilde{d} (\widehat{w},w^{*})$  is close to the population average  $d(\widehat{w},w^{*})$ .
Lemma 4.3. Let  $\{x_i\}_{i=1}^n$  be i.i.d. random variables such that  $x_i \sim \mathcal{D}_x$ . For a sufficiently large constant  $C > 0$ , and for  $n = C \frac{k}{\varepsilon^2} \log \frac{1}{\varepsilon}$  with  $n \geq \frac{k}{2}$ , we have:
$$
\Pr_ {x _ {i} \sim \mathcal {D} _ {x}} \left[ \sup  _ {w \in \mathbb {R} ^ {k}} \left| \widetilde {d} (w, w ^ {*}) - d (w, w ^ {*}) \right| > \varepsilon \right] \leq e ^ {- \Omega (n \varepsilon^ {2})}.
$$
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 \cdot x$ . This implies that for each  $x_i$ , we are able to partition the space of  $w$  with  $O(1 / \varepsilon)$  hyperplanes such that within each cell the TV distance varies by at most  $\varepsilon$ . As we have  $n$  samples and  $w \in \mathbb{R}^k$ , the number of cells induced in  $\mathbb{R}^k$  is  $\propto (n / \varepsilon)^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 = \Theta\left(\frac{k}{\varepsilon^2} \log \frac{1}{\varepsilon}\right)$  and combining Lemma 4.2 with Lemma 4.3 gives  $d(\widehat{w}, w^*) \leq \varepsilon$ .
4.2 ReLU Case
Now consider the single-layer ReLU. We observe  $(x,y)\in \mathbb{R}^k\times \mathbb{R}^d$  such that:
$$
y = \phi \left(W ^ {*} x + \eta\right), \quad \eta \sim \mathcal {N} \left(0, \Sigma^ {*}\right), \tag {7}
$$
where  $\eta \in \mathbb{R}^d$ ,  $W^*$  and  $\phi(\cdot) = \max(\cdot, 0)$  is applied coordinate-wise. The matrices  $W^* \in \mathbb{R}^{d \times k}$  and  $\Sigma^* \in \mathbb{R}^{d \times d}$  are unknown, and we do not observe  $\eta$ . The variable  $x$  is drawn from an arbitrary probability distribution  $\mathcal{D}_x$ , and we make no additional assumptions on  $\mathcal{D}_x$ : this is important, as we will progressively cascade layers, and one should think of  $\mathcal{D}_x$  being the distribution of activations at each layer.
Given a sample  $(x,y)\in \mathbb{R}^k\times \mathbb{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,\Sigma$ , on this sample is given by
$$
\log p _ {W, \Sigma} (y | x) = c - \frac {\log | \Sigma |}{2} + \log \int_ {\substack {t _ {S} \leq 0 \\ t _ {S ^ {c}} = y _ {S ^ {c}}}} \exp \left\{- \frac {(t - W x) ^ {T} \Sigma^ {- 1} (t - W x)}{2} \right\} d t _ {S}. \tag{8}
$$
where  $c$  is a normalization constant which does not depend on  $W$  or  $\Sigma$ . 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 + \eta$  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$ ,  $\Sigma$ .
In this setting, proving an analogue of Theorem 4.1 poses multiple challenges:
- The output  $y$  is  $d$ -dimensional rather than a scalar, and  $\eta$  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  $\Sigma^{*}$ , 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^{*}, \Sigma^{*}$  is difficult: if we see a zero in  $y$ , we do not know its magnitude in  $W^{*}x + \eta$  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\|_{\infty}^{2}}$ .
Nonetheless, we can handle most of these difficulties, and the only assumption we make is that the condition number of  $\Sigma^{*}$  is bounded and known to our estimator.
Assumption 4.4 (Condition number bound). Let  $\lambda_{\mathrm{max}}^{*}, \lambda_{\mathrm{min}}^{*}$  denote the largest and smallest singular values of  $\Sigma^{*}$ . We assume there exists  $\kappa < \infty$  such that  $\frac{\lambda_{\mathrm{max}}^{*}}{\lambda_{\mathrm{min}}^{*}} \leq \kappa$ . We further assume that the value of  $\kappa$  is known to our estimator.
Note that the condition number only allows us to control the correlation between the coordinates of  $W^{*}x + \eta$ . The other challenges introduce by the ReLU, such as the lack of a closed form MLE, the need to estimate  $\Sigma^{*}$  and a mixed density log-likelihood remain.
Under Assumption 4.4, the following theorem shows that the MLE  $\widehat{W},\widehat{\Sigma}$  achieves a small total variation distance. The proof of this theorem is in Appendix C.
Theorem 4.5. Let  $\mathbb{R}_{\kappa}^{d\times d}$  denote the set of positive definite matrices with condition number  $\kappa$ . Given  $n$  samples  $\{(x_i,y_i)\}_{i = 1}^n$  satisfying Assumption 4.4, where  $x_{i}\sim \mathcal{D}_{x}$  i.i.d., and  $y_{i}$  is generated according to (7), let  $\widehat{W},\widehat{\Sigma} := \arg \max_{W\in \mathbb{R}^{d\times k},\Sigma \in \mathbb{R}_{\kappa}^{d\times d}}\frac{1}{n}\sum_{i}\log p_{W,\Sigma}(y_{i}\mid x_{i})$ . Then, for a sufficiently large constant  $C > 0$ ,
$$
n = C \cdot \left(\frac {k d + d ^ {2}}{\varepsilon^ {2}}\right) \log \left(\frac {\kappa k d}{\varepsilon \delta}\right)
$$
samples suffice to ensure that with probability  $1 - \delta$ , we have
$$
d _ {T V} \left(\left(\widehat {W}, \widehat {\Sigma}\right), \left(W ^ {*}, \Sigma^ {*}\right)\right) \leq \varepsilon .
$$