SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Let  $v, z \in \mathcal{Z}^d$  with  $v^{[d] \setminus \{i\}} = z^{[d] \setminus \{i\}}$  be given for some  $i \in [d]$ . Recall the action (8) of Markov kernels  $K^{(i+1)}$  is an expected value with respect to conditional distributions  $\mu^{(i,d]}(dw^{(i,d]} | v^{[i]}$ .
Because  $\nu^d$  has no atoms,  $\nu^d | \mathcal{A}^c$  also has no atoms. Therefore, there is a unique KR-rearrangement  $\widehat{T}$  with  $\widehat{T}_{\sharp} \nu^d = \nu^d | \mathcal{A}^c$ . Then  $\widetilde{T} = T \circ \widehat{T}$  is a KR-rearrangement with
$$
\widetilde {T} _ {\sharp} \nu^ {d} = \mu | \mathcal {B} ^ {c} \tag {37}
$$
by Lemma 9 and we have  $\widetilde{T} (\widehat{v}) = v$ . Lemma 3 implies
$$
\mu^ {(i, d ]} \left(d w ^ {(i, d ]} \mid \mathcal {B} ^ {c}, v ^ {[ i ]}\right) = \widetilde {T} \left(\widehat {v} ^ {[ i ]}, \cdot\right) _ {\sharp} \nu^ {d - i} \tag {38}
$$
An analogous expression holds for the distribution conditioned on  $z$ . We have therefore found two transport functions pushing the reference measure to the respective conditional distributions. By Lemma 5, a coupling of the conditional distributions is then given by
$$
P _ {v, z} ^ {[ i ]} = \left(\widetilde {T} ^ {(i, d)} \left(\widehat {v} ^ {[ i ]}, \cdot\right), \widetilde {T} ^ {(i, d)} \left(\widehat {z} ^ {[ i ]}, \cdot\right)\right) _ {\sharp} \nu^ {d - i}. \tag {39}
$$
Using a change of measure we find
$$
\begin{array}{l} K ^ {(i + 1)} f (v) - K ^ {(i + 1)} f (z) \\ = \int P _ {v, z} ^ {[ i ]} \left(d u ^ {(i, d ]}, d v ^ {(i, d ]}\right) \left(f \left(v ^ {[ i ]} u ^ {(i, d ]}\right) - f \left(z ^ {[ i ]} v ^ {(i, d ]}\right)\right) (40) \\ = \int \left(f \left(v ^ {[ i ]} \widetilde {T} ^ {(i, d)} \left(\widehat {v} ^ {[ i ]}, \tau\right)\right) - f \left(z ^ {[ i ]} \widetilde {T} ^ {(i, d)} \left(\widehat {z} ^ {[ i ]}, \tau\right)\right)\right) \nu^ {d - i} (\tau) (41) \\ \leq \frac {\delta_ {i} (f)}{d} \rho \left(v _ {i}, z _ {i}\right) + \sum_ {j \in (i, d ]} \frac {\delta_ {j} (f)}{d} \int \rho \left(\widetilde {T} ^ {(i, d ]} \left(\widehat {v} ^ {[ i ]}, \tau\right) _ {j}, \widetilde {T} ^ {(i, d ]} \left(\widehat {z} ^ {[ i ]}, \tau\right) _ {j}\right) \nu^ {d - i} (\tau) (42) \\ \leq \frac {\delta_ {i} (f)}{d} \rho \left(v _ {i}, z _ {i}\right) + \sum_ {j \in (i, d ]} \frac {\delta_ {j} (f)}{d} L _ {i j} \rho \left(v _ {i}, z _ {i}\right) (43) \\ \end{array}
$$
which shows
$$
\delta_ {i} \left(K ^ {(i + 1)} f\right) \leq \frac {1}{d} \left(\delta_ {i} (f) + \sum_ {j \in (i, d ]} L _ {i j} \delta_ {j} (f)\right) \tag {44}
$$
for good inputs. We have thus found a Wasserstein matrix  $V^{(i + 1)}$  for  $K^{(i + 1)}$  with entries
$$
V _ {i j} ^ {(i + 1)} = \left\{ \begin{array}{l l} 0 & \text {i f} i > j \\ d ^ {- 1} & \text {i f} i = j \\ d ^ {- 1} L _ {i j} & \text {i f} i <   j \end{array} \right. \tag {45}
$$
in row  $i$  which shows the assertion.
Proof of Theorem 7 For any hypothesis  $h \in \mathcal{H}$ , we have
$$
\begin{array}{l} \mathcal {R} (h) - \mathcal {R} _ {m} (h, \mathcal {D} _ {m}) = \mathbb {E} _ {Z \sim \mu} [ L (h, Z) - \mathcal {R} _ {m} (h, \mathcal {D} _ {m}) ] (46a) \\ = \mathbb {E} _ {Z \sim \mu} \left[ \left(L (h, Z) - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right) \mathbf {1} \{Z \notin \mathcal {B} \} \right] \\ + \mathbb {E} _ {Z \sim \mu} \left[ \left(L (h, Z) - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right) \mathbf {1} \{Z \in \mathcal {B} \} \right] (46b) \\ \leq \mathbb {E} _ {Z \sim \mu} \left[ \left(L (h, Z) - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right) \mathbf {1} \{Z \notin \mathcal {B} \} \right] + \xi (46c) \\ \leq \mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m}) + \xi (46d) \\ \end{array}
$$
where in (46c) we have used that pointwise loss is in  $[0,1]$ . Note that the underlying distribution of the risk  $\mathcal{R}(h)$  is  $\mu$ , while  $\mathcal{D}_m$  are drawn from  $\mu|\mathcal{B}^c$ . The above inequality reconciles this such that a concentration argument for the conditional distribution becomes applicable. For any PAC-Bayes posterior distribution  $\zeta$  and any  $\beta > 0$ , this implies
$$
\begin{array}{l} \mathcal {R} (\zeta) - \mathcal {R} _ {m} (\zeta , \mathcal {D} _ {m}) = \mathbb {E} _ {h \sim \zeta} \mathbb {E} _ {Z \sim \mu} [ L (h, Z) - \mathcal {R} _ {m} (h, \mathcal {D} _ {m}) ] (47a) \\ \leq \mathbb {E} _ {h \sim \zeta} \left[ \mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m}) \right] + \xi (47b) \\ = \frac {1}{\beta} \mathbb {E} _ {h \sim \zeta} \left[ \beta \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right) \right] + \xi (47c) \\ \leq \frac {1}{\beta} \log \mathbb {E} _ {h \sim \pi} \Big [ \exp \big (\beta (\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})) \big) \Big ] \\ + \frac {1}{\beta} \operatorname {K L} [ \zeta : \pi ] + \xi (47d) \\ \end{array}
$$
by Donsker and Varadhan's variational formula [2, Lemma 2.2]. Focusing on the first term, we find
$$
\begin{array}{l} \left. \exp \left(\beta \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right)\right) = \exp \left(\frac {\beta}{m} \sum_ {k \in [ m ]} \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - L (h, Z ^ {(k)})\right)\right) \right. (48a) \\ = \prod_ {k \in [ m ]} \exp \left(\frac {\beta}{m} \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - L (h, Z ^ {(k)})\right)\right) (48b) \\ \end{array}
$$
Each structured datum  $Z^{(k)}$  is drawn independently from  $\mu |\mathcal{B}^c$ . By Proposition 6 there exists a Wasserstein dependency matrix  $\Gamma = \frac{\|\rho\|}{d} D$  for  $\mu |\mathcal{B}^c$  where  $D$  has entries (17). Then
$$
\begin{array}{l} \mathbb {E} _ {\mathcal {D} _ {m} \sim (\mu | \mathcal {B} ^ {c}) ^ {m}} \prod_ {k \in [ m ]} \exp \left(\frac {\beta}{m} \big (\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - L (h, Z ^ {(k)}) \big)\right) \\ = \prod_ {k \in [ m ]} \mathbb {E} _ {Z ^ {(k)} \sim (\mu | \mathcal {B} ^ {c})} \exp \left(\frac {\beta}{m} \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - L (h, Z ^ {(k)})\right)\right) (49a) \\ = \prod_ {k \in [ m ]} \mathbb {E} _ {Z ^ {(k)} \sim \mu | \mathcal {B} ^ {c}} \left[ \exp \left(\frac {\beta}{m} \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - L (h, Z ^ {(k)})\right)\right) \right] (49b) \\ \leq \prod_ {k \in [ m ]} \exp \left(\frac {\beta^ {2}}{8 m ^ {2}} \| \Gamma \delta (\widetilde {L} (h, \cdot)) \| _ {2} ^ {2}\right) \text {b y} (49c) \\ = \exp \left(\frac {\beta^ {2}}{8 m} \| \Gamma \delta (\widetilde {L} (h, \cdot)) \| _ {2} ^ {2}\right) (49d) \\ \leq \exp \left(\frac {\beta^ {2}}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2}\right) (49e) \\ \end{array}
$$
Denote the shorthand
$$
U = \mathbb {E} _ {\mathcal {D} _ {m} \sim (\mu | \mathcal {B} ^ {c}) ^ {m}} \left[ \exp \left(\beta \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m}))\right)\right) \right] \tag {50}
$$
By Markov's inequality it holds
$$
\mathbb {P} _ {\mathcal {D} _ {m} \sim (\mu | \mathcal {B} ^ {c}) ^ {m}} \left[ \exp \left(\beta \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right)\right) \geq \frac {1}{\delta} U \right] \leq \delta \tag {51}
$$
and combining this with (49) we have
$$
\left. \exp \left(\beta \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ L (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right)\right) \leq \frac {1}{\delta} \exp \left(\frac {\beta^ {2}}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2}\right) \right. \tag {52}
$$
with probability at least  $1 - \delta$  over the sample. Using (47) we thus have
$$
\begin{array}{l} \mathcal {R} (\zeta) - \mathcal {R} _ {m} (\zeta , \mathcal {D} _ {m}) \leq \frac {1}{\beta} \left(\log \mathbb {E} _ {h \sim \pi} \left[ \frac {1}{\delta} \exp \left(\frac {\beta^ {2}}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2}\right) \right] + \mathrm {K L} [ \zeta : \pi ]\right) + \xi (53a) \\ = \frac {\beta}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2} + \frac {1}{\beta} \left(\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]\right) + \xi (53b) \\ \end{array}
$$
Ideally, we would minimize the right hand side with respect to  $\beta$ . However, this would mean to have  $\beta$  depend on  $\zeta$  and we thus would not have a uniform bound for all posterior distributions.
Instead, [38] approaches the problem by defining a sequence of constant  $(\delta_j,\beta_j)_{j\in \mathbb{N}_0}$  and bounding the probability that the bound does not hold for any sequence element. Since in the opposite (high-probability) case, the bound holds for all sequence elements, an optimal one can subsequently be chosen dependent on the posterior.
For all  $j\in \mathbb{N}_0$  , define
$$
\delta_ {j} = \delta 2 ^ {- (j + 1)}, \quad \beta_ {j} = 2 ^ {j} \sqrt {\frac {8 m \log \frac {1}{\delta}}{\| \Gamma \widetilde {\delta} \| _ {2} ^ {2}}} \tag {54}
$$
which are independent of  $\zeta$ . Now consider the event  $E_{j}$  that
$$
\left. \exp \left(\beta_ {j} \left(\mathbb {E} _ {Z \sim \mu | \mathcal {B} ^ {c}} [ \ell (h, Z) ] - \mathcal {R} _ {m} (h, \mathcal {D} _ {m})\right)\right) \geq \frac {1}{\delta_ {j}} \exp \left(\frac {\beta_ {j} ^ {2}}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2}\right) \right. \tag {55}
$$
By the above argument leading up to (52), the probability for  $E_{j}$  under a random sample of the conditioned data distribution  $\mu |\mathcal{B}^{c}$  is at most  $\delta_{j}$ . Therefore, the probability that any  $E_{j}$  occurs is bounded by
$$
\mathbb {P} \left(\bigcup_ {j \in \mathbb {N} _ {0}} E _ {j}\right) \leq \sum_ {j \in \mathbb {N} _ {0}} \mathbb {P} \left(E _ {j}\right) \leq \sum_ {j \in \mathbb {N} _ {0}} \delta_ {j} = \delta \tag {56}
$$
Thus, for all posteriors  $\zeta$  with probability at least  $1 - \delta$  none of the events (55) occurs. We may therefore select an index  $j$  dependent on  $\zeta$  to obtain a sharper risk certificate which still holds with probability at least  $1 - \delta$  over the sample conditioned on the good set. For a fixed posterior  $\zeta$ , the optimizer of (53b) would be
$$
\beta^ {*} = \frac {1}{\| \Gamma \widetilde {\delta} \| _ {2}} \sqrt {8 m \left(\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]\right)} \tag {57}
$$
Equating this to (55) and rounding down to the nearest integer gives
$$
j ^ {*} = \left\lfloor \frac {1}{2} \log_ {2} \left(1 + \frac {\operatorname {K L} [ \zeta : \pi ]}{\log \frac {1}{\delta}}\right) \right\rfloor \tag {58}
$$
Denote this number before rounding by  $r$ , i.e.  $j^{*} = \lfloor r\rfloor$ . For any real number  $r$  it holds  $r - 1\leq \lfloor r\rfloor \leq r$ . Therefore
$$
\frac {1}{2} \sqrt {1 + \frac {\operatorname {K L} [ \zeta : \pi ]}{\log \frac {1}{\delta}}} = 2 ^ {r - 1} \leq 2 ^ {j ^ {*}} \leq 2 ^ {r} = \sqrt {1 + \frac {\operatorname {K L} [ \zeta : \pi ]}{\log \frac {1}{\delta}}} \tag {59}
$$
which gives the following bounds on  $u_{j^{*}}$
$$
\frac {1}{2} \sqrt {\frac {8 m \left(\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]\right)}{\| \Gamma \widetilde {\delta} \| _ {2} ^ {2}}} \leq u _ {j *} \leq \sqrt {\frac {8 m \left(\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]\right)}{\| \Gamma \widetilde {\delta} \| _ {2} ^ {2}}} \tag {60}
$$
Likewise, we bound
$$
\begin{array}{l} \operatorname {K L} [ \zeta : \pi ] + \log \frac {1}{\delta_ {j ^ {*}}} = \operatorname {K L} [ \zeta : \pi ] + \log \frac {2}{\bar {\delta}} + j ^ {*} \log 2 (61a) \\ \leq \mathrm {K L} [ \zeta : \pi ] + \log \frac {2}{\delta} + \frac {\log 2}{2} \log_ {2} \left(1 + \frac {\mathrm {K L} [ \zeta : \pi ]}{\log \frac {1}{\delta}}\right) - \log 2 (61b) \\ = \mathrm {K L} [ \zeta : \pi ] + \log \frac {1}{\delta} + \frac {1}{2} \log \left(1 + \frac {\mathrm {K L} [ \zeta : \pi ]}{\log \frac {1}{\delta}}\right) (61c) \\ = \operatorname {K L} [ \zeta : \pi ] + \log \frac {1}{\delta} + \frac {1}{2} \log \left(\log \frac {1}{\delta} + \operatorname {K L} [ \zeta : \pi ]\right) - \frac {1}{2} \log \log \frac {1}{\delta} (61d) \\ \end{array}
$$
The assumption  $\delta \leq \exp(-e^{-1})$  yields  $-\log \log \frac{1}{\delta} \leq 1$  and because  $x + 1 \leq \exp(x)$  for all  $x \in \mathbb{R}$ , we find
$$
\begin{array}{l} \operatorname {K L} [ \zeta : \pi ] + \log \frac {1}{\delta_ {j ^ {*}}} \leq \operatorname {K L} [ \zeta : \pi ] + \log \frac {1}{\delta} + \frac {1}{2} \left(\log \left(\log \frac {1}{\delta} + \operatorname {K L} [ \zeta : \pi ]\right) + 1\right) (62a) \\ \leq \mathrm {K L} [ \zeta : \pi ] + \log \frac {1}{\delta} + \frac {1}{2} (\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]) (62b) \\ = \frac {3}{2} \left(\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]\right) (62c) \\ \end{array}
$$
We can now use the bounds (62c) and (60) in (53b) to bound the expected generalization error
$$
\begin{array}{l} \mathcal {R} (\zeta) - \mathcal {R} _ {m} (\zeta , \mathcal {D} _ {m}) \leq \frac {u _ {j ^ {*}}}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2} + \frac {1}{u _ {j ^ {*}}} \left(\log \frac {1}{\delta_ {j ^ {*}}} + \mathrm {K L} [ \zeta : \pi ]\right) + \xi (63a) \\ \leq \frac {u _ {j ^ {*}}}{8 m} \| \Gamma \widetilde {\delta} \| _ {2} ^ {2} + \frac {3}{2 u _ {j ^ {*}}} \left(\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]\right) + \xi (63b) \\ \leq \frac {1}{2} \| \Gamma \widetilde {\delta} \| _ {2} \sqrt {\frac {\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]}{2 m}} + \frac {3}{2} \| \Gamma \widetilde {\delta} \| _ {2} \sqrt {\frac {\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]}{2 m}} + \xi (63c) \\ = 2 \| \Gamma \widetilde {\delta} \| _ {2} \sqrt {\frac {\log \frac {1}{\delta} + \mathrm {K L} [ \zeta : \pi ]}{2 m}} + \xi (63d) \\ \end{array}
$$
Note that  $\beta^{*}$  would attain the optimal value
$$
\mathcal {R} (\zeta) - \mathcal {R} _ {m} (\zeta , \mathcal {D} _ {m}) \leq \| \Gamma \widetilde {\delta} \| _ {2} \sqrt {\frac {\operatorname {K L} [ \zeta : \pi ] + \log \frac {1}{\delta}}{2 m}} + \xi \tag {64}
$$
which only differs from the above uniform bound by a factor of two. Finally, recall  $\Gamma = \frac{\|\rho\|}{d} D$  where  $D$  has entries (17).
B Additional Lemmata
Lemma 9. Let  $T \colon \Omega \to \Omega$  be a measurable function on a measurable space  $(\Omega, \Sigma)$  and let  $\nu, \mu$  be measures on  $\Omega$  with  $T_{\sharp} \nu = \mu$ . Let  $B \in \Sigma$  be a fixed set with  $\mu(B) > 0$  and  $A = T^{-1}(B)$  its preimage under  $T$ . Then
$$
T _ {\sharp} (\nu | A) = \mu | B. \tag {65}
$$
Proof. Let  $S \in \Sigma$  be arbitrary and let  $\widetilde{\mu} = T_{\sharp}(\nu|A)$ . Then
$$
\widetilde {\mu} (S) = (\nu | A) \left(T ^ {- 1} (S)\right) = \frac {\nu \left(T ^ {- 1} (S) \cap A\right)}{\nu (A)} \tag {66}
$$
$$
(\mu | B) (S) = \frac {\mu (S \cap B)}{\mu (B)} = \frac {\nu \left(T ^ {- 1} (S \cap B)\right)}{\nu (A)} \tag {67}
$$
$$
z \in T ^ {- 1} (S) \cap T ^ {- 1} (B) \Leftrightarrow T (z) \in S \wedge T (z) \in B \Leftrightarrow T (z) \in S \cap B \Leftrightarrow z \in T ^ {- 1} (S \cap B) \tag {68}
$$
thus  $T^{-1}(S) \cap T^{-1}(B) = T^{-1}(S \cap B)$  and consequently  $\widetilde{\mu}(S) = (\mu|B)(S)$ . Since  $S$  was arbitrary, this shows the assertion.
The following theorem exists in various forms in the literature. To make this paper self-contained, we recite the version in [35] which is used to bound moment-generating functions in Proposition 4. Note that we only use the MGF bound (69) in our analysis. However, the concentration inequality (70) also holds analogously under the assumptions of Proposition 4 which may be of independent interest.
Theorem 10 (Azuma-Hoeffding [35, Theorem 4.1]). Let  $(M^{(i)})_{i\in [m]}$  be a martingale difference sequence with respect to a filtration  $(\Sigma_{i})_{i\in [m]}$  of sigma algebras. Suppose that for each  $i\in [m]$  there exist  $\Sigma_{i - 1}$ -measurable random variables  $A^{(i)}$ ,  $B^{(i)}$  such that  $A^{(i)}\leq M^{(i)}\leq B^{(i)}$  almost surely. Then for all  $\lambda \in \mathbb{R}$  it holds that
$$
\mathbb {E} \left[ \exp \left(\lambda \sum_ {i \in [ m ]} M ^ {(i)}\right) \right] \leq \exp \left(\frac {\lambda^ {2}}{8} \sum_ {i \in [ m ]} \| B ^ {(i)} - A ^ {(i)} \| _ {\infty} ^ {2}\right) \tag {69}
$$
and consequently, for any  $t \geq 0$
$$
\mathbb {P} \left(\left| \sum_ {i \in [ m ]} M ^ {(i)} \right| \geq t\right) \leq 2 \exp \left(- \frac {2 t ^ {2}}{\sum_ {i \in [ m ]} \| B ^ {(i)} - A ^ {(i)} \| _ {\infty} ^ {2}}\right). \tag {70}
$$
C Numerical Toy Example
Figure 1: KR rearrangement  $T(z) = (T_{1}(z_{1},z_{2}),T_{2}(z_{1},z_{2}))^{\top}$  transports a uniform reference measure  $\nu^{2}$  on the unit cube  $[0,1]^2$  to a multimodal distribution  $\mu$  (right).
Figure 2: Construction of a good set in toy example. Left:  $L_{12}$  for all inputs. Center: size of an excluded bad set under the data distribution for corresponding largest values of  $L_{12}$ . Right:  $L_{12}$  for good inputs.
To illustrate the concept of transport stability discussed in Proposition 6, as well as the bad set construction of Section 5, we devised a toy example of two-dimensional transport. As a reference distribution, we choose the uniform distribution on the unit cube  $[0,1]^2$ , i.e.  $\mathcal{Z} = [0,1]$ . The metric  $\rho^d$ ,  $d = 2$  is chosen as the  $\ell^1$ -distance  $\rho^2(z,z') = \|z - z'\|_1 = |z_1 - z'_1| + |z_2 - z'_2|$ . The KR transport map has component functions  $T(z_1,z_2) = (T_1(z_1), T_2(z_1,z_2))$ . We make a polynomial ansatz, defining
$$
T _ {1} \left(z _ {1}\right) = \int_ {\left[ 0, z _ {1} \right]} \sum_ {i \in [ 5 ]} \beta_ {i} B _ {i} ^ {5} (\tau) \mathrm {d} \tau , \quad T _ {2} \left(z _ {1}, z _ {2}\right) = \int_ {\left[ 0, z _ {2} \right]} \sum_ {i \in [ 2 ]} \widetilde {\beta} _ {i} \left(z _ {1}\right) B _ {i} ^ {2} (\tau) \mathrm {d} \tau \tag {71}
$$
where  $B_{i}^{m}$  denotes the  $i$ -th Bernstein polynomial of degree  $m$ ,  $\beta_{i} \geq 0$  are parameters and (positive) coefficient functions  $\widetilde{\beta}_{i}(z_{1})$  are again a parameterized combination of Bernstein polynomials (degree 8). Because Bernstein polynomials assume non-negative values on  $[0,1]$ , the components in (71) define a valid KR-rearrangement. The resulting measure transport is illustrated in Figure 1. In order to evaluate the bound of Theorem 7, we need to compute the Lipschitz constant  $L_{12}$  according to equation (16)
$$
\mathbb {E} _ {\tau \sim \nu_ {2}} \left[ \left| T _ {2} \left(z _ {1}, \tau\right) - T _ {2} \left(z _ {1} ^ {\prime}, \tau\right) \right| \right] \leq L _ {1 2} \left| z _ {1} - z _ {1} ^ {\prime} \right| \tag {72}
$$
for all good inputs  $z_{1}, z_{1}^{\prime}$ . Figure 2 (left) shows the values of  $L_{12}$  satisfying (72) for each input  $(z_{1}, z_{1}^{\prime}) \in [0,1]^{2}$ . The sought Lipschitz constant is the largest of these values. In order to avoid regions with large values, Theorem 7 allows for the exclusion of bad inputs. The probability  $\xi$  of this bad set under the data distribution is incurred as a penalty in equation (20). We thus construct bad sets to exclude large values and consider their probability under the data distribution Figure 2 (middle). Finally, the exclusion of such a bad set is shown on the right of Figure 2.