Task-Robust Pre-Training for Worst-Case Downstream Adaptation
Authors: Jianghui Wang
Yang Chen
Xingyu Xie
Cong Fang
Zhouchen Lin
Venue: NeurIPS 2023 poster (2023)
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
Jianghui Wang,' Yang Chen", Xingyu Xie, Cong Fang! Zhouchen Lin
School Intelligence Science and Technology, Peking University
jianghuiwang.ain gmail.com, (yangchen, xyxie, fangcong, zlin) @pku.edu.cn
Abstract
Pre-training has achieved remarkable success when transferred downstream tasks. In machine learning, we care about not only the good performance of model but also its behavior under reasonable shifts of condition. The same philosophy holds when pre-training foundation model. However, the foundation model may not uniformly behave well for 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 pre-training, but semantic features are also required on a downstream task. This paper considers pre- vre-training a model that guarantees a uniformly good performance the downstream tasks. We call this goal as downstream--tn taak robustness. Our method first separates the upstream task into several representative ones and applies simple minimax loss for pre-training. We then design efficient algorithm to solve the minimax loss and prove its convergence in the convex setting. In the experiments, 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.
show both on large-scale natural language processing
Introduction
The rapid development of machine learning is promoting 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 desired tasks using much less data. The hope is to obtain model that serves as infrastructure and is transferable to wide range of tasks. Pre-training plays the role of an engine to acquire the foundation model. Typical pre-training methods upstream tasks. Such pre-trained models can achieve good performance for lot of downstream tasks but may fail in some hard cases. For example, vision pre-trained model for animal and plant recognition may work well for typical characteristics species but fail when identifying mimicry animals and plants; common green mantis can be correctly recognized as a mantis, while an orchid mantis might be falsely classified as an orchid.
minimize the average expected risks of the
In machine learning, one cares about not only the good performance of model but also its behavior under reasonable shifts of conditions. The same philosophy holds for pre-training foundation model. To guarantee uniformly good performance over series of tasks, of pre-training. We call this downstream-lask robustness. The aim to develop a pre-training method train the foundation model that admits good adaptation performance over a series of downstream
has consider the robustness
"Equal contribution.
t Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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, 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 range of tasks and so enhance the safety and reliability of the model.
In recent years, the concept of Distributionally Robust Optimization (DRO) wide attention among theorists and practitioners. Most DRO frameworks training a parameterized model that minimizes the worst-case expectation loss over the data from 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.
66
44
36
has attracted
55
23
22
consider
This paper proposes pre-training method as starting point for downstream- -task robustness. To take 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, language models, we can design upstream tasks by masking different types of words; how 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 simple but practical algorithm called softmax weighted gradient descent to pre-train the model. We prove the algorithm's 's convergence in the ideal convex setting and show its effectiveness in our empirical study.
We consider the application of the framework two experiments Part-of- Speefh masked language models in section [.1 and multi-modal masked image models in section 4.2 We first pre-train 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 all downstream tasks.
We also explain why our framework can benefit downstream-ke task robustness. Specifically, by simpli- fying the model-task relationship, we show fewer samples are needed for the hardest downstream task. The key intuition 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 introduc- tion of the concept of task-robust pretraining, a novel theoretical framework that holds significant potential for future research. (ii) The provision of 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 traditional machine learning task where the data E & follows an underlying distribution P. Given a model, a parameter space θ C Rd, loss function & x Z + R4, the goal is to find the optimal parameter 0* for the model such that 0* = arg mingee Ez-P [e (0) The classic empirical risk minimization (ERM) tackles the problem by first collecting i.i.d. training data {

---

from P and then finding parameter θERM via minimizing the empirical risk:
BRRM = arg min Σ L (0, 2)) lee n1
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 0* in the sense that with high probability least 1 5 < 8 < 1),
EzeP [ (PERM, =)]
min AE
Ez-P [e (0, S,
(2)
when the number of training samples N sufficiently large. We simply denote the training data requirement by (e, 6).
In machine learning, foundation model trained and then adapted for each downstream task. The adaptation process involves initializing downstream model with the pre-trained foundation model's s parameters and then training on the downstream task. Denote the parameter of the foundation model by O Let A be the downstream task space. The goal is to find an initial parameter Ofoundaiion that enables fast adaptations for each downstream task À E A. Different downstream tasks are characterized by different loss functions with a shared data distribution | We use the foundation model's parameters Oroundation as initialization to find an approximately optimal parameter FX,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 Ooonndation denote the number of samples required to find an e-approximately optimal parameter by ERM with probability at least 1 - 5 as Nx (Froundation c, 8). 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.,
0 foundation ==
arg min max N2 (froundation e, 6) @ousdaice XEA
(3)
Directly training for the optimal worst-case initialization is generally infeasible or computationally expensive. Pre-training provides a feasible alternative Opre for 0% (fiundation by training for available surrogate upstream tasks. When the upstream tasks are related to the downstream tasks, the pre- trained 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- rreitriined 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 / as l,- typical choice of the pre-trained parameter is the minimizer of the average expected risk over the T upstream tasks [3 [8 29] i.e.,
T
0 averago == arg min
Σ T
1
B The (0,
θgo
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.,
0 m
arg min max Ex [&g (),
(5)
Bee
[]
as the initial parameter, where [m] denotes the set {1....,m, .m}. We show that & is a better choice than 6 @ivergee in terms of downstream-tan task robustness.
Data distributions vary for different tasks can be modelled by incorporating weighting functions from the data distributions to the loss functions. Consider task where the data distribution P. and dP (z) > 0 for all z E supp(Pt). the expected loss is E.-P, [e(0, =)]. The expected loss can still rewritten as E [0, dis (z)E(0, z).
where (0, 2)
3

---

3 Algorithm
Recall that our minimax pre-training method is to minimize the worst-case expected risks over the upstream tasks, i.e.,
min max Ex>P [C (6,
(6)
REB DE[2]
There is extensive literature on minimax optimization. The minimax optimization algo- rithms can be generally classified into two types: (i) minimization for the maximum function P Le (0, 2)] [12 4 60, 28), and (ii) direct minimax optimization for the objective 37 38]. We introduce a new simple optimization algorithm called softmax weighted gradient descent (Algorithm I that we find is very practical to pre-train the model. The algorithm can be roughly seen as the first type. 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 -tppe
max 1 E
35
59
weights, i.c., Ei= Wo,t (0) VolLa~ [6e (0, 2)], where
exp (aE.. 16, =))
Wa,t (0) = E exp (aEza P (er (0,
and a > is a hyperparameter. (In practice, we use estimations for the expected risks and the gradients minibatch samples.)
The motivation behind softmax weighted gradient descent is use the softmax weighted gradient to approximate the subgradient in the classic subgradient descent for the minimax optimization of ( As the softmax weighted gradient descent algorithm optimizes for the minimax loss directly, can achieve better worst-case loss than other pre-training methods. One advantage of the softmax approximation that 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 each update, it has computational efficiency comparable to gradient descent in deep learning. In contrast, standard minimax algorithms often cost several times gradient oracles 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-wor world tasks than number of benchmark pre- vve-training algorithms.
Algorithm 1 Softmax Weighted Gradient Descent
K-1 softmax hyperparameters {ax ( and an initial parameter 0o: Input: Step sizes {ıı J for k I,.... K do
Compute the softmax weights (Waent (0k- [ as in 7
Update the parameter 02 0--1- 7i Ei= Waz, e (0k-1) VoE. (ee (fk-1, end for
For completeness, we also provide some convergence analysis for our proposed algorithm. We consider relatively basic setting where for all E [T), the loss function E(., =) is convex and L'-Lipschitz continuous for any fixed E Z Intuitively, when the hyperparameter og. is sufficiently
large, the function I Wag.t (6x) Eg-r [&, (0,=)] is good differentiable approximation for the
objective Ez-p [e (6, =)] Softmax weighted gradient descent can be roughly seen as a remedy for non-differentiabiity in subgradient descent, at the expense of controllable approximation errors. We
show in Theorem 3.1 that Algorithm can achieve convergence rate O
if the hyperparameter
ag is large Ö(Vk), This result is comparable to the standard convergence rate O subgradient descent
(+}) of
100 Chapter 3].
Theorem 3.1. Suppose that for all E [T] the loss function Ex(:, convex, L' L'pishhtittz continuous and bounded by B for all 0 E ④ and any fixed z E Z. Denote the optimal solution of (6) 0* and the distance |10g - -6 as Rg- If the step size ak = L'R and the softmax hyperparameter
4vk+1 log 4TB/RH1 for all k O.... K the average 0g of the iteration points in HoL
ak 2
Ho

---

Mask with POS Annotation
Prediction
noum
nown.
The book was written by John
The fook was written by John
hole
The carrot has hole
verd
Why can one do after MBS
former Encoder
Why can one de after MBS persume their audience
They persume
prea.
prep.
We went holiday last
Q uombded
We went to holiday last summer Don't forget take care of baby...
Don't forget take care of baby
CLP How da control my emotions?
QCIP output
R
Trans-
How do control
former Encoder
T2
Tispl
Duplicate
acc1
your horniness.
Task Space
Part-of-Spece Mask BERT. We first sample the datasets from the task space and group them according to different Part-of- -fe-eehhh 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.
Figure 1:
Algorithm
1
ie., 0x K EK- OR. satisfies
2R
max RguP p [61] T]
min max Rgg [f2] S
(8)
lE tE(T)
VK
where f = l, (0x, and 82 = l (0, z).
Remark 3.2. The convexity assumption on the loss functions an oversinplification in deep learning. However, for some cases such 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 {a=] i.c., ag =
Ô Vk to guarantee the convergence rate. our experiments, however, we find that constant softmax hyperparameters,
concretely og for all k O,..., problems. We attribute these phenomena to some properties of deep neural networks, which are left for future exploration.
, K 1, work well for most
Experiments
In this section, we subject our methods to rigorous testing through two experiments, each encompass- ing tasks germane to the fields of Natural Language Processing (NIP) and Computer Vision (CV). A minimalist design approach was adopted for both the models and the tasks to demonstrate the universality 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-Speede Mask BERT model is shown in Figure I PoS- BERT 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 BERT encoder. The loss term with the highest weight selected to enter the optimizer through minimax layer. Finally,
experiments downstream tasks to compare minimax task

---

Table 1: Results GLUE. The "Averages" obtained from GLUE leaderboard. Fl scores
reported for QQP and MRPC. spearman correlations are reported for STS-B, B, Matthews correlations
reported for CoLA, and accuracy scores reported for the other tasks.
Model
Task- Ballncin
MNLI
QOP
ONLI
SST-2
CoLA
SIS-B
MRPC
RTE
Avg.
BERT
84.6
71.2
90.5
93.5
S2.2
85.8
88.9
66.4
79.6
None
84,9
72.6
89.1
90.8
S4.4
83.6
88.1
68.2
79.3
PoS-BERT
Minimax (Ours)
85.6
76.9
88.6
91.3
61.4
84.2
88.2
70.7
81.4 (+1.8
Original image
Masked iepur pathhe
Masked
RGB
Semseg
MAE encoder
Telk Bohut
Figure 2: Multi-Modal Mask MAE: Randomly sampled patches from multiple modalities are projected tokens. Task-specific decoders reconstruct the masked-out patches by first performing cross-attention step from queries to the encoded tokens.
balancing learning algorithm with other methods 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 -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 fine-tuned include MNLI, QQP, QNLI, SST-2, CoLA,, STS-B, MRPC, and RTE. By swapping out the appropriate inputs and outputs, -fr-off-Seeec Mask BERT model many downstream tasks and has unified way handle the tasks that involve single text and text pairs. After setting the masks, we utilize Natural Language Toolkitannotate (NLTK) pseudo-label the words with POS annotations.
4.1.2 Quantitative Result
Quantitative results are presented in Table Our model obtains comparable results on GLUE tasks. PoS-BERT with minimax task-balancing outperforms on half tasks by substantial margin and obtains 1.8% average CoLA, which has the lowest accuracy on BERT, our method gets 9.2% improvement which is significant boost among downstream tasks, Benefiting from task-robust grouping, on QQP and RTE tasks, our method outperforms the original BERT by 5.7 Fl-score and 4.3% accuracy, respectively, Our method also shows superiority on the challenging downstream task. MNLI, by achieving 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 margin of 1 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.
improvement over BERT. As for the most challenging training task,

---

Table 2: Comparison between task-robust method and other task-balancing methods ImageNetIK and ImageNetS50 pre-trainin treininn. 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. (mloU) Depth. (6) Acc. %)
None
92.2
51.9
52.1
Uncertainty
92.6
54.5
70.2
ImageNetS50
GradNorm
93.0
56.5
65.8
DWA 39
93.4
52.7
65.7
Minimax(QUrs)
91.8
61.5
74.1
Uncertainty
82.6
48.9
85.2
Minimax(Ours)
82.3
50.1
85.3
ImageNetIK
Uncertainty
83.3
52.0
86.4
1600
Minimax(Qurs)
83.0
53,2
86.8
Table 3: qualitative comparison between task balancing techniques. T representing the computation cost when no additional task- ballccinnn techniques are employed.
Balance
Balance
Grads Required
No Extra
Method
Magnitude
Learning
Tuning
FLOPs
Motivation
None
Uncertainty
27'
Homoscedastic uncertainty
Gradnorm
V
47
Balunce learning and magnitudes
DWA Minimax (Ours)
37
Balance learning
v
27'
Task robust
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 pre- training, we try to recover each modality from its masked tokens. Each modality is divided into 16x 16 patches and then tokenize the patches with modality-dependerr linear projections. Projected patches are concatenated into a sequence of tokens and given as input to the same transformer encoder. We also add 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 datasets with different scales, ImageNetIK 533 and Ima- geNetS50 (55. to conduct unsupervised training upstream to see whether the minimax pre-training method can help the downstream tasks with poor performance. The classification task 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 the absence of sizeable multi-task dataset with aligned task images [19] 2 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 ImageNet K. We fine-tune our models for 100 epochs and report the top-1 validation accuracy. The result shows 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 specific step, the training weight of the classification tasks will continue decrease.
Semantic segmentation tasks on the NYUv2 dataset. report the mean intersection over the union (mloU) metric. Notice that semantic segmentation is the most challenging task of these downstream transfers. Our method
We further evaluate our models on semantic segmentation tasks
7

---

(N) Inpot
(b) Nome
(c) Uncertainty
(d) Grddrom
(e) DWA
(1 Minimax(Oun)
(2] Ground Irush
Figure 3: Comparative intermediate results. The first and final columns represent the image input and ground truth, respectively, while the intermediary images depict the intermediate results yielded by various task-balancing methodologies. Distinet colors correspond the prediction of different objects. Our approach ensures robustness in downstream tasks.
benefits more than the average loss model from pseudo-labeled modalities 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 improve poorly trained semantic segmentation tasks by increasing the training weight. The following section 5 will explain why simple strategy can significantly help worst-case downstream tasks.
Depth evaluate tasks For depth estimation, we use NYUv2. We report 61 on the NUYv2 test set. showing the percentage of pixels p with error max the less than 1.25 (20]. According to Table 2 My My the accuracy is improved by 3.9% 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, 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 delineates several strategies designed to equilibrate the contribution of each task during the training multi-task network. For qualitative comparison of these methods, refer to Table We appraise these strategies based several criteria [61]. An overview of our examination suggests that our approach achieves 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 semantic segmentation, comparing the performance of various approaches midway through pre- vre-trainig (400 epoch, ImageNetS55). As depicted in Figure 3 our methodology exhibits superior performance, even under conditions of insufficient training.
Explanation
We show why the proposed minimax pre-training method be effective than the average expected risk minimization in some cases. We consider a simplification of the model and the task relationship. Such simplification makes our analysis convenient and intuitive. We assume that for all t E [T',, the function 4(+,²) is j-strongly-convex, L-smooth, and L'-Lipschitz continuous for any fixed E Z and the function (((--) < B for all O and
Note that pre-training on
'In fact, it is sufficient to assume that the expectation E, [.((((,2)) is p-strongly-conve L-smooth, and L' Lipschittz continuous.

---

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.,
T Σ Mcl for all X E A = AT,
Ex
where Ar is the (T' -)-diemnsia probability simplex. We further assume that for each task there exists parameter such that the expected risk of the task is zero, i.e., mingee Ez-p 1c+ (0, [)]=]] - 0 for all E [T].
We first show that in the above setting, the proposed minimax optimization pre-training method can guarantee better worst-case initial expected risk than the minimization method.
Proposition 5.1. Let mma and @ avergge 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
max Esx (0 mare 2)] S max ExP [ (2xx((()eeee (Oaeengge =)]. AEA
(10)
XEA
Remark 5.2. The gap between max EzaP [Ex (0max: =)] and maxxea EzEp [Ex (Overage =)]
can be large. We provide an example in the appendix, where the ratio between maxxeA B [ex (fiverage) 2)] and maxxed EgEP [ex (0mm.))) 2)] 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 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 X, with certain step sizes, the parameters will always be in subset
2 E p [ex (0o. =)]}
O (x()) 0 e 1 |10-0
(11)
where 6g is the initial parameter and 8% = arg mingee Ex~P e (6, 2)].
Proposition 5.3. Suppose that a function f R RR is μ strronnglyy-yconee and -mmoohhhh for all Rd and x² E arg min (xx) Let {xx] )= k= be the sequence generated by gradient descent
with step size 1j >> i.e,, Tik -k-1 - V (x=- for all k E [K-1] 1). If the step size ni S I.
then we have
2 | - x112 (f(zo) - f(x)))
(12)
for all k = -011..., K-1. 1.
By Proposition 5.3 we can deem that the downstream task's 's parameter space is the subset (8)))
Consider the worst sample complexity to find an e-approximately optimal parameter by ERM within the parameter space @x(O)) for downstream task A.
Theorem 5.4. The worst-case sample complexity maxλсл NA (0o, e, 8) with initialization Og satisfies
8dB² 2 1+
16L' V
2
8
2 log
max N₂ (fo, c,) S xEA
log
max E (ex (0o. =)1 AEA
+
(13)
2 8
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 initializa- tion Og is O (log max E--P [E> (fo. 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- initial expected risk is logarithmic 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.

---

Conclusion
This paper introduces the concept of downstream-tan 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-tast robustness holds significant promise for the reliable and safe deployment of Al infrastructure.
Acknowledgments
C. Fang and Z. Lin were supported by National Key R&D Program of China (2022200660301, 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. PCI,2011 A12) an Qualcomm.
10

---

References
[1] Sara Atito Ali Ahmed, Muhammad Awais, and Josef Kittler. Sit: Self-supervised vision transformer. Computing Research Repository, 2021.
[2] Roman Bachmann, David Mizrahi, Andrei Atanov, and Amir Zamir. Multimae: Multi-modal multi-task masked autoencoders. arXiv preprint arXiv:2204001678, 2022.
[3] Alexei Baevski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. data2vec: A general framework for sulf-supe learning in speech, vision and language. Computing Research Repository, 2022.
[4] Adil M Bagirov, Bülent Karasözen, and Meral Sezer. Discrete gradient method: derivative-fiee method for nonsmooth optimization. Journal of Optimization Theory and Applications, 2008
[5] Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image trans- formers. In International Conference on Learning Representations (ICLR), 2021.
[6] Aharon Ben- 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.
[7] Steven Bird, Ewan Klein, and Edward Loper. Natural language processing with Python: analyzing text with the natural language toolkit. "O'Reilly Media, Inc.", ", 2009.
[8] Ali Borji. A categorical archive of chatgpt failures. arXiv preprint arXiv arXiv:2302.03494 2023.
[9] 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.
[10] Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning. 2015.
[11] 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.
[12] Jacques Antonin Chatelon, Donald W Hearn, and Timothy J Lowe. A subgradient algorithm for certain minimax and minisum problems, Mathematical Programming, 1978.
[13] 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.
[14] 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.
[15] Tianlong Chen, Sijia Liu, Shiyu Chang, Yu Cheng, Lisa Amini, and Zhangyang Wang. Adver- sarial robustness: From self-supervised pre-training to fine-tuning. In Conference on Computer Vision and Pattern Recognition (CVPR), pages 699-708, 2020.
[16] Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gra- dient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning (ICML), 2018.
[17] 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.
[18] 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 Linguisties: Human Language Technologies (NAACL.HHLT), 2019.

---

[19] Carl Doersch and Andrew Zisserman. Multi-task self supervvide visual learning. In International Conference on Computer Vision (ICCV), 2017.
[20] Carl Doersch and Andrew Zisserman. Multi-task self-supervised visual learning. In International Conference on Computer Vision (ICCV), 2017.
[21] 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.
[22] John C Duchi, Peter W Glynn, and Hongseok Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 2021.
[23] John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statisties, 2021.
[24] Alaaeldin El-Nouby, Gautier Izacard, Hugo Touvron, Ivan Laptev, Hervé Jégou, and Edouard Grave. Are large-scale datasets necessary for self-supe pre-training? Computing Research Repository, 2021.
[25] Shang-Hur Gao, Zhong- Yu Li, Ming-Hsuan Yang, Ming- Ming Cheng, Junwei Han, and Philip Torr. Large-scale unsupervised semantic segmentation. arXiv preprint 2Xvi:2006.00144 2021.
[26] Golnaz Ghiasi, Yin Cui, Aravind Srinivas, Rui Qian, Tsung- Yi Lin, Ekin D Cubuk, Quoe 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.
[27] 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.
[28) Warren Hare and Mason Macklem. Derivative- free optimization methods for finite minimax problems. Optimization Methods and Software, 2013.
[29] 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.
[30] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems (NeurIPS), 29, 2016.
[31] 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.
[32] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems (NeurlPS). 2018.
[33] Ziyu Jiang, Tianlong Chen, Ting Chen, and Zhangyang Wang. Robust pre-training by adversarial contrastive learning. Advances in Neural Information Processing Systems (NeurIPS). 33:16199- 16210, 2020.
[34] 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
[35] Galina M Korpelevich. The extragradient method for finding saddle points and other problems. Matecon, 1976.
[36] 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.
12

---

[37] Tianyi Lin, Chi Jin, and Michael Jordan. On gradient descent ascent for nonconvex-comee concave minimax problems. In International Conference on Machine Learning (ICML), PMLR, 2020.
[38] Tianyi Lin, Chi Jin, and Michael Jordan. Near-optimal algorithms for minimax optimization. In Conference on Learning Theory. PMLR, 2020.
[39] 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.
[40] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. International Conference on Learning Representations (ICLR), 2019.
[41] Uya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations (ICLR), 2017.
[42] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019.
[43] Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg S Corrado, and Jeff Dean. Distributed repre- sentations of words and phrases and their compositionality. Advances in Neural Information Processing Systems (NeurIPS), 26, 2013.
[44] Hongseok Namkoong and John C Duchi. Variance-based regularization with convex objectives. Advances in Neural Information Processing Systems (NeurIPS), 30, 2017.
[45] Yurii Nesterov and Laura Scrimali. Solving strongly monotone variational and quasi- variational inequalities. Discrete and Continuous Dynamical Systems. 2011.
[46] OpenAL. Gpt-4 technical report, 2023.
[47] Oscar Ovvedo-Trespalaccon 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 obtain common safety-related information and advice. Available at SSRN 4346827, 2023.
[48] Jason Phang, Thibault Févry, and Samuel R. Bowman. Sentence encoders stilts: Supplemen- tary training on intermediate labeled-data tasks. Computing Research Repository, 2018.
[49] Alec Radford, Karthik Narasimhan, tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. OpenAI, 2018
[50] Alec Radford and Jeffrey Wu. Language models are unsupervised multitask learners. OpenAI. 2019.
[51] 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.
[52] 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.
[53] 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 (UCV), 11:3)2211-252, 2015.
[54] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy I Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. International Conference on Learning Representations (ICLR), 2020.
[55] Soroosh Shafieezadeh Abadeh, Peyman M Mohajerin Esfahani, and Daniel Kuhn. Distribution- ally robust logistic regression. Advances in Neural Information Processing Systems (NeurIPS),
28, 2015.
13

---

[56] Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus. Indoor segmentation and support inference from rgbd images. European Conference on Computer Vision (ECCV). 2012.
[57] Wilson L Taylor. "cloze procedure": A new tool for measuring readability. Journalism quarterly, 1953.
[58] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397) 2023.
[59] Paul Tseng, On linear convergence of iterative methods for the variational inequality problem. Journal of Computational and Applied Mathematics, 1995.
[60] Paul Tseng. Accelerated proximal gradient methods for convex optimization. Technical report, Technical report, University of Washington, Seattle, 2008., 2008
[61] Simon Vandenhende, Stamatios Georgoulis, Wouter Van Gansbeke, Mare Proesmans, Dengxin Dai, and Lue Van Gool. Multi-task learning for dense prediction tasks: survey. IEEE transactions on pattern analysis and machine intelligence. 2021.
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[63] Roman Vershynin. digh-dimmesiooal probability: An introduction with applications in data
science. Cambridge university press, 2018.
[64] 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.
[65] 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.
[66] 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.
[67] 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.
[68] Brian D Ziebart, Andrew Maas, J Andrew Bagnell, Anind Dey, et al. Maximum en- tropy inverse reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2008
14

---

Supplementary Materials
A Related Work
Minimax Optimization in Deep Learning. Minimax optimization has wide application in deep learning, c.g., adversarial robustness [71 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 [T3] 33 Our work proposes minimax optimization procedure for pre-training but with 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) first proposed by [57] as a Cloze task. Adapted as a novel pre-training task, MLM and its autoregressive counterpart ,e88- BERT 18 GPT 49 50 [ and T5 [51], highly successful methods to deal with NLP problems. ML. 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 excellently so that various downstream tasks can utilize the pre-trained representations, In particular, BERT 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.
scale
Masked image modeling. Encouraged by transformers, which have gradually become primary architecture for generic language understanding, VIT [21] later illusion the potential of adopting pure transformer 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 [ 14 21 24 29, 66), discrete tokens [5] 67). and (deep) features E 655. 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
B.1.1 Proof of Theorem 3.1
For notation simplicity, we define that
fe(0) Ez-P (Ce =)]:
Fo(0)= Σ Wa.t (0x) Ezmp 16a (0, z)]
F(0) max fe(0), (E(T]
espiaa...n1..00... where W.,t r (0) Eis exp(aE.-pppp (4y
Our proof consists of two parts. The first part is show how well (ox))) approximates F(Ox) with our choice of the softmax hyperparameter a. The second part to analyze the dynamics of the algorithm and the total convergence rate.
15

---

We first show that Fx(Os) 2 F(Ox) + Ro for on 4yk+1
log
47Byk+1
Define Tk,)(0)
2v++I
NoLL"
Ro
ft E [2) fo (0k) 2 F(Ox) - e). If ax 2 c2
log
TB for an > 0, we have
Fx(Ok)
Wor (0x) f((Ox)
Wan,t (0x) fe(0s)
1 (0x)
EIET... (0x) exp (e(x))
(F(0k) (x)
Er= exp fa (0x))
E EigTer, (8x) exp (a)f((((()))
(F(Ot) Ex)
Σ EvETua (0a) exp (e
1 1 T T'epp (ax (F())) - -k))
2
(F(0)) - cK.
exp (agF((:))
[1 T Texp (-anEx)] (F(Ox) Ex)
F(Bs) 2ck-
The inequality A is due to the definition of (0). The inequality B is because EogTan (0,) exp (0x)) I T exp (as (F(0))-es)) (x)) by the definition of (0,c(0) and there exists t E Tx.ex (0) such that fü; F(Ox). which further implies EveTe.,, (0x) exp (a 2 exp (a+F((t)) The inequality C is because the hyperparameter as 2 cN cn
log TB Specifically, let Ro
<s 4vk+I and on for all k = 0,.... 1.
R/
4vK+I R.L log ATB correspondingly, we have Fs 2 F(0x) 2vk+1
R.
We then analyze the dynamics of the algorithm and give the total convergence rate. For k O,.... K- 1, it holds that
(A) (0x) I Fx((") + (VF)(0)), -
(14)
B Fx(00)= 2n (104:1 0x + +10 0-11 -110 - -01)3)
where l" E arg max F(0). The inequality is due the convexity of (e)) The equal- ity B is due to the update steps Ox+1 O2 (.(a)) in Algorithm I and the fact (a, b) (|all + |bb|2 - le - b.
Plugging the approximation error (0x) F(Ow) +
Rot' 2Vk+1 and the inequality Fa(0) < F(0) for
O....,K K 1 and all 0 e θ into [14), we have
F(Ox) < F(0") + 2n (10k-1 0xl12 + |10/ - 8¹122 -1 - 00(²) +
Ro
(15)
2 +I
16

---

Taking the average over lo = O,...... K 1, we further have
1 K-1 RoLl
F(B+) - F(0))
1
10 0xll + 1100 -0"12--110
1²1
K Σ
2n
+
S <F(0)) + 21
K
10R+1 0a11² + 1100 - -0-13)
+
K-1
ROL'
K
2 +
k=U
S F(00) + 1
Σ
10 0x11 + 1100 0²1)?
Ra T!
21
VK
(8)
2 P(O') + KL
R
R
2
2n
VK
@ F(0") + 2R
vK
(16)
The inequality A is due to the fact Σ VK
1 < 2VK. The inequality B is because for k O, K 1, the function (0) is L'pLishhiitz continuous, which implies ||0x+1 11
n²2 Vv 2 17² L The equality Cs due to our choice for the step sizes, i.c., 77k: = D'KK for all & = O..... K 1.
By the convexity of F(0), we have F(Ox) S t EKK F(Ox). 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 Ex and the task space A, the equation
max
Eg~ [ex (0, max EgsP [Le (0, xEλ
1E[2]
holds for all e e.
By the definition of 0r max: we have
max ΛEλ
Eg [eA (0* (0mx:: 2)] max Ean [Le (0 2)] :E(T]
max EzpP P [Le (Oiverages 2) :E(T]
max Bgn [ex (8 (Oiverage+ 2)
xEΛ
which is the result to prove.
B.2.2 Proof of Proposition 5.3
Proposition 5.3 a standard result of gradient descent for strongly-convex and smooth functions. We
include the proof here for completeness. By the convexity of f(z) and the choice of the step size we have
f(#x+1) S f(xx) + (Vf(xx), Fx+1 -xx) + Lf || xx11"
L |
=f(xx)
2
< fxxx,
which means the objective values are nonincreasing and implies f(xx) S ((0) for all k E [K].
By the strsonyn convexity at the point x", it holds for all x E IRd that
2 (f(x) f(x)))
|| x= 112 < 14
(17)
17

---

Combining 11 and the sequence (f(xx) [[(=x))[= decreasing, we obtain
2 (s(zo) - [(x))), for all l: 0, 1...,, K 1,
| - x²11² v /
as desired.
B.2.3 Proof of Theorem
The following proposition characterizes the sample complexity to find an e-approximately optimal parameter by ERM for downstream task À within the parameter space O (0)). Theorem 5.4 follows directly from Proposition B.I by considering the worst-case downstream task À E A. The remaining is to prove Proposition B.
Proposition B.1. For a given task À E A and a parameter space (x let the parameter 0% E x be the minimizer in the the empirical for N, i..d. samples 1 a distribution P. ie., 0% arg mitlgee, ex (0a) with probability at least 1 if if
Nx Ein Ex(0,z)) zi). The parameter 0%, is e-approximately optimal
8dB²
16L'
2
8B 2
2 8
N₂
log 1
- E*
+
log
(18)
22
where E- = Esup lex (fo, 2)]
Denote Ez-P P [e^ (0, =)] as (1( and N E Ex 2i) as (0). First, we note
Pr (fx 05 (x7) 26)
([Ax (0:) A (0x)
[A (02) f (02)) + [1A00))
Pr
(4)
([x (0:) A (8:) + [A (0)) fa (82) 20)
Pr
(19)
(0) A (0) 00)22 2
Pr
sup
(00)
(0) fa (8) 2 2
Pr
sup (0a)
The inequality A is because S (0:) (x) S by the definition of 0z.
We derive the upper bound by covering numbers. We only consider Euclidean space for simplicity. Definition B.2 (Covering numbers [63] Chapter 4)). Consider subset S Rd and let >> 0. A subset N C S is called an e-net of S if every point in S is within distance of some points of AV i.e. for all æ E S, there exists x6 A such that | xo S The smallest possible cardinality of an e-net of S is called the covering number of S and is denoted C(S, e), e),iie. i.e.
C(S,c) = min {IW] |W c- of 5}. Consider an e'ntet Aeett((e,)))) (V)) e') of (00) where = 3I7
= By the definition of e'nntt, we have
sup
/A (0) fa (0) S sup
[A (0) Aa (0)
+ 4"
e
(20)
(0))
(E))))(()).c))
Combining (99 and
20
we obtain
Pr
fa
0%
- fa (0 26)
< Pr
sup (E)))
fa (0) - fa (0) 2 4
(21)
<C(()(B)).c)
sup (E))
Pr (1 (0) fa (0) 4 5)
We leverage the upper bounds of covering numbers of balls 63 Chapter 4].
18

---

Lemma B.3. The covering number of a ball of radius R, denoted as BR. in R satisfies
C(Bn. - (2h +1 (2++1)
2R
By Lemma B.3 we have
2R C(Ox(B)),
e +
(22)
where Rx
2 B 1ex (fu, 2)]. V By Hoeffording's ss inequality, for each 9 E (()( c'), we have
Pr (1xx - (1) fx
9) S 2xex
nc²
(23)
813
Plugging 22 and 23 into
have
Pr (Ix (0)) Sx (0 2 1
2Rx +1) exp
N2a²
(24)
8B²
By 24 when the number of samples N₃ satisfies
SdB²
16L'
8B2
log
1
- Exnp 1tA (0o, 2]]
+
log
M
22
8'
we have Pr fAx (02) fA (02) 26) 2d.
C Example for Remark 5.2
Consider an extreme example where R, EewP [0,,2
A(6 1)2 where < A T
and EZ~P (0, 2)] = 8 for t = .,...,77 T. Then we have 6 further implies that
/ and 6
ivverage A+T It
A
A
mmax
Ez-p (ex (0 =)]
XEA
(VA 2
A(T 1)2
max EznP [ex (Oiverage: )) XEA
average
(A + T 1)2"
If A = T 1, we have
Larcuue 2 (VT -1)) 1)²
which is as large O
D Training details
D.1 Part-of-Speech Mask BERT Training Setting
In this work, we denote the number of layers (i.c.. Transformer blocks) as L, the hidden size as H. and the number of self-attention heads A. For comparison purposes, we primarily report results on two models with the same size: PoS- BERT BASE and BERT BERTBASE (0=s((==12 H=768, A=12, Total 100Mm The model is trained with Adam [42] by setting 31 0.9, ,3 0.999, and L2 weight decay of 0.01. The learning rate is warmed up over the first 10K steps peak value of le-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
1e-6,
19

---

Hyperparam
Part-of-Speech Mask BERT
Number of Layers
Hidden size
768
Attention heads
12
Attention heads size
64
Hyperparam
GLUE
Dropout
0.1
Learning Rate
2c-5
Warmup steps
10K
Batch Size
32
Peak Learning Rate
2e-4
Weight Decay
0.1
Batch Size
256
Learning Rate Decay
Linear
Weight Decay
0.01
Warmup Ratio
0.06
Max Steps
1000K
Learning Rate Decay
Linear
Table 5: Hyperparameters for fine-
Adam
le-6
tuning Part-of-Speech Mask BERT GLUE.
Adam 31
0.9
Adam 32
0.999
Gradient Clipping
0.0
Table 4: Hyperparameters for pre-training Part-of-Speech
Mask BERT.
D.2 Multi-Modal Mask MAE Training Setting
We use Vit-B 21 with patch size of 16x 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 K and 800 epochs ImageNetS50. We choose AdamW as the optimizer with a base learning rate of le-4 and weight decay of 0.05. We first warm up the learning rate with 40 epochs and then decay it with cosine decay [4 We set the batch size to 2048 and trained the models using AxA1000 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 224x224 pixels and apply a random horizontal flip with probability of 0.5. The hyperparameters for experiments shown Table and 7
Hyperparam
[None, GradNorm, DWA)
Uncertainty
Minimax
Batch Size
2048
2048
2048
Learning Rate
8e-4
8e-4
8e-4
Min Learning Rate
le-6
le-6
le-6
Weight Decay
0.05
0.05
0.05
Adamw
1e-8
1e-8
1e-8
Adamw 31
0.9
0.9
0.9
Adamw 82
0.95
0.95
0.95
Epoch
[800]
[400, 800, 1600)
(400, 800, 1600)
Warm up Epoch
40
40
Learning Rate Schedule
cosine decay 98
cosine decay 98
cosine decay 98
Non-masked tokens
Input resolution Augmentation
224x224
224x224
224x224
RanzemResizeeCrop
RandomResizeCCopp
RandomResizzeCrup
Dropout
0.0
0.0
0.0
Patch Size
16
16
16
Table 6: Hyperparameters for pre-training Multi-Modal Mask MAE. We only pre-train 800 epochs on ImageNeiS50, and pre-train both 400 and 1600 epochs on ImageNet
20

---

Classification
Semantic Segmentation
Depth
Hyperparam
ImageNetIK
ImageNetS50
ImageNetS50
NYUv2
NYUv2
Epoch
100
100
100
100
2000
Warm up Epoch
5
5
20
20
100
Batch Size
1024
1024
1024
1024
2048
Learning Rate
4e-3
4e-3
le-4
le-4
le-4
Min Learning Rate
le-6
le-6
le-6
Ic-6
0
Weight Decay
0.05
0.05
0.05
0.05
le-4
Adamw B1
0.9
0.9
0.9
0.9
0.9
Adamw B2
0.999
0.999
0.999
0.999
0.999
Layer Decay
0.65
0.65
0.75
0.75
0.75
Patch Size
16
16
16
16
16
Drop path
0.1
0.1
0.1
0.1
/
LR Schedule
cosine decay
cosine decay 224x224
cosine decay 224x224
cosine decay
cosine decay
Input resolution Augmentation
224x224
224 2244
256x256
Rand(9, 0.5)
Rand(9, 0.5)
LSJ
LSJ
LSJ
Table 7: Hyperparameters for fine-tuning Multi-Modal Mask MAE 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 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.
21


Reviews:
[{"review_id": "TShrYJFucZ", "reviewer": "Reviewer_sSwq", "summary": "This 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.", "questions": "1. 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?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. 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\n", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback on our work. We truly value your constructive comments and have taken steps to address your concerns.\n\n1. **Discussion with some multi-task learning methods**: We acknowledge the importance of including a discussion on multi-task learning methods. In the updated manuscript, we will compare our approach with the methods you highlighted: “Gradient Surgery for Multi-Task Learning, NeurIPS 20” and “Task-customized Masked Autoencoder via Mixture of Cluster-conditional Experts, ICLR 23”, providing a richer context for our contribution. For [1], they focus more on the gradient direction to ensure the model performs well on each task. This aligns with the Gradnorm method we used for comparison. The downside of this approach is its high computational complexity, and it doesn't guarantee robust performance for downstream tasks. [2] trains each expert using only semantically relevant images by leveraging cluster-conditional gates. The difference from our approach is that we emphasize optimizing the training process to enhance the performance of downstream tasks.\n\n[1] Gradient Surgery for Multi-Task Learning, NeurlPS 20\n\n[2] Task-customized Masked Autoencoder via Mixture of Cluster-conditional Experts, ICLR 23\n\n2. **Pre-training model with each pre-training task**: Acknowledging your concerns, we've added experiments that directly address performance in worst-case downstream tasks. Notably, since downstream task initial inputs are contingent upon the RGB modality, its absence would hinder task progression. If the upstream task does not incorporate any RGB modality input, the downstream task will be challenging to accomplish. Our results, which all include RGB modality input, affirm that the worst-case downstream task can effectively bolster downstream tasks.\n\n    |  | Cls. | Semseg. | Depth. |\n    | --- | --- | --- | --- |\n    | RGB | 82.9 | 42.7 | 80.8 |\n    | Semseg+RGB | **83.1** | **51.0** | 83.4 |\n    | Depth+RGB | 82.6 | 46.2 | **86.1** |\n\n3. **Clarifications for Table 1**: The 'None' row in Table 1 signifies the original BERT without loss balancing. We'll make this clearer in the revised manuscript. Also, a comparison between average scores and ours will highlight the distinctions.\n\n    |  | MNLI | QQP | QNLI | SST-2 | CoLA | STS-B | MRPC | RTE |\n    | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n    | Average | 84.9 | 72.6 | **89.1** | 90.8 | 54.4 | 83.6 | 88.1 | 68.2 |\n    | Ours | **85.6** | **76.9** | 88.6 | **91.3** | **61.4** | **84.2** | **88.2** | **70.7** | \n\n4. **Comparison with Multi-task Baselines on the PoS-BERT Setting**: We concur with your suggestion and will integrate comparisons with multiple multi-task baselines in the PoS-BERT setting in our updated manuscript.\n\nIn conclusion, we have earnestly addressed your feedback and believe our revisions will clarify and underscore the significance of our approach. We remain committed to refining our manuscript, aiming for clarity and meaningful contributions. Thank you once again for your insights."}, {"review_id": "jYNkCsmjMD", "reviewer": "Reviewer_Jr2o", "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.", "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", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Firstly, thank you for your thoughtful feedback on our paper. I would like to delve deeply into the central theme of our work: the critical importance of robustness. Let us clarify and address your concerns:\n\n**Consistent Improvement**：\n\n1. Achieving consistent improvements in downstream tasks solely through pre-training and finetuning is often not feasible in many robust scenarios. For instance, in adversarial learning, robustness often accompanies a decline in task performance. Similarly, in domain adaptation, sometimes improving generalization to a new domain may result in a decrease in the performance of the source domain [1]. Another example is in multi-task learning where optimizing for a secondary task might interfere with the primary task performance[2].\n\n  [1] Domain-Adaptive Neural Networks for Object Detection, ICCV 2019\n\n  [2] Multi-task learning with deep neural networks: a survey\n\n2. Our experiments demonstrate a **holistic** improvement in our downstream tasks: other tasks' performance has hardly decreased, meanwhile tasks that initially performed poorly have shown more significant enhancement. \n\nIn Table 1, our method outperforms on half of the tasks by a substantial margin and obtains a **1.8% average score** improvement. Concerning the most challenging training task, CoLA, which has the lowest accuracy on BERT, our method attains a 9.2% improvement, marking a significant boost among downstream tasks. Due to task-robust grouping, on the QQP and RTE tasks, our method surpasses the original BERT by 5.7 F1-score and 4.3% accuracy, respectively. Compensating for better performance on more challenging tasks, our method concedes minor correctness on some downstream tasks that already transfer well. On QNLI, SST-2, STS-B, and MRPC, our results fall below that of the original BERT model by margins of 1% accuracy, 1.9% accuracy, 1.6 Spearson correlation, and 0.7 F1-score.\n\nIn Table 2, our method astonishingly achieves 9% and 22% improvements in the Semseg and Depth downstream tasks, respectively. In contrast, there's only a 0.4% performance drop in the classification task. The **average** uplift in downstream tasks reaches 10.2%. The overall performance of our downstream tasks is on the rise. Compared to the tasks that improved, the degree of decline in the declining tasks is negligible.\n\n3. From the data in Table 2, it can be observed that as training becomes sufficiently comprehensive and the duration extends, the tasks that initially showed a decline eventually start to improve. This suggests that given more training time, the model's innate potential to adapt and recalibrate can be harnessed further, eventually boosting those tasks which initially lagged behind.\n\n4. Our paper places a significant emphasis on robustness. In the realm of machine learning, and especially in pre-trained models, a model's ability to consistently perform across a vast spectrum of tasks is invaluable. This robustness isn't merely about excelling in frequently encountered tasks but ensuring minimal performance degradation in edge cases or less conventional scenarios. This balance between performance in general and edge cases is essential for the broad applicability and reliability of models in real-world scenarios.\n\nWe fully acknowledge the concerns you've raised about the potential trade-offs. But we genuinely believe that in the grander scheme of things, ensuring consistent and robust performance across the board is of paramount importance, potentially outweighing the marginal gains in isolated tasks.\n\nOnce again, I appreciate the depth of your engagement with our work. We will strive to make our stance and the significance of our approach clearer in the subsequent revisions of the manuscript."}, {"review_id": "KxXzT4WduJ", "reviewer": "Reviewer_hh6P", "summary": "The 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.", "questions": "- **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", "limitations": "I 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.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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", "weaknesses": "- **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed feedback on our submission. We highly appreciate your insights and have carefully considered each point you raised. Let us clarify and address your concerns:\n\n**Weeknesses**\n\n1. **Weak Comparison to multi-task learning optimization methods**: \n\n  We acknowledge that including a broader range of comparisons would provide more clarity on the efficacy of our approach. We are indeed familiar with the papers you've mentioned, and in retrospect, agree that it would have been beneficial to compare with methods like PCGrad/CaGrad/GradDrop. Your insights into other task-balancing methods have been enlightening. While the methods we initially compared encompass a broad spectrum of strategies, we acknowledge the potential value of additional comparisons. We are currently in the process of conducting supplementary experiments in this direction and anticipate sharing the results soon.\n\n  You brought up an essential point about simply tuning the weights of each task loss as scalar hyperparameters. This approach indeed has similarities with the uncertainty method we employed for comparison, a point underscored in the article \"Do Current Multi-Task Optimization Methods in Deep Learning Even Help?\". This paper, as you mentioned, reiterated this view and referenced the uncertainty-based paper. Using uncertainty for a benchmark is a robust representation, and the methods we selected for comparison are emblematic of the broader approaches in the field. However, to satiate any residual curiosity, we provide supplementary experiments treating the weight solely as a hyperparameter:\n\n  | | Cls. | Semseg. | Depth. |\n  | --- | ---  | ---     | ---    |\n  | new baseline | **92.7** | 55.2    | 66.8   |\n  | Ours | 91.8 | **61.5** | **74.1** |\n\n  As the results attest, our methodology retains its superiority.\n\n2. **Unclear assumptions on upstream and downstream tasks**: \n\n  Our initial intention was to show how our method functions when upstream and downstream tasks are closely related, as we believe this is a common practical scenario. However, we acknowledge that for general applicability, we need to explicitly state and validate this assumption. \n\n  As you rightly pointed out in Section 5, the convex combination assumption is pivotal to our theoretical analysis. This framework allows us to effectively model the nexus between upstream and downstream tasks and harness the shared patterns and feature representations that benefit both tasks. Though it might appear idealistic, this assumption is grounded in reality. It resonates with the foundational concepts in multi-task learning, where task synergies are typically quantified through shared representations. Recognizing the dynamic nature of real-world applications, we agree that they may not always perfectly adhere to this assumption. To accommodate more diverse task scenarios, this assumption can be eased by integrating extra approximation errors.\n\n  We apologize for any ambiguity that may have arisen. To clarify, our upstream tasks are focused on different modalities of mask recovery, while the downstream tasks target specific computer vision tasks, such as classification and semantic segmentation. They are distinct, and we'll ensure this distinction is more pronounced in the revised manuscript.\n\n**Questions**\n\n1.**Section 4.2**\n\nThank you for bringing up this point. We understand where the confusion might have arisen, and we'd like to clarify the relationship between the upstream and downstream tasks in our computer vision example.\n\nThe upstream tasks in our computer vision example differ from the downstream tasks not only in terms of input data but also in the nature of the tasks themselves. The upstream tasks encompass three different domains of masked image recovery, whereas the downstream tasks involve classification, semantic segmentation, and depth estimation. Some of these are classification tasks, while others are autoregressive tasks. Our method is designed to harness the shared information and structures from the upstream tasks for the benefit of the downstream tasks. While it might bear resemblance to a multi-task scenario, such a setting is customary in the realm of pre-training [1, 2, 3, 4]. We apologize for any confusion caused and appreciate your highlighting this matter.\n\n[1] Multi-Task Self-Training for Learning General Representations, ICCV 2021\n\n[2] Variational mixture-of-experts autoencoders for multimodal deep generative models. 2019\n\n[3] Factors of influence for transfer learning across diverse appearance domains and task types, TPAMI 2021\n\n[4] Mid-level visual representations improve generalization and sample efficiency for learning visuomotor policies, 2018\n\n2.**Suggestions on Writing:**\n\n  We appreciate the feedback on specific lines and sections. It's evident that we could improve our presentation to make our contributions and methodologies clearer:\n\n  - We will relocate and rephrase the discussion on the link between expected risk and sample complexity, ensuring its relevance is apparent.\n  - The claim on line 259 will be revised to more accurately reflect our findings from Figure 3.\n  - We will introduce a more structured related work section as you've suggested, to help readers position our work in the broader context.\n\n  In summary, we are grateful for your feedback as it offers us a clear direction to refine and bolster our work. Ensuring the relevance and robustness of our methodology is paramount to us."}, {"review_id": "Sj0xkk5Trs", "reviewer": "Reviewer_UtAg", "summary": "The 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. ", "questions": "1. 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? ", "limitations": "The 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. ", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. 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. ", "weaknesses": "The 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.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your detailed feedback. We are particularly grateful for your insightful comments which have helped us identify key areas for enhancement in our manuscript. We would like to address the specific points you raised:\n\n1. **Convergence Rate of the Minimax Optimization Algorithm**:\n\nYou rightly observed the convergence rate of our proposed algorithm, which is $O(1/\\sqrt{K})$. We'd like to clarify that, theoretically, this rate is optimal for the setting we addressed in Theorem 3.1, which is a linear-nonsmooth-convex minimax optimization problem [1]. Despite the fact that the rate $O(1/\\sqrt{K})$ is optimal in our setting, it is possible to achieve a faster convergence rate with stronger assumptions or extra oracles. We leave further investigation for future work. We fully recognize the importance of placing our results within the broader context of established results, and we will include this comparison in the paper to provide a more comprehensive view. Your suggestion to provide convergence comparisons, whether theoretical or empirical, is highly appreciated, and we will make sure to incorporate this in our next iteration.\n\n   [1] Sébastien Bubeck, Convex Optimization: Algorithms and Complexity, Chapter 3, Theorem 3.13 and its proof\n\n2. **Choice of Representative Upstream Tasks**:\n\nThe basic idea is to choose upstream tasks to be prototypical or informative to leverage shared information and structures for downstream tasks. In practice, upstream tasks can be chosen by heuristics, domain-specific knowledge (e.g., our work and [2]), or automatically finding underlying patterns or relationships between tasks via algorithms such as clustering [3]. Thanks for your suggestion. We leave this study as a future work.\n\n   [2] Shiori Sagawa et al., Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization, ICLR 2020.\n\n   [3] Yaqiang Yao et al., Robust Task Grouping with Representative Tasks for Clustered Multi-Task Learning, KDD 2019.\n\n3. **Example on Line 279**:\n\nWe sincerely apologize for this oversight. The intended example for line 279 can be accessed in the PDF attached to the global rebuttal (under the 'Author Rebuttal' section). It appears that we mistakenly omitted this example from the appendix, even though we referenced it in the main body of our manuscript. Your meticulous attention to the content has allowed us to identify this gap, and we are truly grateful. In light of this, we will ensure the inclusion of the example in the revised version of the paper, positioning it appropriately within the appendix. We appreciate your patience and understanding as we work towards rectifying this oversight.\n\nOnce again, thank you for your constructive feedback, which has been instrumental in guiding our revisions. We are committed to enhancing our paper based on your observations and look forward to presenting a more polished and comprehensive version."}]

Author rebuttal:
We attach PDF files here.

Meta-review: The paper investigates the intriguing problem of ensuring performance guarantees for worst-case downstream applications through transfer learning via pre-training. To achieve this goal, it introduces a minimax loss for pre-training and develops an efficient algorithm tailored to the task. Experimental results in both vision and language domains demonstrate the effectiveness of the proposed model. Four reviewers participated in the evaluation process, resulting in four weak accept ratings and one borderline reject rating. To address the negative concerns, the final revision should include additional discussions on multi-task learning, provide a clearer explanation of the connection between upstream and downstream tasks, and offer a better illustration of the proposed objective.
Decision: Accept (poster)