SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Definition A.1 (Frechet variance). The Frechet variance is defined as
$$
\Psi_ {\mathbf {w}} (V) := \sum_ {i \in [ K ]} \mathbf {w} _ {i} d ^ {2} (V, V _ {i}).
$$
Additionally, we will require a technical assumption related to the sensitivity of the Fréchet variance to different node choices.
Assumption A.2 (Frechet variance is  $\frac{1}{\alpha}$ -bi-Lipschitz). For a metric space defined by a graph  $G = (\mathcal{V}, \mathcal{E})$ , the Frechet variance is  $K$ -bi-Lipschitz if there exists a  $K \geq 1$  such that
$$
\frac {1}{K} d ^ {2} (V, \tilde {V}) \leq | \Psi_ {\mathbf {w}} (V) - \Psi_ {\mathbf {w}} (\tilde {V}) | \leq K d ^ {2} (V, \tilde {V})
$$
for all  $V, \tilde{V} \in \mathcal{V}$ , and a fixed  $\mathbf{w} \in \Delta^{K-1}$ . For our purposes, such a  $K$  always exists: consider setting  $K = \mathrm{diam}(G)^2 \max_{V_1, V_2 \in \mathcal{V}} |\Psi_{\mathbf{w}}(V_1) - \Psi_{\mathbf{w}}(V_2)|$ . However, this is a very conservative bound that holds for all graphs that we consider. Instead, we assume access to the largest  $\alpha = \frac{1}{K} \leq 1$  such that
$$
\alpha d ^ {2} (V, \tilde {V}) \leq | \Psi_ {\mathbf {w}} (V) - \Psi_ {\mathbf {w}} (\tilde {V}) |,
$$
which may be problem dependent.
Theorem A.3 (LOKI sample complexity). Let  $\mathcal{V}$  be a set of points on the  $d$  dimensional 2-norm ball of radius  $R$ , and let  $\Lambda \subseteq \mathcal{V}$  be the set of  $K$  observed classes. Assume that  $\Lambda$  forms a  $2R / (\sqrt{d} K - 1)$ -net under the Euclidean distance. Assume that training examples are generated by drawing  $n$  samples from the following process: draw  $x \sim \mathcal{N}(y, I)$  where  $y \sim \mathrm{Unif}(\Lambda)$ , and at test time, draw  $x \sim \mathcal{N}(y, I)$  where  $y \sim \mathrm{Unif}(\mathcal{V})$ . Assume that we estimate a Gaussian mixture model with  $K$  components (each having identity covariance) on the training set and obtain probability estimates  $\hat{\mathbb{P}}(y_i | x)$  for  $i \in [K]$  for a sample  $(x, y_*)$  from the test distribution. Then with high probability, under the following model,
where  $d$  is the dimensionality of the input and  $\alpha$  is our parameter under Assumption A.2.
Proof. We begin by detailing the data-generating process.
At training time, our underlying data-generating process,  $\mathcal{D}_{\mathrm{train}}$ , is as follows:
- We begin with a  $\frac{2R}{\sqrt{dK - 1}}$ -net,  $\mathcal{V}$ , on the  $d$  dimensional 2-norm ball of radius  $R$  under the Euclidean distance, and let  $\Lambda \subseteq \mathcal{V}$
Draw  $y\sim \mathrm{Unif}(\mathcal{V})$
- Discard draws if  $y \notin \Lambda$ .
Draw  $x\sim \mathcal{N}(y,I)$
At test time, we do not discard draws and allow for classes not in  $\Lambda$  and use the following data generating process,  $\mathcal{D}_{\mathrm{test}}$ :
- We begin with a  $\frac{2R}{\sqrt{dK - 1}}$ -net,  $\mathcal{V}$ , on the  $d$  dimensional 2-norm ball of radius  $R$  under the Euclidean distance, and let  $\Lambda \subseteq \mathcal{V}$
- Draw  $y \sim \operatorname{Unif}(\mathcal{Y})$ .
Draw  $x\sim \mathcal{N}(y,I)$
Given a labeled training dataset  $D = \{(x_{i},y_{i})\}_{i = 1}^{n}$  containing  $n$  points drawn from  $\mathcal{D}_{\mathrm{train}}$  with  $|\Lambda| = K$  distinct classes, we would like to fit a  $K$ -component Gaussian mixture model with identity covariance. We first perform mean estimation of each of the classes separately using the median-of-means estimator [15, 27, 24]. Using the estimator yields the following parameter estimation bound [27, 4]:
$$
\left| \left| y _ {i} - \hat {y} _ {i} \right| \right| _ {2} \leq O \left(\sqrt {\frac {d \log K / \delta}{n}}\right)
$$
with probability  $1 - \delta$
Next, we consider the relationships between four quantities:  $\Psi_{\mathbb{P}}(y_{*}), \Psi_{\mathbb{P}}(\hat{y}_{*}), \Psi_{\widehat{\mathbb{P}}}(\hat{y}_{*}), \Psi_{\widehat{\mathbb{P}}}(\hat{y}_{*})$ , where  $\mathbb{P}$  is the vector of probabilities from the true Gaussian mixture,  $\widehat{\mathbb{P}}$  is the vector of probabilities from the estimated model,  $y_{*} \in m_{\Lambda}(\mathbb{P})$  is the target class, and  $\hat{y}_{*} \in m_{\Lambda}(\widehat{\mathbb{P}})$  is the predicted class. While it is problem-dependent as to whether  $\Psi_{\mathbb{P}}(y_{*}) \leq \Psi_{\widehat{\mathbb{P}}}(\hat{y}_{*})$  or  $\Psi_{\mathbb{P}}(y_{*}) > \Psi_{\widehat{\mathbb{P}}}(\hat{y}_{*})$ , a similar argument holds for both cases. So without loss of generality, we assume that  $\Psi_{\mathbb{P}}(y_{*}) \leq \Psi_{\widehat{\mathbb{P}}}(\hat{y}_{*})$ . Then by the definition of the Fréchet mean, the following inequalities hold:
$$
\Psi_ {\mathbb {P}} (y _ {*}) \leq \Psi_ {\widehat {\mathbb {P}}} (\widehat {y} _ {*}) \leq \Psi_ {\widehat {\mathbb {P}}} (y _ {*}),
$$
$$
\Psi_ {\widehat {\mathbb {P}}} \left(\widehat {y} _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \leq \Psi_ {\widehat {\mathbb {P}}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right). \tag {2}
$$
We proceed by obtaining upper and lower bounds of the existing bounds in Equation 2. First, we will obtain an upper bound on  $\Psi_{\widehat{\mathbb{P}}}(\boldsymbol {y}_*) - \Psi_{\mathbb{P}}(\boldsymbol {y}_*)$
$$
\begin{array}{l} | \Psi_ {\widehat {\mathbb {P}}} (y _ {*}) - \Psi_ {\mathbb {P}} (y _ {*}) | = \left| \sum_ {i \in [ K ]} (\widehat {\mathbb {P}} - \mathbb {P}) d ^ {2} (y _ {*}, V _ {i}) \right| \\ = \left| \sum_ {i \in [ K ]} \left(\widehat {\mathbb {P}} \left(y _ {i} | x\right) - \mathbb {P} \left(y _ {i} | x\right)\right) \left| \left| y _ {*} - y _ {i} \right| \right| _ {2} ^ {2} \right| \\ = \left| \sum_ {i \in [ K ]} \left(\frac {\widehat {\mathbb {P}} (x \mid y _ {i}) \mathbb {P} (y _ {i})}{\sum_ {j \in [ K ]} \widehat {\mathbb {P}} (x \mid y _ {j}) \mathbb {P} (y _ {j})} - \frac {\mathbb {P} (x \mid y _ {i}) \mathbb {P} (y _ {i})}{\sum_ {j \in [ K ]} \mathbb {P} (x \mid y _ {j}) \mathbb {P} (y _ {j})}\right) \| y _ {*} - y _ {i} \| _ {2} ^ {2} \right| \tag {3} \\ \end{array}
$$
$$
\begin{array}{l} = \left| \sum_ {i \in [ K ]} \left(\frac {\exp \left\{- \frac {1}{2} | | x - \hat {y} _ {i} | | _ {2} ^ {2} \right\} \frac {1}{K}}{\sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} | | x - \hat {y} _ {j} | | _ {2} ^ {2} \right\} \frac {1}{K}} \right. \right. \\ \left. - \frac {\exp \left\{- \frac {1}{2} \| x - y _ {i} \| _ {2} ^ {2} \right\} \frac {1}{K}}{\sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} \| x - y _ {j} \| _ {2} ^ {2} \right\} \frac {1}{K}}\right) \| y _ {*} - y _ {i} \| _ {2} ^ {2}. \\ = \left| \sum_ {i \in [ K ]} \left(\frac {\exp \left\{- \frac {1}{2} | | x - \hat {y} _ {i} | | _ {2} ^ {2} \right\}}{\sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} | | x - \hat {y} _ {j} | | _ {2} ^ {2} \right\}} \right. \right. \\ \left. - \frac {\exp \left\{- \frac {1}{2} \| x - y _ {i} \| _ {2} ^ {2} \right\}}{\sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} \| x - y _ {j} \| _ {2} ^ {2} \right\}}\right) \| y _ {*} - y _ {i} \| _ {2} ^ {2}. \tag {4} \\ \end{array}
$$
For notational convenience, we define the following:
$$
\begin{array}{l} a _ {i} := \exp \left\{- \frac {1}{2} | | x - y _ {i} | | _ {2} ^ {2} \right\}, \\ \hat {a} _ {i} := \exp \left\{- \frac {1}{2} | | x - \hat {y} _ {i} | | _ {2} ^ {2} \right\}, \\ b := \sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} | | x - y _ {j} | | _ {2} ^ {2} \right\}, \\ \hat {b} := \sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} \| x - \hat {y} _ {j} \| _ {2} ^ {2} \right\}, a n d \\ c _ {i} := \left\| \hat {y} _ {*} - y _ {i} \right\| _ {2} ^ {2}. \\ \end{array}
$$
Then (4) becomes:
$$
\begin{array}{l} \left| \sum_ {i \in [ K ]} \left(\frac {\hat {a} _ {i}}{\hat {b}} - \frac {a _ {i}}{b}\right) c _ {i} \right| = \left| \sum_ {i \in [ K ]} \left(\frac {\hat {a} _ {i}}{\hat {b}} - \frac {\hat {a} _ {i}}{b} + \frac {\hat {a} _ {i}}{b} - \frac {a _ {i}}{b}\right) c _ {i} \right| \\ = \left| \sum_ {i \in [ K ]} \left(\frac {\hat {a} _ {i} - a _ {i}}{b} + \hat {a} _ {i} \left(\frac {1}{\hat {b}} - \frac {1}{b}\right)\right) c _ {i} \right| \\ \leq \left| \sum_ {i \in [ K ]} \left(\frac {\hat {a} _ {i} - a _ {i}}{b}\right) c _ {i} \right| + \left| \sum_ {i \in [ K ]} \left(\hat {a} _ {i} \left(\frac {1}{\hat {b}} - \frac {1}{b}\right)\right) c _ {i} \right| \\ = \left| \sum_ {i \in [ K ]} \frac {a _ {i}}{b} \left(\frac {\hat {a} _ {i}}{a _ {i}} - 1\right) c _ {i} \right| + \left| \frac {b - \hat {b}}{b} \sum_ {i \in [ K ]} \frac {\hat {a} _ {i}}{\hat {b}} c _ {i} \right| \\ = \left| \sum_ {i \in [ K ]} \frac {a _ {i}}{b} \left(\frac {\hat {a} _ {i}}{a _ {i}} - 1\right) c _ {i} \right| + \left| \left(\sum_ {i \in [ K ]} \frac {a _ {i}}{b} \left(\frac {\hat {a} _ {i}}{a _ {i}} - 1\right)\right) \left(\sum_ {i \in [ K ]} \frac {\hat {a} _ {i}}{\hat {b}} c _ {i}\right) \right|. \tag {5} \\ \end{array}
$$
Now we define the following:  $L := ||x - y_z||_2^2$  with  $z \in \arg \min_j ||x - y_j||_2^2$  is the smallest distance to a class mean, and  $L + E_i := ||x - y_i||_2^2$  with  $E_i > 0$ . Similarly, define  $\hat{L} := ||x - \hat{y}_z||_2^2$  with  $z \in \arg \min_j ||x - \hat{y}_j||_2^2$  is the smallest distance to an estimated class mean, and  $\hat{L} + \hat{E}_i := ||x - \hat{y}_i||_2^2$  with  $\hat{E}_i > 0$ . Finally, we define  $T := \sqrt{\hat{L}} + \sqrt{\hat{L} + \hat{E}_i} + \sqrt{L} + \sqrt{L + E_i}$ .
Then we can bound the parts separately:
For  $i\neq z$  , we have
$$
\begin{array}{l} \frac {a _ {i}}{b} = \left(\frac {\exp \left\{- \frac {1}{2} | | x - y _ {i} | | _ {2} ^ {2} \right\}}{\sum_ {j \in [ K ]} \exp \left\{- \frac {1}{2} | | x - y _ {j} | | _ {2} ^ {2} \right\}}\right) \\ \leq \left(\frac {\exp \{- \frac {1}{2} (L + E _ {i}) \}}{\exp \{- \frac {1}{2} L \} + \exp \{- \frac {1}{2} (L + E _ {i}) \}}\right) \\ = \frac {1}{\exp \left\{\frac {1}{2} E _ {i} \right\}} \tag {6} \\ \end{array}
$$
and in the case of  $i = z$ , we bound  $\frac{a_z}{b} \leq 1$ . So overall, we have  $\frac{a_i}{b} \leq \frac{1}{\exp\{1_{i \neq z} \frac{1}{2} E_i\}}$  for all  $i \in [K]$ .
$$
\begin{array}{l} c _ {i} = \left\| y _ {*} - y _ {i} \right\| _ {2} ^ {2} \\ \leq 2 \| x - y _ {i} \| _ {2} ^ {2} + 2 \| x - y _ {*} \| _ {2} ^ {2} \\ = 2 \left(L + E _ {i} + \left\| x - y _ {*} \right\| _ {2} ^ {2}\right). \tag {7} \\ \end{array}
$$
$$
\begin{array}{l} \frac {\hat {a} _ {i}}{a _ {i}} - 1 = \exp \left\{\frac {1}{2} | | x - y _ {i} | | _ {2} ^ {2} - \frac {1}{2} | | x - \hat {y} _ {i} | | _ {2} ^ {2} \right\} - 1 \\ \approx 1 + \frac {1}{2} | | x - y _ {i} | | _ {2} ^ {2} - \frac {1}{2} | | x - \hat {y} _ {i} | | _ {2} ^ {2} - 1 \\ = \frac {1}{2} \left\langle \hat {y} _ {i} - y _ {i}, 2 x - y _ {i} - \hat {y} _ {i} \right\rangle (8) \\ \leq \frac {1}{2} | | \hat {y} _ {i} - y _ {i} | | \cdot | | 2 x - y _ {i} - \hat {y} _ {i} | | \\ \leq \left| \left| \hat {y} _ {i} - y _ {i} \right| \right| \left(\left| \left| x - y _ {i} \right| \right| + \left| \left| x - \hat {y} _ {i} \right| \right|\right) \\ = \left\| \hat {y} _ {i} - y _ {i} \right\| \left(\left\| x - y _ {i} \right\| + \left\| x - y _ {i} + y _ {i} - \hat {y} _ {i} \right\|\right) \\ \leq \left| \left| \hat {y} _ {i} - y _ {i} \right| \right| \left(2 \left| \left| x - y _ {i} \right| \right| + \left| \left| y _ {i} - \hat {y} _ {i} \right| \right|\right) \\ = \left\| \hat {y} _ {i} - y _ {i} \right\| \left(2 \sqrt {L + E _ {i}} + \left\| y _ {i} - \hat {y} _ {i} \right\|\right). (9) \\ \end{array}
$$
Next, we must control  $E_{i} - \hat{E}_{i}$  in order to obtain the bound for  $\frac{\hat{a}_i}{\hat{b}}$ .
$$
\begin{array}{l} E _ {i} - \hat {E} _ {i} = \left(\left\| x - y _ {i} \right\| ^ {2} - L\right) - \left(\left\| x - \hat {y} _ {i} \right\| ^ {2} - \hat {L}\right) \\ = \left\| x - y _ {i} \right\| ^ {2} - \left\| x - \hat {y} _ {i} \right\| ^ {2} + \left\| x - \hat {y} _ {z} \right\| ^ {2} - \left\| x - y _ {z} \right\| ^ {2} \\ = \left\langle \hat {y} _ {i} - y _ {i}, 2 x - y _ {i} - \hat {y} _ {i} \right\rangle + \left\langle y _ {z} - \hat {y} _ {z}, 2 x - \hat {y} _ {z} - y _ {z} \right\rangle \\ \leq \left\| \hat {y} _ {i} - y _ {i} \right\| \cdot \left\| 2 x - y _ {i} - \hat {y} _ {i} \right\| + \left\| y _ {z} - \hat {y} _ {z} \right\| \cdot \left\| 2 x - \hat {y} _ {z} - y _ {z} \right\| \\ \leq \left\| y _ {i} - \hat {y} _ {i} \right\| \left(\left\| x - \hat {y} _ {i} \right\| + \left\| x - y _ {i} \right\|\right) + \left\| \hat {y} _ {z} - y _ {z} \right\| \left(\left\| x - y _ {z} \right\| + \left\| x - \hat {y} _ {z} \right\|\right) \\ \leq c \sqrt {\frac {d \log K / \delta}{n}} \left(\left| \left| x - \hat {y} _ {i} \right| \right| + \left| \left| x - y _ {i} \right| \right| + \left| \left| x - y _ {z} \right| \right| + \left| \left| x - \hat {y} _ {z} \right| \right|\right) \\ = c \sqrt {\frac {d \log K / \delta}{n}} \left(\sqrt {\hat {L}} + \sqrt {\hat {L} + \hat {E} _ {i}} + \sqrt {L} + \sqrt {L + E _ {i}}\right) \\ \end{array}
$$
$$
\hat {E} _ {i} \geq \max  \left\{0, E _ {i} - c T \sqrt {\frac {d \log K / \delta}{n}} \right\}.
$$
Using (6), we obtain
$$
\begin{array}{l} \frac {\hat {a} _ {i}}{\hat {b}} \leq \frac {1}{\exp \left\{\mathbf {1} _ {i \neq z} \frac {1}{2} \hat {E} _ {i} \right\}} \\ \leq \frac {1}{\exp \left\{\mathbf {1} _ {i \neq z} \max  \left\{0 , \frac {1}{2} E _ {i} - \frac {c T}{2} \sqrt {\frac {d \log K / \delta}{n}} \right\} \right\}} \tag {10} \\ \end{array}
$$
Plugging (6), (7), (9), and (10) into (5), we obtain
$$
\left| \sum_ {i \in [ K ]} \frac {a _ {i}}{b} \left(\frac {\hat {a} _ {i}}{a _ {i}} - 1\right) c _ {i} \right| + \left| \left(\sum_ {i \in [ K ]} \frac {a _ {i}}{b} \left(\frac {\hat {a} _ {i}}{a _ {i}} - 1\right)\right) \left(\sum_ {i \in [ K ]} \frac {\hat {a} _ {i}}{\hat {b}} c _ {i}\right) \right|
$$
$$
\begin{array}{l} \leq \sum_ {i \in [ K ]} \frac {| | \hat {y} _ {i} - y _ {i} | | \left(2 \sqrt {L + E _ {i}} + | | y _ {i} - \hat {y} _ {i} |\right)}{\exp \left\{\mathbf {1} _ {i \neq z} \frac {1}{2} E _ {i} \right\}} 2 (L + E _ {i} + | | x - y _ {*} | | _ {2} ^ {2}) \\ + \left(\sum_ {i \in [ K ]} \frac {| | \hat {y} _ {i} - y _ {i} | | (2 \sqrt {L + E _ {i}} + | | y _ {i} - \hat {y} _ {i} | |)}{\exp \left\{\mathbf {1} _ {i \neq z} \frac {1}{2} E _ {i} \right\}}\right) \\ \cdot \left(\sum_ {i \in [ K ]} \frac {2 \left(L + E _ {i} + \left| \left| x - y _ {*} \right| \right| _ {2} ^ {2}\right)}{\exp \left\{\mathbf {1} _ {i \neq z} \max  \left\{0 , \frac {1}{2} E _ {i} - \frac {c T}{2} \sqrt {\frac {d \log K / \delta}{n}} \right\} \right\}}\right) \\ \leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \sum_ {i \in [ K ]} \frac {L + E _ {i} + | | x - y _ {*} | | _ {2} ^ {2}}{\exp \left\{\mathbf {1} _ {i \neq z} \frac {1}{2} E _ {i} \right\}}\right) \\ \end{array}
$$
Next, recalling the fact that our class means form an  $\varepsilon$ -net on the radius- $R$  ball, we use Lemma 5.2 from [40] to bound  $L$  as  $L \leq \frac{2R}{\sqrt{d}K - 1}$ .
$$
\leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \sum_ {i \in [ K ]} \frac {\frac {R}{\sqrt {d} K} + E _ {i} + | | x - y _ {*} | | _ {2} ^ {2}}{\exp \left\{\mathbf {1} _ {i \neq z} \frac {1}{2} E _ {i} \right\}}\right). \tag {11}
$$
Next, we will consider cases on  $E_{i}$ . We will consider the cases in which  $E_{i} < \log R^{2}$  and  $E_{i} \geq \log R^{2}$ .
We will begin with the case in which  $E_{i} < \log R^{2}$ , then Equation 11 becomes:
$$
\leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \left(\sum_ {i \in [ K ]} \frac {R}{\sqrt {d} K} + | | x - y _ {*} | | _ {2} ^ {2}\right)\right).
$$
Next, the case in which  $E_{i} \geq \log R^{2}$ , then Equation 11 becomes:
$$
\begin{array}{l} \leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \sum_ {i \in [ K ]} \frac {\frac {R}{\sqrt {d} K} + 2 \log R + | | x - y _ {*} | | _ {2} ^ {2}}{R}\right) \\ \leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \left(\sum_ {i \in [ K ]} \frac {1}{\sqrt {d} K} + | | x - y _ {*} | | _ {2} ^ {2}\right)\right). \\ \end{array}
$$
Setting  $M \leq K$  to be the number of terms for which  $E_{i} < \log R^{2}$  holds and by combining the two bounds, we obtain the following:
$$
\leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \left(\left(\frac {M R}{\sqrt {d} K} + \frac {K - M}{\sqrt {d} K}\right) + K | | x - y _ {*} | | _ {2} ^ {2}\right)\right).
$$
Ultimately, our bound is
$$
\left| \Psi_ {\widehat {\mathbb {P}}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \right| \leq O \left(\sqrt {\frac {d \log K / \delta}{n}} \left(\left(\frac {M R}{\sqrt {d} K} + \frac {K - M}{\sqrt {d} K}\right) + K \| x - y _ {*} \| _ {2} ^ {2}\right)\right). \tag {12}
$$
Next, we will obtain a lower bound on  $\Psi_{\widehat{\mathbb{P}}}(\hat{y}_{*}) - \Psi_{\mathbb{P}}(y_{*})$
$$
\begin{array}{l} \left| \Psi_ {\widehat {\mathbb {P}}} \left(\widehat {y} _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \right| \geq \left| \Psi_ {\mathbb {P}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(\widehat {y} _ {*}\right) \right| - \left| \Psi_ {\widehat {\mathbb {P}}} \left(\widehat {y} _ {*}\right) - \Psi_ {\mathbb {P}} \left(\widehat {y} _ {*}\right) \right| \quad t r i a n g l e \\ \geq \alpha d ^ {2} \left(y _ {*}, \hat {y} _ {*}\right) - \left| \Psi_ {\widehat {\mathbb {P}}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \right| \quad \text {A s s u m p t i o n A . 2}. \\ \end{array}
$$
Combining both of these bounds with Equation 2, we obtain
$$
\begin{array}{l} \alpha d ^ {2} \left(y _ {*}, \hat {y} _ {*}\right) - \left| \Psi_ {\widehat {\mathbb {P}}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \right| \leq \Psi_ {\widehat {\mathbb {P}}} \left(\hat {y} _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \leq \Psi_ {\widehat {\mathbb {P}}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \\ \leq \left| \Psi_ {\mathbb {P}} \left(y _ {*}\right) - \Psi_ {\mathbb {P}} \left(y _ {*}\right) \right| \\ \Rightarrow \alpha d ^ {2} (y _ {*}, \hat {y} _ {*}) \leq 2 | \Psi_ {\hat {\mathbb {P}}} (y _ {*}) - \Psi_ {\mathbb {P}} (y _ {*}) | \\ \Rightarrow d ^ {2} (y _ {*}, \hat {y} _ {*}) \leq O \left(\frac {1}{\alpha} \sqrt {\frac {d \log K / \delta}{n}} \left(\frac {M R}{\sqrt {d} K} + \frac {K - M}{\sqrt {d} K} + K | | x - y _ {*} | | _ {2} ^ {2}\right)\right) \\ \Rightarrow \mathbb {E} _ {(x, y _ {*}) \sim \mathcal {D} _ {\text {t e s t}}} [ d ^ {2} (y _ {*}, \hat {y} _ {*}) ] \leq O \left(\frac {1}{\alpha} \sqrt {\frac {d \log K / \delta}{n}} \left(\frac {M R}{\sqrt {d} K} + \frac {K - M}{\sqrt {d} K} + d K\right)\right), \\ \end{array}
$$
Next, we obtain a bound on  $M$ . Since  $M \leq K$  is the number of terms for which  $E_{i} < \log R^{2}$ , we have that  $M = VK$ , where  $V \in [0,1]$  is the ratio of the volumes of the ball for which  $E_{i} < \log R^{2}$  and the ball containing the entire set of classes  $\mathcal{V}$ . We have
$$
\begin{array}{l} E _ {i} = \left\| x - y _ {i} \right\| _ {2} ^ {2} - \left\| x - y _ {z} \right\| _ {2} ^ {2} <   \log R ^ {2} \\ \left\| x - y _ {i} \right\| _ {2} ^ {2} <   \log R ^ {2} + L \\ \leq \log R ^ {2} + \frac {2 R}{\sqrt {d} K - 1} \leq R ^ {2} \\ \Rightarrow \sqrt {\log R ^ {2} + \frac {2 R}{\sqrt {d} K - 1}} \leq R, \\ \end{array}
$$
which is an upper bound of the radius of the ball for which  $E_{i} < \log R^{2}$  is true. Now we construct  $V$ :
$$
V = \left(\frac {\sqrt {\log R ^ {2} + \frac {2 R}{\sqrt {d} K - 1}}}{R}\right) ^ {d}.
$$
Combining this with our error bound, we have
$$
\begin{array}{l} \mathbb {E} _ {(x, y _ {*}) \sim \mathcal {D} _ {\text {t e s t}}} [ d ^ {2} (y _ {*}, \hat {y} _ {*}) ] \leq \\ O \left(\frac {1}{\alpha} \sqrt {\frac {d \log K / \delta}{n}} \left(\frac {\left(\frac {\sqrt {\log R ^ {2} + \frac {2 R}{\sqrt {d} K - 1}}}{R}\right) ^ {d} R}{\sqrt {d}} + \frac {1 - \left(\frac {\sqrt {\log R ^ {2} + \frac {2 R}{\sqrt {d} K - 1}}}{R}\right) ^ {d}}{\sqrt {d}} + d K\right)\right). \\ \end{array}
$$
However, we can choose  $K$  to weaken our dependence on  $R$ .
$$
\begin{array}{l} V = \left(\frac {\sqrt {\log R ^ {2} + \frac {2 R}{\sqrt {d K} - 1}}}{R}\right) ^ {d} \leq \frac {1}{R} \\ \Rightarrow K \geq \frac {2}{\sqrt {d} \left(R ^ {1 - \frac {2}{d}} - \frac {\log R}{R}\right)} + 1. \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} _ {(x, y _ {*}) \sim \mathcal {D} _ {\text {t e s t}}} [ d ^ {2} (y _ {*}, \hat {y} _ {*}) ] \leq O \left(\frac {1}{\alpha} \sqrt {\frac {d \log K / \delta}{n}} \left(\frac {1}{\sqrt {d}} + \frac {1 - \frac {1}{R}}{\sqrt {d}} + \frac {\sqrt {d}}{\left(R ^ {1 - \frac {2}{d}} - \frac {\log R}{R}\right)} + d\right)\right) \\ \leq O \left(\frac {1}{\alpha} \sqrt {\frac {d \log K / \delta}{n}} \left(\frac {1}{\sqrt {d}} + \frac {\sqrt {d}}{\left(R ^ {1 - \frac {2}{d}} - \frac {\log R}{R}\right)} + d\right)\right) \\ \leq O \left(\frac {d}{\alpha} \sqrt {\frac {\log K / \delta}{n}} \left(\frac {1}{\left(R ^ {1 - \frac {2}{d}} - \frac {\log R}{R}\right)} + \sqrt {d}\right)\right) \\ \end{array}
$$
Discussion of Theorem A.3. The above bound contains several standard quantities, including the dimensionality of the inputs  $d$ , the radius  $R$  of the ball from which labels are drawn, the number of classes and samples used to estimate the model  $K$  and  $n$  respectively, and  $\delta$ , which is used to control the probability that the bound holds. Naturally, the bound scales up with higher  $d$  and  $K$ , and improves with an increased number of samples  $n$ . The dependence on  $R$  is also related to its dependence on  $d$  — so long as we have that  $K \geq 2(\sqrt{d}(R^{1 - \frac{2}{d}} - \log R / R))^{-1} + 1$  and for  $d \geq 2$ ,  $R$  increases, which improves the bound. The bound includes a metric space dependent quantity  $\alpha$ , which is related to the amount that the Fréchet variance can change subject to changes in its argument, i.e., which node is considered.
A.2 Proof of Theorem A.4 (minimum locus cover for trees)
Theorem A.4 (Minimum locus cover for trees). The minimum locus cover for a tree graph  $T$  is  $\Lambda = \text{Leaves}(T)$ .
Proof. We would like to show to show two things:
1. any node can be resolved by an appropriate construction of  $\mathbf{w}$  using only the leaves and
2. removing any leaf makes that leaf node unreachable, no matter what the other nodes are in  $y$ -i.e., we must use at least all of the leaf nodes.
If both of these conditions hold (i.e., that the leaves form a locus cover and that they are the smallest such set), then the leaves of any tree form a minimum locus cover. To set up the problem, we begin with some undirected tree,  $T$ , whose nodes are our set of labels:  $T = (\mathcal{V}, E)$  where  $\Lambda = \operatorname{Leaves}(T) \subseteq \mathcal{V}$  is the set of leaf nodes. Moreover, let  $\Lambda$  be a tuple of the leaf nodes. We will start by proving that  $\Lambda$  is a locus cover. We will show this by cases on  $v \in \mathcal{V}$ , the node that we would like to resolve:
1. If  $v \in \Lambda$ , i.e.  $v$  is a leaf node, then setting  $\boldsymbol{w}_i = 1\{\boldsymbol{\Lambda}_i = v\}$  yields
$$
\begin{array}{l} m_{\boldsymbol{\Lambda}}(\boldsymbol {w}) = \operatorname *{arg  min}_{y\in \mathcal{Y}}\sum_{i = 1}^{K}\mathbf{1}\{\boldsymbol{\Lambda}_{i} = v\} d^{2}(y,\boldsymbol{\Lambda}_{i}) \\ = \operatorname *{arg  min}_{y\in \mathcal{Y}}d^{2}(y,v) = \{v\} . \\ \end{array}
$$
2. If  $v \notin \Lambda$ , we have a bit more work to do. Since  $v$  is an internal node, consider any pair of leaves,  $l_1 \neq l_2$  such that  $v$  is along the (unique) path between  $l_1$  and  $l_2$ :  $v \in \Gamma(l_1, l_2)$ . Then
$$
\text {s e t} \boldsymbol {w} _ {i} = \left\{ \begin{array}{l l} \frac {d (v , l _ {2})}{d (l _ {1} , l _ {2})} \text {i f} \boldsymbol {\Lambda} _ {i} = l _ {1}, \\ \frac {d (v , l _ {1})}{d (l _ {1} , l _ {2})} \text {i f} \boldsymbol {\Lambda} _ {i} = l _ {2}, \\ 0 \text {o t h e r w i s e} \end{array} \right. \quad \text {y i e l d s}
$$
$$
\begin{array}{l} m_{\boldsymbol{\Lambda}}(\boldsymbol {w}) = \operatorname *{arg  min}_{y\in \mathcal{Y}}\frac{d(v,l_{2})d^{2}(y,l_{1}) + d(v,l_{1})d^{2}(y,l_{2})}{d(l_{1},l_{2})} \\ = \operatorname * {a r g   m i n} _ {y \in \mathcal {Y}} d (v, l _ {2}) d ^ {2} (y, l _ {1}) + d (v, l _ {1}) d ^ {2} (y, l _ {2}) \\ = \{v \}. \\ \end{array}
$$