SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\frac {\operatorname * {P r} (t = k | (\theta_ {0} , \theta_ {1} , \ldots , \theta_ {M} , y))}{\operatorname * {P r} (t = m | (\theta_ {0} , \theta_ {1} , \ldots , \theta_ {M} , y))} = \frac {\exp (g (\theta_ {k} , y))}{\exp (g (\theta_ {m} , y))}.
$$
Marginalizing out all other components, we obtain
$$
\frac {\Pr (t = k \mid (\theta_ {k} , \theta_ {m} , y))}{\Pr (t = m \mid (\theta_ {k} , \theta_ {m} , y))} = \frac {\exp (g (\theta_ {k} , y))}{\exp (g (\theta_ {m} , y))}. \tag {18}
$$
Matching the restriction (18) with the true population (17), the restricted projection is attainable if and only if
$$
\exp (g (\theta_ {t}, y) = p (\theta_ {t}, y) / q (\theta_ {t} | y),
$$
so that the optimal classifier needs to be
$$
\Pr (t | \theta_ {1}, \dots , \theta_ {K}, y) = \frac {p (\theta_ {t} , y) / q (\theta_ {t} | y)}{\sum_ {k = 1} ^ {K} p (\theta_ {k} , y) / q (\theta_ {k} | y)}.
$$
It happens that this MCMC restricted optimal classifier matches the IID optimal classifier (10). It follows from the proof in Theorem 9 that the classification divergence using the restricted classifier is still  $D_4(p,q)$ , as if  $\{\tilde{\theta}_m\}_{m=1}^M$  are IID samples from  $q(\theta|y)$ .
C.4 Large sample limit and rate as  $M\to \infty$  (Theorem 3)
In the multivariate Example 4, from each draw  $\theta, y, \tilde{\theta}_{1:M}$ , we generate  $M + 1$  examples from  $K \coloneqq M + 1$  classes; one label from each. We will use the index starting from 1 in this subsection. For the  $k$ -th example,  $t_k = k$ , and the permutation for classifier reads  $\phi_k$  is a vector including  $y$  and  $K$  copies of  $\theta$ , where the  $k$ -th copy is the prior draw, and the remaining  $\theta$  are from  $q(|y)$ . Slightly abused the notation, in this subsection, we call this long feature vector as  $\theta_{1:K}$ ; it is a vector in the  $\Theta^{M + 1}$  space. We now prove Theorem 3.
Proof. First, we write the true conditional label probability in this process:
$$
\pi (t | \theta_ {1: K}, y) = \frac {p \left(\theta_ {t} | y\right) \prod_ {j \neq t} p \left(\theta_ {j} | y\right)}{\sum_ {t ^ {\prime}} p \left(\theta_ {t ^ {\prime}} | y\right) \prod_ {j \neq t ^ {\prime}} p \left(\theta_ {j} | y\right)} = \frac {p \left(\theta_ {t} | y\right) / q \left(\theta_ {t} | y\right)}{\sum_ {j} p \left(\theta_ {j} | y\right) / q \left(\theta_ {j} | y\right)}.
$$
Plug it as the classifier, we obtain the optimal ELPD or negative cross entropy:  $\mathrm{ELPD} = \mathbb{E}\log \pi (t|\theta_{1:K},y)$ ,
$$
\begin{array}{l} \mathrm {E L P D} = \mathbb {E} \frac {p \left(\theta_ {t} | y\right) / q \left(\theta_ {t} | y\right)}{\sum_ {k = 1} ^ {K} p \left(\theta_ {k} | y\right) / q \left(\theta_ {k} | y\right)} \\ = \mathbb {E} \log \frac {p \left(\theta_ {t} \mid y\right)}{q \left(\theta_ {t} \mid y\right)} - \mathbb {E} \log \sum_ {k} p \left(\theta_ {k} \mid y\right) / q \left(\theta_ {k}\right). \\ \end{array}
$$
The first term above is simply
$$
\mathbb {E} \log \frac {p (\theta_ {t} | y)}{q (\theta_ {t} | y)} = \mathrm {K L} (p (\theta | y) | | q (\theta | y))
$$
According to our definition (4), the divergence is ELPD offset by an entropy term,
$$
D _ {4} := \operatorname {E L P D} + \log K.
$$
We now derive the limit of  $D_4 - \mathrm{KL}(p(\theta | y) || q(\theta | y))$  when  $M \to \infty$  (or equivalently  $K = M + 1 \to \infty$ )
$$
\begin{array}{l} \Delta := D _ {4} - \operatorname {K L} (p (\theta | y) | | q (\theta | y)) \\ = \operatorname {E L P D} + \log K - \operatorname {K L} (p (\theta | y) | | q (\theta | y)) \\ = \log K - \mathbb {E} \log \left(\sum_ {k} \frac {p \left(\theta_ {k} | y\right)}{q \left(\theta_ {k} | y\right)}\right) \\ = - \mathbb {E} \log \left(\frac {1}{K} \sum_ {k = 1} ^ {K} \frac {p \left(\theta_ {k} | y\right)}{q \left(\theta_ {k} | y\right)}\right) \\ \end{array}
$$
Given any label value  $1 \leq \leq K$ ,  $\theta_t \sim p(\cdot | y)$ , and all the remaining  $\theta_j \sim q(\cdot | y)$  for  $j \neq t$ .
$$
X _ {k} = \frac {1}{K} \sum_ {k = 1} ^ {K} \frac {p \left(\theta_ {k} \mid y\right)}{q \left(\theta_ {k} \mid y\right)} = \frac {1}{K} \frac {p \left(\theta_ {t} \mid y\right)}{q \left(\theta_ {t} \mid y\right)} + \frac {1}{K} \sum_ {k \neq t} \frac {p \left(\theta_ {t} \mid y\right)}{q \left(\theta_ {t} \mid y\right)}.
$$
The first term
$$
\frac {1}{K} \sum_ {k = 1} ^ {K} \frac {p \left(\theta_ {t} | y\right)}{q \left(\theta_ {t} | y\right)} \rightarrow 0.
$$
The second term, call it  $\Delta_2$  is the mean of IID means as  $\theta_j \sim q(\cdot | y)$  for  $j \neq t$ , the law of large number yields
$$
\Delta_ {2} := \frac {1}{K} \sum_ {k \neq t} ^ {K} \frac {p \left(\theta_ {t} | y\right)}{q \left(\theta_ {t} | y\right)}, \quad \Delta_ {2} \rightarrow \mathbb {E} _ {x \sim q (x | y)} \frac {p (x | y)}{q (x | y)} = 1.
$$
$$
\Delta \rightarrow 0, \text {a s} K \rightarrow \infty .
$$
$$
D _ {4} - \operatorname {K L} (p (\theta | y) | | q (\theta | y)) \rightarrow 0,
$$
which finished the first part of the proof.
Now let's derive the rate. Under regularity conditions, for example, the variance of density ratio is bounded, i.e., there exists a constant  $C < \infty$ , such that for all  $y$
$$
\operatorname {V a r} _ {\theta_ {t} \sim q (\theta_ {t} | y)} \left(\frac {p (\theta_ {t} | y)}{q (\theta_ {t} | y)}\right) <   C,
$$
then CLT holds, such that the second term above has a normal limit,
$$
\sqrt {K} \left(\Delta_ {2} - 1\right)\rightarrow \text {n o r m a l} (0, \sigma^ {2}), \quad \text {i n}
$$
$$
\begin{array}{l} \sigma^ {2} = \operatorname {V a r} _ {q} \left(\frac {p \left(\theta_ {t} | y\right)}{q \left(\theta_ {t} | y\right)}\right) \\ = \mathbb {E} _ {\theta , y \sim q (\theta , y)} \left(\frac {p (\theta | y)}{q (\theta | y)} - 1\right) ^ {2} \\ = \mathbb {E} _ {y} \mathbb {E} _ {\theta \sim q (\theta | y)} \left(\frac {p (\theta | y)}{q (\theta | y)} - 1\right) ^ {2} \\ = \chi^ {2} (p (\theta | y) \mid | q (\theta | y)), \\ \end{array}
$$
which is the definition of the conditional chi-squared divergence.
Consider a Taylor series expansion,  $\log (1 + x) = x - \frac{1}{2} x^2 + o(x^2)$ . Using the Delta method to express the log function and the Slutsky theorem to ignore the zero term, we get
$$
K \mathbb {E} \log (\Delta_ {2}) \rightarrow - \frac {\sigma^ {2}}{2}.
$$
Plug this in the definition of  $\Delta$  we obtain the desired convergence rate,
$$
D _ {4} = \mathrm {K L} (p (\theta | y) | | q (\theta | y)) - \frac {1}{2 M} \chi^ {2} (p (\theta | y) | | q (\theta | y)) + o (1 / M).
$$
This proves the claims in Theorem 3.
Figure 9: Our designed permutation vs the naive permutation.
Original batch from one run
label t	features φ
(θ, y)	1
(θ1, y)	0
(θ2, y)	0
...	...
(θM, y)	0
Table 2: Batch permutation of labels
a batch of permuted labels
label t	another batch of permutation
label t
0	0
1	0
0	0
...	...
0	1
C.5 Valid hypothesis testing (Theorem 4)
It is not trivial to have a valid permutation test since we do not have IID examples. One naive permutation is to permute all the labels  $t$  (across batches  $i$ ). The resulting permutation testing is not uniform under the null (Figure 9).
In contrast, when we permute the label, we only ask for within-batch permeation (permute the index  $\{t_1,\dots ,t_L\}$  in each batch). See Table 2 for an illustration. Let's prove this permutation is valid.
Proof. Under the null,  $p(\theta | y) = q(\theta | y)$  almost everywhere. According to the label mapping  $\Phi \in \mathbb{F}$ , label  $t$  is independent of  $\phi$  given  $y$ . We first show that label  $t$  is independent of  $\phi$ .
In general conditional independence does not lead to unconditional independence. But because here we design the generating process by  $\pi(y, t, \phi) = \pi_Y(y)\pi_t(t)\pi(\phi|y, \theta)$ . Under the null we have  $\pi(y, t, \phi) = \pi_t(t)(\pi_Y(y)\pi(\phi|y))$ . Hence  $t$  needs to be independent of  $\phi$ .
For any classifier  $c$ , because  $c(\phi)$  is a function of  $\phi$ ,  $c(\phi)$  and  $t$  are also independent. Let  $\pi_{\Phi}(\phi)$  be the marginal distribution of  $\phi$  in  $\pi$ .
Now we are computing the cross entropy (with respect to the population  $\pi$ ).  $\mathrm{LPD} = \sum_{j=1}^{N} U(c(\phi_n), t_n)$  where  $U$  is the log score  $U(P, x) = \log P(x)$ . It is a random variable because we have a finite validation-set. Now we are conducting a permutation of label  $t$ , the permuted label  $\tilde{t}$  is an independent draw from the same marginal distribution of  $t$ ,  $\pi_t(t)$ . Because of the independence,  $\mathbb{E}_{\phi, t} U(c(\phi), t) = \mathbb{E}_t \mathbb{E}_{\phi} U(c(\phi), t)$ , hence the permuted  $\mathrm{LPD}_b \stackrel{d}{=} \mathrm{LPD}$ , where  $\mathrm{LPD}_b$  is the computed LPD in the  $b$ -th permutation. Therefore  $\operatorname*{Pr}(\mathrm{LPD} \leq x) = \operatorname*{Pr}(\mathrm{LPD}_b \leq x)$ . In expectation (with an infinitely amount of repeated random permutation), this  $p$ -value will be uniform on  $[0, 1]$ .
In practice, the computed p-value from finite permutation can only take values on a finite set  $\{0,1 / B,\ldots ,(B - 1) / B,1\}$ . More precisely, under the null hypothesis, this permutation test p-value is discreetly-uniform on this set. This is because for any  $0\leq m\leq B$
$$
\operatorname * {P r} (p = m / B) = \int_ {0} ^ {1} \binom {B} {m} p ^ {m} (1 - p) ^ {B - m} d p = 1 / (B + 1), \forall m.
$$
Hence, the permutation  $p$ -value based on  $B$  random permutations is uniformly distributed on the set  $\{0, 1 / B, \ldots, (B - 1) / B, 1\}$ .
C.6 Why does classifier help test? (Theorem 5)
Let's recap the binary label-generating process:
1. Sample  $y \sim$  marginal  $p(y)$ ,
2. Sample  $k \sim$  Bernoulli(w), where  ${w}_{1} = 1/M,{w}_{2} = \left( {M - 1}\right) /M$  .
3. Sample  $\theta \sim \pi_k(\theta |z)$ , where  $\pi_0 = p(\theta |y)$ ,  $\pi_1 = q(\theta |y)$ .
4. Return  $(z,k,x)$
Theorem 5 state that the optimal classifier has a sufficiency propriety in that (a) let  $\hat{c}$  be the probability of label 1 in the optimal classifier as per (7), and let  $\pi_c^p$  and  $\pi_c^q$  be the one-dimensional distribution of this  $\hat{c}(\phi)$  when  $(\theta, y)$  is sampled from  $p(\theta, y)$  or from  $p(y) q(\theta | y)$  respectively, then (i) Conditional on the summary statistic  $\hat{c}$ , the label  $t$  is independent of all features  $\phi = (\theta, y)$ . (ii) There is no loss of information in divergence as the joint divergence is the same as the projected divergence,  $D_1(p, q) = D_1(\pi_c^p, \pi_c^q)$ .
Proof. There are three random variables,  $\theta, y,$  and  $t$ . The optimal classifier  $\hat{c}$  is the probability in this true joint:
$$
\hat {c} (\theta , y) = \Pr (t = 1 | (\theta , y)).
$$
To show sufficiency or conditional independence, all we need to show is that, conditional on any given value of  $\hat{c}$ ,  $\operatorname*{Pr}(t = 1 | (\theta, y), c)$  does not depend on  $(\theta, y)$  (in other words,  $\operatorname*{Pr}(t = 1 | (\theta, y), c)$  is a function of  $c$  only). This becomes obvious as
$$
\Pr (t = 1 | (\theta , y), c (\theta , y)) = \Pr (t = 1 | (\theta , y)) = \hat {c} (\theta , y).
$$
Now we prove that there is "no loss" in divergence.
$$
D _ {1} (p, q) = w \operatorname {K L} (p (\theta , y) | | r (\theta , y)) + (1 - w) \operatorname {K L} (q (\theta , y) | | r (\theta , y))
$$
We express the first term in  $\hat{c}$
$$
\begin{array}{l} \operatorname {K L} (p (\theta , y) \mid | w p (\theta , y) + (1 - w) q (\theta , y)) \\ = \operatorname {K L} (\pi (\theta , y | t = 0) | | \pi (\theta , y)) \\ = \operatorname {K L} (\pi (\hat {c} | t = 0) \mid \mid \pi (\hat {c})) - \operatorname {K L} (\pi (\theta , y | \hat {c}, t = 0) \mid \mid \pi (\theta , y | \hat {c})) \\ \end{array}
$$
This steps uses the chain rule of KL divergence:  $\mathrm{KL}[p(x,y)\mid q(x,y)] = \mathrm{KL}[p(x)\mid q(x)] + \mathrm{KL}[p(y\mid x)\mid q(y\mid x)]$
Using conditional independence:
$$
\pi (\theta , y | \hat {c}, t = 0) = \pi (\theta , y | \hat {c}, t = 1)
$$
Hence  $\mathrm{KL}(\pi (\theta ,y|\hat{c},t = 0)\mid ||\pi (\theta ,y|\hat{c})) = 0$  . Therefore,
$$
\operatorname {K L} (p (\theta , y) | | r (\theta , y)) = \operatorname {K L} (\pi (\hat {c} | t = 0) | | \pi (\hat {c}))
$$
where  $\pi (c) = \pi (\hat{c} |t = 0) + (1 - w)\pi (\hat{c} |t = 0)$
$$
\operatorname {K L} (q (\theta , y) | | r (\theta , y)) = \operatorname {K L} (\pi (\hat {c} | t = 1) | | \pi (\hat {c}))
$$
This proves  $D_{1}(p,q) = D_{1}(\pi_{c}^{p},\pi_{c}^{q})$
C.7 The maximum discriminative generator (Theorem 2)
We save the proof of Theorem 2 in the end for its length.
The generator  $\phi$  contains a few degrees of freedom: the number of classes  $K$ , the number of examples  $L$ , and how to design and the label-feature pairs. In the binary labeling:  $t_1 = 1, \phi_1 = (\theta_1, y)$  and the remaining  $t_k = 0, \phi_k = (\theta_k, y)$  for  $2 \leq k \leq M + 1$ . The multi-class  $\Phi^*$  assigns labels  $1: K$  as
$$
\Phi^ {*}: t _ {k} = k, \phi_ {k} = \left(\operatorname {P e r m} _ {1 \rightarrow k} \left(\theta_ {1}, \dots , \theta_ {M}\right), y\right), 1 \leq k \leq K = M + 1. \tag {19}
$$
Before the main proof that the multiclass permutation creates the largest divergence, let's first convince that the multi-class classification produced higher divergence than the binary one.
In the binary classification with  $M = 1$  (one draw from  $p$  and one draw from  $q$ )
$$
D _ {1} (p, q) = \frac {1}{2} \operatorname {K L} \left(p (\theta | y), \frac {p (\theta | y) + q (\theta | y)}{2}\right) + \frac {1}{2} \operatorname {K L} \left(q (\theta | y), \frac {p (\theta | y) + q (\theta | y)}{2}\right).
$$
In the multi-class classification with  $M = 1$
$$
D _ {4} (p, q) = \mathrm {K L} \left(p (\theta^ {1} | y) q (\theta^ {2} | y), \frac {p (\theta^ {1} | y) q (\theta^ {2} | y) + q (\theta^ {1} | y) p (\theta^ {2} | y)}{2}\right).
$$
Use the joint  $\mathrm{KL} >$  marginal KL, we have
$$
D _ {4} (p, q) \geq \operatorname {K L} \left(p (\theta^ {1} | y), \frac {p (\theta^ {1} | y) + q (\theta^ {1} | y)}{2}\right).
$$
Likewise,
$$
D _ {4} (p, q) \geq \operatorname {K L} \left(q \left(\theta^ {2} | y\right), \frac {p \left(\theta^ {2} | y\right) + q \left(\theta^ {2} | y\right)}{2}\right).
$$
Add these two lines we obtain
$$
D _ {4} (p, q) \geq D _ {1} (p, q).
$$
To prove that this multi-class permutation produced the uniformly largest divergence (across  $M, K$ ,  $p, q$ ) optimal, we organize the proof into lemmas 5 to 9. For notation brevity, we denote  $\hat{M} \coloneqq M + 1$  in these lemmas to avoid using index  $M + 1$ .
Lemma 10. For an arbitrary integer  $L$ , any given output space  $\mathcal{V}$ , and any input space  $\mathcal{X}$  that has at least  $L$  elements, if there are two functions mapping  $\mathcal{X}^L$  to  $\mathcal{V}$
$$
f _ {1}, f _ {2}: (x _ {1}, \dots ,, x _ {L}) \mapsto y \in \mathcal {Y}.
$$
satisfying the following propriety:
- for any probability distribution  $\pi$  on  $\mathcal{X}$ , when  $x_{1},\ldots ,x_{n}$  are  $L$  iid random variables with law  $\pi$ ,  $f_{1}(x_{1},\ldots ,x_{L})$  has the same distribution as  $f_{2}(x_{1},\ldots ,x_{L})$ ,
then there must exist a permutation of  $1:L$ , denoted by  $\sigma(1:L)$ , such that
$$
f _ {2} \left(x _ {1}, \dots , x _ {L}\right) = f _ {1} \left(x _ {\sigma (1)}, \dots , x _ {\sigma (L)}\right).
$$
Proof. For any  $L$  distinct values  $a_1, \ldots, a_L$ ,  $a_i \in \mathcal{X}$ , let  $\pi$  be a mixture distribution of  $L$  delta functions:
$$
\pi = \sum_ {m = 1} ^ {L} \delta (a _ {m}) p _ {m},
$$
where the  $m$ -th mixture probability is
$$
p _ {m} = C \left(\frac {1}{2}\right) ^ {L ^ {m - 1}}, C ^ {- 1} = \sum_ {m = 1} ^ {L} \left(\frac {1}{2}\right) ^ {L ^ {m - 1}}.
$$
$C$  is chosen such that  $\sum_{m=1}^{L} p_m = 1$ .
Now that  $x_{1}, \ldots, x_{L}$  are  $L$  IID random variables from this  $\pi$ ,  $f_{1}(x_{1}, \ldots, x_{L})$  is also a mixture of delta functions, For any sequence of input indices  $(u_{1}, \ldots, u_{L}) \in \{1 : L\}^{L}$ ,
$$
\begin{array}{l} \Pr (f _ {1} (x _ {1}, \dots , x _ {L}) = f _ {1} \left(a _ {u _ {1}}, \dots , a _ {u _ {L}}\right)) = \prod_ {m = 1} ^ {L} \left(\left(C / 2 ^ {L ^ {m - 1}}\right) ^ {\sum_ {j = 1} ^ {L} 1 \left(u _ {j} = m\right)}\right) \\ = C ^ {L} \left(\frac {1}{2}\right) ^ {\sum_ {m = 1} ^ {L} \left(\sum_ {j = 1} ^ {L} 1 \left(u _ {j} = m\right) L ^ {m - 1}\right)}, \tag {20} \\ \end{array}
$$
in which the power index can be written as
$$
\left(\sum_ {j = 1} ^ {L} 1 (u _ {j} = 1), \dots , \sum_ {j = 1} ^ {L} 1 (u _ {j} = L)\right) _ {L} := \sum_ {m = 1} ^ {L} \left(\sum_ {j = 1} ^ {L} 1 (u _ {j} = m) L ^ {m - 1}\right)
$$
as an  $L$ -decimal-integer.
Next, we study the law of  $f_{2}(x_{1},\ldots ,x_{L})$ :
$$
\Pr (f _ {2} (x _ {1}, \dots , x _ {L}) = f _ {2} (a _ {1}, \dots , a _ {L})) = C ^ {L} (\frac {1}{2}) ^ {(1, 1, \dots , 1) _ {L}}.
$$
Because  $f_{2}(x_{1},\ldots ,x_{L})$  and  $f_{1}(x_{1},\ldots ,x_{L})$  have the same distribution,  $f_{2}(a_{1},\ldots ,a_{L})$  needs to match the value at which  $f_{1}(x_{1},\ldots ,x_{L})$  has probability  $C^L (\frac{1}{2})^{(1,1,\dots,1)_L}$  to attain. Comparing with (20), this probability is only attained when  $\sum_{j = 1}^{L}1(u_{j} = m) = 1,\forall m$ . That is, there exists a  $\sigma$ , a permutation of  $1:L$ , such that  $u_{1},\ldots ,u_{L} = \sigma (1,2,\ldots ,L)$ .
Matching the value of  $f_{2}(x_{1},\ldots ,x_{L})$  we obtain
$$
f _ {2} \left(a _ {1}, \dots , a _ {L}\right) = f _ {1} \left(a _ {\sigma (1)}, \dots , a _ {\sigma (L)}\right).
$$
Because the choice of vector  $(a_{1},\ldots ,a_{L})$  is arbitrary, we have
Because augmentation increases divergence, for the purpose of finding the largest divergence, to produce the largest divergence, we only need to consider an augmented generator that includes all  $y$
$$
\Phi^ {a u g}: (y, \theta_ {1, \dots , L}) \mapsto \left\{\left(\left(\phi_ {1}, y\right), t _ {1}\right), \dots , \left(\left(\phi_ {K}, y\right), t _ {K}\right) \right\}.
$$
It is enough to consider the generator of which  $\phi_{k} = \phi_{k}(\theta_{1},\dots ,\theta_{L})$  are  $K$  functions of  $(\theta_{1},\ldots ,\theta_{L})$
Lemma 11. For any augmented generator  $\Phi^{arg}$  satisfies  $\mathbb{F}$ , the null, there must exist  $K$  permutations  $\sigma_{1}(1:L),\ldots ,\sigma_{K}(1:(L))$ , with the convention  $\sigma_1(1:L) = 1:L$ , such that
$$
\phi_ {k} \left(\theta_ {1}, \dots , \theta_ {L}\right) = \phi_ {1} \left(\theta_ {\sigma_ {k} (1: (L))}\right).
$$
Proof. Use Lemma (10) for  $(K - 1)$  times.
Lemma 12. For any augmented generator  $\Phi^{arg}$  satisfies  $\mathbb{F}$ , all feature-label generator can be replaced by a permutation,  $\phi_k(\theta_1,\ldots \theta_L) = (\theta_{\sigma_k(1)},\ldots \theta_{\sigma_k(L)})$ , while the divergence does not decrease.
Proof. From the previous lemma,
$$
\phi_ {k} \left(\theta_ {1}, \dots , \theta_ {L}\right) = \phi_ {1} \left(\theta_ {\sigma_ {k} (1: L)}\right).
$$
The augmented feature is now  $(\phi_{1}(\theta_{\sigma_{k}(1:L)}),y)$ , a transformation of  $(\theta_{\sigma_k(1:L)}),y)$ . Using the raw feature  $(\theta_{\sigma_k(1:L)}),y)$  keeps the divergence non-decreasing.
Now that we only want to consider permutation-based generators: there exists a  $K$ , and  $K$  permutations  $\sigma_{k}(1:L)$ .
$$
\Phi^ {a u g}: (y, \theta_ {1, \dots , L}) \mapsto \left\{\left(\left(\phi_ {1}, y\right), t _ {1}\right), \dots , \left(\left(\phi_ {k}, y\right), t _ {k}\right) \right\}.
$$
$$
\phi_ {k} \left(\theta_ {1}, \dots , \theta_ {L}\right) = \theta_ {\sigma_ {k} (1: L)}.
$$
Given a  $p \neq q$ , any permutations  $\theta_{\sigma_i(1:L)}$  contains one copy from  $p$  and  $(L - 1)$  copies from  $q$ . It suffices to only consider those permutations  $\theta_{\sigma_i(1:L)}$  whose distributions are unique.
Lemma 13. Given any  $p \neq q$  and  $L$  fixed, assuming all  $\theta_{\sigma_i(1:L)}$  has different distributions, then  $D(q,p,\Phi^{arg})$  is an increasing function on  $K$ .
Proof. Using that joint KL is always not smaller than the KL divergence of sub-coordinates.
Lemma 14. Among all permutations  $\sigma(\theta_1, \ldots, \theta_L)$ , the maximum number of distinct distributions are  $K = L$ .
Proof. The total number of permutations is  $L!$ . Because  $\theta_{2:L}$  are IID given  $y$ , the permutation of index  $2: L$  does not change the distribution. When  $p \neq q$ , the total number of distributionally distinct permutations is  $L! / (L - 1)! = L = M + 1$ .
It is clear that the proposed multi-class permutation  $\Phi^{*}$  attains this maximum number of distinct distributions, which proves its optimality among all generators from all previous lemmas, thereby the proof of our Theorem 2.