SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Bastian Boll
Image & Pattern Analysis Group
Heidelberg University
bastian.boll@iwr.uni-heidelberg.de
Christoph Schnorr
schnoerr@math.uni-heidelberg.de
In structured prediction, target objects have rich internal structure which does not factorize into independent components and violates common i.i.d. assumptions. This challenge becomes apparent through the exponentially large output space in applications such as image segmentation or scene graph generation. We present a novel PAC-Bayesian risk bound for structured prediction wherein the rate of generalization scales not only with the number of structured examples but also with their size. The underlying assumption, conforming to ongoing research on generative models, is that data are generated by the Knothe-Rosenblatt rearrangement of a factorizing reference measure. This allows to explicitly distill the structure between random output variables into a Wasserstein dependency matrix. Our work makes a preliminary step towards leveraging powerful generative models to establish generalization bounds for discriminative downstream tasks in the challenging setting of structured prediction.
Structured prediction is the task of predicting an output which itself contains internal structure. As an example, consider the problem of image segmentation. The output to be predicted is an assignment of semantic classes to each image pixel. However, the segmentation problem is not merely a pixelwise classification, because each pixel is not independently assigned a semantic class. If two pixels are adjacent in the image plane, it is more likely that they belong to the same class than to different ones. A way to remedy this problem is to enumerate all possible segmentations of an image and treat the problem as classification. However, the output space of this classification is now exponentially large. This exponential (in the number of pixels) size of the output space indicates a much more difficult problem than the unstructured case of classification. In fact, it also presents an immediate challenge for any statistical learning theory which requires assumptions on the number of output classes [44].
From a practical perspective, structured prediction problems present a challenge of label acquisition. For instance, dense manual segmentation of an image is significantly more labour intensive than manual classification. In turn, one would expect that pixelwise segmentation of an image contains much richer information to be exploited in supervised learning. However, this is not represented in typical statistical learning theory which assumes all data to be independently drawn from the true underlying distribution. In the extreme case of only a single structured example being available for training, these statistical learning theories can not make any meaningful statement on generalization.
Addressing this point in particular, [38] presents an analysis of the dependency structure in the output and proves a risk certificate which decays in both the number of structured examples  $m$  and their size  $d$ . Referring back to the example of image segmentation, this amounts to a high-probability bound on the fraction of misslabeled pixels which decays with the number of labeled pixels observed during training as opposed to merely the number of segmented images.
At the core of statistical learning theory lies the concentration of measure phenomenon which posits that a stable function of a large number of weakly dependent random variables will take values close to its mean [37, 9]. This is relevant because model risk, the expected loss on unseen data, is the mean of empirical risk under the draw of the sample. Learning theories can thus be built on concentration of measure results. In the present work, we focus on PAC-Bayesian learning theory [51, 43, 42]. For an overview of PAC-Bayesian theory we refer to [13, 27, 1].
In addition to a concentration result such as a moment-generating function (MGF) bound, PAC-Bayesian arguments employ a change of measure, typically via Donsker and Varadhan's variational formula. In particular, stochastic predictors, i.e. distributions over a hypothesis class of models are considered. A core objective is to construct generalization bounds which hold uniformly over all stochastic predictors called PAC-Bayes posteriors. Model complexity is then measured as relative entropy to a reference stochastic predictor called PAC-Bayes prior. As this terminology alludes to, the PAC-Bayes posterior is informed by more data than the PAC-Bayes prior. Note however, that PAC-Bayesian theory generalizes Bayesian theory because prior and posterior are not typically connected via likelihood.
PAC-Bayesian theory has gained considerable attention in recent years due to the demonstration of non-vacuous risk certificates in deep learning [23]. Since then, a line of research has succeeded in tightening the risk bounds of deep classifiers [48, 47, 16]. In addition, multiple authors have studied ways to weaken underlying assumptions of bounded loss [29, 28] and i.i.d. data [3].
The majority of recent works in this field focuses on classification or regression. Structured prediction has received comparatively little recent attention, an exception being the work [11] which provides a PAC-Bayesian perspective on the implicit loss embedding framework for structured prediction [15].
Here, we continue a line of research started by [39, 40, 38] which aims to construct PAC-Bayesian risk bounds for structured prediction that account for generalization from a single (but large) structured datum. Instrumental to their analysis is the stability of inference and quantified dependence in the data distribution. The latter is expressed in terms of  $\vartheta$ -mixing coefficients, the total variation distance between data distributions conditioned on fixed values of a subset of variables. For structured prediction with Hamming loss, a coupling of such conditional measures can be constructed [24] such that  $\vartheta$ -mixing coefficients yield an upper-bound that allows to invoke concentration of measure arguments [36]. The result is an MGF bound which the authors employ in a subsequent PAC-Bayesian construction - achieving generalization from a single datum.
The underlying assumption of these previous works is that data are generated by a Markov random field (MRF). This model assumption is somewhat limiting because Markov properties, certain conditional independences, likely do not hold for many real-world data distributions. In addition, MRFs are difficult to work with computationally. Exact inference in MRFs is NP-hard [57] and thus learning, which often contains inference as a subproblem, presents significant computational roadblocks. Even once an MRF has been found which represents data reasonably well, numerical evaluation of the PAC-Bayesian risk certificate proposed in [38] will again be computationally difficult.
Nevertheless, we share the sentiment of previous authors that some assumption on the data-generating process is required in structured prediction. This is because conditional data distributions, the distribution of data conditioned on a fixed set of values for a subset of variables, are central to establishing concentration of measure via the martingale method [35]. Consider again the example of image segmentation. Once we have fixed a sufficiently large number of pixels to arbitrary values (and class labels), even a large dataset will not contain an abundance of data which match these values and thus provide statistical power to learn the conditional distribution. This problem is well-known in conditional density estimation [56].
1.3 Contribution, Organization
In this work, we propose to instead assume a triangular and monotone transport, a Knothe-Rosenblatt (KR) rearrangement [32, 49, 12, 8, 41] of a reference measure as data model. This choice is attractive for multiple reasons. First, any data distribution which does not contain atoms can be represented
uniquely in this way [6] which should suffice to represent many distributions of practical interest. With regard to conditional distributions, the KR-rearrangement has the convenient property that conditioning on a fixed value for a subset of variables can again be represented by KR-rearrangement. We will use this property in our construction of coupling measures between conditional distributions.
- We present a novel PAC-Bayesian risk bound for structured prediction wherein the rate of generalization scales not only with the number of structured examples but also with their size.
- Based on data generated by KR-rearrangement of a tractable reference measure, we distill relevant structure of the data distribution into a Wasserstein dependency matrix. Our analysis hinges on state-of-the-art results in concentration theory [35] which serve to bound moment-generating functions by properties of the Wasserstein dependency matrix. We subsequently invoke a PAC-Bayesian argument to derive the desired risk certificate.
- We also propose to leverage a construction of bad input data as a computational tool to find entries of the Wasserstein dependency matrix.
We stress the fact that many established approaches to generative modelling can be seen as instances of measure transport. For instance, it includes normalizing flows [54, 53, 33, 46, 50], diffusion models [52, 30], generative adversarial networks and variational autoencoders [10, 26]. While most measure transport models which currently enjoy empirical success are not KR-rearrangements, we hope that the methods presented here can lay the foundation of leveraging powerful generative models to build risk certificates for discriminative downstream tasks.
Organization The central concepts of concentration theory and measure transport are described in the preliminary Sections 2 and 3. The main results are presented in Section 4. In Section 5 we present a first discussion of computational aspects related to the presented framework. The paper closes on a discussion of limitations in Section 6 and a conclusion in Section 7.
Basic Notation For any  $d \in \mathbb{N}$ , denote  $[d] = \{1, \ldots, d\}$ . If  $z \in \mathcal{Z}^d$  is a vector, we refer to the subvector of entries with index in a set  $I \subseteq [d]$  as  $z^I$ . In particular, index sets of interest will be half-open and closed intervals  $(i, d] \subseteq [d]$  and  $[i, d] \subseteq [d]$ . Analogously, we will index the output of vector-valued functions  $f^I$  and marginal measures  $\mu^I$ . For a set  $\mathcal{B} \subseteq \mathcal{Z}^d$ , we denote its complement in  $\mathcal{Z}^d$  by  $\mathcal{B}^c = \mathcal{Z}^d \setminus \mathcal{B}$  and for a measure  $\mu$  on  $\mathcal{Z}^d$ , we denote the conditional measure given  $\mathcal{B}^c$  as  $\mu | \mathcal{B}^c$ . If  $Z$  is a random variable with distribution  $\mu$  on  $\mathcal{Z}^d$  and  $I, J \subseteq [d]$  are disjoint index sets with  $I \cup J = [d]$ , we denote the conditional law of  $Z^I$  given  $Z^J = z^J$  as  $\mu (dz^I | z^J)$ .
2 Concentration of Measure and Generalization
Let  $\mathcal{X}$  denote an input space and  $\mathcal{Y}$  denote an output space. Let  $\mu$  be a distribution on  $\mathcal{Z}^d = (\mathcal{X} \times \mathcal{Y})^d$ . There are two restrictions inherent to this setup. First, an input is always paired with an output and thus the number of inputs needs to match the number of outputs. Second, all structured data will be drawn from  $\mu$  and thus the size of each structured datum will be the same. Otherwise,  $\mathcal{X}$  and  $\mathcal{Y}$  can in principle be arbitrary sets which admit metrics. For concreteness, think of  $\mathcal{X} = [0,1] \subseteq \mathbb{R}$  as being a set of gray values and  $\mathcal{Y} = \mathbb{R}$  containing signed distances from a semantic boundary [45] in an image with  $d$  pixels. In this case,  $\mathcal{Z}^d$  contains all binary segmentations of grayvalue images.
The goal of learning is to find parameters  $\theta$  which define a predictor  $\phi_{\theta} \colon \mathcal{X}^d \to \mathcal{Y}^d$  such that the risk
$$
\mathcal {R} (\theta) = \mathbb {E} _ {(X, Y) \sim \mu} [ L (\phi_ {\theta} (X), Y) ] \tag {1}
$$
with respect to some loss function  $L\colon \mathcal{Y}^d\times \mathcal{Y}^d\to \mathbb{R}$  is minimized. However, the true risk (1) is typically intractable because  $\mu$  is unknown. A related tractable quantity is the empirical risk
$$
\mathcal {R} _ {m} (\theta , \mathcal {D} _ {m}) = \frac {1}{m} \sum_ {k \in [ m ]} L \left(\phi_ {\theta} \left(X ^ {(k)}\right), Y ^ {(k)}\right) \tag {2}
$$
of a sample  $\mathcal{D}_m = (X^{(k)},Y^{(k)})_{k\in [m]}$  drawn from  $\mu^m$ . We further assume that the loss of structured outputs is the mean of bounded pointwise loss  $\ell \colon \mathcal{V}\times \mathcal{V}\to [0,1]$
$$
L \left(\widetilde {y} ^ {(k)}, y ^ {(k)}\right) = \frac {1}{d} \sum_ {i \in [ d ]} \ell \left(\widetilde {y} _ {i} ^ {(k)}, y _ {i} ^ {(k)}\right). \tag {3}
$$
In PAC-Bayesian constructions, we consider stochastic predictors  $\zeta$ , i.e. measures on a hypothesis space  $\mathcal{H}$  of predictors  $\phi_{\theta} \colon \mathcal{X}^d \to \mathcal{Y}^d$  as identified with measures on the underlying parameter space from which  $\theta$  is selected. We then take the above notions of risk and empirical risk in expectation over parameter draws
$$
\mathcal {R} (\zeta) = \mathbb {E} _ {\theta \sim \zeta} [ \mathcal {R} (\theta) ], \quad \mathcal {R} _ {m} (\zeta , \mathcal {D} _ {m}) = \mathbb {E} _ {\theta \sim \zeta} \left[ \mathcal {R} _ {m} (\theta , \mathcal {D} _ {m}) \right]. \tag {4}
$$
The expected value of empirical risk with respect to the sample is the true risk. For this reason, a central tool for the study of generalization is the concentration of measure phenomenon. Informally, it states that a stable function of many weakly dependent random variables concentrates on its mean. We will invoke a line of reasoning put forward in [35] and propose a novel approach to structured prediction based on the measure-transport framework outlined in Section 1. To this end, we first define the following formal notions of stability and dependence.
Let  $\rho$  be a metric such that  $\mathcal{Z}$  has finite diameter
$$
\| \rho \| = \sup  _ {\xi , \xi^ {\prime} \in \mathcal {Z}} \rho (\xi , \xi^ {\prime}) <   \infty \tag {5}
$$
and let  $\rho^d (z,z') = \sum_{i\in [d]}\rho (z_i,z_i')$  denote the corresponding product metric on  $\mathcal{Z}^d$
Definition 1 (Local oscillation). Let  $f\colon \mathcal{Z}^d\to \mathbb{R}$  be Lipschitz with respect to  $\rho^d$ . Then the quantities
$$
\delta_ {i} (f) = \sup  _ {z, z ^ {\prime} \in \mathcal {Z} ^ {d}, z _ {[ d ] \backslash \{i \}} ^ {\prime} = z _ {[ d ] \backslash \{i \}}} \frac {| f (z) - f \left(z ^ {\prime}\right) |}{\rho \left(z _ {i} , z _ {i} ^ {\prime}\right)}, \quad i \in [ d ] \tag {6}
$$
are called the local oscillations of  $f$ .
The vector of local oscillations gives a granular account of stability. In order to discuss interdependence of data in a probability space  $(\mathcal{Z}^d,\mu ,\Sigma)$ , define the Markov kernels
$$
K ^ {(i)} (z, d w) = \delta_ {z ^ {[ i - 1 ]}} \left(d w ^ {[ i - 1 ]}\right) \otimes \mu^ {[ i, d ]} \left(d w ^ {[ i, d ]} | z ^ {[ i - 1 ]}\right), \quad i \in [ d ] \tag {7}
$$
as well as  $K^{(d + 1)}(z, dw) = \delta_z(dw)$  and their action on functions
$$
K ^ {(i)} f (z) = \int f (y) K ^ {(i)} (z, d w) = \int f \left(z ^ {[ i - 1 ]} w ^ {[ i, d ]}\right) \mu^ {[ i, d ]} \left(d w ^ {[ i, d ]} \mid z ^ {[ i - 1 ]}\right) \tag {8}
$$
where in the edge case  $i = 1$ , the condition on  $z^{[i-1]} = z^{\{\}}$  is removed. Here  $K^{(i)}(z, dw)$  is a Borel measure for every  $z$  and (8) computes the expected value of  $f$  at  $z$ , conditioned on the fixed realization of the subvector  $z^{[i-1]}$ . It turns out that the effect of the kernel (7) on local oscillations serves to quantify dependence of data with joint distribution  $\mu$ .
Definition 2 (Wasserstein matrix). For  $i \in [d + 1]$ , let  $K^{(i)}$  denote the Markov kernel (8). A matrix  $V^{(i)} \in \mathbb{R}_{\geq 0}^{d \times d}$  is called a Wasserstein matrix [25] for  $K^{(i)}$ , if
$$
\delta_ {k} \left(K ^ {(i)} f\right) \leq \sum_ {j \in [ d ]} V _ {k j} ^ {(i)} \delta_ {j} (f), \quad \forall k \in [ d ] \tag {9}
$$
for any function  $f\colon \mathcal{Z}^d\to \mathbb{R}$  which is Lipschitz with respect to  $\rho^d$
The two concepts defined above will be used in Section 4 to construct a moment-generating function bound via the martingale method.
3 Triangular Measure Transport
Suppose a structured output is composed of  $d > 0$  unstructured data in a space  $\mathcal{Z}$ . Then the target measure  $\mu$  of interest is a measure on  $\mathcal{Z}^d$  which does not factorize into simpler distributions. A popular method of representing complex joint distributions of interdependent random variables is to define a map  $T\colon \mathcal{Z}^d\to \mathcal{Z}^d$  which transports a tractable factorizing reference measure  $\nu^{d}$  to the target measure  $\mu$ , i.e.  $T_{\sharp}\nu^{d} = \mu$ . This abstract framework encompasses many generative models such as normalizing flows [54, 53, 33, 46, 50], diffusion models [52, 30], generative adversarial networks and variational autoencoders [10, 26]. Here, we focus on transport maps  $T$  which are monotone and triangular in the sense that  $T(z)_i$  only depends on the inputs  $z^{[i]}$  and each  $T(z^{[i-1]},\cdot)_i$  is an increasing function. Such a map is called a Knothe-Rosenblatt (KR) rearrangement [32, 49, 12, 8, 41]. If both  $\nu^{d}$  and  $\mu$  have no atoms then the KR rearrangement exists and is unique [6]. In particular, normal distribution  $\nu^{d}$  and any absolutely continuous (with respect to the Lebesgue measure) distribution  $\mu$  meet these criteria. The KR rearrangement has the useful property that certain conditional distributions have a simple representation.
Lemma 3 (Lemma 1 of [41]). Let  $T \colon \mathcal{Z}^d \to \mathcal{Z}^d$  be the KR-rearrangement which satisfies  $T_{\sharp}\nu^{d} = \mu$ . For arbitrary  $i \in [d]$ , let  $z^{[i]} \in \mathcal{Z}^i$  be fixed. Then
$$
\mu \left(d w ^ {(i, d)} \mid z ^ {[ i ]}\right) = T ^ {(i, d)} \left(\bar {z} ^ {[ i ]}, \cdot\right) _ {\sharp} \nu^ {d - i} \tag {10}
$$
where  $\overline{z}^{[i]}$  is the unique element of  $\mathcal{Z}^i$  such that  $T^{[i]}(\overline{z}^{[i]}) = z^{[i]}$ .
Numerical realization of KR rearrangements has recently received attention [5] and more broadly, a variety of triangular transport architectures exists [20, 21]. However, we do not focus on numerical considerations in the present theoretical work.
One may wonder if data generated by KR-rearrangement implicitly restricts the choice of possible input and output spaces since the monotonicity requirement on  $T$  can only be satisfied if the underlying set  $\mathcal{Z}$  is ordered. However, many feature spaces of practical interest are still permissible for what follows. In particular, both  $\mathcal{X}$  and  $\mathcal{Y}$  may be Euclidean spaces or hypercubes. In keeping with the introductory image segmentation example, suppose  $\mathcal{X} = [0,1]^3$  contains RGB color values and  $\mathcal{Y} = [-1,1]$  contains signed distance from a semantic boundary in the image plane. Then we can interpret  $\mathcal{Z}^d = ([0,1]^3 \times [-1,1])^d$  as a product of compact intervals in  $\mathbb{R}^{4d}$  which is clearly permissible as underlying space for KR-rearrangement. All presented results hold irrespective of whether  $\mathcal{Z}$  contains such internal dimensions as long as a natural ordering of the set  $\mathcal{Z}$  is induced.
4 PAC-Bayesian Risk Certificate
In this section we present a novel PAC-Bayesian risk bound for structured prediction which combines three main ingredients.
(1) A concentration of measure theorem for dependent data (Theorem 4) which builds on the notion of a Wasserstein dependency matrix;
(2) a simple construction of coupling measures between conditional distributions (Lemma 5) which serves to represent the Wasserstein dependency matrix;
(3) a PAC-Bayesian argument (Theorem 7) employing Donsker-Varadhan's variational formula in concert with concentration of measure results.
The first theorem summarizes key results from [35] on the concentration of measure phenomenon for dependent random variables. We have slightly generalized by augmenting the underlying Doob martingale construction with the inclusion of a set  $\mathcal{B}$  of bad inputs. For inputs in this set, data stability requirements do not necessarily hold. We call the complement  $\mathcal{B}^c = \mathcal{Z}^d\setminus \mathcal{B}$  the set of good inputs. The concept of good and bad inputs as well as related proof techniques were originally proposed by [38]. Here, we incorporate them into the more general concentration of measure formalism of [35]. A full proof of the following theorem is deferred to Appendix A.
Theorem 4 (Moment-generating function (MGF) bound for good inputs). Let  $\mathcal{B} \subseteq \mathcal{Z}^d$  be a measurable set of bad inputs. Suppose for each  $i \in [d + 1]$ ,  $V^{(i)}$  is a Wasserstein matrix for the
Markov kernel  $K^{(i)}$  defined in (7) on the set of good inputs, that is
$$
\delta_ {k} \left(K ^ {(i)} \widetilde {f}\right) \leq \sum_ {j \in [ d ]} V _ {k j} ^ {(i)} \delta_ {j} (\widetilde {f}), \quad \forall k \in [ d ] \tag {11}
$$
for all Lipschitz (with respect to  $\rho^d$ ) functions  $\widetilde{f} \colon \mathcal{B}^c \to \mathbb{R}$ . Define the Wasserstein dependency matrix
$$
\Gamma \in \mathbb {R} ^ {d \times d}, \quad \Gamma_ {i j} = \| \rho \| V _ {i j} ^ {(i + 1)} \tag {12}
$$
Then for all Lipschitz functions  $f\colon \mathcal{Z}^d\to \mathbb{R}$ , the following MGF bound holds
$$
\mathbb {E} _ {z \sim \mu | \mathcal {B} ^ {c}} \left[ \exp \left(\lambda (f (z) - \mathbb {E} _ {\mu | \mathcal {B} ^ {c}} f)\right) \right] \leq \exp \left(\frac {\lambda^ {2}}{8} \| \Gamma \delta (f) \| _ {2} ^ {2}\right). \tag {13}
$$
An upper bound on the moment generating function will be used in the PAC-Bayesian argument concluding this section. The function  $f$  in question will be the loss of a structured datum  $z$ . Regarding (13), our goal is to bound the norm  $\| \Gamma \delta(f) \|_2^2$  through properties of the data distribution. We will use the fact that data is represented by measure transport to establish such a bound after the following preparatory lemma.
Lemma 5 (Coupling from transport). Let  $\nu^d$  be a reference measure on  $\mathcal{Z}^d$  and  $F, G\colon \mathcal{Z}^d \to \mathcal{Z}^d$  be measurable maps. Define the map  $(F, G)$  by
$$
(F, G) \colon \mathcal {Z} ^ {d} \rightarrow \mathcal {Z} ^ {d} \times \mathcal {Z} ^ {d}, \quad z \mapsto (F (z), G (z)) \tag {14}
$$
Then  $(F,G)_{\sharp}\nu^{d}$  is a coupling of  $F_{\sharp}\nu^{d}$  and  $G_{\sharp}\nu^{d}$ .
Proof. Let  $A \subseteq \mathcal{Z}^d$  be measurable, then
$$
(F, G) _ {\sharp} \nu^ {d} (A, \mathcal {Z} ^ {d}) = \nu^ {d} \left(\left(F, G\right) ^ {- 1} (A, \mathcal {Z} ^ {d})\right) = \nu^ {d} \left(F ^ {- 1} (A)\right) = F _ {\sharp} \nu^ {d} (A) \tag {15}
$$
which shows that  $F_{\sharp}\nu^{d}$  is the first marginal of  $(F,G)_{\sharp}\nu^{d}$ . An analogous argument for the second marginal shows the assertion.
By assuming  $\mu$  to be represented via KR-rearrangement of a factorizing reference measure, Lemma 3 gives an explicit representation of KR-rearrangement for conditional distributions. From there, we invoke Lemma 5 to construct a coupling between conditional distributions and subsequently follow a line of reasoning put forward in [35] to explicitly construct Wasserstein matrices for the kernels (7) which yield a bound on (13) by Theorem 4. This leads to the following proposition. A full proof is deferred to Appendix A.