SlowGuess/NAACL_2022_OCR huggingface dataset:
Chengyue Gong $^{1}$ , Xiaocong Du $^{2}$ , Dhruv Choudhary $^{2}$ , Bhargav Bhushanam $^{2}$ , Qiang Liu $^{1}$ , Arun Kejariwal $^{2}$
<sup>1</sup> University of Texas at Austin, <sup>2</sup> Meta Platforms, Inc. {cygong, lqiang} @cs.utexas.edu {xiaocongdu, choudharydhruv, bbhushanam, akejariwal} @fb.com
Learning embedding layers (for classes, words, items, etc.) is a key component of lots of applications, ranging from natural language processing, recommendation systems to electronic health records, etc. However, the frequency of real-world items follows a long-tail distribution in these applications, causing naive training methods perform poorly on the rare items. A line of previous works address this problem by transferring the knowledge from the frequent items to rare items by introducing an auxiliary transfer loss. However, when defined improperly, the transfer loss may introduce harmful biases and deteriorate the performance.
In this work, we propose a harmless transfer learning framework that limits the impact of the potential biases in both the definition and optimization of the transfer loss. On the definition side, we reduce the bias in transfer loss by focusing on the items to which information from high-frequency items can be efficiently transferred. On the optimization side, we leverage a lexicographic optimization framework to efficiently incorporate the information of the transfer loss without hurting the minimization of the main prediction loss function. Our method serves as a plug-in module and significantly boosts the performance on a variety of NLP and recommendation system tasks.
Since the advent of the deep learning era, learnable embedding layers for categorical and discrete features are basic modules for neural networks in many fields, such as natural language processing (NLP) (Vaswani et al., 2017; Devlin et al., 2018; Bahdanau et al., 2014, e.g.), recommendation systems (Zhang et al., 2019; Guo et al., 2017, e.g.), electronic health record (e.g. Kalyan and Sangeetha, 2020; Qian et al., 2017; Choi et al., 2018). Discrete items are first mapped to continuous representations in the feature space through
the embedding layers and then processed by other neural modules.
However, a key challenge in learning item embeddings is due to the long tail phenomenon (Clauset et al., 2009): the frequencies of different items are usually extremely imbalanced and there often exists a large number of rare items that appear only a small number of times. As a result, it is very difficult to properly estimate the embeddings of rare items (e.g. Yin et al., 2020; Peng et al., 2019; Gao et al., 2019; Gong et al., 2018; Li et al., 2020a). The rare items often suffer from the under-fitting issue because they appear infrequently in the data (Sennrich et al., 2015; Provilkov et al., 2019).
To solve this long tail problem, transfer learning techniques are proposed (e.g. Gong et al., 2018; Li et al., 2020b; Yin et al., 2020; Chen et al., 2020) by propagating useful information from high-frequency items to rare items. The idea of these methods is to minimize the linear combination of the main prediction loss  $\ell_{\mathrm{main}}$  with an auxiliary transfer loss  $\ell_{\mathrm{transfer}}$  which encourages that the embeddings of the high-frequency and rare items follow a similar distribution, and hence allows us to the regularize the training of the rare items using the information from the common items. The basic assumption of the existing transfer learning techniques is that the distribution of the underlying embedding of the high-frequency and rare items are similar, which is necessarily false in practice. As a result, the transfer loss  $\ell_{\mathrm{transfer}}$  introduces potentially harmful bias and impacts the overall performance.
The goal of this work is to limit the potential harmful information in the transfer learning pipeline. We achieve this following two complementary directions:
1) Improving the definition of the transfer loss  $\ell_{\mathrm{transfer}}$  so that it inherently introduces less bias.
2) Adaptively balancing the combination coefficient of  $\ell_{\mathrm{main}}$  and  $\ell_{\mathrm{transfer}}$ , so that the transferring
Findings of the Association for Computational Linguistics: NAACL 2022, pages 504 - 516 July 10-15, 2022 ©2022 Association for Computational Linguistics
information is efficiently incorporated without being conflicting with the main loss  $\ell_{\mathrm{main}}$
To improve the definition of  $\ell_{\mathrm{transfer}}$ , we leverage the co-appearance information to identify the rare items that are expected to be embedded similarly to the high-frequency items; by applying the transfer loss only between them, we significantly reduce the potential harmful bias. To adaptively balance  $\ell_{\mathrm{main}}$  and  $\ell_{\mathrm{transfer}}$ , we formulate the transfer learning as a lexicographic (lexico) optimization, in which we minimize  $\ell_{\mathrm{transfer}}$  subject to that the minimum of  $\ell_{\mathrm{main}}$  is achieved. This allows us to incorporate the information of  $\ell_{\mathrm{transfer}}$ , without interfering the minimization of the main loss  $\ell_{\mathrm{main}}$ . We provide a simple algorithm for solving the lexico optimization by extending the dynamic barrier algorithm of (Gong et al., 2021) to incorporate the item-wise structure in our problem.
We verify the performance of our method in multiple applications, from language model, machine translation, named entity recognition to click-through rate (CTR) prediction. Compared to different baseline transfer learning methods, our method achieves better or comparable results without tuning the hyper-parameters in all the applications.
2 Problem Set: Learning with Long Tail
In the following, we first formulate the problem and then introduce the details about our method. We start by introducing the framework of transfer learning for rare item embeddings. Assume we have a dataset  $\mathcal{D}$  regarding a set of discrete items  $\mathcal{I}$ , such as the words in NLP and the users and products in recommendation systems. A typical deep learning model consists of two following parts:
1) a set of continuous embedding vectors  $e = \{e_i \in \mathbb{R}^m : i \in \mathcal{I}\}$ , which map the discrete items in  $\mathcal{I}$  into the  $\mathbb{R}^m$  Euclidean space;
2) a deep neural network  $f_{\theta}$ , which takes the continuous embedding of the inputs and makes the desirable predictions. Assume  $f$  is indexed by a trainable parameter  $\theta \in \mathbb{R}^{m'}$ .
Both  $e$  and  $\theta$  are parameters that we learn using the training data. In practice, the items  $\mathcal{I}$  tend to exhibit a long tail distribution, consisting of both high-frequency, informative items that appear many times in the dataset, and also a large set of rare items that only appear a small number of times. As a result, the embedding vectors of the rare items may not be well estimated and hence deteriorate the performance.
To address the long tail problem, a number of works have been proposed to train the parameters  $\{e,\theta \}$  by regularizing the main prediction loss with a transfer loss that propagate information from the frequent items to the rare items:
$$
\min  _ {\{e, \theta \}} \left\{\mathcal {L} (e, \theta) := \ell_ {\text {m a i n}} (e, \theta) + \lambda \ell_ {\text {t r a n s f e r}} (e) \right\}, \tag {1}
$$
where  $\lambda > 0$  is a combination coefficient; the  $\ell_{\mathrm{main}}$  is the main prediction loss on the training data based on model  $f_{\theta}$  and the embedding  $e$ , which is usually cross-entropy loss, KL divergence, and  $\ell_{\mathrm{transfer}}$  is a transfer loss, designed to transfer the knowledge from the frequent items (the source domain) to the rare items (target domain). Denote by  $S$  and  $\mathcal{T}$  the items in the source and target domain, respectively. The transfer loss is defined as
$$
\ell_ {\text {t r a n s f e r}} (e) = \operatorname {D i s t} \left(\left\{e _ {i} \right\} _ {i \in \mathcal {S}}, \left\{e _ {i} \right\} _ {i \in \mathcal {T}}\right), \tag {2}
$$
where measures the discrepancy between the empirical distributions of the source and target items  $\{e_i\}_{i\in S}$  and  $\{e_i\}_{i\in T}$ . The discrepancy measure  $Dist(\cdot)$  can be defined in various ways, such as Maximum Mean Discrepancy (MMD) (Chen et al., 2020), adversarial networks (Gong et al., 2018; Yin et al., 2020), moment matching (Peng et al., 2019), Wasserstein distance (Xu et al., 2018) Additional techniques, such as meta learning (Zhu et al., 2021) and gradient alignment (Li et al., 2020b) have also been explored.
In practice, the source domain  $S$  and target domain  $\mathcal{T}$  is chosen based on the item frequency with various heuristics. For example, Gong et al. (2018) and Li et al. (2020b) uses the Pareto Principle (Dunford et al., 2014) (or 80/20 rule), which places the top  $20\%$  high-frequency items in  $S$  and the rest in  $\mathcal{T}$ . Further, Yin et al. (2020) define the target domain as the items with frequency lower than a dataset-dependent threshold.
3 Harmless Transferring for Rare Items
The basic assumption that underpins the transfer loss  $\ell_{\mathrm{transfer}}$  in (2) is that the distribution of the high-frequency items  $e_S$  and rare items  $e_I$  is similar. However, this is not necessarily true in practice. There could be rare items that find no similar high-frequency items and should be embedded differently. When this happens, the transfer loss  $\ell_{\mathrm{transfer}}$  may provide biased and harmful information and hence deteriorate the overall performance.
In this work, we propose two techniques to limit the potential harmful information in the transfer
learning pipeline, by improving both the definition and the optimization of  $\ell_{\mathrm{transfer}}$
1) In Section 3.1, we improve the definition of  $\ell_{\mathrm{transfer}}$  in (2) to reduce potential harmful bias. This is made possible by filtering out the irrelevant items from the target domain  $\mathcal{T}$  based on the coappearance information between the items.
2) In Section 3.2, we further limits the potential harmful impact of  $\ell_{\mathrm{transfer}}$  by replacing the regularized loss in (1) with a lexicographic (lexico) optimization which minimizes  $\ell_{\mathrm{transfer}}$  subject to that the minimum of  $\ell_{\mathrm{main}}$  is achieved. This allows us to prioritize the optimization of  $\ell_{\mathrm{main}}$  when the gradient direction  $\ell_{\mathrm{main}}$  and  $\ell_{\mathrm{transfer}}$  are conflicting with each other. We provide a simple algorithm for solving the lexico optimization by extending the dynamic barrier algorithm of (Gong et al., 2021) to incorporate the item-wise structure in our problem.
Figure 1: A demonstration of our method. 1) For tail items, we set the target domain as items which have similar (the dashed line denotes the large similarity) frequent items (the orange circles), and remove tail items with no similar frequent items (the orange triangle), 2) We dynamically control the coefficient between the main loss and transfer loss in a way that efficiently incorporates the transfer loss information without hurting the optimization of the main loss.
3.1 Selective Transfer Loss
In previous works, source  $S$  and target  $\mathcal{T}$  domains are decided solely based on frequency information, which can not ensure that their embedding vectors  $e_S$  and  $e_T$  have a similar distribution. This may introduce harmful information into  $\ell_{\mathrm{transfer}}$ . We address this problem by incorporating co-appearance information between the items, so that target domain  $\mathcal{T}$  only include the rare items that find some high-frequency items.
Specifically, starting with a pair of  $S$  and  $\mathcal{T}$  using standard method (say, based on frequency cutoff), we refine  $\mathcal{T}$  to only include rare items that find similar items in  $S$ . To do so, we characterize each item  $i \in \mathcal{I}$  by its co-appearance with the high-frequency
items. Specifically, let  $s_i = [s_{ij}]_{j \in S}$ , where  $s_{ij}$  denotes the number of times when item  $i$  and  $j$  co-appear in an n-gram window. Without further information, the co-appearance vector  $s_i$  provides a naive representation of item  $i$ . Then, we filter the target domain  $\mathcal{T}$  by removing the rare items that has low similarity with every high frequency item, yielding a more selective target domain:
$$
\mathcal {T} _ {\eta} ^ {\prime} = \mathcal {T} \cap \{i \in \mathcal {I}: \max  _ {j \in \mathcal {S}} \cos (s _ {i}, s _ {j}) \geq \eta \}, \tag {3}
$$
$\cos (\cdot ,\cdot)$  denotes the cosine similarity between two vectors, and  $\eta$  is a given threshold. We then replace  $\mathcal{T}$  with  $\mathcal{T}_{\eta}^{\prime}$  in the definition of  $\ell_{\mathrm{transfer}}$ , that is,
$$
\ell_ {\text {t r a n s f e r}} (e) = \operatorname {D i s t} \left(\left\{e _ {i} \right\} _ {i \in \mathcal {S}}, \left\{e _ {i} \right\} _ {i \in \mathcal {T} _ {\eta} ^ {\prime}}\right).
$$
In this way, we filter out the irrelevant items from the target domain, and hence reduce the potential biases in  $\ell_{\mathrm{transfer}}$ .
3.2 Harmless Transfer via Lexico Optimization
Despite the improved definition above, the transfer loss  $\ell_{\mathrm{transfer}}$  may still contain harmful bias that will be inherited by the final result as we minimize the regularized loss function in (1). As a result, the choice of the regularization coefficient  $\lambda$  in (1) becomes critically important: a large  $\lambda$  may overemphasize the transfer loss and amplify the harmful biases, while a small  $\lambda$  may yield insufficient transferring. The optimal choice of  $\lambda$  depends on both  $\ell_{\mathrm{main}}$  and  $\ell_{\mathrm{transfer}}$  and need to be selected case by case for individual problems using grid search or other hyper-parameter optimization approaches, which yields high computational cost.
To incorporate the benefit of the transfer loss without hurting the optimization of the main loss, we propose to estimate the parameters  $\{e,\theta \}$  by solving a lexicographic optimization problem (e.g., Dempe and Zemkoho, 2020; Gong et al., 2021):
$$
\min  _ {e} \ell_ {\text {t r a n s f e r}} (e), \text {s . t .} \{e, \theta \} \in \arg \min  \ell_ {\text {m a i n}}, \tag {4}
$$
where  $\arg \min \ell_{\mathrm{main}}$  denotes the set of (local) minimum points of  $\ell_{\mathrm{main}}$ . This means that we prioritize the optimization of the main loss, and minimize the  $\ell_{\mathrm{transfer}}$  only within the optimum set of  $\ell_{\mathrm{main}}$ .
A key fact that underpins the formulation in (4) is that the modern neural networks that we use in practice are almost always over-parameterized, and hence the optimum set of  $\ell_{\text{main}}$  is not a set of
isolated points, but connected manifolds consisting of an infinite number of points. Therefore, by searching in the solution manifold the point that minimizes  $\ell_{\mathrm{transfer}}$ , we gain improvement on the transfer loss without hurting the main loss.
Dynamic Barrier Gradient Descent We adopt the dynamic barrier gradient descent (DBGD) algorithm by (Gong et al., 2021) to solve the lexicographic optimization problem in (4). DBGD iteratively updates the parameters by
$$
(e, \theta) \leftarrow (e, \theta) - \epsilon (\mu_ {e}, \mu_ {\theta}),
$$
where  $\epsilon$  is a step size and  $\mu = (\mu_{e},\mu_{\theta})$  is the update direction for parameters  $(e,\theta)$  given by solving the following optimization problem:
$$
\begin{array}{l} \mu = \underset {\nu} {\operatorname {a r g m i n}} \| \nabla_ {(e, \theta)} \ell_ {\text {t r a n s f e r}} - \nu \| ^ {2} \tag {5} \\ s. t. \quad \left\langle \nabla_ {(e, \theta)} \ell_ {\text {m a i n}}, \nu \right\rangle \geq \phi , \\ \end{array}
$$
where  $\phi$  is a non-negative barrier function that equals zero only when we reach the minimum of  $\ell_{\mathrm{main}}$ . For example, for  $\ell_{\mathrm{main}} \geq 0$ , we can take
$$
\phi = \alpha \min (\| \nabla_ {(e, \theta)} \ell_ {\mathrm {m a i n}} \| ^ {2}, \ell_ {\mathrm {m a i n}})
$$
as suggested by (Gong et al., 2021), where  $\alpha > 0$  is a positive coefficient. The idea of (5) is as follows:
1. By enforcing that  $\left\langle \nabla_{(e,\theta)}\ell_{\mathrm{main}},\mu \right\rangle \geq \phi \geq 0$  we ensure that we always monotonically minimize  $\ell_{\mathrm{main}}$  (when step size  $\epsilon$  is sufficiently small), as  $\ell_{\mathrm{main}}$  is the main loss function.
2. By minimizing  $\| \nabla_{(e,\theta)}\ell_{\mathrm{transfer}} - \mu \| ^2$  under the inner product constraint, we  $\ell_{\mathrm{transfer}}$  is minimized as much as possible without hurting the optimization of  $\ell_{\mathrm{main}}$ .
As shown in (Gong et al., 2021), the problem in Eq. 5 yields a simple closed form solution:
$$
\mu = \nabla_ {(e, \theta)} \ell_ {\text {t r a n s f e r}} + \lambda \nabla_ {(e, \theta)} \ell_ {\text {m a i n}}, \tag {6}
$$
where  $\lambda$  is a non-negative coefficient defined by
$$
\lambda = \max  \left(\frac {\phi - \nabla_ {(e , \theta)} \ell_ {\text {t r a n s f e r}} ^ {\top} \nabla_ {(e , \theta)} \ell_ {\text {m a i n}}}{\| \nabla_ {(e , \theta)} \ell_ {\text {m a i n}} \| ^ {2}}, \quad 0\right). \tag {7}
$$
Eq (6) can be viewed as the gradient of the regularized loss (2). The key difference is that the  $\lambda$  in (6) is decided adaptively by formula (7) in each iteration of the algorithm based on the inner product of the gradients, rather than pre-determined in the beginning.
3.2.1 Item-wise Dynamic Barrier Descent
To best exploit the special structure of our problem, we propose to modify the off-the-shelf algorithm above yield better results.
Update of  $\theta$  Because the transfer loss  $\ell_{\mathrm{transfer}}(e)$  is independent with the neural network parameter  $\theta$ , that is,  $\nabla_{\theta}\ell_{\mathrm{transfer}}(e) = 0$ , the update on  $\theta$  follows the standard gradient:  $\mu_{\theta} = \lambda \nabla_{\theta}\ell_{\mathrm{main}}$ . It finds its handy to modify it to be  $\mu_{\theta} = \nabla_{\theta}\ell_{\mathrm{main}}$ , that is, we update  $\theta$  by standard gradient descent on  $\ell_{\mathrm{main}}$ .
Update of Embedding Vectors  $e$  Another critical special property of our problem is that  $e = \{e_i\}_{i \in \mathcal{I}}$  consists of the embedding vectors of all the items. It would be desirable to fine-grainedly control the co-efficient for different items; doing so can also decrease the computational cost since we only need to calculate the coefficient for the items that appear in the training mini-batch at each iteration.
Specifically, we update each embedding vector  $e_i$  by  $e_i \gets e_i - \epsilon \mu_{e_i}$ , where  $\mu_{e_i}$  is decided by
$$
\begin{array}{l} \mu_ {e _ {i}} = \underset {\nu_ {i}} {\mathrm {a r g m i n}} | | \nabla_ {e _ {i}} \ell_ {\mathrm {t r a n s f e r}} - \nu_ {i} | | ^ {2} \\ s. t. \quad \left\langle \nabla_ {e _ {i}} \ell_ {\text {m i n}}, v _ {i} \right\rangle \geq \phi_ {i}, \\ \end{array}
$$
where  $\phi_{i}$  is an item-wise control barrier associated with  $e_i$ :  $\phi_{i} = \alpha \min (\| \nabla_{e_{i}}\ell_{\mathrm{main}}\|^{2},\ell_{\mathrm{main}})$  (we simply take  $\alpha = 2$  in practice). Similar to (6), we can show that  $\mu_{e_i}$  is a linear combination of  $\nabla_{e_i}\ell_{\mathrm{main}}$  and  $\nabla_{e_i}\ell_{\mathrm{transfer}}$ , with an item-wise coefficient  $\lambda_{i}$ .
See Algorithm 1 for details of the update rule.
Memory And Time Cost Compared with the standard gradient descent on the main loss  $\ell_{\mathrm{main}}$ , the main additional computational cost that we introduce is calculating the gradient of  $\ell_{\mathrm{transfer}}$ , which is much faster than that of  $\ell_{\mathrm{main}}$  because it only involves embedding vectors  $e$ , and does not need to backpropagate on the network  $f_{\theta}$ . In practice, we find that we introduce an additional  $5\%$  to  $10\%$  time cost. For memory cost, we require to store the gradient for  $\ell_{\mathrm{transfer}}$  at each iteration. Consider that we store the gradient information for items in one given mini-batch instead of the whole embedding layer, the additional memory cost is light.
Transfer Learning for Long-Tail Items Transfer learning techniques are proposed (e.g. Gong et al., 2018; Gao et al., 2019; Li et al., 2020b; Yin
Algorithm 1: Controllable Item-wise Optimization for Embeddings
Denote  $\theta_t, e_t$  and  $\{\epsilon_t\}$  as the parameters, embeddings and learning rate at  $t$ -th iteration. for Iteration  $t$  do
$$
\begin{array}{l} \theta_ {t + 1} \longleftarrow \theta_ {t} - \epsilon_ {t} \nabla_ {\theta_ {t}} \ell_ {\text {m a i n}} \left(e _ {t}, \theta_ {t}\right), \\ e _ {t + 1, i} \leftarrow e _ {t, i} - \epsilon_ {t} \left(\nabla_ {e _ {t, i}} \ell_ {\text {t r a n s f e r}} (e _ {t}) + \lambda_ {t, i} \nabla_ {e _ {t, i}} \ell_ {\text {m a i n}} (e _ {t}, \theta_ {t})\right), \quad \forall \text {i t e m s} i \\ \end{array}
$$
$$
\begin{array}{l} \lambda_ {t, i} = \max \left(\frac {\phi_ {t , i} - \nabla_ {e _ {t , i}} \ell_ {\mathrm {t r a n s f e r}} (e _ {t}) ^ {\top} \nabla_ {e _ {t , i}} \ell_ {\mathrm {m a i n}} (e _ {t} , \theta_ {t})}{\| \nabla_ {e _ {t , i}} \ell_ {\mathrm {m a i n}} (e _ {t} , \theta_ {t}) \| ^ {2}}, 0\right), \\ \phi_ {t, i} = \alpha \min  (\| \nabla_ {e _ {t, i}} \ell_ {\operatorname * {m a i n}} (e _ {t}) \| ^ {2}, \ell_ {\operatorname * {m a i n}} (e _ {t})). \\ \end{array}
$$
et al., 2020; Chen et al., 2020) to leverage useful information from high-frequency items to low-frequency items. These methods use 1) different domain knowledge to split the  $S$  and  $\mathcal{T}$  domain, 2) different distance functions to transfer the knowledge, and 3) different training pipelines. We have discussed the first two directions above. For training pipelines, some (e.g. Gong et al., 2018; Li et al., 2020b; Yin et al., 2020) jointly train the head and tail item embeddings while others first train the head item embeddings and then train the rare item embeddings (Kang et al., 2019). Improving long-tail item representations is not only an important topic for discrete input items but also useful for long-tail classification problems (e.g. Tang et al., 2020; Kang et al., 2019; Tan et al., 2020). The Softmax layer can be viewed as an embedding layer in which each dimension stands for an embedding vector for one class.
Lexicographic optimization Lexicographic optimization is a problem traditionally studied in operation research and economics (Ehrgott, 1998; Boggs and Tolle, 1995; Lewis and Gale, 1994). Its applications in deep learning tasks are demonstrated in Gong et al. (2021) with the dynamic barrier algorithm. Similar methods that trade-off multiple objective functions can be found in multi-objective optimization methods, e.g. PCGrad (Yu et al., 2020) and its variants (e.g. Lin et al., 2020; Javaloy and Valera, 2021; Liu et al., 2021a).
We answer the following two key questions through empirical studies: 1) could our target domain selection strategy outperform the simple frequency-based methods? 2) could the lexicographic optimization approach outperform the simple linear
combination method defined in (2)?
We test our method on a variety of tasks, including natural language modeling, named entity recognition, machine translation, and recommendation systems. In all experiments, we set the value of  $\eta$  to be the value of 0.2-quantile of all the similarity scores. It means that we remove  $20\%$  words in the rough target domain which only considered frequency. For algorithm 1, we always set  $\alpha = 2$  and use the default optimizer (e.g., Adam) in each application to set the step size  $\epsilon_{t}$ . In all experiments, we average the results over three trials. We do not report the variance in the table since the variance is small (e.g.  $< 0.1$ ) in many cases. We refer readers to the appendix for full tables with the error bar.
For notation, we denote by 'select' our target domain selection strategy, and 'lexico' the lexicographic optimization algorithm 1, and 'ours' the combination of two techniques.