Title: Task-Robust Pre-Training for Worst-Case Downstream Adaptation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Pre-training

Robustness

Multi-task learning
Abstract: 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.

Task-Robust Pre-Training for Worst-Case Downstream Adaptation

Abstract

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.

1 Introduction

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 ∈ Z follows an underlying distribution P . Given a model, a parameter space Θ ⊂ R d , a loss function ℓ : Θ × Z ↦→ R + , the goal is to find the optimal parameter θ ∗ for the model such that θ ∗ = arg min θ ∈ Θ E z ∼ P [ ℓ ( θ, z )] . The classic empirical risk minimization (ERM) tackles the problem by first collecting i.i.d. training data { z i } N i =1

from P and then finding a parameter ˆ θ ERM via minimizing the empirical risk:

In statistical learning theory, it is well-known that under mild conditions (such as the VC dimension of the model is bounded above), ˆ θ ERM is a good approximation of θ ∗ in the sense that with high probability at least 1 -δ ( 0 < δ ≪ 1 ),

when the number of training samples N is sufficiently large. We simply denote the training data requirement by N ( ϵ, δ ) .

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 θ foundation. Let Λ be the downstream task space. The goal is to find an initial parameter θ foundation that enables fast adaptations for each downstream task λ ∈ Λ . Different downstream tasks are characterized by different loss functions with a shared data distribution 3 . We use the foundation model's parameters θ foundation as initialization to find an approximately optimal parameter ˆ θ λ, 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 λ and the model initialized by θ foundation, denote the number of samples required to find an ϵ -approximately optimal parameter by ERMwith probability at least 1 -δ as N λ ( θ foundation , ϵ, δ ) . 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.,

Directly training for the optimal worst-case initialization is generally infeasible or computationally expensive. Pre-training provides a feasible alternative θ ∗ pre-train for θ ∗ 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 ℓ 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.,

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.,

as the initial parameter, where [ m ] denotes the set { 1 , . . . , m } . We show that θ ∗ max is a better choice than θ ∗ average in terms of downstream-task robustness.

3 Algorithm

Recall that our minimax pre-training method is to minimize the worst-case expected risks over the upstream tasks, i.e.,

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 ∈ [ T ] E z ∼ P [ ℓ t ( θ, 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 θ along the direction of the gradient weighted by softmax-type weights, i.e., ∑ T t =1 w α,t ( θ ) ∇ θ E z ∼ P [ ℓ t ( θ, z )] , where

and α > 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.

For completeness, we also provide some convergence analysis for our proposed algorithm. We consider a relatively basic setting where for all t ∈ [ T ] , the loss function ℓ ( · , z ) is convex and L ′ -Lipschitz continuous for any fixed z ∈ Z Intuitively, when the hyperparameter α k is sufficiently large, the function ∑ T t =1 w α k ,t ( θ k ) E z ∼ P [ ℓ t ( θ, z )] is a good differentiable approximation for the objective E z ∼ P [ ℓ t ( θ, 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 ( 1 √ K ) if the hyperparameter

α k is as large as ˜ O ( √ k ) . This result is comparable to the standard convergence rate O ( 1 √ K ) of subgradient descent [10, Chapter 3].

Theorem 3.1. Suppose that for all t ∈ [ T ] the loss function ℓ t ( · , z ) is convex, L ′ -Lipschitz continuous and bounded by B for all θ ∈ Θ and any fixed z ∈ Z . Denote the optimal solution of (6) as θ ∗ and the distance ∥ θ 0 -θ ∗ ∥ as R 0 . If the step size η k = η = R 0 L ′ √ K and the softmax hyperparameter α k ≥ 4 √ k +1 R 0 L ′ log 4 TB √ k +1 R 0 L ′ for all k = 0 , . . . , K -1 , the average ¯ θ K of the iteration points in

The

book

was written by John ...

The

carrot

has a hole ...

Why can one

do

after MBS ...

They

persume

their audience ...

We went

on

holiday last summer ...

Don't forget to take care

of

baby ...

How do I control my

emotions?

How do you control

your horniness.

opti-

maizer

Duplicate

!

Trans-

former

Encoder

Trans-

former

Encoder

"

#

"

$

"%

The

book

was written by John ...

The

carrot

has a hole ...

Why can one

do

after MBS ...

They

persume

their audience ...

We went

to

holiday last summer ...

Don't forget to take care

of

baby ...

verb.

prep.

QQP tasks

QQP output

pretrained

Mask with POS Annotation

Prediction

&[()*]

&#

&,

&[*-.]

&#

/

&0

/

1

2 #

2 ,

2[*-.]

2 #

/

2 0

/

&[()*]

&#

&,

&[*-.]

&0

/

&#

/

1

2 #

2 ,

2[*-.]

2 #

/

2 0

/

3

Task Space

...

...

...

...

...

...

...

verb.

prep.

...

...

...

...

Algorithm 1, i.e., ¯ θ K = 1 ∑ K -1 θ k

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 { α k } K -1 k =0 , i.e., α k = ˜ O ( √ k ) to guarantee the convergence rate. In our experiments, however, we find that constant softmax hyperparameters, or more concretely α k = 1 for all k = 0 , . . . , K -1 , work well for most problems. We attribute these phenomena to some properties of deep neural networks, which are left for future exploration.

4 Experiments

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

MAE

encoder

Decoder

Decoder

Decoder

Task-Robust

Loss

...

...

...

Original images

RGB

Semseg

Depth

Masked input patches

Masked targets

backward

dirt-merged

grass-merged

SSEL

0.

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.

4.2 CV Scenario: Multi-Modal Mask MAE

4.2.1 Model and Settings

Architectures An overview of MM-MAE is shown in Figure 4.1.2. Multi-Modal Mask MAE contains three encoders, each of which processes different modalities of one image. During pretraining, we try to recover each modality from its masked tokens. Each modality is divided into 16 × 16 patches and then tokenize the patches with modality-dependent linear projections. Projected patches are concatenated into a sequence of tokens and given as input to the same transformer encoder. We also add a global token with 2D sine-cosine positional embeddings. Each task owns a specialized decoder, and the computational cost of decoders scales linearly with the number of tasks.

Tasks and Datasets We select two datasets with different scales, ImageNet1K [53] and ImageNetS50 [25], to conduct unsupervised training upstream to see whether the minimax pre-training method can help the downstream tasks with poor performance. The classification task is evaluated on the validation part of the original dataset, while the semantic segmentation and depth estimation tasks are validated on the NYUv2 dataset [56] by fine-tuning. Due to the absence of a sizeable multi-task dataset with aligned task images [19, 2] we generate pseudo-labels on ImageNet and ImageNetS50 with GPT-3 and Mask2Former.

4.2.2 Quantitative Result

Classification tasks The quantitative results are presented in Table 2. We evaluate our models and baseline by fine-tuning them on the supervised ImageNetS50 and ImageNet1K. We fine-tune our models for 100 epochs and report the top-1 validation accuracy. The result shows a tiny gap between our method and the average method in the classification task. Classification tasks are regarded as the least challenging task category of the three. Cause different downstream tasks have different optimal parameter requirements, this gap is unavoidable. After the pre-training of the model reaches a specific step, the training weight of the classification tasks will continue to decrease.

Semantic segmentation tasks We further evaluate our models on semantic segmentation tasks on the NYUv2 dataset. We report the mean intersection over the union (mIoU) metric. Notice that semantic segmentation is the most challenging task of these downstream transfers. Our method

(a) Input

(b) None

(c) Uncertainty

(d) Gradnorm

(e) DWA

(f) Minimax(Ours)

(g) Ground truth

benefits more than the average loss model from pseudo-labeled modalities as input. In particular, the correctness is improved by 9.6% on ImageNetS50 pre-training. With the progress of model training, our task-robust loss forces the model to improve poorly trained semantic segmentation tasks by increasing the training weight. The following section 5 will explain why a simple strategy can significantly help worst-case downstream tasks.

Depth evaluate tasks For depth estimation, we use NYUv2 . We report δ 1 on the NUYv2 test set, showing the percentage of pixels p with error max { ˆ y p y p , y p ˆ y p } less than 1.25 [20]. According to Table 2, the accuracy is improved by 3.9% on ImageNetS50 pre-training with the help of the downstream-task robustness loss function. The depth estimation task in the same data volume has a higher tolerance for prediction errors per pixel than semantic segmentation. However, it is still more complicated than the classification task, which only predicts the image once. After the classification task is well-trained, the depth estimation task will benefit from our strategy in the subsequent training.

4.3 Qualitative Comparison

Table 2 delineates several strategies designed to equilibrate the contribution of each task during the training of a multi-task network. For a qualitative comparison of these methods, refer to Table 3. We appraise these strategies based on several criteria [61]. An overview of our examination suggests that our approach achieves a synergistic blend of simplicity, efficiency, and effectiveness.

To facilitate a more intuitive comparison of the differences in results throughout the training process, we undertook downstream tasks in semantic segmentation, comparing the performance of various approaches midway through pre-training (400 epoch, ImageNetS50). As depicted in Figure 3, our methodology exhibits superior performance, even under conditions of insufficient training.

5 Explanation

We show why the proposed minimax pre-training method can be more effective than the average expected risk minimization in some cases. We consider a simplification of the model and the task relationship. Such a simplification makes our analysis convenient and intuitive. We assume that for all t ∈ [ T ] , the function ℓ t ( · , z ) is µ -strongly-convex, L -smooth, and L ′ -Lipschitz continuous for any fixed z ∈ Z and the function ℓ t ( · , · ) ≤ B for all θ ∈ Θ and z ∈ Z 4 . Note that pre-training on

irrelevant upstream tasks does little help to the downstream tasks in general. Here, we only discuss the case where the upstream tasks and the downstream tasks are closed related. We ideally assume that the loss functions of the downstream tasks are convex combinations of the loss functions of the upstream tasks, i.e.,

where ∆ T is the ( T -1) -dimensional probability simplex. We further assume that for each task there exists a parameter such that the expected risk of the task is zero, i.e., min θ ∈ Θ E z ∼ P [ ℓ t ( θ, z )] = 0 for all t ∈ [ T ] .

We first show that in the above setting, the proposed minimax optimization pre-training method can guarantee a better worst-case initial expected risk than the minimization method.

Proposition 5.1. Let θ ∗ max and θ ∗ average be the pre-trained parameters obtained by minimizing the maximal expected risk and the average expected risk, respectively. Then for the worst-case expected risks of the downstream tasks, we have

Remark 5.2 . The gap between max λ ∈ Λ E z ∼ P [ ℓ λ ( θ ∗ max , z )] and max λ ∈ Λ E z ∼ P [ ℓ λ ( θ ∗ average , z )] can be large. We provide an example in the appendix, where the ratio between max λ ∈ Λ E z ∼ P [ ℓ λ ( θ ∗ average , z )] and max λ ∈ Λ E z ∼ P [ ℓ λ ( θ ∗ max , z )] is as large as O ( T ) .

We then illustrate that a good initialization can serve as an implicit regularization. We suppose that the downstream tasks are trained with gradient descent. (For stochastic gradients with bounded variances, the analysis below also holds for sufficiently small step sizes, within neglectable approximation errors.) For the downstream task λ , with certain step sizes, the parameters will always be in a subset

where θ 0 is the initial parameter and θ ∗ λ = arg min θ ∈ Θ E z ∼ P [ ℓ λ ( θ, z )] .

Proposition 5.3. Suppose that a function f : R d ↦→ R is µ f -strongly-convex and L f -smooth for all x ∈ R d and x ∗ ∈ arg min x ∈ R d f ( x ) . Let { x k } K -1 k =0 be the sequence generated by gradient descent with a step size η > 0 , i.e., x k = x k -1 -η ∇ f ( x k -1 ) for all k ∈ [ K -1] . If the step size η ≤ 1 L f , then we have

for all k = 0 , 1 , . . . , K -1 .

By Proposition 5.3, we can deem that the downstream task's parameter space is the subset Θ λ ( θ 0 ) .

Consider the worst sample complexity to find an ϵ -approximately optimal parameter by ERM within the parameter space Θ λ ( θ 0 ) for a downstream task λ ∈ Λ .

Theorem 5.4. The worst-case sample complexity max λ ∈ Λ N λ ( θ 0 , ϵ, δ ) with initialization θ 0 satisfies

Theorem 5.4 characterizes the upper bound of the worst-case sample complexity of downstream tasks. If we regard ϵ and δ as constants, the worst-case sample complexity with respect to the initialization θ 0 is O (log max λ ∈ Λ E z ∼ P [ ℓ λ ( θ 0 , z )]) . Combined with (10), Theorem 5.4 demonstrates that the proposed minimax pre-training procedure implies tighter sample complexity than the average minimization pre-training procedure in the worst case. Though the dependency on the worst-case initial expected risk is logarithmic in the upper bound analysis, we find that the initialization can have an evident effect on the generalization of the downstream tasks in practice. We claim that the upper bound for general cases may not be tight for our deep learning applications. Special structures in applications might lead to tighter bounds for generalization errors, which remains for further study.

6 Conclusion

This paper introduces the concept of downstream-task robustness for pre-training, aiming to improve the performance of foundation models across various downstream tasks. As models such as ChatGPT become more prevalent, safety and consistent performance are increasingly important. Our proposed minimax loss for pre-training, validated through extensive experiments, offers a potential solution to enhance the robustness and safety of such models. In the future, we will explore grouping the upstream tasks adaptively. We would say though still in its early stages, the study of downstream-task robustness holds significant promise for the reliable and safe deployment of AI infrastructure.

7 Acknowledgments

C. Fang and Z. Lin were supported by National Key R&D Program of China (2022ZD0160301). C. Fang was also supported by the NSF China (No. 62376008) and Wudao Foundation. Z. Lin was also supported by the NSF China (No. 62276004), the major key project of PCL, China (No. PCL2021A12) an Qualcomm.

References

Sara Atito Ali Ahmed, Muhammad Awais, and Josef Kittler. Sit: Self-supervised vision transformer. Computing Research Repository, 2021.

Roman Bachmann, David Mizrahi, Andrei Atanov, and Amir Zamir. Multimae: Multi-modal multi-task masked autoencoders. arXiv preprint arXiv:2204.01678, 2022.

Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. data2vec: A general framework for self-supervised learning in speech, vision and language. Computing Research Repository, 2022.

Adil M Bagirov, Bülent Karasözen, and Meral Sezer. Discrete gradient method: derivative-free method for nonsmooth optimization. Journal of Optimization Theory and Applications, 2008.

Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image transformers. In International Conference on Learning Representations (ICLR), 2021.

Aharon Ben-Tal, Dick Den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 2013.

Steven Bird, Ewan Klein, and Edward Loper. Natural language processing with Python: analyzing text with the natural language toolkit. "O'Reilly Media, Inc.", 2009.

Ali Borji. A categorical archive of chatgpt failures. arXiv preprint arXiv:2302.03494, 2023.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems (NeurIPS), 2020.

Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 2015.

Nicholas Carlini, Anish Athalye, Nicolas Papernot, Wieland Brendel, Jonas Rauber, Dimitris Tsipras, Ian J. Goodfellow, Aleksander Madry, and Alexey Kurakin. On evaluating adversarial robustness. Computing Research Repository, 2019.

Jacques Antonin Chatelon, Donald W Hearn, and Timothy J Lowe. A subgradient algorithm for certain minimax and minisum problems. Mathematical Programming, 1978.

Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu, Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu, and Wen Gao. Pre-trained image processing transformer. In Conference on Computer Vision and Pattern Recognition (CVPR), 2021.

Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning (ICML), pages 1691-1703. PMLR, 2020.

Tianlong Chen, Sijia Liu, Shiyu Chang, Yu Cheng, Lisa Amini, and Zhangyang Wang. Adversarial robustness: From self-supervised pre-training to fine-tuning. In Conference on Computer Vision and Pattern Recognition (CVPR), pages 699-708, 2020.

Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning (ICML), 2018.

Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified adversarial robustness via randomized smoothing. In International Conference on Machine Learning (ICML), pages 1310-1320. PMLR, 2019.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), 2019.

Carl Doersch and Andrew Zisserman. Multi-task self-supervised visual learning. In International Conference on Computer Vision (ICCV), 2017.

Carl Doersch and Andrew Zisserman. Multi-task self-supervised visual learning. In International Conference on Computer Vision (ICCV), 2017.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.

John C Duchi, Peter W Glynn, and Hongseok Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 2021.

John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics, 2021.

Alaaeldin El-Nouby, Gautier Izacard, Hugo Touvron, Ivan Laptev, Hervé Jégou, and Edouard Grave. Are large-scale datasets necessary for self-supervised pre-training? Computing Research Repository, 2021.

Shang-Hua Gao, Zhong-Yu Li, Ming-Hsuan Yang, Ming-Ming Cheng, Junwei Han, and Philip Torr. Large-scale unsupervised semantic segmentation. arXiv preprint arXiv:2106.03149, 2021.

Golnaz Ghiasi, Yin Cui, Aravind Srinivas, Rui Qian, Tsung-Yi Lin, Ekin D Cubuk, Quoc V Le, and Barret Zoph. Simple copy-paste is a strong data augmentation method for instance segmentation. In Conference on Computer Vision and Pattern Recognition (CVPR), pages 2918-2928, 2021.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 2020.

Warren Hare and Mason Macklem. Derivative-free optimization methods for finite minimax problems. Optimization Methods and Software, 2013.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Conference on Computer Vision and Pattern Recognition (CVPR), pages 16000-16009, 2022.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems (NeurIPS), 29, 2016.

Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European Conference on Computer Vision (ECCV), pages 646-661. Springer, 2016.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems (NeurIPS), 2018.

Ziyu Jiang, Tianlong Chen, Ting Chen, and Zhangyang Wang. Robust pre-training by adversarial contrastive learning. Advances in Neural Information Processing Systems (NeurIPS), 33:1619916210, 2020.

Alex Kendall, Yarin Gal, and Roberto Cipolla. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2018.

Galina M Korpelevich. The extragradient method for finding saddle points and other problems. Matecon, 1976.

Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems (NeurIPS), 2020.

Tianyi Lin, Chi Jin, and Michael Jordan. On gradient descent ascent for nonconvex-concave minimax problems. In International Conference on Machine Learning (ICML). PMLR, 2020.

Tianyi Lin, Chi Jin, and Michael I Jordan. Near-optimal algorithms for minimax optimization. In Conference on Learning Theory. PMLR, 2020.

Shikun Liu, Edward Johns, and Andrew J Davison. End-to-end multi-task learning with attention. In Conference on Computer Vision and Pattern Recognition (CVPR), 2019.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. International Conference on Learning Representations (ICLR), 2019.

Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations (ICLR), 2017.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019.

Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. Advances in Neural Information Processing Systems (NeurIPS), 26, 2013.

Hongseok Namkoong and John C Duchi. Variance-based regularization with convex objectives. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.

Yurii Nesterov and Laura Scrimali. Solving strongly monotone variational and quasi-variational inequalities. Discrete and Continuous Dynamical Systems, 2011.

OpenAI. Gpt-4 technical report, 2023.

Oscar Oviedo-Trespalacios, Amy E Peden, Thomas Cole-Hunter, Arianna Costantini, Milad Haghani, Sage Kelly, Helma Torkamaan, Amina Tariq, James David Albert Newton, Timothy Gallagher, et al. The risks of using chatgpt to obtain common safety-related information and advice. Available at SSRN 4346827, 2023.

Jason Phang, Thibault Févry, and Samuel R. Bowman. Sentence encoders on stilts: Supplementary training on intermediate labeled-data tasks. Computing Research Repository, 2018.

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. OpenAI, 2018.

Alec Radford and Jeffrey Wu. Language models are unsupervised multitask learners. OpenAI, 2019.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 2020.

Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems (NeurIPS), 34:8583-8595, 2021.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision (IJCV), 115(3):211-252, 2015.

Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. International Conference on Learning Representations (ICLR), 2020.

Soroosh Shafieezadeh Abadeh, Peyman M Mohajerin Esfahani, and Daniel Kuhn. Distributionally robust logistic regression. Advances in Neural Information Processing Systems (NeurIPS), 28, 2015.

Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus. Indoor segmentation and support inference from rgbd images. In European Conference on Computer Vision (ECCV), 2012.

Wilson L Taylor. 'cloze procedure': A new tool for measuring readability. Journalism quarterly, 1953.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.

Paul Tseng. On linear convergence of iterative methods for the variational inequality problem. Journal of Computational and Applied Mathematics, 1995.

Paul Tseng. Accelerated proximal gradient methods for convex optimization. Technical report, Technical report, University of Washington, Seattle, 2008., 2008.

Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Marc Proesmans, Dengxin Dai, and Luc Van Gool. Multi-task learning for dense prediction tasks: A survey. IEEE transactions on pattern analysis and machine intelligence, 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science. Cambridge university press, 2018.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Annual Conference on Empirical Methods in Natural Language Processing (EMNLP), 2018.

Chen Wei, Haoqi Fan, Saining Xie, Chao-Yuan Wu, Alan Yuille, and Christoph Feichtenhofer. Masked feature prediction for self-supervised visual pre-training. In Conference on Computer Vision and Pattern Recognition (CVPR), 2022.

Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In Conference on Computer Vision and Pattern Recognition (CVPR), pages 9653-9663, 2022.

Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. Image bert pre-training with online tokenizer. In International Conference on Learning Representations (ICLR), 2021.

Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2008.

A Related Work

Minimax Optimization in Deep Learning. Minimax optimization has wide application in deep learning, e.g., adversarial robustness [11, 17], distributional robustness [22, 36, 54], adversarial generative models [27], imitation learning [68, 30]. In the field of pre-training, previous work has employed minimax optimization for adversarial robustness [15, 33]. Our work proposes a minimax optimization procedure for pre-training but with a different formulation. While the previous works aims at adversarial robustness, our work focus on the worst-case generalization of downstream tasks.

Masked language modeling. Masked language modeling (MLM) was first proposed by [57] as a Cloze task. Adapted as a novel pre-training task, MLM and its autoregressive counterpart , e.g. , BERT [18], GPT [49, 50, 9], and T5 [51], are highly successful methods to deal with NLP problems. MLM first masks out some tokens from the input sentences and then trains the model to retrieve the missing context from the rest of the tokens. These methods have been demonstrated to scale excellently so that various downstream tasks can utilize the pre-trained representations. In particular, BERT is constructed based on the transformer [62] model. After preparing the input samples, an embedding layer and a stack of Transformer layers are followed to conduct the bi-directional semantic modeling. We exploit BERT, the most typical masked language model, as the backbone model to process our ablation experiments.

Masked image modeling. Encouraged by transformers, which have gradually become a primary architecture for generic language understanding, ViT [21] later illusion the potential of adopting a pure transformer in image tasks. To generalize better for vision tasks and motivated by the success of BERT [18] in NLP, many recent works propose various masked image prediction methods for pre-training vision models in a self-supervised way. These methods reconstruct the target such as pixels [1, 14, 21, 24, 29, 66], discrete tokens [5, 67], and (deep) features [3, 65]. Notably, the masked autoencoder (MAE) [29] adopts an asymmetric design to allow the large encoder to operate only on unmasked patches and is followed by a lightweight decoder to reconstruct the complete signal from the latent representation along with mask tokens. MultiMAE [2] leverages the efficiency of the MAE approach and extends it to multi-modal and multitask settings. Based on MultiMAE, we apply our approach to increase the transfer capability for downstream tasks.

B Proofs

B.1 Convergence Rate of Algorithm 1

B.1.1 Proof of Theorem 3.1

For notation simplicity, we define that

Our proof consists of two parts. The first part is to show how well F k ( θ k ) approximates F ( θ k ) with our choice of the softmax hyperparameter α . The second part is to analyze the dynamics of the algorithm and the total convergence rate.

Supplementary Materials

We first show that F k ( θ k ) ≥ F ( θ k ) + R 0 L ′ 2 √ k +1 for α k ≥ 4 √ k +1 R 0 L ′ log 4 TB √ k +1 R 0 L ′ . Define T k,ϵ ( θ ) = { t ∈ [ T ] | f t ( θ k ) ≥ F ( θ k ) -ϵ } . If α k ≥ 1 ϵ k log TB ϵ k for an ϵ k > 0 , we have

The inequality A is due to the definition of T k,ϵ ( θ ) . The inequality B is because ∑ t ̸∈ T k,ϵ k ( θ k ) exp( α k f t ( θ k )) ≤ T exp( α k ( F ( θ k ) -ϵ k )) by the definition of T k,ϵ ( θ ) and there exists t ∗ k ∈ T k,ϵ k ( θ ) such that f t ∗ k = F ( θ k ) , which further implies ∑ t ′ ∈ T k,ϵ k ( θ k ) exp( α k f t ′ ( θ k )) ≥ exp( α k F ( θ k )) . The inequality C is because the hyperparameter α k ≥ 1 ϵ k log TB ϵ k . Specifically, let ϵ k = R 0 L ′ 4 √ k +1 and α k ≥ 4 √ k +1 R 0 L ′ log 4 TB √ k +1 R 0 L ′ correspondingly, we have F k ( θ k ) ≥ F ( θ k ) + R 0 L ′ 2 √ k +1 for all k = 0 , . . . , K -1 .

We then analyze the dynamics of the algorithm and give the total convergence rate. For k = 0 , . . . , K -1 , it holds that

where θ ∗ ∈ arg max θ ∈ Θ F ( θ ) . The inequality A is due to the convexity of F k ( θ ) . The equality B is due to the update steps θ k +1 = θ k -η ∇ F k ( θ k ) in Algorithm 1 and the fact ⟨ a, b ⟩ = 1 2 ( ∥ a ∥ 2 + ∥ b ∥ 2 -∥ a -b ∥ 2 ) .

Plugging the approximation error F k ( θ k ) ≥ F ( θ k ) + R 0 L ′ 2 √ k +1 and the inequality F k ( θ ) ≤ F ( θ ) for k = 0 , . . . , K -1 and all θ ∈ Θ into (14), we have

Taking the average over k = 0 , . . . , K -1 , we further have

The inequality A is due to the fact ∑ K k =1 1 √ k < 2 √ K . The inequality B is because for k = 0 , . . . , K -1 , the function F k ( θ ) is L ′ -Lipschitz continuous, which implies ∥ θ k +1 -θ k ∥ 2 = η 2 ∥∇ F k ( θ k ) ∥ 2 ≤ η 2 L ′ 2 . The equality C is due to our choice for the step sizes, i.e., η k = η = R 0 L ′ √ K for all k = 0 , . . . , K -1 .

By the convexity of F ( θ ) , we have F ( ¯ θ K ) ≤ 1 K ∑ K -1 k =0 F ( θ k ) . Combined with (16), we attain the desired result.

B.2 Analysis for the Minimax Pre-training Method

B.2.1 Proof of Proposition 5.1

By the assumption on the downstream task losses ℓ λ and the task space Λ , the equation

holds for all θ ∈ Θ .

By the definition of θ ∗ max , we have

which is the result to prove.

B.2.2 Proof of Proposition 5.3

Proposition 5.3 is a standard result of gradient descent for strongly-convex and smooth functions. We include the proof here for completeness. By the convexity of f ( x ) and the choice of the step size η we have

which means the objective values are nonincreasing and implies f ( x k ) ≤ f ( x 0 ) for all k ∈ [ K ] . By the µ f -strongly convexity at the point x ∗ , it holds for all x ∈ R d that

Combining (17) and the sequence { f ( x k ) } K k =0 decreasing, we obtain

as desired.

B.2.3 Proof of Theorem 5.4

The following proposition characterizes the sample complexity to find an ϵ -approximately optimal parameter by ERM for a downstream task λ within the parameter space Θ λ ( θ 0 ) . Theorem 5.4 follows directly from Proposition B.1 by considering the worst-case downstream task λ ∈ Λ . The remaining is to prove Proposition B.1.

Proposition B.1. For a given task λ ∈ Λ and a parameter space Θ λ ( θ 0 ) , let the parameter ˆ θ ∗ λ ∈ Θ λ ( θ 0 ) be the minimizer in of the empirical risk for N λ i.i.d. samples { z i } N i =1 from a distribution P , i.e., ˆ θ ∗ λ = arg min θ ∈ Θ λ ( θ 0 ) = 1 N λ ∑ N λ i =1 ℓ λ ( θ, z i ) . The parameter ˆ θ ∗ λ is ϵ -approximately optimal with probability at least 1 -δ if

where E ∗ = E z ∼ P [ ℓ λ ( θ 0 , z )] .

The inequality A is because ˆ f λ ( ˆ θ ∗ λ ) -ˆ f λ ( θ ∗ λ ) ≤ 0 by the definition of ˆ θ ∗ λ .

We derive the upper bound by covering numbers. We only consider Euclidean space for simplicity. Definition B.2 (Covering numbers [63, Chapter 4]) . Consider a subset S ⊂ R d and let ϵ > 0 . A subset N ⊂ S is called an ϵ -net of S if every point in S is within distance ϵ of some points of N , i.e., for all x ∈ S , there exists x 0 ∈ N such that ∥ x -x 0 ∥ ≤ ϵ . The smallest possible cardinality of an ϵ -net of S is called the covering number of S and is denoted C ( S, ϵ ) , i.e.,

Consider an ϵ ′ -net N (Θ λ ( θ 0 ) , ϵ ′ ) of Θ λ ( θ 0 ) where ϵ ′ = ϵ 8 L ′ . By the definition of ϵ ′ -net, we have

Combining (19) and (20), we obtain

We leverage the upper bounds of covering numbers of balls [63, Chapter 4].

Lemma B.3. The covering number of a ball of radius R , denoted as B R , in R d satisfies

By Hoeffording's inequality, for each θ ∈ N (Θ λ ( θ 0 ) , ϵ ′ ) , we have

Plugging (22) and (23) into (21), we have

By (24), when the number of samples N λ satisfies

C Example for Remark 5.2

Consider an extreme example where Θ = R , E z ∼ P [ ℓ 1 ( θ, z )] = A ( θ -1) 2 where 1 < A < T and E z ∼ P [ ℓ t ( θ, z )] = θ 2 for t = 2 , . . . , T . Then we have θ ∗ max = √ A √ A -1 and θ ∗ average = A A + T -1 . It further implies that

If A = T -1 , we have ℓ average ℓ max ≥ 1 4 ( √ T -1 -1 ) 2 , which is as large as O ( T ) .

D Training details

D.1 Part-of-Speech Mask BERT Training Setting

In this work, we denote the number of layers (i.e., Transformer blocks) as L , the hidden size as H , and the number of self-attention heads as A . For comparison purposes, we primarily report results on two models with the same size: PoSBERT BASE and BERT BASE ( L =12, H =768, A =12, Total Parameters=110M). The model is trained with AdamW [42] by setting β 1 = 0.9, β 2 = 0.999, ϵ = 1e-6, and L 2 weight decay of 0.01. The learning rate is warmed up over the first 10K steps to a peak value of 1e-4, then linearly decayed. We duplicate training data ten times to avoid using the same mask for each training instance in every epoch so that each sequence is masked in 10 different ways over the 40 training epochs. Thus, each training sequence was seen with the same mask four times. The hyperparameters for experiments are shown as Table 4 and Table 5.

D.2 Multi-Modal Mask MAE Training Setting

We use Vit-B [21] with a patch size of 16 × 16 pixels as the backbone for our MAE experiments, and estimate the model's performance under different pre-training epochs, i.e., 400 and 1,600 epochs on ImageNet1K and 800 epochs on ImageNetS50. We choose AdamW as the optimizer with a base learning rate of 1e-4 and weight decay of 0.05. We first warm up the learning rate with 40 epochs and then decay it with cosine decay [41]. We set the batch size to 2048 and trained the models using 8 × A100 GPUs with automatic mixed precision enabled. Our data augmentations are straightforward. We randomly crop the images, setting the random scale between 0.2 and 1.0 and the random aspect ratio between 0.75 and 1.33. Afterward, we resize the crops to 224×224 pixels and apply a random horizontal flip with a probability of 0.5. The hyperparameters for experiments are shown as Table 6 and Table 7.

Table 7: Hyperparameters for fine-tuning Multi-Modal Mask MAE on various downtasks. The augmentation strategy LSJ is large scale jittering [26]. And we use drop path [31] in classification and semantic segmentation tasks.

E Limitation

Contemporary pre-training models are consistently enlarging in size. However, due to limitations associated with computational power and the non-disclosure tendency of large-scale models, we were unable to conduct our experimentation directly on ultra-large models such as LLaMA [58], GPT3 [9], and V-MoE [52]. Notwithstanding, we have validated our hypothesis on two commonly encountered domains and model frameworks, thus illustrating the extensive applicability of our proposed methodology.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors proposed a concept called task-robust pre-training, which means pre-training a model that guarantees a uniformly good performance over the downstream tasks. A minimax loss and a corresponding algorithm to solve it is provided with convergence analyses in the convex setting. Experiments on large-scale natural language processing and computer vision datasets demonstrated the performance, and additional theoretical explanations were provided to address the benefits of the proposed method.\n\nSTRENGTHS:\n1. The concept of task-robust pre-training is novel. It can be seen as a generalization of distributionally robust optimization. The softmax weighted gradient descent method is also a new minimax optimization algorithm suited for the new concept. \n2. Theoretical analyses on the convergence rate is provided. \n3. Experiment results are impressive, considering the complexity of the tasks. \n4. The paper is well organized and clearly written.\n\nWEAKNESSES:\nThe convergence rate for the new minimax optimization algorithm is $O(1/\\sqrt{K})$, which is comparable to that of subgradient descent, and it is known that subgradient descent can be unbearably slow in practice. It would be nice if convergence comparisons between existing methods and the new method can be provided, either theoretically or empirically.\n\nQUESTIONS:\n1. In Section 5, under a simplified setting, the loss functions of the downstream tasks are assumed to be convex combinations of the loss functions of the upstream tasks. For more general tasks (other than NLP or CV), how should representative upstream tasks be chosen? Does it depend on heuristics? \n2. Where is the example mentioned on line 279 of the paper?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations in both the main paper and the appendix. There does not appear to be any potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the paper of downstream adaptation from a large pretrained \"foundation\" model with a focus on robutness: Some downstream tasks may be harder to transfer to than others. To tackle this, the paper proposes a pre-training strategy to encourage uniformly good performance across all downstream tasks.\n\nThe key component of the pretraining strategy is the min-max objective of **Equation 5**: Given a set of upstream tasks, rather than minimizing the sum of losses across all tasks during pretraining, Equation 5 proposes to minimize the maximum loss across all tasks (worst-case scenario). To avoid the non-differentiable max operation,  the softmax approximation is used in **Equation 7**: In summary, the model is pretrained using a weighted average of each upstream task loss: the weights form a probability distribution and are defined as the softmax version of the losses for each task (with a temperature hyperparameter $\\alpha$.)\n\nThe proposed pretraining strategy is then evaluated on a NLP and computer vision setting, based on the resulting downstream tasks' performance. It is mainly compared to the standard strategy of optimizing the sum of upstream task losses, as well as some previous task balancing methods.\n\nSTRENGTHS:\n- The paper tackles an interesting problem that is relevant to practical usage: Finetuning one large, pretrained multi-task model to a relevant downstream task is a very common application.\n\n- The paper experiments on two different modalities: text and images\n\n- The paper provides some theoretical justifications (generalization bound) for the proposed training objective\n\nWEAKNESSES:\n- **Weak Comparison to multi-task learning optimization methods:** Table 1 only compares the proposed method to the standard average weighting of the task losses. While more complete, the comparison in Table 2 is limited to only few task balancing methods (uncertainty, Gradnorm, DWA) ; the paper would benefit from comparing to more recent methods, e.g. PCGrad/CaGrad/GradDrop. In addition some papers such as *\"Do Current Multi-Task Optimization Methods in Deep Learning Even Work ?, NeurIPS 2022\"* and *\"In Defense of Unit Scalarization, NeurIPS 2022\"* suggest that simply tuning the weights of each task loss as scalar hyperparameters yields a much stronger baseline than the usual \"average weighted multi-task\" and can even perform on par to other more advanced task balancing methods: This would be a fairer baseline than the \"None\" method presented in experiments. \n\n- **Unclear assumptions on upstream and downstream tasks.** Both Section 5 and the experiments in Section 4.2 seem to rely on the fact that the upstream and downstream tasks are very closely related. I think this assumption should appear much earlier in the text and in clearer terms. For instance *line 107* of the introduction only mention one example of related upstream/downstrean tasks (masked image modeling and classification/detection). However we have clearer / stricter assumptions in the rest of the text: in **Section 4.2**, upstream tasks and downstream tasks are identical. And in **section 5** it is expected that the downstream task losses should be convex combinations of some of the upstream task losses.\n\nQUESTIONS:\n- **Section 4.2** Do I understand correctly that the in the computer vision example, the upstream tasks exactly match the downstream tasks  (just with different input data) ? This seems a bit different than the setting described at the beginning of the paper as it is closer to a standard multi-task scenario; while in the introduction/setting section there does not seem to be any explicit assumption made between the upstream and downstream tasks\n\n\n- **Suggestions on writing** I found some parts of the writing either not well substantiated or unclear on what they want to convey. For instance:\n   - lines 86-103 states classical machine learning results linking expected risk and sample complexity. However this seems only relevant to Section 5, thus seems a bit ouf of place in the introduction \n  -  line 259: \"As depicted in Figure 3, our methodology exhibits superior performance even under conditions of insufficient training\" -> This seems a too strong statement to derive from Figure 3\n\n- **Suggestion:** The paper would benefit from a more structured independent related work section, in particular comparing to multi-task learning literature\n\nLIMITATIONS:\nI did not find very explicitly stated limitations, but the paper does recognize that the proposed generalization bound in Theorem 5.4 may not be applicable/tight enough for deep learning models.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
In order to improve the robustness of the pre-trained model on downstream tasks, the authors propose a simple optimization algorithm, softmax weighted gradient descent, to minimize the worst-case expected risk of upstream tasks.

STRENGTHS:
The idea of the article is simple and easy to understand, the proof is sufficient, and the experimental part is also very sufficient.

WEAKNESSES:
In the experimental part, the method does not seem to show a consistent improvement. For example, as shown in Table 1, although the author mentioned that the model has significantly improved performance on many more challenging tasks, it has worse performance than the previous model on some downstream tasks that have performed well. See Table 2 for the same reason. Could this be improved with some tweaks for consistency?

QUESTIONS:
According to the weekness I mentioned above, my question is whether such a strategy is a trade-off in the performance between the best case and the worst case on the downstream task, and cannot achieve consistency improvement?

LIMITATIONS:
See the questions mentioned before
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces the concept of worst-case adaptation, as each pre-training task may transfer differently to various downstream tasks. In order to address this issue, the authors propose a novel approach using a minmax loss and weighted gradient descent to optimize for the worst-case loss among all pre-training tasks. The proposed approach is supported by detailed theoretical analysis and experiments, which demonstrate its effectiveness in improving performance on a range of downstream tasks.\n\nSTRENGTHS:\n1. Personally I really like the perspective of worst-case downstream adaptation for pre-training models,  and view this as an important problem. The proposed approach of minimizing the worst pre-training task using a proxy and reweighing technique on the gradient is found to be effective and reasonable. The authors' motivation and method are both sound, which suggests that this work could have important implications for improving the performance of pre-trained models on downstream tasks. \n2. Detailed theoretical analysis is provided.\n3. The writing is good and easy to follow.\n\nWEAKNESSES:\n1. Discussion with some multi-task learning methods is needed. \n\n[1] Gradient Surgery for Multi-Task Learning, NeurlPS 20\n\nThey check the gradient direction to make sure the model do well in each task.\n\n[2] Task-customized Masked Autoencoder via Mixture of Cluster-conditional Experts, ICLR 23\n\nThey propose a moe architecture to ease the negative transfer from pre-training to fine-tuning, thus achieving better performance on all downstream datasets.\n\n2. I have some concerns about the way the effectiveness of the proposed method is presented in the experiments. From my perspective, the experiments only show a trade-off compared to the average expected risk, rather than a clear improvement in worst-case downstream adaptation. \nOne suggestion is to collect the worst results of models on many downstream tasks, where the models are pre-trained with each pre-training task only, the average expected risk, and the worst-case risk. By comparing the performance of these models, it would be possible to demonstrate a clear improvement in worst-case downstream adaptation.\n\nQUESTIONS:\n1. For Table 1, does the ‘None’ row mean the original BERT or the PoS-BERT but without your minimax loss? Comparison between /theta_{average}^{\\*} (Equ. 4) and /theta_{max}^{\\*} (your method, Equ. 5) is needed.\n2. Why not comparing several multi-task baselines on the PoS-BERT setting?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}