SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hong Chen $^{1}$ , Xin Wang $^{1,2*}$ , Yuwei Zhou $^{1}$ , Yijian Qin $^{1}$ , Chaoyu Guan $^{1}$ , Wenwu Zhu $^{1,2*}$ $^{1}$ Department of Computer Science and Technology, Tsinghua University  
 $^{2}$ Beijing National Research Center for Information Science and Technology, Tsinghua  
{h-chen20,zhou-yw21,qinyj19,guancy19}@mails.tsinghua.edu.cn  
{xin_wang,wwzhu}@tsinghua.edu.cn
Current auxiliary learning methods mainly adopt the methodology of reweighing losses for the manually collected auxiliary data and tasks. However, these methods heavily rely on domain knowledge during data collection, which may be hardly available in reality. Therefore, current methods will become less effective and even do harm to the primary task when unhelpful auxiliary data and tasks are employed. To tackle the problem, we propose a joint data-task generation framework for auxiliary learning (DTG-AuxL), which can bring benefits to the primary task by generating the new auxiliary data and task in a joint manner. The proposed DTG-AuxL framework contains a joint generator and a bi-level optimization strategy. Specifically, the joint generator contains a feature generator and a label generator, which are designed to be applicable and expressive for various auxiliary learning scenarios. The bi-level optimization strategy optimizes the joint generator and the task learning model, where the joint generator is effectively optimized in the upper level via the implicit gradient from the primary loss and the explicit gradient of our proposed instance regularization, while the task learning model is optimized in the lower level by the generated data and task. Extensive experiments show that our proposed DTG-AuxL framework consistently outperforms existing methods in various auxiliary learning scenarios, particularly when the manually collected auxiliary data and tasks are unhelpful.
Auxiliary learning aims to improve the model generalization ability on the primary task with the help of related auxiliary tasks [1, 2]. This learning paradigm has been widely adopted and has shown its effectiveness in various areas, like image classification [3], recommendation [4] and reinforcement learning [5, 6]. Different auxiliary tasks are often chosen manually according to the primary task, e.g., [7] utilize the task of visual attribute classification to help the fine-grained bird classification and [4] improve the click conversion rate prediction with the help of click-through rate prediction task.
Most existing works utilize the auxiliary information by first reweighing the losses of the auxiliary data and tasks, and then use the sum of the weighted losses together with the primary loss to optimize the task learning model. The weights are employed to balance the primary loss and the auxiliary losses to avoid negative auxiliary transfer, which are tuned with HPO tools [4, 3, 8]. More recent works [9, 6, 10, 1, 11] propose to dynamically weigh different auxiliary losses during training.
However, the existing methods require that there exists beneficial information in the auxiliary data and tasks, so that the beneficial loss terms can be selected through the loss reweighing process. This condition cannot always be satisfied because whether the auxiliary data or task is beneficial depends on many factors including the chosen auxiliary task, the scale of primary task dataset and the selected
*Corresponding Authors.
learning model for the tasks [1, 11], making it difficult to manually collect the beneficial auxiliary data and task via prior knowledge. Therefore, existing approaches may adopt useless auxiliary information and finally do harm to the primary task when the involved auxiliary data and tasks are improperly collected, as observed in [1, 12]. Although [2] propose to generate fine-grained auxiliary classification tasks, this method can be only applied to the classification primary task. Additionally, they only generate the new task without considering new data, while the data-level information is pointed out to be an important factor in auxiliary learning [11].
To address the problem, in this paper, we propose to simultaneously generate the auxiliary data and task for auxiliary learning in a joint manner. However, there are three challenges for the joint generation. First, it is challenging to design a generic framework that can accommodate various tasks with different inputs. This is because the types of data and tasks are quite diversified in different auxiliary learning scenarios, e.g., the primary task of image classification takes visual images as input and outputs categorical labels for classification [7], while the primary task of rating prediction for recommendation takes tabular data as input and outputs numeric labels for regression [4]. Second, it is challenging to develop a generation framework that is expressive enough to produce beneficial data and tasks for the primary task. Finally, to guarantee the jointly generated auxiliary data and task are beneficial to primary task, how to effectively optimize the parameters in the generation framework is a challenging problem as well.
To tackle these challenges, we propose a joint Data-Task Generation framework for Auxiliary Learning (DTG-AuxL), which involves a joint generator and a bi-level optimization strategy. Specifically, the joint generator consists of a feature generator and a label generator, which generates the new auxiliary data and new task based on the existing manually collected data and tasks. The data generation process is conducted in the feature space so that it can accommodate various input types, while the label generator architecture is inspired by the model in recommendation which can accommodate both categorical labels for classification and numeric labels for regression. Moreover, we introduce nonlinear interaction terms in the joint generator, making it more expressive to produce beneficial auxiliary data and tasks. To effectively optimize the joint generator and the task learning model, we propose the bi-level optimization strategy with instance regularization. In the lower optimization, the task learning model is optimized by the generated auxiliary data and task. In the upper optimization for the joint generator, we not only utilize the implicit gradient from the primary loss but also the explicit gradient from the proposed instance regularization, to avoid label generation mode collapse. Extensive experiments show that DTG-AuxL outperforms existing methods in various auxiliary learning scenarios, especially when the manually collected auxiliary data and task are unhelpful to the primary task. We summarize our contributions as follows,
- We propose to simultaneously generate auxiliary data and tasks in a joint manner in auxiliary learning for the first time, to the best of our knowledge.
- We propose DTG-AuxL, a joint data-task generation framework applicable in various auxiliary learning scenarios, containing a joint generator and a bi-level optimization strategy.
- Extensive experiments in various auxiliary learning scenarios demonstrate the superiority of our proposed DTG-AuxL framework. We further analyze when and how DTG-AuxL works to bring performance boost.
Auxiliary Learning The most widely adopted way in auxiliary learning is to combine the loss of the primary task and the losses of the auxiliary tasks in a linear way [3, 7, 4, 5], where the linear weights are tuned manually or with HPO methods. More recent works [9, 6, 10] propose to automatically assign dynamic weights to the auxiliary losses. [9] propose to assign weights to each loss based on the cosine gradient similarity between the primary loss and each auxiliary loss. Later work [10] aims to make the weighted auxiliary gradient sum closest to the gradient of the primary loss. [13, 1] utilize the bi-level optimization strategy to learn the auxiliary weights, where [1] even propose to combine the losses not only limited to a linear form, but also a nonlinear form. [11, 14] point out that only considering the task-level information is not sufficient, so they jointly consider the weights of different tasks and different data samples within the same task through a joint selector. As aforementioned, these reweighing methods will easily fail to bring improvement when the chosen auxiliary tasks and data are unhelpful. There are also works [2, 1] that generate
Optimization
Figure 1: The overall DTG-AuxL framework. In the model design part, the joint generator contains a feature generator and a label generator. The feature generator utilizes the primary and auxiliary features to generate the new auxiliary feature  $f_{g,j}$ , whose label  $\hat{y}_{g,j}$  is generated by the label generator by combining the information of all the auxiliary and primary labels. The optimization part shows that we optimize the task learning model and the joint generator in an alternating bi-level manner.
fine-grained classification auxiliary tasks for the primary classification task. However, they can only be applied to a classification problem. Additionally, they only generate auxiliary tasks without new data, limiting their performance especially when the data of the primary task is inadequate, which is an often encountered scenario in auxiliary learning [9, 1].
Multi-task learning Multi-task learning aims to share information among tasks to improve model performance. However, the goal of multi-task learning is to obtain good performance for all the tasks, while auxiliary learning only focuses on the primary task. Multi-task learning methods can be categorized into three parts [15]: multi-task architecture design [16, 17], multi-task optimization [18, 19] and multi-task relationship learning [20], where the multi-task optimization methods involve techniques for optimizing several losses, like loss reweighing [18] and gradient modulation [21].
The overall DTG-AuxL framework is shown in Figure 1. Next, we give the problem formulation, describe the designs of the joint generator, and present the proposed bi-level optimization strategy.
3.1 Preliminaries and Problem Formulation
In auxiliary learning, we have one primary task  $T_{p}$ , and totally  $K$  auxiliary tasks  $\{T_{ai}\}_{i=1}^{K}$ . Each of these tasks has its corresponding training dataset, including the primary dataset  $D_{p} = \{(x_{p,j}, y_{p,j})\}_{j=1}^{|D_{p}|}$ , and the dataset for each auxiliary task  $D_{ai} = \{(x_{ai,j}, y_{ai,j})\}_{j=1}^{|D_{ai}|}$ , where  $| \cdot |$  denotes the data sample number of the dataset. If the auxiliary tasks share the same input with the primary task, which is a widely encountered and studied scenario in previous works [11, 1, 9, 10], we have  $x_{ai,j} = x_{p,j}$  and  $|D_{ai}| = |D_{p}|$ . Besides the datasets, we have a task learning model parameterized by  $\theta$  which is used to learn all the tasks together. There is also a validation dataset  $D_{v}$  which is used to evaluate the model performance on the primary task. With these notations, the widely adopted training objective of auxiliary learning is:
$$
\mathcal {L} _ {t} (\theta) = \mathcal {L} _ {p} \left(D _ {p}; \theta\right) + \sum_ {i = 1} ^ {K} w _ {i} \cdot \mathcal {L} _ {a i} \left(D _ {a i}; \theta\right), \tag {1}
$$
where  $\mathcal{L}_p$  and  $\mathcal{L}_{ai}$  indicate the loss functions of the primary and each auxiliary task. Current reweighing methods focus on how to decide  $w_{i}$  so that the task learning model  $\theta$  can achieve the best
performance on  $D_v$ . However, as aforementioned, the terms  $\mathcal{L}_{ai}(D_{ai},\theta)$  are defined by the manually collected auxiliary data and tasks, with no guarantee to bring benefits for the primary task. Therefore, we propose to simultaneously generate the new beneficial auxiliary data and task in a joint manner, with the following training objective:
$$
\mathcal {L} _ {t} (\theta , \phi) = \mathcal {L} _ {p} (D _ {p}; \theta) + \sum_ {i = 1} ^ {K} w _ {i} \mathcal {L} _ {a i} (D _ {a i}; \theta) + w _ {g} \mathcal {L} _ {g} (D _ {g} (\phi); \theta), \tag {2}
$$
where the last term is the loss on the generated auxiliary data and task.  $D_g(\phi) = \{x_{g,j},\hat{y}_{g,j}\}_{j = 1}^{|D_g|}$  is the generated dataset which contains the new data together with the new label defining the new task, and  $\phi$  denotes the parameters of the joint generator that are used to generate  $D_{g}$ . Since we only care about the performance of the primary task in auxiliary learning, we keep  $\mathcal{L}_g$  the same loss function as  $\mathcal{L}_p$ , which can be cross entropy or MSE loss, etc. Note that we still keep the original auxiliary loss terms in our training objective, so that it can accommodate the scenario where the existing manually collected auxiliary data and tasks are beneficial. The task weights  $w_{i}$  and  $w_{g}$  are also learnable, which will be optimized together with the generator parameters, and we uniformly denote them all as  $\phi$  for convenience. Next, we discuss the details of the generator and how we optimize  $\phi$  and  $\theta$ .
3.2 Joint Generator Design
The joint generator involves a feature generator to generate new features and a label generator to generate a new task for the new feature, as shown in the Model Design of Figure 1.
3.2.1 Feature Generator
Since we expect that our data generator can tackle different types of data, a neat and elegant solution is to conduct the generation process in the feature space. In the feature space, different types of input data are transformed to vectors, so we can tackle these vectors in a unified way. Specifically, for the input from different tasks  $[x_{p,j}, x_{a1,j}, \dots, x_{aK,j}]$ , the task learning model will first map them to the feature space with an encoder which is generally noted as "backbone", and then use different task-specific heads to tackle each of the tasks. We generate the new features based on the features extracted via the backbone. In another word, the input of the feature generator is a feature list  $[f_{p,j}, f_{a1,j}, \dots, f_{aK,j}]$ , where  $f_{p,j} = f_{\text{backbone}}(x_{p,j}; \theta)$  is a  $d$ -dimension feature. We next describe how the new features are generated.
Feature linear term. A natural way to generate new features is to combine different features with linear masks. We assign an individual feature mask to each of the  $K + 1$  features, i.e., we have a mask list  $[m_p, m_{a1}, \dots, m_{aK}]$ , where each mask is a  $d$ -dimension vector. Then, we denote the subscript set for the task IDs  $\{p, a1, \dots, aK\}$  as  $S$ , and the linear combination is conducted as follows:
$$
f _ {g l, j} = \sum_ {i \in S} \hat {m} _ {i} * f _ {i, j}, \hat {m} _ {i} [ k ] = \frac {e ^ {m _ {i} [ k ]}}{\sum_ {j \in S} e ^ {m _ {j} [ k ]}}, k = 1 \dots d, \tag {3}
$$
where  $gl$  indicates "generated linearly", * means the element-wise multiplication, and  $\hat{m}_i$  is the normalized mask.  $m_i[k]$  is the  $k^{th}$  element of the mask  $m_i$ , and is normalized by softmax with the elements in the same dimension in other masks. This linear combination with normalized mask works as feature selection from all the input features, and then combines them into a new one.
Feature nonlinear term. To make the generated features more expressive, we introduce an MLP (Multi-Layer-Perceptron) to capture the nonlinear feature interaction:
$$
f _ {g n, j} = M L P _ {F} \left(\left[ f _ {p, j}; f _ {a 1, j}; \dots ; f _ {a K, j} \right]\right), \tag {4}
$$
where  $gn$  indicates "generated nonlinearly". All the features are concatenated together and the nonlinear feature interactions are modeled by an MLP, whose output dimension is  $d$ . Finally, the generated feature is the sum of the linear and nonlinear term, i.e.,  $f_{g,j} = f_{gl,j} + f_{gn,j}$ .
3.2.2 Label Generator
The label generator aims to generate a proper label for the feature produced by the feature generator. We use the labels of all the input data as the input of the label generator to make the generated label expressive enough, i.e., the input of the label generator is a list of labels  $[y_{p,j}, y_{a1,j}, \dots, y_{aK,j}]$ .
However, this label list may contain both numeric labels (if the corresponding task is regression) and categorical labels (if the corresponding task is classification). How to utilize different types of label information to generate a new label is the key problem. Inspired by the CTR (click-through rate) model in recommendation [22, 23], which predicts the user's preference towards an item also based on both their numeric and categorical features, we design the label generator with both linear terms and deep nonlinear terms similar to those in the CTR model.
Label linear term. The linear term models the direct and independent relationship between each input label and the generated label. We simply keep the dimension of the generated label the same as the primary label  $y_{p,j}$  and more adaptive ways to choose the dimension can be an interesting future work. Here, we assume that the dimension of  $y_{p,j}$  is  $m$ , where  $m = 1$  if the primary task  $T_{p}$  is regression, or  $m$  equals to the number of the total categories if  $T_{p}$  is classification. If the primary task is regression, the generated label is also a scalar value for regression, while if the primary task is classification, the generated label is an  $m$ -dimension probability distribution vector. Specifically, for each task  $\mathrm{ID} i \in S$ , if  $T_{i}$  is a classification task and it has totally  $d_{i}$  categories, there is an embedding table  $E_{i}$  with dimension  $d_{i} \times m$  for this task. We directly map the label  $y_{i,j}$  in task  $T_{i}$  to its  $m$ -dimension embedding space through the embedding table:
$$
y _ {g l, i, j} = E _ {i} \left(y _ {i, j}\right). \tag {5}
$$
If  $T_{i}$  is a regression task, we follow [23] to map  $y_{i,j}$  to its  $m$ -dimension embedding. Specifically, we also have an embedding table  $E_{i}$  for this task  $T_{i}$ , whose dimension is  $d_{i} \times m$ , where  $d_{i} = H$  is a hyper-parameter which is fixed to 10 in our experiments. Then  $y_{i,j}$  is mapped into the  $m$ -dimension embedding space as follows:
$$
c _ {i, j} = \operatorname {S o f t m a x} \left(\operatorname {L i n e a r} \left(y _ {i, j}\right)\right), c _ {i, j} \in R ^ {H}, y _ {g l, i, j} = \sum_ {k = 1} ^ {H} c _ {i, j} [ k ] E _ {i} [ k ], \tag {6}
$$
where the numeric label  $y_{i,j}$  is linearly transformed to a  $H$ -dimension vector  $c_{i,j}$ , which is used to weigh all the embeddings in the embedding table to obtain the final embedding. Then the final label linear term is obtained by the sum of all the input label embeddings  $y_{gl,j} = \sum_{i\in S}y_{gl,i,j}$ .
Label nonlinear term. In addition to the label linear term capturing how each input label independently influences the final generated label, we also propose a nonlinear term to model the influence of more complex label interactions on the generated label. Specifically, we introduce another group of embedding tables for all the input tasks  $\{EN_{i}\}_{i\in S}$ , where the ways to obtain the embeddings of the categorical and numeric labels are the same as eq. (5) and eq. (6), respectively. The dimension of  $EN_{i}$  is  $d_{i}\times m_{n}$ , where  $d_{i}$  equals to the number of total categories if  $T_{i}$  is classification.  $d_{i}$  equals to  $H$  (i.e., 10) for regression task, and  $m_{n}$  is a hyper-parameter. With these embedding tables, we can map all the input label list  $[y_{p,j},y_{a1,j},\dots ,y_{aK,j}]$  to an embedding list  $[e_{p,j},e_{a1,j},\dots ,e_{aK,j}]$ . Then the nonlinear label interactions are captured through an MLP:
$$
y _ {g n, j} = M L P _ {L} \left(\left[ e _ {p, j}; e _ {a 1, j}; \dots ; e _ {a K, j} \right]\right). \tag {7}
$$
The final generated label for the generated feature  $f_{g,j}$  is the sum of the linear and nonlinear terms,
$$
y _ {g, j} = y _ {g l, j} + y _ {g n, j}. \tag {8}
$$
Label bias term. Besides the linear and nonlinear terms, we also introduce a label bias term that guides the generated label to possess similar semantic meaning to the label from the target primary task,  $y_{p,j}$ . As such, we add  $y_{p,j}$  as the label bias term as follows,
$$
\hat {y} _ {g, j} = \alpha y _ {p, j} + (1 - \alpha) * \operatorname {n o r m} \left(y _ {g, j}\right), \tag {9}
$$
where  $\alpha \in (0,1)$  is a learnable parameter initialized as 0.5 and norm() is the normalization function. If the primary task is classification, norm() will be softmax() and  $y_{p,j}$  will be converted to the one-hot form. If the primary task is regression with range  $(a,b)$ , then norm() is set to be  $(b - a)\text{sigmoid}(\cdot) + a$ . This bias term makes the generated label lie in the same space as the original label, enabling us to more conveniently explore its semantic meaning, which is also verified to improve model performance in our experiments.
3.2.3 Discussion about the Share Input Scenario
The proposed generator can generate the new feature by combining the features from the primary and auxiliary tasks. However, in the most widely studied auxiliary learning scenario [9, 1, 11], all the
tasks share the same input data. We only have the training dataset  $\{(x_{p,j},y_{p,j},y_{a1,j},\dots ,y_{aK,j})\}_{j = 1}^{|D_p|}$ . Therefore, we cannot obtain the auxiliary data  $[x_{a1,j},\dots ,x_{aK,j}]$  from the auxiliary tasks to conduct the feature generation. To tackle this problem, we obtain the new auxiliary data by randomly sampling another data sample in the dataset, which is easy to implement with randomly shuffling the training batch to match another sample for the original sample. Specifically, for a sample  $(x_{p,j},y_{p,j},y_{a1,j},\dots ,y_{aK,j})$ , we consider another sample  $(x_{p,j2},y_{p,j2},y_{a1,j2},\dots ,y_{aK,j2})$  in the dataset to provide new auxiliary data information. The input of the feature generator is  $[f_{p,j},f_{p,j2}]$ , which are the features extracted by the backbone with  $[x_{p,j},x_{p,j2}]$  as input. When generating new labels, we need to combine all the existing labels of the two samples, i.e., the input of the label generator is  $[y_{p,j},y_{a1,j},\dots ,y_{aK,j},y_{p,j2},y_{a1,j2},\dots ,y_{aK,j2}]$ , and then the label generation process is the same as before. Note that there are two small details: i) During the label embedding process,  $y_{ai,j}$  and  $y_{ai,j2}$  share the same embedding table, because they both belong to the same task. ii) For the bias term in eq. (9), we now have two labels from the primary task,  $y_{p,j}$  and  $y_{p,j2}$ , and then eq. (9) can be adjusted to:
$$
\hat {y} _ {g, j} = \alpha \left(\beta y _ {p, j} + (1 - \beta) y _ {p, j 2}\right) + (1 - \alpha) * \operatorname {n o r m} \left(y _ {g, j}\right), \tag {10}
$$
where  $\beta \in (0,1)$  is a learnable parameter of the generator. Till now, we have described the designs of the joint generator, which is applicable in various auxiliary learning scenarios. However, the generator has several parameters to be optimized, like the masks and MLPs. We denote all the learnable parameters in the generator together with the task weights  $w_{i}$  in eq. (2) as  $\phi$ . Next, we will elaborate how we optimize the task learning model parameters  $\theta$  and the generator parameters  $\phi$ .
3.3 Optimization Strategy
Bi-level optimization. In our whole framework, the task learning model parameters  $\theta$  are expected to minimize the loss of all the selected and generated tasks, while the generator parameters  $\phi$  aim to make  $\theta$  achieve the best performance on the primary task. These two different goals give rise to the following bi-level optimization problem:
$$
\phi^ {*} = \arg \min  _ {\phi} \mathcal {L} _ {p} (\theta^ {*} (\phi); D _ {v}),
$$
$$
s. t. \theta^ {*} (\phi) = \arg \min  _ {\theta} \mathcal {L} _ {t} (\theta , \phi), \tag {11}
$$