SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jianghui Wang\*, Yang Chen\*, Xingyu Xie, Cong Fang† Zhouchen Lin†
jianghuiwang.ai@gmail.com, {yangchen, xyxie, fangcong, zlin}@pu.edu.cn
Pre-training has achieved remarkable success when transferred to downstream tasks. In machine learning, we care about not only the good performance of a model but also its behavior under reasonable shifts of condition. The same philosophy holds when pre-training a foundation model. However, the foundation model may not uniformly behave well for a series of related downstream tasks. This happens, for example, when conducting mask recovery regression where the recovery ability or the training instances diverge like pattern features are extracted dominantly on pre-training, but semantic features are also required on a downstream task. This paper considers pre-training a model that guarantees a uniformly good performance over the downstream tasks. We call this goal as downstream-task robustness. Our method first separates the upstream task into several representative ones and applies a simple minimax loss for pre-training. We then design an efficient algorithm to solve the minimax loss and prove its convergence in the convex setting. In the experiments, we show both on large-scale natural language processing and computer vision datasets our method increases the metrics on worse-case downstream tasks. Additionally, some theoretical explanations for why our loss is beneficial are provided. Specifically, we show fewer samples are inherently required for the most challenging downstream task in some cases.
The rapid development of machine learning is promoting a shift in the learning paradigm, where one first trains a very large model, often called a foundation model, with massive data, and then adapts it to desired tasks using much less data. The hope is to obtain a model that serves as an infrastructure and is transferable to a wide range of tasks. Pre-training plays the role of an engine to acquire the foundation model. Typical pre-training methods are to minimize the average expected risks of the upstream tasks. Such pre-trained models can achieve good performance for a lot of downstream tasks but may fail in some hard cases. For example, a vision pre-trained model for animal and plant recognition may work well for typical characteristics species but fail when identifying mimicry animals and plants; a common green mantis can be correctly recognized as a mantis, while an orchid mantis might be falsely classified as an orchid.
In machine learning, one cares about not only the good performance of a model but also its behavior under reasonable shifts of conditions. The same philosophy holds for pre-training a foundation model. To guarantee a uniformly good performance over a series of tasks, one has to consider the robustness of pre-training. We call this downstream-task robustness. The aim is to develop a pre-training method to train the foundation model that admits a good adaptation performance over a series of downstream
tasks. It is crucial to achieving the downstream-task robustness for pre-training: (i) safety is critical for some applications, such as deep learning systems in medicine and finance; (ii) our goal of the foundation model requires reliably good performance on all downstream tasks.
In recent times, popular large-scale models like ChatGPT [46] have also faced safety issues, sparking discussions among various stakeholders [8]. The concerns largely stem from the potential misuse or unintended behavior of the model in real-world applications. Some argue that the vast and diverse knowledge base of these models, coupled with their ability to generate human-like text, could be leveraged for malicious purposes. Others worry about the potential of the model to generate inappropriate or harmful content [47].
Addressing these safety concerns is crucial, particularly in the context of pre-training foundation models that are intended for downstream tasks. An initial way to mitigate these safety issues is to consider downstream-task robustness. By focusing on downstream-task robustness, we expect that the models are resilient to perturbations in the input data, thereby reducing their susceptibility to adversarial attacks or misuse. This approach can help in maintaining consistent performance across a range of tasks and so enhance the safety and reliability of the model.
In recent years, the concept of Distributionally Robust Optimization (DRO) [6, 44, 36] has attracted wide attention among theorists and practitioners. Most DRO frameworks [55, 23, 22] consider training a parameterized model that minimizes the worst-case expectation loss over the data from a family of probability distributions. Downstream-task robustness can be considered a generalization of DRO. The destination of downstream-task robustness is to guarantee good worst-case performance for a series of downstream adaptations.
This paper proposes a pre-training method as a starting point for downstream-task robustness. To take a step forward, our method considers learning several upstream tasks. The choice of how to design the upstream tasks allows us to incorporate prior knowledge of the domain and problems. For example, in language models, we can design upstream tasks by masking different types of words; how we generate such upstream tasks by grouping samples reflects our prior knowledge of the natural language. Then instead of minimizing the average expected risk of the upstream tasks, we minimize the worst-case expected risk. We also introduce a simple but practical algorithm called softmax weighted gradient descent to pre-train the model. We prove the algorithm's convergence in the ideal convex setting and show its effectiveness in our empirical study.
We consider the application of the framework in two experiments — Part-of-Speech masked language models in section 4.1 and multi-modal masked image models in section 4.2. We first pre-train a foundation model with the proposed task-robust pre-training method on multiple upstream tasks generated by different masks and adapt the foundation model for downstream tasks. Compared with the average expected risk minimization, our method achieves better worst-case performance and comparable average performance on all downstream tasks.
We also explain why our framework can benefit downstream-task robustness. Specifically, by simplifying the model-task relationship, we show fewer samples are needed for the hardest downstream task. The key intuition is that proper worse-case training for upstream tasks leads to an initiation close to the solution for the worst-case downstream tasks, thus reducing the downstream burden.
The contributions of our study can be primarily encapsulated within three key areas: (i) The introduction of the concept of task-robust pretraining, a novel theoretical framework that holds significant potential for future research. (ii) The provision of a simple yet efficacious method for task-robust pretraining, accompanied by a comprehensive exposition of its theoretical feasibility. (iii) A series of empirical validations across multiple domains, substantiating the effectiveness of our methodology.
2 Setup and Methodology
Consider a traditional machine learning task where the data  $z \in \mathbb{Z}$  follows an underlying distribution  $P$ . Given a model, a parameter space  $\Theta \subset \mathbb{R}^d$ , a loss function  $\ell : \Theta \times \mathbb{Z} \mapsto \mathbb{R}_+$ , the goal is to find the optimal parameter  $\theta^*$  for the model such that  $\theta^* = \arg \min_{\theta \in \Theta} \mathbb{E}_{z \sim P}[\ell(\theta, z)]$ . The classic empirical risk minimization (ERM) tackles the problem by first collecting i.i.d. training data  $\{z_i\}_{i=1}^N$
from  $P$  and then finding a parameter  $\hat{\theta}_{\mathrm{ERM}}$  via minimizing the empirical risk:
$$
\hat {\theta} _ {\mathrm {E R M}} := \underset {\theta \in \Theta} {\arg \min } \frac {1}{n} \sum_ {i = 1} ^ {N} \ell (\theta , z _ {i}). \tag {1}
$$
In statistical learning theory, it is well-known that under mild conditions (such as the VC dimension of the model is bounded above),  $\hat{\theta}_{\mathrm{ERM}}$  is a good approximation of  $\theta^{*}$  in the sense that with high probability at least  $1 - \delta$  ( $0 < \delta \ll 1$ ),
$$
\mathbb {E} _ {z \sim P} \left[ \ell \left(\hat {\theta} _ {\mathrm {E R M}}, z\right) \right] - \min  _ {\theta \in \Theta} \mathbb {E} _ {z \sim P} [ \ell (\theta , z) ] \leq \epsilon , \tag {2}
$$
when the number of training samples  $N$  is sufficiently large. We simply denote the training data requirement by  $N(\epsilon, \delta)$ .
In machine learning, a foundation model is trained and then adapted for each downstream task. The adaptation process involves initializing a downstream model with the pre-trained foundation model's parameters and then training on the downstream task. Denote the parameter of the foundation model by  $\theta_{\mathrm{foundation}}$ . Let  $\Lambda$  be the downstream task space. The goal is to find an initial parameter  $\theta_{\mathrm{foundation}}$  that enables fast adaptations for each downstream task  $\lambda \in \Lambda$ . Different downstream tasks are characterized by different loss functions with a shared data distribution<sup>3</sup>. We use the foundation model's parameters  $\theta_{\mathrm{foundation}}$  as initialization to find an approximately optimal parameter  $\hat{\theta}_{\lambda,\mathrm{ERM}}$  by ERM. We study the sample complexity required to guarantee a good approximate solution with high probability while considering the effect of initialization. For the task  $\lambda$  and the model initialized by  $\theta_{\mathrm{foundation}}$ , denote the number of samples required to find an  $\epsilon$ -approximately optimal parameter by ERM with probability at least  $1 - \delta$  as  $N_{\lambda}$  ( $\theta_{\mathrm{foundation}}, \epsilon, \delta$ ). The aim is to find the optimal initialization that minimizes the worst-case sample complexity required to find an approximately optimal parameter for all tasks, i.e.,
$$
\theta_ {\text {f o u n d a t i o n}} ^ {*} := \underset {\theta_ {\text {f o u n d a t i o n} \in \Theta}} {\arg \min } \underset {\lambda \in \Lambda} {\max } N _ {\lambda} \left(\theta_ {\text {f o u n d a t i o n}}, \epsilon , \delta\right). \tag {3}
$$
Directly training for the optimal worst-case initialization is generally infeasible or computationally expensive. Pre-training provides a feasible alternative  $\theta_{\mathrm{pre - train}}^*$  for  $\theta_{\mathrm{foundation}}^*$  by training for available surrogate upstream tasks. When the upstream tasks are related to the downstream tasks, the pretrained parameter can lead to lower initial expected risks and accelerate downstream training. For example, if we pre-train a model on generated upstream tasks of reconstructing images corrupted by different masks, the pre-trained model can learn some prior knowledge for general vision tasks; with the pre-trained parameter as the initialization, we can accelerate the training process of downstream vision tasks such as image classification and object detection [13]. Consider there are  $T$  representative upstream tasks. Denote the loss function of the task  $t$  as  $\ell_t$ . A typical choice of the pre-trained parameter is the minimizer of the average expected risk over the  $T$  upstream tasks [43, 18, 29], i.e.,
$$
\theta_ {\text {a v e r a g e}} ^ {*} := \underset {\theta \in \Theta} {\arg \min } \frac {1}{T} \sum_ {t = 1} ^ {T} \mathbb {E} _ {z \sim P} \left[ \ell_ {t} (\theta , z) \right]. \tag {4}
$$
However, minimizing the average expected risk over upstream tasks may neglect extreme cases and lead to limited benefit for some downstream tasks.
To alleviate the aforementioned issue of the average expected risk minimization, we propose to use the minimizer of the worst-case expected risks over the upstream tasks, i.e.,
$$
\theta_ {\max } ^ {*} := \underset {\theta \in \Theta} {\arg \min } \underset {t \in [ T ]} {\max } \mathbb {E} _ {z \sim P} \left[ \ell_ {t} (\theta , z) \right], \tag {5}
$$
as the initial parameter, where  $[m]$  denotes the set  $\{1,\dots ,m\}$ . We show that  $\theta_{\mathrm{max}}^{*}$  is a better choice than  $\theta_{\mathrm{average}}^{*}$  in terms of downstream-task robustness.
3Data distributions vary for different tasks can be modelled by incorporating weighting functions from the data distributions to the loss functions. Consider a task where the data distribution is  $P_{t}$  and  $\frac{dP_t}{dP}(z) > 0$  for all  $z \in \mathrm{supp}(P_t)$ , the expected loss is  $\mathbb{E}_{z \sim P_t}[\ell(\theta, z)]$ . The expected loss can still be rewritten as  $\mathbb{E}_{z \sim P}[\ell_t(\theta, z)]$ , where  $\ell_t(\theta, z) = \frac{dP_t}{dP}(z)\ell(\theta, z)$ .
Recall that our minimax pre-training method is to minimize the worst-case expected risks over the upstream tasks, i.e.,
$$
\underset {\theta \in \Theta} {\min } \max  _ {t \in [ T ]} \mathbb {E} _ {z \sim P} \left[ \ell_ {t} (\theta , z) \right] \tag {6}
$$
There is extensive literature on minimax optimization. The minimax optimization algorithms can be generally classified into two types: (i) minimization for the maximum function  $\max_{t\in [T]}\mathbb{E}_{z\sim P}[\ell_t(\theta ,z)]$  [12, 4, 60, 28], and (ii) direct minimax optimization for the objective [35, 59, 45, 37, 38]. We introduce a new simple optimization algorithm called softmax weighted gradient descent (Algorithm 1) that we find is very practical to pre-train the model. The algorithm can be roughly seen as the first type. It is an adaptation of the classic subgradient descent for minimizing the maximum function to enable its practical use in pre-training. Concretely, in one update, we take a descent step at the current point  $\theta$  along the direction of the gradient weighted by softmax-type weights, i.e.,  $\sum_{t = 1}^{T}w_{\alpha ,t}(\theta)\nabla_{\theta}\mathbb{E}_{z\sim P}[\ell_t(\theta ,z)]$ , where
$$
w _ {\alpha , t} (\theta) := \frac {\exp \left(\alpha \mathbb {E} _ {z \sim P} \left[ \ell_ {t} (\theta , z) \right]\right)}{\sum_ {t ^ {\prime} = 1} ^ {T} \exp \left(\alpha \mathbb {E} _ {z \sim P} \left[ \ell_ {t ^ {\prime}} (\theta , z) \right]\right)}, \tag {7}
$$
and  $\alpha > 0$  is a hyperparameter. (In practice, we use estimations for the expected risks and the gradients on minibatch samples.)
The motivation behind softmax weighted gradient descent is to use the softmax weighted gradient to approximate the subgradient in the classic subgradient descent for the minimax optimization of (6). As the softmax weighted gradient descent algorithm optimizes for the minimax loss directly, it can achieve better worst-case loss than other pre-training methods. One advantage of the softmax approximation is that it avoids the non-differentiability caused by the maximum operator via the softmax approximation, making the algorithm easily implementable for pre-training applications. Also, as the softmax weighted gradient descent includes only single weighted gradient step in each update, it has computational efficiency comparable to gradient descent in deep learning. In contrast, standard minimax algorithms often cost several times gradient oracles in single step. Moreover, our algorithm can be directly combined with commonly-used optimization tricks in deep learning, such as momentum and adaptive learning rates. In our experiments, we observe that the algorithm with a simple implementation achieves better worst-case errors in various real-world tasks than a number of benchmark pre-training algorithms.
Algorithm 1 Softmax Weighted Gradient Descent
Input: Step sizes  $\{\eta_k\}_{k=1}^{K-1}$ , softmax hyperparameters  $\{\alpha_k\}_{k=0}^{K-1}$  and an initial parameter  $\theta_0$ ; for  $k = 1, \dots, K - 1$  do
Compute the softmax weights  $\{w_{\alpha_k,t}(\theta_{k - 1})\}_{t = 1}^T$  as in (7);
Update the parameter as  $\theta_{k}\gets \theta_{k - 1} - \eta_{k}\sum_{t = 1}^{T}w_{\alpha_{k},t}\left(\theta_{k - 1}\right)\nabla_{\theta}\mathbb{E}_{z\sim P}\left[\ell_{t}\left(\theta_{k - 1},z\right)\right]$
For completeness, we also provide some convergence analysis for our proposed algorithm. We consider a relatively basic setting where for all  $t \in [T]$ , the loss function  $\ell(\cdot, z)$  is convex and  $L'$ -Lipschitz continuous for any fixed  $z \in \mathbb{Z}$ . Intuitively, when the hyperparameter  $\alpha_k$  is sufficiently large, the function  $\sum_{t=1}^{T} w_{\alpha_k, t}(\theta_k) \mathbb{E}_{z \sim P}[\ell_t(\theta, z)]$  is a good differentiable approximation for the objective  $\mathbb{E}_{z \sim P}[\ell_t(\theta, z)]$ . Softmax weighted gradient descent can be roughly seen as a remedy for non-differentiability in subgradient descent, at the expense of controllable approximation errors. We show in Theorem 3.1 that Algorithm 1 can achieve a convergence rate  $O\left(\frac{1}{\sqrt{K}}\right)$  if the hyperparameter  $\alpha_k$  is as large as  $\tilde{O}(\sqrt{k})$ . This result is comparable to the standard convergence rate  $O\left(\frac{1}{\sqrt{K}}\right)$  of subgradient descent [10, Chapter 3].
Theorem 3.1. Suppose that for all  $t \in [T]$  the loss function  $\ell_t(\cdot, z)$  is convex,  $L'$ -Lipschitz continuous and bounded by  $B$  for all  $\theta \in \Theta$  and any fixed  $z \in \mathbb{Z}$ . Denote the optimal solution of (6) as  $\theta^*$  and the distance  $\|\theta_0 - \theta^*\|$  as  $R_0$ . If the step size  $\eta_k = \eta = \frac{R_0}{L' \sqrt{K}}$  and the softmax hyperparameter  $\alpha_k \geq \frac{4\sqrt{k + 1}}{R_0 L'} \log \frac{4TB\sqrt{k + 1}}{R_0 L'}$  for all  $k = 0, \dots, K - 1$ , the average  $\bar{\theta}_K$  of the iteration points in
Figure 1: Part-of-Speech Mask BERT. We first sample the datasets from the task space and group them according to different Part-of-Speech types, and then recover the predicted sentence by a BERT encoder. The optimizer selects the most challenging task and then updates the model's weight through a minimax layer.
Algorithm 1, i.e.,  $\bar{\theta}_K = \frac{1}{K}\sum_{k=0}^{K-1}\theta_k$ , satisfies
$$
\max  _ {t \in [ T ]} \mathbb {E} _ {z \sim P} \left[ \ell_ {1} \right] - \min  _ {\theta \in \Theta} \max  _ {t \in [ T ]} \mathbb {E} _ {z \sim P} \left[ \ell_ {2} \right] \leq \frac {2 R _ {0} L ^ {\prime}}{\sqrt {K}}, \tag {8}
$$
where  $\ell_1 = \ell_t\left(\bar{\theta}_K,z\right)$  and  $\ell_{2} = \ell_{t}(\theta ,z)$
Remark 3.2. The convexity assumption on the loss functions is an oversimplification in deep learning. However, for some cases such as the neural tangent kernel [32], the deep neural networks exhibit properties similar to convexity. In our experiments with non-convex models, we also observe that the algorithm behaves well.
Remark 3.3. The above analysis requires increasing softmax hyperparameters  $\{\alpha_{k}\}_{k = 0}^{K - 1}$ , i.e.,  $\alpha_{k} = \tilde{O}\left(\sqrt{k}\right)$  to guarantee the convergence rate. In our experiments, however, we find that constant softmax hyperparameters, or more concretely  $\alpha_{k} = 1$  for all  $k = 0,\dots ,K - 1$ , work well for most problems. We attribute these phenomena to some properties of deep neural networks, which are left for future exploration.
In this section, we subject our methods to rigorous testing through two experiments, each encompassing tasks germane to the fields of Natural Language Processing (NLP) and Computer Vision (CV). A minimalist design approach was adopted for both the models and the tasks to demonstrate the universality of our design across a broad spectrum of model tasks. We extended the functionalities of BERT and MAE, thereby constructing Part-of-Speech Mask BERT (PoS-BERT) and Multi-Modal Mask MAE (MM-MAE), respectively.
4.1 NLP Scenario: Part-of-Speech Mask BERT
4.1.1 Model and Settings
Architectures An overview of the Part-of-Speech Mask BERT model is shown in Figure 1. PoSBERT model first samples the datasets from the task space and groups them according to different Part-of-Speech types. The loss function term is calculated separately for each data group entering the BERT encoder. The loss term with the highest weight is selected to enter the optimizer through a minimax layer. Finally, we run experiments on downstream tasks to compare our minimax task
Table 1: Results on GLUE. The "Averages" are obtained from GLUE leaderboard. F1 scores are reported for QQP and MRPC, , spearman correlations are reported for STS-B, Matthews correlations are reported for CoLA, and accuracy scores are reported for the other tasks.
Model	Task-Balancing	MNLI	QQP	QNLI	SST-2	CoLA	STS-B	MRPC	RTE	Avg.
BERT	-	84.6	71.2	90.5	93.5	52.2	85.8	88.9	66.4	79.6
PoS-BERT	None	84.9	72.6	89.1	90.8	54.4	83.6	88.1	68.2	79.3
Minimax (Ours)	85.6	76.9	88.6	91.3	61.4	84.2	88.2	70.7	81.4 (+1.8)
Figure 2: Multi-Modal Mask MAE: Randomly sampled patches from multiple modalities are projected to tokens. Task-specific decoders reconstruct the masked-out patches by first performing a cross-attention step from queries to the encoded tokens.
balancing learning algorithm with other methods to verify our theory. We follow the common practice to design the feature representations for masked language modeling and next-sentence prediction.
Tasks and Datasets During pre-training, Part-of-Speech Mask BERT has two objectives: masked language modeling and next-sentence prediction. We define 9 task categories that recover different parts of speech-type words on masked language modeling: (1)verb, (2) noun, (3) adjective, (4) determiner, (5) adverb, (6) pronoun, (7) preposition, (8) conjunction, (9) interjection. Following previous work [18, 48, 40], we evaluate our pre-trained models on downstream tasks using the GLUE [64] benchmarks. Downstream tasks we fine-tuned include MNLI, QQP, QNLI, SST-2, CoLA, STS-B, MRPC, and RTE. By swapping out the appropriate inputs and outputs, Part-of-Speech Mask BERT can model many downstream tasks and has a unified way to handle the tasks that involve single text and text pairs. After setting the masks, we utilize Natural Language Toolkitannotate (NLTK) [7] to pseudo-label the words with POS annotations.
4.1.2 Quantitative Result
Quantitative results are presented in Table 1. Our model obtains comparable results on GLUE tasks. PoS-BERT with minimax task-balancing outperforms on half tasks by a substantial margin and obtains a  $1.8\%$  average score improvement over BERT. As for the most challenging training task, CoLA, which has the lowest accuracy on BERT, our method gets a  $9.2\%$  improvement which is a significant boost among downstream tasks. Benefiting from task-robust grouping, on QQP and RTE tasks, our method outperforms the original BERT by 5.7 F1-score and  $4.3\%$  accuracy, respectively. Our method also shows superiority on the challenging downstream task, MNLI, by achieving a  $1.1\%$  higher matched accuracy. As compensation for working better on the more challenging tasks, our method loses little correctness on some downstream tasks that already transfer well. On QNLI, SST-2, STS-B, and MRPC, our results are lower than that of the original BERT model by a margin of  $1\%$  accuracy,  $1.9\%$  accuracy, 1.6 spearson correlation, and 0.7 F1-score. The empirical result shows that, with our task-robust pre-training strategy, the downstream tasks perform on the whole, especially those tricky tasks. We expect future work to further improve these results by incorporating more sophisticated multi-task and grouping procedures.
Table 2: Comparison between task-robust method and other task-balancing methods on ImageNet1K and ImageNetS50 pre-training. Our methodology demonstrates superior performance across the majority of downstream tasks, particularly excelling in the most challenging among them.
Pre-training	Downstream
Data	Epoch	Task-Balancing	Cls. (Top-1 Acc. %)	Semseg. (mIoU)	Depth. (δ1 Acc. %)
ImageNetS50	800	None	92.2	51.9	52.1
Uncertainty [34]	92.6	54.5	70.2
GradNorm [16]	93.0	56.5	65.8
DWA [39]	93.4	52.7	65.7
Minimax(Ours)	91.8	61.5	74.1
ImageNet1K	400	Uncertainty	82.6	48.9	85.2
Minimax(Ours)	82.3	50.1	85.3
1600	Uncertainty	83.3	52.0	86.4
Minimax(Ours)	83.0	53.2	86.8