SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ruicheng Xian $^{1\dagger}$  Honglei Zhuang $^{2}$  Zhen Qin $^{2}$  Hamed Zamani $^{3\dagger}$  Jing Lu $^{2}$  Ji Ma $^{2}$  Kai Hui $^{2}$  Han Zhao $^{1}$  Xuanhui Wang $^{2}$  Michael Bendersky $^{2}$
<sup>1</sup>University of Illinois Urbana-Champaign
{rxian2,hanzhao}@illinois.edu
{hlz,zhenqin,ljwinnie,maji,kaihuibj,xuanhui,bemike}@google.com
<sup>3</sup>University of Massachusetts Amherst
zamani@cs.umass.edu
Domain adaptation aims to transfer the knowledge learned on (data-rich) source domains to (low-resource) target domains, and a popular method is invariant representation learning, which matches and aligns the data distributions on the feature space. Although this method is studied extensively and applied on classification and regression problems, its adoption on ranking problems is sporadic, and the few existing implementations lack theoretical justifications. This paper revisits invariant representation learning for ranking. Upon reviewing prior work, we found that they implement what we call item-level alignment, which aligns the distributions of the items being ranked from all lists in aggregate but ignores their list structure. However, the list structure should be leveraged, because it is intrinsic to ranking problems where the data and the metrics are defined and computed on lists, not the items by themselves. To close this discrepancy, we propose list-level alignment—learning domain-invariant representations at the higher level of lists. The benefits are twofold: it leads to the first domain adaptation generalization bound for ranking, in turn providing theoretical support for the proposed method, and it achieves better empirical transfer performance for unsupervised domain adaptation on ranking tasks, including passage reranking.
Learning to rank applies machine learning techniques to solve ranking problems that are at the core of many everyday products and applications, including search engines and recommendation systems [34]. The availability of ever-increasing amounts of training data and advances in modeling are constantly refreshing the state-of-the-art of many ranking tasks [40]. Yet, the training of practical and effective ranking models on tasks with little to no annotated data remains a challenge [57]. A popular transfer learning framework for this problem is domain adaptation [41]: given source domains with labeled data that are relevant to the target domain of interest, domain adaptation methods optimize the model on the source domain and make the knowledge transferable by utilizing (unlabeled) target data.
For domain adaptation, there are task-specific and general-purpose (unsupervised) methods. E.g., for text ranking in information retrieval where we only have access to target documents without any training query or relevance annotation, a method is to use generative language models to synthesize
$^{\dagger}$ Work performed while at Google Research.
Lists of items  $(q - d$  pairs) to be ranked
Figure 1: Item-level and list-level alignment on a ranking problem. Each list contains  $\ell$  items to be ranked. Whereas item-level alignment aligns the distributions of items but ignores the list structure, list-level alignment preserves this structure and aligns the distributions of lists.
Item-level alignment
Distributions of items (list structure is ignored)
List-level alignment
Distributions of lists
relevant queries for the unannotated documents, then train the ranking model on the synthesized query-document pairs [37, 54, 65]. For general-purpose methods, which are applicable to all kinds of tasks, perhaps the most well-known is domain-invariant representation learning [38, 17]. While it is studied extensively on classification and regression problems [71, 73], applications of invariant representation learning on ranking problems are sporadic [12, 58, 68], and the existing implementations lack theoretical justifications, which casts doubt on their applicability to learning to rank.
In this paper, we revisit the method of learning invariant representations for domain adaptation on ranking problems. We found that, from our review of prior work [12, 58, 68], existing implementations of the method all perform what we call item-level alignment (Fig. 1 and Section 3.2.1), which aligns the distributions of the items being ranked (e.g., query-document pairs) aggregated from all lists and discards their list structure (as in, e.g., q-d pairs with the same query all belong to the same list). However, both the data and the metrics for ranking are defined and computed on lists, not the items by themselves, so intuitively the list structure should be leveraged in the feature learning process as it is intrinsic to the problem. To close this discrepancy, we propose and analyze list-level alignment, which aligns the distributions of the lists and preserves their list structure (Section 3.2.2).
The conceptual leap from item-level to list-level alignment is significant. First, by analyzing list-level representation invariance, we establish the first domain adaptation generalization bound for ranking (Section 4), in turn providing theoretical support for our proposed method as well as a foundation for future studies. Second, list-level alignment provides empirical benefits to domain adaptation on ranking problems, which we demonstrate on two ranking tasks, passage reranking (Section 5) and Yahoo! LETOR web search ranking (Appendix D), where it outperforms zero-shot learning, item-level alignment, as well as a method specific to text ranking based on query synthesis [37].
We define a ranking problem by a joint distribution  $\mu$  of length- $\ell$  lists and their nonnegative scores, denoted by  $X = (X_{1},\dots ,X_{\ell})\in \mathcal{X}$  and  $Y = (Y_{1},\dots ,Y_{\ell})\in \mathbb{R}_{\geq 0}^{\ell}$ . For example, in information retrieval, the items in each list are query-document pairs of the same query, and the scores indicate the relevance of the documents to the query. Note that the lists are permutation-invariant, because switching the order of the items (and the corresponding scores) does not change the content of the list; concretely, let  $S_{\ell}$  denotes the set of permutations on  $[\ell ]\coloneqq \{1,2,\dots ,\ell \}$ , then permutation invariance means that  $\mu ((x_1,\dots ,x_\ell),(y_1,\dots ,y_\ell)) = \mu ((x_{\pi_1},\dots ,x_{\pi_\ell}),(y_{\pi_1},\dots ,y_{\pi_\ell}))$  for all  $(x,y)$  and  $\pi \in S_{\ell}$ . We assume that the scores are a function of the list,  $Y = y(X)$ , so the problem can be equivalently given by the marginal distribution  $\mu^{X}$  of lists and a scoring function  $y:\mathcal{X}\to \mathbb{R}_{\geq 0}^{\ell}$ .
Learning to rank aims to learn a ranker that takes lists  $x$  as input and outputs rank assignments  $r \in S_{\ell}$ , where  $r_i \in [\ell]$  is the predicted rank of item  $i$ , such that  $r$  recovers the descending order of the scores  $y$ , i.e.,  $y_i > y_j \iff r_i < r_j$  for all  $i, j$ . The more common setup, which we will adopt, is
<sup>1</sup>Although here the input lists are written to be decomposable into the  $\ell$  items they contain, generally and more abstractly, they need not be so.
to train a scorer,  $f: \mathcal{X} \rightarrow \mathbb{R}^{\ell}$ , with the rank assignments obtained by first computing the ranking scores  $s = f(x) \in \mathbb{R}^{\ell}$ , where  $s_i$  is the predicted score of item  $i$ , then taking the descending order of  $s$  (alternatively, the rank assignments can be generated probabilistically; see Section 4).
To quantitatively measure the quality of the predicted ranks, we use (listwise) ranking metrics,  $u: S_{\ell} \times \mathbb{R}_{\geq 0}^{\ell} \to \mathbb{R}_{\geq 0}$ , which computes a utility score via comparing the predicted rank assignments to the ground-truth scores of the list. Common metrics in information retrieval include reciprocal rank and normalized discounted cumulative gain [63, 26]:
Definition 2.1 (Reciprocal Rank). If the ground-truth scores are binary, i.e.,  $y \in \{0,1\}^{\ell}$ , then the reciprocal rank (RR) of a predicted rank assignments  $r \in S_{\ell}$  is
$$
\operatorname {R R} (r, y) = \max  \left(\left\{r _ {i} ^ {- 1}: 1 \leq i \leq \ell , y _ {i} = 1 \right\} \cup \{0 \}\right).
$$
The average RR of a ranking model  $f^{\prime}:\mathcal{X}\to S_{\ell}$  on a dataset,  $\mathbb{E}_{(X,Y)\sim \mu}[\mathrm{RR}(f'(X),Y)]$ , is referred to as the mean reciprocal rank (MRR).
Definition 2.2 (NDCG). The discounted cumulative gain (DCG) and the normalized DCG (with identity gain function) of a predicted rank assignments  $r \in S_{\ell}$  are
$$
\operatorname {D C G} (r, y) = \sum_ {i = 1} ^ {\ell} \frac {y _ {i}}{\log \left(r _ {i} + 1\right)}, \quad \text {a n d} \quad \operatorname {N D C G} (r, y) = \frac {\operatorname {D C G} (r , y)}{\operatorname {I D C G} (y)},
$$
where  $\mathrm{IDCG}(y) = \max_{r' \in S_\ell} \mathrm{DCG}(r', y)$ , called the ideal DCG, is the maximum DCG of a list, which is attained by the descending order of  $y$ .
Domain Adaptation. In this learning setting, we have a source domain  $(\mu_S^X, y_S)$  and a target domain  $(\mu_T^X, y_T)$  with different data distributions, and the goal is to learn a good ranking model for the target domain by leveraging all available resources: whereas access to source domain labeled training data is always assumed, we may only be given unlabeled data for the target domain (this scenario is called unsupervised domain adaptation). A popular method for this adaptation is domain-invariant representation learning, which matches and aligns the source and target domain data distributions on the feature space so that their distributions appear similar to the model.
3 Learning Domain-Invariant Representations for Ranking
We begin by reviewing invariant representation learning for domain adaptation and the optimization technique of adversarial training, then describe and compare two instantiations of this framework to ranking problems in Section 3.2—item-level alignment, which is implemented in prior work, and our proposed list-level alignment.
For representation learning, we train composite models of the form  $f = h \circ g$ , where  $g: \mathcal{X} \to \mathcal{Z}$  is a shared feature map and  $h$  is a task head on the shared features. As an example, if the model  $f$  is end-to-end implemented as an  $m$ -layer multilayer perceptron (MLP), then we could treat the first  $(m - 1)$  layers as  $g$  and the last as  $h$ .
And, recall the definitions of Lipschitz continuity and Wasserstein distance [15]:
Definition 3.1 (Lipschitz Continuity). A function  $f: \mathcal{X} \to \mathcal{Y}$  from metric space  $(\mathcal{X}, d_{\mathcal{X}})$  to  $(\mathcal{Y}, d_{\mathcal{Y}})$  is  $L$ -Lipschitz, denoted by  $f \in \mathrm{Lip}(L)$ , if  $d_{\mathcal{Y}}(f(x), f(x')) \leq L d_{\mathcal{X}}(x, x')$  for all  $x, x' \in \mathcal{X}$ .
Definition 3.2 (Wasserstein Distance). The Wasserstein-1 distance between probability measures  $p, q$  on metric space  $\mathcal{X}$  is denoted and given by  $W_{1}(p,q) = \sup_{\psi : \mathcal{X} \to \mathbb{R}, \psi \in \mathrm{Lip}(1)} \int_{\mathcal{X}} \psi(x)(p(x) - q(x)) \, \mathrm{d}x$ .
3.1 Invariant Representation Learning via Adversarial Training
Given a source domain  $\mu_S$  and a target domain  $\mu_T$ , invariant representation learning aims to learn a shared feature map  $g$  s.t. the distributions are aligned on the feature space, i.e.,  $D(\mu_S^Z, \mu_T^Z) \approx 0$ , where  $D$  is a divergence measure or probability metric, e.g., Wasserstein distance, and the source feature distribution  $\mu_S^Z$  (analogously for target  $\mu_T^Z$ ) is defined to be the distribution of  $Z = g(X)$ ,  $X \sim \mu_S^X$ . The idea is that a composite model with an aligned feature representation is transferable between
2For adapting from multiple source domains, see [71].
domains—this is supported by the following generalization bound on classification problems [52, 5]. We will establish a similar result for ranking in Section 4.
Theorem 3.3. Let binary classification problems on a source and a target domain be given by joint distributions  $\mu_S, \mu_T$  of inputs and labels,  $(X,Y) \in \mathcal{X} \times \{0,1\}$ . Define the error rate of a predictor  $f: \mathcal{X} \to [0,1]$  on  $\mu$  by  $\mathcal{R}(f) = \mathbb{E}_{(X,Y)} \sim_{\mu}[f(X) \mathbb{1}[Y \neq 1] + (1 - f(X)) \mathbb{1}[Y \neq 0]]$ , where  $\mathbb{1}[\cdot]$  is the indicator function.
Let  $\mathcal{H} \subset [0,1]^{\mathcal{Z}}$  be an  $L$ -Lipschitz class of prediction heads, and for any  $g \in \mathcal{G}$ , define the minimum joint error rate by  $\lambda_g^* = \min_{h'} (\mathcal{R}_S(h' \circ g) + \mathcal{R}_T(h' \circ g))$ , then for all  $h \in \mathcal{H}$ ,
$$
\mathcal {R} _ {T} (h \circ g) \leq \mathcal {R} _ {S} (h \circ g) + 2 L W _ {1} \left(\mu_ {S} ^ {Z}, \mu_ {T} ^ {Z}\right) + \lambda_ {g} ^ {*}.
$$
This result bounds the target domain risk of the composite model  $h \circ g$  by its source task risk, plus the divergence of the feature distributions and the optimal joint risk. It therefore suggests that a good model for the target could be obtained by learning an informative domain-invariant feature map  $g$  (s.t.  $\lambda_g^*$  and  $W_1(\mu_S^Z, \mu_T^Z) \approx 0$ ) and training  $h$  to minimize the source task risk  $\mathcal{R}_S$ .
Invariant representation learning algorithms capture the aforementioned objectives with the joint training objective of  $\min_{h,g}(\mathcal{L}_S(h\circ g) + \lambda D(\mu_S^Z,\mu_T^Z))$  [35, 17, 13], where  $\mathcal{L}_S$  is the source task loss, and  $\lambda >0$  is a hyperparameter that controls the strength of invariant representation learning. This objective learns domain-invariant features by optimizing  $g$  to minimize the distributional discrepancy, and minimizing source risk at the same time s.t. the learned features are useful for the task. Since it does not require labeled data but only unlabeled ones for estimating  $W_{1}(\mu_{S}^{Z},\mu_{T}^{Z})$ , it is applicable to unsupervised domain adaptation. If labeled target data are available, they can be incorporated by, e.g., adding a target loss  $\mathcal{L}_T$ , which will help keep  $\lambda_g^*$  low.
Adversarial Training. When the divergence measure  $D$  is simple and admits closed-form expressions, e.g., maximum-mean discrepancy [19], the objective above can be minimized directly [35], but they are typically too weak at modeling complex distributions that arise in real-world problems. To achieve feature alignment under stronger measures, e.g., JS divergence or Wasserstein distance, a well-known approach is adversarial training [18, 17, 2], which reformulates the above objective into
$$
\mathcal {L} _ {\text {j o i n t}} (h, g) = \min  _ {h \in \mathcal {H}, g \in \mathcal {G}} \left(\mathcal {L} _ {S} (h \circ g) - \lambda \min  _ {f _ {\text {a d}} \in \mathcal {F} _ {\text {a d}}} \mathcal {L} _ {\text {a d}} (g, f _ {\text {a d}})\right),
$$
and optimizes it using gradient descent-ascent with a gradient reversal layer added on top of  $g$  in the adversarial component.
The adversarial component,  $-\min_{f_{\mathrm{ad}}} \mathcal{L}_{\mathrm{ad}}(g, f_{\mathrm{ad}})$ , can be shown to upper bound the divergence between  $\mu_S^Z, \mu_T^Z$ . It involves an adversary  $f_{\mathrm{ad}}: \mathcal{Z} \to \mathbb{R}$  (parameterized by neural networks) and an adversarial loss function  $\ell_{\mathrm{ad}}: \mathbb{R} \times \{0, 1\} \to \mathbb{R}$ , whose inputs are the output of  $f_{\mathrm{ad}}$  and the domain identity  $a$  (we set target domain to  $a = 1$ ). The adversarial loss is computed over both domains:
$$
\mathcal {L} _ {\mathrm {a d}} (g, f _ {\mathrm {a d}}) := \mathbb {E} _ {X \sim \mu_ {S} ^ {X}} [ \ell_ {\mathrm {a d}} (f _ {\mathrm {a d}} \circ g (X), 0) ] + \mathbb {E} _ {X \sim \mu_ {T} ^ {X}} [ \ell_ {\mathrm {a d}} (f _ {\mathrm {a d}} \circ g (X), 1) ].
$$
If we choose the 0-1 loss,  $\ell_{\mathrm{ad}}(\hat{a}, a) = (1 - a)\mathbb{1}[\hat{a} \geq 0] + a\mathbb{1}[\hat{a} < 0]$ , then the adversarial component upper bounds  $W_1(\mu_S^Z, \mu_T^Z)$  (see Proposition A.1); here,  $f_{\mathrm{ad}}$  acts as a domain discriminator for distinguishing the domain identity, and  $\mathcal{L}_{\mathrm{ad}}$  is the balanced classification error of  $f_{\mathrm{ad}}$ . For training, the 0-1 loss is replaced by a surrogate loss, e.g., the logistic loss [18],  $\ell_{\mathrm{ad}}(\hat{a}, a) = \log(1 + e^{(1 - 2a)\hat{a}})$ .
3.2 Invariant Representation Learning for Ranking
This section describes and compares two instantiations of the invariant representation learning framework above for ranking: item-level alignment, and our proposed list-level alignment. The key difference between them is the choice of  $\mu^{Z}$  whose divergence  $W_{1}(\mu_{S}^{Z},\mu_{T}^{Z})$  is to be minimized.
Model Setup. We consider a composite scoring model  $f = h \circ g$  where the feature map  $g$  is s.t. given an input list  $x = (x_{1},\dots ,x_{\ell})$ , it outputs a list of  $k$ -dimensional feature vectors,  $z = g(x) = (v_{1},\dots ,v_{\ell}) \in \mathbb{R}^{\ell \times k}$ , where  $v_{i} \in \mathbb{R}^{k}$  corresponds to item  $i$ . Ranking scores are then
This definition assumes that the output class is probabilistic according to  $\mathbb{P}(\widehat{Y} = 1\mid X = x) = f(x)$
Figure 2: List-level alignment is a stronger requirement than item-level alignment, which, in addition to pairing target domain items to source items, it also pairs target lists to source lists.
obtained by, e.g., projecting each  $v_{i}$  to  $\mathbb{R}$  with a (shared) linear layer. This setup, which we will adopt in our experiments, is common with listwise ranking models and in many ranking systems [10]: in neural text ranking, each feature vector can be the embedding of the input text computed by a language model [39, 21, 76].
3.2.1 Item-Level Alignment
In item-level alignment (ItemDA), the distributions of the  $k$ -dimensional feature vectors from all lists in aggregate are aligned [12, 58, 68], i.e.,  $\mu_{S}^{Z,\mathrm{item}} \approx \mu_{T}^{Z,\mathrm{item}}$ , where
$$
\mu^ {Z, \text {i t e m}} (v) := \mathbb {P} _ {Z \sim \mu^ {Z}} (v \in Z) = \mathbb {P} _ {X \sim \mu^ {X}} (v \in g (X)), \quad \operatorname {s u p p} (\mu^ {Z, \text {i t e m}}) \subseteq \mathbb {R} ^ {k}.
$$
In other words (when both domain have the same number of lists), it finds a matching between the items of the source and target domains that pairs each target item with a source one (Fig. 2). Note that the list structure is not retained in this definition, because one cannot necessarily tell whether two items drawn from the bag of feature vectors,  $v, v' \sim \mu^{Z,\mathrm{item}}$ , belong to the same list.
To implement item-level alignment, the discriminator  $f_{\mathrm{ad}}$  (usually an MLP) is set up to take individual feature vectors  $v \in \mathbb{R}^k$  as input. In each forward-pass, the feature map  $g$  computes a batch of feature vector lists,  $\{z_i\}_{i \in [b]} = \{(v_{i1}, \dots, v_{il})\}_{i \in [b]}$ , then the discriminator takes the items batched from all lists (this step discards the list structure),  $\{v_{ij}\}_{i \in [b], j \in [\ell]}$ , and predicts their domain identities.
Item-level alignment is identical to invariant representation learning implementations for classification and regression, e.g., DANN [17], which also operate on "items", since neither the data nor the metrics in those problem settings have any explicit structural assumptions. However, ranking problems are inherently defined with list structures: the inputs to the model are organized by lists, and ranking metrics are evaluated on lists rather than the items by themselves. This discrepancy casts doubt on the applicability of item-level alignment to ranking problems, and moreover, it is unclear whether domain adaptation for ranking via item-level alignment is justified from a theoretical perspective.
3.2.2 List-Level Alignment
Closing the discrepancy discussed above, we propose list-level alignment (ListDA), which aligns the distributions of the lists of feature vectors on the  $\mathbb{R}^{\ell \times k}$  space, i.e.,  $\mu_S^{Z,\mathrm{list}}\approx \mu_T^{Z,\mathrm{list}}$  , where
$$
\mu^ {Z, \text {l i s t}} (z) := \mathbb {P} _ {Z \sim \mu^ {Z}} (z) = \mathbb {P} _ {X \sim \mu^ {X}} (z = g (X)), \quad \operatorname {s u p p} (\mu^ {Z, \text {l i s t}}) \subseteq \mathbb {R} ^ {\ell \times k}.
$$
This means finding a matching that not only pairs all source and target items, but also pairs items from the same target list to items from the same source list. Therefore, list-level alignment is a stronger requirement than item-level alignment,  $\mu_{S}^{Z,\mathrm{list}} = \mu_{T}^{Z,\mathrm{list}} \Longrightarrow \mu_{S}^{Z,\mathrm{item}} = \mu_{T}^{Z,\mathrm{item}}$ , while the converse is generally not true (see Fig. 2 for a picture, or Example A.2)!
To implement list-level alignment, the discriminator has to make predictions on lists,  $z \in \mathbb{R}^{\ell \times k}$ , i.e., in each forward-pass it would be fed a batch of feature vector lists,  $\{z_i\}_{i \in [b]} = \{(v_{i1}, \dots, v_{il})\}_{i \in [b]}$  and output  $b$  predictions. For this, a possible design choice is to flatten each feature vector list into a long  $(\ell k)$ -dimensional vector and use an MLP as  $f_{\mathrm{ad}}$ . But recall from Section 2 that the
input lists are permutation-invariant, and so are the feature vector lists. Since the above setup does not incorporate this property and an exponential amount of samples is required to see all possible permutations, this would be inefficient to optimize; yet, the optimality of  $f_{\mathrm{ad}}$  is essential to the success of adversarial training. So instead, we use transformers with mean-pooling less positional encoding as the discriminator [60], which is permutation-invariant to the input.
Compared to item-level alignment, our list-level alignment is supported by a domain adaptation generalization bound (Theorem 4.7, established in the next section), and achieves better transfer performance in our Section 5 and Appendix D evaluations. These results suggest that the list structure is essential for an effective domain adaptation on ranking problems.
4 Generalization Bound for Ranking with List-Level Alignment
Based on the list-level alignment proposed above, we establish a domain adaptation generalization bound for ranking by analyzing list-level representation invariance. The discussions in this section consider learning a composite scoring model,  $f = h \circ g: \mathcal{X} \to \mathbb{R}^{\ell}$ , but need not assume that the list feature representation  $z = g(x)$  can be decomposed into item-level feature representations as in the model setup of Section 3.2.2, only that it resides in a metric space  $\mathcal{Z}$ .
Our result can be viewed as an instantiation of the framework established originally for classification by Ben-David et al. [5]. But since ranking is a very different task with its unique technical challenges, our analysis differs from those in prior work [52]. We will introduce appropriate assumptions for handling these difficulties leading up to the main Theorem 4.7, which the readers may skip to.
The first hurdle is the discontinuity of the sorting operation. So rather than taking the descending order of the predicted scores to obtain the rank assignments, we generate them probabilistically using a Plackett-Luce model [43, 36] with the exponentiated scores as its parameters [10, 20]. This makes the computation of the utility (Lipschitz) continuous w.r.t. the raw output scores of the model.
Definition 4.1 (Plackett-Luce Model). A Plackett-Luce (P-L) model with parameters  $w \in \mathbb{R}_{>0}^{\ell}$  specifies a distribution over  $S_{\ell}$ , whose probability mass function  $p_w$  is
$$
p _ {w} (r) = \prod_ {i = 1} ^ {\ell} \frac {w _ {I (r) _ {i}}}{\sum_ {j = i} ^ {\ell} w _ {I (r) _ {j}}}, \quad \forall r \in S _ {\ell},
$$
where  $I(r)_i$  is the index of the item with rank  $i$ ,  $r_{I(r)_i} = i$ .
Assumption 4.2. Given predicted scores  $s = f(x) \in \mathbb{R}^{\ell}$ , we generate the rank assignments probabilistic from a P-L model by  $R \sim p_{\exp(s)} \in S_{\ell}$ , where  $\exp$  is taken coordinate-wise.