SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
f _ {n} (W, P) - f _ {n} ^ {\prime} (W, P) = \frac {1}{n} \sum_ {i = 1} ^ {2 n} s _ {i} \cdot d _ {T V} \left(p _ {W, P} \left(y | x _ {i}\right), p _ {W ^ {*}, P ^ {*}} \left(y | x _ {i}\right)\right).
$$
For a fixed  $W, P$  and  $x_{1},\ldots ,x_{2n}$ , the random variables  $(s_1,\dots,s_{2n})$  are a permutation distribution, so negatively associated. Then the variables  $s_i\cdot 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:
$$
\Pr \left[ \left| f _ {n} (W, P) - f _ {n} ^ {\prime} (W, P) \right| > \varepsilon \right] <   e ^ {- n \varepsilon^ {2} / 2}, \tag {43}
$$
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  $\Sigma^{-1}$ . In the following,  $\lambda_{\max}(P)$  denotes max eigenvalue of a matrix  $P$ ,  $\lambda_{\min}(P)$  denotes the min eigenvalue, and  $\lambda_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  $\lambda \in [L\cdot \lambda_{\min}(P^{*}),U\cdot \lambda_{\max}(P^{*})]$ , and for each  $\lambda$ , we will construct a net over matrices that have max eigenvalue  $\leq \lambda$ .
Now consider  $\lambda > 0$  that lies in the following discrete set:
$$
\left. \left\{\lambda_ {\min } \left(P ^ {*}\right) 2 ^ {j}, j \in \lceil \log_ {2} \left(\kappa \frac {U}{L}\right) \rceil \right\} \right.
$$
Following definition C.12, let  $\Omega_{\lambda}$  denote the subset of positive definite matrices in  $\mathbb{R}^{d\times d}$  that have condition number  $\kappa$  and max-eigenvalue in  $[\frac{\lambda}{2},\lambda ]$ . Similarly, following Definition C.12, let  $\widetilde{\Omega}_{\lambda ,\beta}$  denote the gridded version of  $\Omega_{\lambda}$ , where entries in the matrix are multiples of  $\lambda \beta$ .
For any  $P\in \Omega_{\lambda}$  , let  $\widetilde{P}\in \widetilde{\Omega}_{\lambda ,\beta}$  be the matrix obtained by rounding down every element in  $P$
By the Data Processing Inequality, for any  $W \in \mathbb{R}^{d \times k}$ , we have
$$
d _ {T V} \Big (p _ {W, P} (y | x), p _ {W, \widetilde {P}} (y | x) \Big) \leq d _ {T V} (\mathcal {N} (W x; P), \mathcal {N} (W x; \widetilde {P})).
$$
By Lemma C.18, we can upper bound the RHS of the above inequality by
$$
d _ {T V} (\mathcal {N} (W x; P), \mathcal {N} (W x; \widetilde {P})) \leq O (d ^ {2} \beta \kappa).
$$
$$
\beta = O \Big (\frac {\varepsilon}{d ^ {2} \kappa} \Big),
$$
we have a partition of size  $O\left(\left(d^{2}\kappa /\varepsilon\right)^{d^{2}}\right)$  per  $\lambda$  such that:
$$
d _ {T V} \left(p _ {W, P} (y | x), p _ {W, \widetilde {P}} (y | x)\right) \leq O (\varepsilon).
$$
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
$$
\left| f _ {n} (W, P) - f _ {n} ^ {\prime} (W, P) \right| \leq \left| f _ {n} (W, P) - f _ {n} (W, \widetilde {P}) \right| + \left| f _ {n} (W, \widetilde {P}) - f _ {n} ^ {\prime} (W, \widetilde {P}) \right| + \left| f _ {n} ^ {\prime} (W, \widetilde {P}) - f _ {n} ^ {\prime} (W, P) \right|.
$$
Using the cover over  $P$ , the first and last term on the RHS are  $O(\varepsilon)$ . This gives
$$
\left| f _ {n} (W, P) - f _ {n} ^ {\prime} (W, P) \right| \leq O (\varepsilon) + \left| f _ {n} (W, \widetilde {P}) - f _ {n} ^ {\prime} (W, \widetilde {P}) \right|. \tag {44}
$$
For a fixed  $\widetilde{W},\widetilde{P}$ , we will have (43) using a Chernoff bound. Since  $\widetilde{P}$  is already finite, we will now construct a net over  $W$  for each  $\widetilde{P}$ .
It is sufficient to bound  $d_{TV}(p_{W,\widetilde{P}}(y|x_i), p_{W',\widetilde{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,
$$
\left| d _ {T V} \left(p _ {W, \bar {P}} (y | x _ {i}), p _ {W ^ {*}, P ^ {*}} (y | x _ {i})\right) - d _ {T V} \left(p _ {W ^ {\prime}, \bar {P}} (y | x _ {i}), p _ {W ^ {*}, P ^ {*}} (y | x _ {i}) \right| \leq O (\varepsilon). \right.
$$
for all  $i \in [2n]$ . It is also sufficient to bound  $d_{TV}(p_{W,\widetilde{P}}(y|x_i), p_{W',\widetilde{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  $\mathcal{O}\left(\frac{d}{\epsilon^{3/2}} \sqrt{\log\left(\frac{2d}{\epsilon}\right)}\right)$  per row of  $W$  such that for any  $W, W'$  in a cell, either
$$
d _ {T V} \left(p _ {W, \widetilde {P}} \left(y \mid x _ {i}\right), p _ {W ^ {*}, P ^ {*}} \left(y \mid x _ {i}\right)\right) \geq d _ {T V} \left(p _ {W, \widetilde {P}} \left(y _ {j} \mid x _ {i}\right), p _ {W ^ {*}, P ^ {*}} \left(y _ {j} \mid x _ {i}\right) \geq 1 - \varepsilon \right.
$$
$$
d _ {T V} \left(p _ {W, \widetilde {P}} (y | x _ {i}), p _ {W ^ {\prime}, \widetilde {P}} (y | x _ {i}) \leq \epsilon \right.
$$
$$
d _ {T V} \left( \right.p _ {W, \tilde {P}} \left(y \mid x _ {i}\right), p _ {W ^ {\prime}, \tilde {P}} \left(y \mid x _ {i}\right) \leq \sum_ {j} d _ {T V} \left(p _ {W, \tilde {P}} \left(y \mid x _ {i}\right) _ {j} \mid z _ {i}, p _ {W ^ {\prime}, \tilde {P}} \left(y _ {j} \mid x _ {i}\right)\right) \leq \epsilon / d.
$$
Therefore, for each  $i$  regardless of the value of  $W_{j}^{*}z_{i}$  there are at most  $\mathcal{O}\left(\frac{d}{\epsilon^{3 / 2}}\sqrt{\log\left(\frac{2d}{\epsilon}\right)}\right)$  partitioning hyperplanes.
We then take the intersection of all  $2n$  partitions (for each data point  $z_{i}$ ). The cells of this partition are defined by  $2n$  sets of  $\mathcal{O}\left(\frac{d}{\epsilon^{3/2}} \sqrt{\log\left(\frac{2d}{\epsilon}\right)}\right)$  parallel hyperplanes. Since  $z \in \mathbb{R}^k$ , the number of cells is at most  $\mathcal{O}\left(\left(\frac{nd}{\epsilon^{3/2}} \sqrt{\log\left(\frac{2d}{\epsilon}\right)}\right)^k\right)$ .
Hence the total number of cells for  $d$  rows is at most  $\mathcal{O}\left(\left(\frac{nd}{\epsilon^{3/2}}\sqrt{\log\left(\frac{2d}{\epsilon}\right)}\right)^{dk}\right)$ .
Putting everything together. Finally, for any  $W \in \mathbb{R}^d$  let  $\widetilde{W} \in N_{\widetilde{P}}$  be the representative of its cell. Recall that each representative  $\widetilde{P}$  of  $P$  induces a different cover  $N_{\widetilde{P}}$  over  $W$ . Let  $N$  be the net over the precision matrices  $P$ .
By definition of the cells,
$$
\left| d _ {T V} \left(p _ {W, \widetilde {P}} \left(y \mid x _ {i}\right), p _ {W ^ {*}, P ^ {*}} \left(y \mid x _ {i}\right)\right) - d _ {T V} \left(p _ {\widetilde {W}, \widetilde {P}} \left(y \mid x _ {i}\right), p _ {W ^ {*}, \Sigma^ {*}} \left(y \mid x _ {i}\right)\right) \right| <   O (\varepsilon).
$$
for all  $i \in [2n]$ . Thus
$$
\left| \left(f _ {n} ^ {\prime} (W, \widetilde {P}) - f _ {n} (W, P)\right) - \left(f _ {n} ^ {\prime} (\widetilde {W}, \widetilde {P}) - f _ {n} (\widetilde {W}, P)\right) \right| \leq O (\varepsilon).
$$
$$
\begin{array}{l} \Pr \left[ \sup  _ {W \in \mathbb {R} ^ {d \times k}, P \in \mathbb {R} ^ {d \times d}} \left| f _ {n} ^ {\prime} (W, P) - f _ {n} (W, P) \right| > \varepsilon \right] \\ \leq \Pr \left[ \max  _ {w \in N _ {P}, P \in N} \left| f _ {n} ^ {\prime} (W, P) - f _ {n} (W, P) \right| > \frac {\varepsilon}{4} \right] \\ \leq e ^ {\log | N | + \log | N _ {P} | - \left(\frac {\varepsilon}{4}\right) ^ {2} n / 2} \\ \end{array}
$$
As there are  $\log \kappa \frac{U}{L}$  partitions over  $P$  (corresponding to the maximum possible eigenvalue of  $P$ ), each with  $(O\left(\frac{d^2\kappa}{\varepsilon}\right))^{d^2}$  elements, we have
$$
\log | N | \lesssim d ^ {2} \log \left(\frac {d ^ {2} \kappa}{\varepsilon}\right) + \log \log \left(\kappa \frac {U}{L}\right).
$$
and each cover  $N_P$  over  $W$  has size
$$
\log | N _ {P} | = 2 k d \log \left(\frac {d}{\epsilon^ {3 / 2}} \sqrt {\log \left(\frac {2 d}{\epsilon}\right)}\right).
$$
$$
\operatorname * {P r} [ \sup _ {W, P} f _ {n} ^ {\prime} (W, P) - f _ {n} (W, P) > \varepsilon ] <   e ^ {- \Omega (\varepsilon^ {2} n)}.
$$
Lemma C.19. Let  $y = \phi(\mu^{*} + \eta_{\sigma^{*}})$  where  $\mu^{*}, \sigma^{*}$  are fixed, and  $y_{\mu, \sigma} = \phi(\mu + \eta_{\sigma})$ . We partition the space  $\mathbb{R}$  of  $\mu$  s.t. for  $\mu, \mu'$  in a cell, either
$$
d _ {T V} \left(p _ {\mu , \sigma} (y), p _ {\mu^ {\prime}, \sigma} (y)\right) \leq \epsilon / 2 d.
$$
$$
d _ {T V} \left(p _ {\mu , \sigma} (y), p _ {\mu^ {*}, \sigma^ {*}} (y)\right) \geq 1 - \epsilon \quad a n d \quad d _ {T V} \left(p _ {\mu^ {\prime}, \sigma} (y), p _ {\mu^ {*}, \sigma^ {*}} (y)\right) \geq 1 - \epsilon .
$$
Then the number of cells is at most  $\mathcal{O}\left(\frac{d}{\epsilon^{3/2}} \sqrt{\log\left(\frac{2d}{\epsilon}\right)}\right)$ .
Proof. In one dimension
$$
d _ {T V} \left(p _ {\mu , \sigma} (y), p _ {\mu^ {*}, \sigma^ {*}} (y)\right) = d _ {T V} \left(p _ {c \mu , c \sigma} (y), p _ {c \mu^ {*}, c \sigma^ {*}} (y)\right)
$$
where  $c$  is a constant. So, we can assume WLOG that  $\sigma^{*} = 1$ . The number of cells in the grid only depends on  $\sigma / \sigma^{*}$ .
Now, we show that, regardless of the value of  $\mu^{*}$  we only need to make a grid on a segment of length at most  $3\max (\sigma ,1)\sqrt{\log(2d / \epsilon)}$  . This is because for any  $\mu$  outside the ranges specified below the  $d_{TV}(p_{\mu ,\sigma}(y),p_{\mu^*,\sigma^*}(y))\geq 1 - \epsilon$
- If  $\mu^{*} \leq -\sqrt{\log(2d / \epsilon)}$  and for any  $\mu$  such that  $\mu \geq \sigma \sqrt{\log(2d / \epsilon)}$ , the  $d_{TV}(p_{\mu^{*},\sigma^{*}}(y), p_{\mu,\sigma}(y)) \geq 1 - \epsilon$ .
- If  $0 \geq \mu^{*} \geq -\sqrt{\log(2d / \epsilon)}$  and for any  $\mu$  s.t.  $\mu \geq \max(\sigma, 1)\sqrt{\log(2d / \epsilon)}$ , the  $d_{TV}(p_{\mu^{*},\sigma^{*}}(y), p_{\mu,\sigma}(y))$  is the same as in the linear case and since,  $\mu - \mu^{*} \geq \max(\sigma, 1)\sqrt{\log(2d / \epsilon)}$ , the  $d_{TV}(p_{\mu^{*},\sigma^{*}}(y), p_{\mu,\sigma}(y)) \geq 1 - \epsilon$ .
- If  $0 \leq \mu^{*} \leq \sqrt{\log(2d / \epsilon)}$ , for any  $\mu$  s.t.  $\mu - \mu^{*} \geq \max(\sigma, 1) \sqrt{\log(2d / \epsilon)}$ , the
- If  $\mu^{*} \geq \sqrt{\log(2d / \epsilon)}$  then for any  $\mu$  s.t.  $\mu - \mu^{*} \geq \max(\sigma, 1) \sqrt{\log(2d / \epsilon)}$  using the same argument as above, we have, the  $d_{TV}(p_{\mu^{*}, \sigma^{*}}(y), p_{\mu, \sigma}(y)) \geq 1 - \epsilon$ . Moreover, this is also true for  $\mu$  s.t.  $-\sigma \sqrt{\log(2d / \epsilon)} \leq \mu \leq \mu^{*} - \max(\sigma, 1) \sqrt{\log(2d / \epsilon)}$ . Therefore, in this case, we have an additional cell.
In addition to the above, for any  $\mu, \mu' \leq -\sigma \sqrt{\log(2d / \epsilon)}$ , the  $d_{TV}(p_{\mu,\sigma}(y), p_{\mu',\sigma}(y)) \leq \epsilon / 2d$  since both  $y_{\mu,\sigma}, y_{\mu',\sigma}$  are only non-zero with probability at most  $\epsilon / 2d$ . Therefore, for all the above cases, we only need to partition a segment of length at most  $3\max(\sigma, 1)\sqrt{\log(2d / \epsilon)}$
Moreover, for  $\sigma$  sufficiently small we can do better. We only need to partition a space of  $\sigma \sqrt{\log(2d / \epsilon)}$ . This is primarily because when  $\sigma$  sufficiently small, for any  $\mu$  in the linear case we have that  $d_{TV}(p_{\mu ,\sigma}(y),p_{\mu^*,\sigma^*}(y))\geq 1 - \epsilon$ .
It is easy to see that  $d_{TV}(p_{\mu, \sigma}(y), p_{0,1}(y)) \geq d_{TV}(p_{0, \sigma}(y), p_{0,1}(y))$ . The PDFs of  $\mathcal{N}(0, \sigma)$  and  $\mathcal{N}(0, 1)$  intersect at  $x = \pm \sigma \sqrt{\frac{\log(1 / \sigma^2)}{1 - \sigma^2}}$ . To show that  $d_{TV}(p_{0, \sigma}(y), p_{0,1}(y)) \geq 1 - \epsilon$ , it is now sufficient to show that
$$
1 - 2 \Phi (- | x | / \sigma) \geq 1 - \epsilon \quad \text {a n d} \quad 1 - 2 \phi (- | x |) \leq \epsilon ,
$$
where  $\Phi (x)$  is the CDF of the standard normal distribution. By using classical bounds on  $\Phi (x)$ , we have that
$$
\Phi (- | x | / \sigma) \leq \frac {\exp^ {- x ^ {2} / 2 \sigma^ {2}}}{| x | / \sigma} = \frac {\exp^ {- \frac {\log (1 / \sigma^ {2})}{2 (1 - \sigma^ {2})}}}{| x | / \sigma}
$$
which is  $\leq \epsilon /2$  if  $\sigma^2\leq \epsilon /2$  . And,
$$
\Phi (- | x |) \geq \frac {| x | \exp^ {- x ^ {2} / 2}}{x ^ {2} + 1}
$$
which is  $\geq (1 - \epsilon) / 2$  if
When  $\mu^{*} \leq -\sqrt{\log(2d / \epsilon)}$  the same argument as above shows that if  $\mu > \sigma \sqrt{\log(2d / \epsilon)}$  then the  $d_{TV}(p_{\mu^{*},\sigma^{*}}(y), p_{\mu,\sigma}(y)) \geq$  the difference in the probabilities at 0 which is bigger than  $1 - \epsilon$ . We consider the case where  $\mu^{*} \geq -\sqrt{\log(2d / \epsilon)}$ .
Since, when  $\sigma$  is small, for any  $\mu$  in the linear case the TV distance is large it is sufficient to have  $\mu$  large enough so that the intersection of the PDFs are positive and we are in the linear case.
The point of intersection assuming  $\sigma^{*} = 1$  is given by
$$
x = \frac {\mu_ {1} - \mu_ {2} / \sigma^ {2} \pm \sqrt {(\mu_ {1} - \mu_ {2}) ^ {2} / \sigma^ {2} + \left(\frac {1}{\sigma^ {2}} - 1\right) \log (\sigma^ {2})}}{\frac {1}{\sigma^ {2}} - 1}
$$
which is positive whenever
$$
\mu_ {2} \geq \sigma \sqrt {\log (1 / \sigma^ {2} + \mu_ {1} ^ {2}} \geq \sigma \sqrt {\log (2 d / \epsilon}.
$$
For the rest of the space, we partition  $\mu \in \mathbb{R}^k$  s.t. for any  $\mu ,\mu^{\prime}$  in a cell,
$$
| \mu - \mu^ {\prime} | \leq \sigma \epsilon / 2 d.
$$
This implies that for any  $\mu, \mu'$  in a cell, either,
$$
d _ {T V} \left(p _ {\mu , \sigma} (y), p _ {\mu^ {\prime}, \sigma} (y)\right) \leq \epsilon / 2 d
$$
$$
d _ {T V} \left(p _ {\mu , \sigma} (y), p _ {\mu^ {*}, \sigma^ {*}} (y)\right) \geq 1 - \epsilon \text {a n d} d _ {T V} \left(p _ {\mu^ {\prime}, \sigma} (y), p _ {\mu^ {*}, \sigma^ {*}} (y)\right) \geq 1 - \epsilon .
$$
Then the number of cells is the max of  $\mathcal{O}(d\sqrt{\log(2d / \epsilon)} /\epsilon)$  (when  $\sigma$  small) or  $\mathcal{O}(d\sqrt{\log(2d / \epsilon)} /\sigma \epsilon)\leq d\sqrt{\log(2d / \epsilon)} /\epsilon^{3 / 2}$  (when  $\sigma$  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
$$
d _ {T V} \left(X ^ {j}, \hat {X} ^ {j}\right) \leq \epsilon / 2 \quad \text {a n d}
$$
$$
d _ {T V} (\phi (\hat {W} _ {M L E} ^ {j} X _ {j} + \eta_ {j}), \phi (W ^ {j} X _ {j} + \eta_ {j})) \leq \epsilon / 2
$$
$$
\begin{array}{l} d _ {T V} (\phi (\hat {W} _ {M L E} ^ {j} \hat {X} _ {j} + \eta_ {j}), \phi (W ^ {j} X _ {j} + \eta_ {j})) \\ \leq d _ {T V} \left(\phi \left(\hat {W} _ {M L E} ^ {j} \hat {X} _ {j} + \eta_ {j}\right), \phi \left(\hat {W} _ {M L E} ^ {j} X _ {j} + \eta_ {j}\right)\right) \\ + d _ {T V} (\phi (\hat {W} _ {M L E} ^ {j} X _ {j} + \eta_ {j}), \phi (W ^ {j} X _ {j} + \eta_ {j})) \\ \leq \epsilon \\ \end{array}
$$
where we use the fact that  $d_{TV}(f(X), f(Y)) \leq d_{TV}(X, Y)$ .
(a) Total Variation Upper Bound vs  $d$
(b) Total Variation vs.  $n$  for Different  $\mathcal{D}_x$
Figure 4: (a) Plots Pinsker's upper bound on the TV distance as  $d$  gets large. We set  $\Sigma^{*} = I_{d}$  and  $W^{*} = \mathbb{1}_{d\times 1}$ , thus setting input dimension  $k = 1$ .  $n = 5000$  samples are taken. As we might expect, the upper bound is increasing in  $d$ . Each point is determined by 2000 samples. (b) Plot of TV vs.  $n$  for additional distributions of  $x$ . All three distributions follow roughly the same trend, each point is determined by 2500 samples.
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 experiments, 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  $\sigma^{*} = 1$  and the  $w^{*} = \mathbb{1}_{k\times 1}$ . In each case the MLE is solved via gradient descent with backtracking line search, and we check a first order condition  $\| \nabla_{w,\sigma}\log p_{w,\sigma}((y\mid x))\| ^2 < \delta = 10^{-3}$  as the exit condition. We verify that increasing or decreasing  $\delta$  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_{\hat{w},\hat{\sigma}}(y\mid x),p_{w,\sigma}(y\mid 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  $\Sigma^{*} = I_{d}$  and take let  $W^{*} = b\mathbb{1}_{d\times 1}$ , where  $b$  will vary across our experiments. We set  $n = 10000$ . In Figure 3b we set  $n = 5000$  and adjust  $\Sigma$  such that one diagonal entry is  $\kappa^{1/2}$ , and the other is  $\kappa^{-1/2}$ , making the total condition number  $\kappa$ .
In both of these experiments, we restrict the MLE computation to be over diagonal  $\Sigma$  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:
$$
\begin{array}{l} d_{TV}\Bigl((\widehat{W},\widehat{\Sigma}),(W^{*},\Sigma^{*})\Bigr) = \frac{1}{2}\int_{\mathbb{R}^{d}_{\geq 0}}\biggl|p_{\widehat{W},\widehat{\Sigma}}(y\mid x) - p_{W^{*},\Sigma^{*}}(y\mid x)\biggr|dy \\ = \frac {1}{2} \int_ {\mathbb {R} _ {\geq 0} ^ {3}} \left| \prod_ {i = 1} ^ {d} p _ {\widehat {W}, \widehat {\Sigma}} (y _ {i} \mid x) - \prod_ {i = 1} ^ {d} p _ {W ^ {*}, \Sigma^ {*}} (y _ {i} \mid x) \right| d y, \\ \end{array}
$$
where  $y_{i}$  is the  $i^{\text{th}}$  element of  $y$ . Though at first glance 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 non-negative orthant cone  $\mathbb{R}_{\geq 0}^{3}$  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:
$$
\begin{array}{l} d _ {T V} \left(\left(\widehat {W}, \widehat {\Sigma}\right), \left(W ^ {*}, \Sigma^ {*}\right)\right) = \\ \frac {1}{2} \sum_ {S ^ {\prime} \in 2 ^ {[ d ]}} \int_ {\mathbb {R} _ {\geq 0} ^ {| S ^ {\prime} |}} \left| \prod_ {i \in S ^ {\prime}} p _ {\widehat {W}, \widehat {\Sigma}} \left(y _ {i} \mid x\right) \prod_ {i \in \left(S ^ {\prime}\right) ^ {c}} \Phi \left(\widehat {\Sigma} _ {i i} ^ {- 1 / 2} \widehat {W} _ {i}\right) - \prod_ {i \in S ^ {\prime}} p _ {W ^ {*}, \Sigma^ {*}} \left(y _ {i} \mid x\right) \prod_ {i \in \left(S ^ {\prime}\right) ^ {c}} \Phi \left(\Sigma_ {i i} ^ {- 1 / 2} W _ {i}\right) \right| d y _ {S ^ {\prime}}. \tag {45} \\ \end{array}
$$
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  $\sigma^{*}$  and some  $W^{*} \in \mathbb{R}^{1,k}$  and describe how to compute the likelihood function. We defer the proof of the
log-concavity to the following section, which covers the more general case with  $d \geq 1$  When we re-parameterize as  $u = -W / \sigma$ ,  $v = 1 / \sigma$ , the likelihood function is written as:
$$
f _ {u, v} (y) = - \frac {1}{2} \sum_ {i \in S ^ {\prime}} \left(v y _ {i} - u \cdot x _ {i}\right) ^ {2} + \left| S ^ {\prime} \right| \log (v) + \sum_ {i \in S} \log \Phi (- u \cdot x _ {i}) \tag {46}
$$
where in this case we let  $S = \{i \mid y_i = 0\}$ , where  $y_i$  is the  $i^{\text{th}}$  sample in the set  $\{y_i, x_i\}_{i=1}^n$ . 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  $\eta$  follows a similar approach.
Numerical concerns. In (46), the term  $\log \Phi(-u \cdot x_i)$  presents some numerical concerns when  $u \cdot x_i \gg 0$  if we naively compute  $\Phi(-u \cdot 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:
$$
\log \Phi (- x) = - \log m (x) - \frac {1}{2} \log (2 \pi) - \frac {1}{2} x ^ {2}, x > 0.
$$
Since  $m(x)$  changes relatively slowly in  $x$  compared to  $\Phi(-x)$ , this greatly improves numerical stability.
F.2 Multidimensional Case
In the multi-dimensional case, we will generally use the more standard natural parameters:
$$
U := \frac {\Sigma^ {- 1}}{2},
$$
$$
v := - \Sigma^ {- 1} W x.
$$