TapWeight: Reweighting Pretraining Objectives for Task-Adaptive Pretraining

## Abstract

Abstract Large-scale general domain pretraining followed by downstream-specific finetuning has become a predominant paradigm in machine learning. However, discrepancies between the pretraining and target domains can still lead to performance degradation in certain cases, underscoring the need for task-adaptive continued pretraining (TAP). TAP methods typically involve continued pretraining on task-specific unlabeled datasets or introducing additional unsupervised learning objectives to enhance model capabilities. While many TAP methods perform continued pretraining with multiple pretraining objectives, they often determine the tradeoff parameters between objectives manually, resulting in suboptimal outcomes and higher computational costs. In this paper, we propose TapWeight, a task-adaptive pretraining framework which automatically determines the optimal importance of each pretraining objective based on downstream feedback. TapWeight reweights each pretraining objective by solving a multi-level optimization problem. We applied TapWeight to both molecular property prediction and natural language understanding tasks, significantly surpassing baseline methods. Experimental results validate the effectiveness and generalizability of TapWeight. Our code is publicly available at https://anonymous.4open.science/r/TapWeight-9A2E .

## 1 Introduction

Foundation models pretrained on large-scale general domain corpora have achieved state-of-the-art performance across a wide range of tasks (He et al., ; Devlin et al., ; Brown et al., ). These models, which capture general knowledge for specific modalities such as text or images through unsupervised learning, are typically adapted to downstream tasks via finetuning. However, when there is a domain discrepancy between the pretraining corpus and the target task, direct finetuning of the pretrained model often fails to deliver optimal results (Lee et al., ; Chen et al., ; Xie et al., ). To address this challenge, downstream task-adaptive continued pretraining, or task-adaptive pretraining (TAP), has been introduced. TAP bridges this gap by introducing an additional continued pretraining stage between general domain pretraining and task specific finetuning. For example,  Gururangan et al. () conducts task-adaptive pretraining by performing unsupervised learning on the unlabeled data of the downstream task.  Wu et al. () introduces an additional perturbation masking objective during continued pretraining of a BERT model (Devlin et al., ), enhancing its performance on dialogue understanding tasks.

Among these, many existing task-adaptive pretraining methods consist of multiple pretraining objectives (Wu et al., ; Gao et al., ; Cui et al., ), making it challenging to determine the relative importance of each objective. Some TAP methods assign equal weight to each pretraining objective (Lee et al., ; Wu et al., ), disregarding their varying impact on downstream performance. For instance,  Gao et al. () shows that pretraining BERT with a contrastive learning (CL) objective results in better downstream performance on semantic textual similarity (STS) datasets than using masked language modeling (MLM) loss, indicating that the CL objective is more important than the MLM objective for these tasks. Other approaches attempt to manually tune the importance ratios through hyperparameter search (Gao et al., ), which often results in suboptimal performance and increased computational costs. This issue becomes particularly severe when the number of pretraining objectives is large, such as with the task-adaptive pretraining of a popular molecular model Imagemol, which involves 5 distinct pretraining objectives (Zeng et al., ).

To address the aforementioned challenges, we propose a novel framework, TapWeight, designed to learn the optimal tradeoff parameters between various pretraining objectives during task-adaptive pretraining. The goal is to learn these optimal tradeoff parameters such that the pretrained model, after finetuning on a downstream task, achieves the best downstream task performance. Our approach involves a three-level optimization framework to learn these parameters. In the first stage, we perform task-adaptive pretraining using initial tradeoff parameters, denoted as $\lambda$. These parameters are kept fixed during this stage and will be updated in subsequent stages. The resulting pretrained model is thus a function of $\lambda$. In the second stage, the pretrained model from above stage is finetuned on the training split of the downstream dataset. Consequently, the finetuned model becomes an implicit function of the tradeoff parameters. In the third stage, the finetuned model is evaluated on the validation split of the downstream dataset, and the tradeoff parameters $\lambda$ are optimized by minimizing the validation loss. This end-to-end process allows the three stages to dynamically influence one another, forming an integrated framework that optimizes task-adaptive pretraining process and enhances downstream task performance. Moreover, TapWeight is applicable to any pretrained model with multiple pretraining objectives, regardless of data modality or downstream task type, demonstrating superior generalizability compared to existing task-adaptive pretraining methods (Nishida et al., ; Cui et al., ). Figure 1 illustrates the complete framework of TapWeight.

Figure: Figure 1: An Overview of TapWeight. In the first stage, the model undergoes multi-objective pretraining with fixed tradeoff ratios between objectives. In the second stage, the pretrained model is finetuned on the training split of the downstream dataset. In the third stage, the finetuned model is evaluated on the validation split of the downstream dataset to compute a loss, and the trainable tradeoff parameters fixed in the first stage are learned by minimizing this validation loss.
Refer to caption: x1.png

We apply TapWeight for task-adaptive pretraining of both a molecule representation model, Imagemol (Zeng et al., ), and a language model, RoBERTa (Liu et al., ). Evaluating its performance across 13 molecular property prediction datasets and 8 natural language understanding tasks, TapWeight significantly outperforms baseline methods. The superior performance of TapWeight highlights its effectiveness and generalizability. Our contribution can be summarized as follows:

- •
We propose TapWeight, an approach that automatically searches for the tradeoff parameters across multiple pretraining objectives and performs reweighted task-adaptive pretraining. TapWeight is formulated within a multi-level optimization (MLO) framework. We employ an efficient gradient descent algorithm to solve the MLO problem, obtaining the optimal tradeoff parameters for multiple pretraining objectives. Our implementation of TapWeight is publicly available.
- •
We apply TapWeight for task-adaptive pretraining of a molecule representation model and a language model. Extensive experiments on 13 downstream datasets in molecular property prediction and 8 datasets in natural language understanding underscore its effectiveness and generalizability.

## 2 Related Works

### 2.1 Domain / Task Adaptive Pretraining

To bridge the gap between general domain pretraining and downstream tasks in a specific domain, domain-adaptive pretraining (DAP) and task-adaptive pretraining (TAP) have been introduced (Gururangan et al., ). DAP performs continued pretraining on a large, unlabeled corpus from a similar domain as the downstream task. For example, BioBERT continues to pretrain a BERT model on a large-scale biomedical corpus, enhancing its performance on a variety of biomedical text mining tasks (Lee et al., ). Similarly, LegalBERT continues to pretrain a BERT model on legal documents to improve performance on legal NLP tasks (Chalkidis et al., ), while SciBERT leverages a large multi-domain corpus of scientific publications for further pretraining, enhancing its effectiveness on scientific NLP tasks (Beltagy et al., ). More recently, MEDITRON performs continued pretraining of a Llama-2 model with 80 billion parameters on text in medical domain, showing significant performance gains on major medical benchmarks (Chen et al., ).

Although DAP significantly improves model performance on downstream tasks, it needs a large corpus of unlabeled data in a specific domain, which is not always available. To address this limitation, multiple task-adaptive pretraining (TAP) methods have emerged, which do not rely on additional domain-specific corpora beyond the downstream dataset itself. TAP methods can also be viewed as a novel finetuning process, where standard finetuning is preceded by low-cost continued pretraining. For instance, TAPT performs continued pretraining directly on the unlabeled training split of the downstream dataset (Gururangan et al., ). TAPTER first trains new word embeddings using the unlabeled training split of the downstream dataset, and then use these embeddings for continued pretraining of the model (Nishida et al., ). SimCSE introduces an additional constrastive learning loss in addition to the original masked language modelling loss to further pretrain a RoBERTa model, specifically enhancing its capability on standard semantic textual similarity tasks (Gao et al., ). While existing TAP methods are effective, they are typically tailored to specific downstream tasks or data modalities (Wu et al., ; Cui et al., ). In contrast, TapWeight is applicable to any pretrained model with multiple pretraining objectives, underscoring its broad generalizability.

### 2.2 Multi-level Optimization

Many machine learning tasks can be formulated as multi-level optimization (MLO) problems, such as neural architecture search (Liu et al., ; Chen et al., ; Xu et al., ), meta learning (Finn et al., ; Rajeswaran et al., ; Zhang et al., ), and hyperparameter optimization (Lorraine et al., ; Lorraine & Duvenaud, ; Mackay et al., ). MLO problems consist of multiple levels of optimization problems that are mutually dependent, making it challenging for common automatic differentiation algorithms to handle them. To tackle this challenge, multiple algorithms (Lorraine et al., ; Liu et al., ; Rajeswaran et al., ) and libraries (Choe et al., ; ) have been proposed to efficiently compute gradients in MLO problems.

Recently, MLO techniques have been widely adopted in data reweighting and task reweighting. In these methods, the weights of data or tasks are often treated as hyperparameters and optimized in the upper levels of MLO problems. For example, MetaWeightNet learns an explicit weighting function for each data point to maximize the performance on a small amount of unbiased meta-data (Shu et al., ). DoGE optimizes weights for each data domain using a small proxy model to guide the pretraining of larger models (Fan et al., ). MetaWeighting learns tradeoff parameters for each task in multi-task learning to minimize generalization loss (Mao et al., ). Our method also falls within this category, with a specific focus on reweighting pretraining objectives for downstream task-adaptive continued pretraining.

## 3 Method

### 3.1 Overview

Given $n$ continued pretraining objectives $\mathcal{T}_{1},\mathcal{T}_{2},...\mathcal{T}_{n}$ and their corresponding training losses $\mathcal{L}_{1},\mathcal{L}_{2},...\mathcal{L}_{n}$, we formulate the multi-objective continued pretraining loss $\mathcal{L}_{pt}$ as:

$$ $\mathcal{L}_{pt}(\theta,\lambda,\mathcal{D}_{pt})=\sum_{i=1}^{n}\lambda_{i} \mathcal{L}_{i}(\theta,\mathcal{D}_{pt})$ (1) $$

where $\mathcal{D}_{pt}$ is the unsupervised pretraining dataset, $\theta$ denotes the pretraining model parameters, and $\lambda_{i}$ is the tradeoff parameter for each pretraining objective. We denote the target downstream task as $\mathcal{D}_{ft}$ and split it into $\mathcal{D}_{tr},\mathcal{D}_{val}$ and $\mathcal{D}_{ts}$, which are training, validation and test splits respectively.

In our framework, TapWeight, we aim to automatically search for the optimal tradeoff weights $\lambda=\{\lambda_{1},...,\lambda_{n}\}$, so that the pretrained model will achieve highest performance on $\mathcal{D}_{val}$ after finetuned on a downstream dataset $\mathcal{D}_{ft}$. To achieve this, our method consists of three end-to-end stages. In the first stage, we perform continued pretraining of the model, with tradeoff weights tentatively fixed. In the second stage, we conduct finetuning of the pretrained model on the training split of the downstream dataset. In the third stage, we compute a loss by applying the finetuned model on the validation split of the downstream dataset, and optimize the tradeoff parameters by minimizing this loss. We next formally define these three stages under a multi-level optimization framework.

### 3.2 TapWeight Framework

##### Level I

In the first level, we aim to perform continued pretraining for the model. Formally, the optimization problem (OP) is to optimize the model weights $\theta$ to minimize the multi-objective pretraining loss $\mathcal{L}_{pt}$ on a unlabeled dataset $\mathcal{D}_{pt}$:

$$ $\displaystyle\theta^{*}(\lambda)=\operatorname*{arg\,min}_{\theta}\mathcal{L}_ {pt}(\theta,\lambda,\mathcal{D}_{pt})$ (2) $$

Since the optimal solution $\theta^{*}$ to this problem depends on the value of the tradeoff parameter, it is an implicit function of $\lambda$, denoted as $\theta^{*}(\lambda)$.

##### Level II

In the second level, we aim to finetune the pretrained model with optimal parameters $\theta^{*}$ obtained from previous level on the downstream dataset. However, formulating the optimization problem here with the same set of parameters $\theta$ as the lower level imposes high computation and memory burdens, as it requires differentiating through the whole gradient update trajectory in the lower level (Rajeswaran et al., ). Optimizing distinct sets of parameters at different levels enables the use of implicit differentiation methods, which significantly reduces computational costs, as detailed in Section 3.3. Therefore, we create a model with new parameters $\omega$ that are different from those in the pretrained model, but with a regularization loss $\mathcal{R}$ between $\omega$ and $\theta$ to encourage them to be close. This proximal constraint casts strong dependence between $\omega^{*}$ and $\theta^{*}$, closely resembling the real finetuning process. Formally, the OP in this level is to optimize $\omega$ by minimizing the weighted summation of finetuning loss $\mathcal{L}_{tr}$ and the proximal regularization loss $\mathcal{R}$:

$$ $\displaystyle\omega^{*}(\theta^{*}(\lambda))$ $\displaystyle=\operatorname*{arg\,min}_{\omega}\mathcal{L}_{tr}(\omega, \mathcal{D}_{tr})+\gamma\mathcal{R}(\omega,\theta^{*}(\lambda))$ (3) $$

where $\mathcal{D}_{tr}$ is the training split of the downstream dataset, and $\gamma$ is a tradeoff hyperparameter to balance the finetuning loss and regularization loss. In practice, we select the mean squared error (MSE) loss as the regularization loss $\mathcal{R}$. The optimal solution of $\omega$ in this level is a function of $\theta^{*}$ due to the loss term $\mathcal{R}$, which is in turn a function of $\lambda$, denoted as $\omega^{*}(\theta^{*}(\lambda))$.

##### Level III

In the third level, we aim to search for the optimal tradeoff parameters $\lambda^{*}$ between pretraining objectives. Formally, the OP in this level is to optimize $\lambda$ to minimize the validation loss $\mathcal{L}_{val}$:

$$ $\displaystyle\min_{\lambda}\mathcal{L}_{val}(\omega^{*}(\lambda),\mathcal{D}_{ val})$ (4) $$

where $\mathcal{D}_{val}$ is the validation split of the downstream dataset.

##### Multi-level Optimization Framework

In this way, we formulate a three-level optimization problem with OPs in different levels mutually dependent on each other:

$$ $\displaystyle\min_{\lambda}\mathcal{L}_{val}(\omega^{*}(\lambda),\mathcal{D}_{ val})$ (5) $\displaystyle s.t.\quad\omega^{*}(\theta^{*}(\lambda))$ $\displaystyle=\operatorname*{arg\,min}_{\omega}\mathcal{L}_{tr}(\omega, \mathcal{D}_{tr})+\gamma\mathcal{R}(\omega,\theta^{*}(\lambda))$ $\displaystyle\theta^{*}(\lambda)$ $\displaystyle=\operatorname*{arg\,min}_{\theta}\mathcal{L}_{pt}(\theta,\lambda ,\mathcal{D}_{pt})$ $$

By solving this multi-level optimization problem, we are able to reweight each continued pretraining objective based on feedback from validation performance on downstream tasks. In practice, both $\theta$ and $\omega$ in Equation 5 are initialized with model weights from general-domain pretraining.

### 3.3 Optimization Algorithm

In this section, we illustrate the algorithm we use to efficiently approximate the gradient of loss $\mathcal{L}_{val}$ in the third level with respect to the tradeoff parameter $\lambda$. This full derivative $\frac{d\mathcal{L}_{val}}{d\lambda}$ can be computed with the following equation using chain rule:

$$ $\displaystyle\frac{d\mathcal{L}_{val}}{d\lambda}={\color[rgb]{0.5,0.5,0}\frac{ \partial\mathcal{L}_{val}}{\partial\omega^{*}}}\times{\color[rgb]{1,0,0}\frac{ \partial\omega^{*}}{\partial\theta^{*}}}\times{\color[rgb]{1,0,0}\frac{ \partial\theta^{*}}{\partial\lambda}}$ (6) $$

In the right hand side of Equation 6, the olive term, a partial derivative vector, can be directly computed with popular automatic differentiation libraries, such as Pytorch (Paszke et al., ). However, directly computing the two red terms, which are best-response Jacobian matrices, can be computationally prohibitive due to the lack of analytical solutions to these optimization problems. Inspired by previous works (Lorraine et al., ; Zhang et al., ), we use Implicit Function Theorem (IFT) based methods to approximate the best-response Jacobian matrices. We include more details of IFT based gradient computation method in Appendix A. In this way, we are able to compute both red terms in Equation 6 efficiently, thereby obtaining the gradient of $\mathcal{L}_{val}$ with respect to $\lambda$. We then optimize the tradeoff parameter $\lambda$ with gradient descent. The complete algorithm is implemented using the Betty library (Choe et al., ; ).

**Table 1: Results of molecular property prediction on 8 classification tasks in MoleculeNet benchmark, in terms of AUROC. The best results are shown in bold.**
| Method | BACE | BBBP | ClinTox | Sider | Tox21 | ToxCast | HIV | MUV | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset Size | 1,513 | 2,039 | 1,478 | 1,427 | 7,831 | 8,575 | 41,127 | 93,087 |  |
| AttrMask | 77.2 | 70.2 | 68.6 | 60.4 | 74.2 | 62.5 | 74.3 | 73.9 | 70.2 |
| ContextPred | 78.6 | 71.2 | 73.7 | 59.3 | 73.3 | 62.8 | 75.8 | 72.5 | 70.9 |
| GraphMVP | 76.8 | 68.5 | 79.0 | 62.3 | 74.5 | 62.7 | 74.8 | 75.0 | 71.7 |
| Imagemol | 80.1 | 67.3 | 78.5 | 63.6 | 76.5 | 65.4 | 75.6 | 78.4 | 73.2 |
| TapWeight (ours) | 83.1 | 71.2 | 81.3 | 64.5 | 77.0 | 66.1 | 78.4 | 80.5 | 75.3 |

## 4 Experiments

### 4.1 Molecular Property Prediction

In this section, we use TapWeight for task-adaptive pretraining of molecular image models and validate the effectiveness of our framework on the downstream task of molecular property prediction.

#### 4.1.1 Preliminary

Given a large unlabeled molecular dataset $\mathcal{D}=\{x_{i}\}_{1\leq i\leq n}$ containing millions of molecules, we define a multi-objective continued pretraining loss inspired by Imagemol (Zeng et al., ):

$$ $\displaystyle\mathcal{L}(x)=\lambda_{1}\mathcal{L}_{mg1}(x)+\lambda_{2} \mathcal{L}_{mg2}(x)+\lambda_{3}\mathcal{L}_{mg3}(x)+\lambda_{4}\mathcal{L}_{ jpp}(x)+\lambda_{5}\mathcal{L}_{mcl}(x)$ (7) $$

where $x$ represents a molecular image, and $\lambda=\{\lambda_{i}\}_{1\leq i\leq 5}$ are tradeoff parameters. $\mathcal{L}_{mg1}$, $\mathcal{L}_{mg2}$, and $\mathcal{L}_{mg3}$ are MACCS key (Durant et al., ) clustering-based classification losses with different number of clusters. $\mathcal{L}_{jpp}$ is a jigsaw puzzle prediction loss, where the model solves a jigsaw puzzle on the same molecular image. $\mathcal{L}_{mcl}$ is a mask-based contrastive learning loss, which generates constrastive pairs by masking molecular images. Details of these pretraining objectives can be found in Appendix B.1.

The multi-objective loss $\mathcal{L}$ is optimized on the complete unlabeled dataset $\mathcal{D}$ to train a molecular image encoder. The learnt encoder can be further finetuned on downstream datasets for various molecular tasks. Existing approaches typically set the tradeoff parameters $\lambda$ equally across different pretraining objectives, overlooking the varying contributions of each objective to specific downstream tasks (Zeng et al., ). We address this challenge by applying TapWeight framework for continued pretraining of the molecular image encoder.

#### 4.1.2 Experimental Settings

We perform continued pretraining of a pretrained Imagemol model on a dataset $\mathcal{D}$, consisting of 1 million molecules from PubChem (Kim et al., ). For downstream tasks, we employ the MoleculeNet benchmark, which includes 8 classification datasets focused on predicting biophysical and physiological properties essential for drug discovery (Wu et al., ). We generate the training, validation and test split of these downstream datasets by applying scaffold splitting with an 8:1:1 ratio. We use AUROC as the evaluation metric for all classification datasets, MAE for Qm7 and Qm9 datasets, and RMSE for all other regression datasets. In addition to Imagemol, we benchmark against Graph Neural Network (GNN)-based molecular property prediction methods, including pretraining approaches such as attribute masking, context prediction Hu et al. (), and GraphMVP (Liu et al., ). The pretrained molecular image encoder is based on a ResNet18 model, with the final classification layer removed (He et al., ). We set the number of clusters for the loss terms $\mathcal{L}_{mg1}$, $\mathcal{L}_{mg2}$, and $\mathcal{L}_{mg3}$ to 100, 1,000, and 10,000, respectively. During the continued pretraining, we set the unrolling step in the MLO framework to be 1. We use the SGD optimizer with a step learning rate scheduler across all three optimization levels. All experiments are conducted on NVIDIA A100 GPUs. More detailed descriptions are provided in the Appendix for the datasets (B.2), baselines (B.3), and hyperparameter settings (B.4).

#### 4.1.3 Results

**Table 2: Results of molecular property prediction on 5 regression tasks in MoleculeNet benchmark. The best results are shown in bold.**
| Method | Freesolv | Esol | Lipo | Qm7 | Qm9 |
| --- | --- | --- | --- | --- | --- |
| Dataset Size | 642 | 1,128 | 4,200 | 6,830 | 133,885 |
| AttrMask | 2.95 | 1.37 | 0.81 | 161.7 | 5.03 |
| ContextPred | 3.01 | 1.35 | 0.83 | 153.2 | 4.95 |
| GraphMVP | 2.21 | 1.13 | 0.79 | 134.5 | 4.76 |
| Imagemol | 3.04 | 1.11 | 0.76 | 141.0 | 4.52 |
| TapWeight (ours) | 1.91 | 1.06 | 0.76 | 126.0 | 4.28 |

Table 1 show the results of various methods across 8 molecular property classification tasks from MoleculeNet benchmark. Our method outperforms all baseline methods on all 8 datasets, showcasing the effectiveness of our method. On average, our method achieves an AUROC of 75.3, compared to 73.2 for the Imagemol model without continued pretraining. Similarly, Table 2 displays the results for 5 regression tasks in the MoleculeNet benchmark, where our method once again surpasses all baselines on each task. Experimental results validate the effectiveness of our method on both classification and regression tasks. Specifically, the superior performance of our method over Imagemol validates the necessity of downstream-guided continued pretraining following general pretraining. Notably, our method consistently outperforms baseline approaches regardless of the size of the finetuning dataset, demonstrating the robustness of our approach. It is worth mentioning that TAPT (Gururangan et al., ) is not applicable to this task, as clustering-based losses, such as $\mathcal{L}_{mg3}$, are not well-suited for direct application on small unlabeled datasets where the number of data points is smaller than the predefined number of clusters. In contrast, TapWeight does not face such limitations, demonstrating its generalizability.

### 4.2 Natural Language Understanding

In this section, we validate the effectiveness of TapWeight for continued pretraining of a masked language model (MLM) with its application to natural language understanding tasks.

#### 4.2.1 Prelimimary

Given a large-scale raw-text dataset $\mathcal{D}=\{x_{i}\}_{1\leq i\leq n}$ consisting of millions of sentences, we define the following continued pretraining loss:

$$ $\displaystyle\mathcal{L}(x)=\lambda_{1}\mathcal{L}_{mlm}(x)+\lambda_{2} \mathcal{L}_{cl}(x)+\lambda_{3}\mathcal{L}_{sop}(x)$ (8) $$

where $x$ is a sentence, and $\lambda=\{\lambda_{i}\}_{1\leq i\leq 3}$ are tradeoff parameters. $\mathcal{L}_{mlm}$ represents the masked language model loss, which involves randomly masking tokens in the input sentences and predicting these masked tokens (Devlin et al., ). $\mathcal{L}_{cl}$ denotes the contrastive learning loss, where an input sentence is used to predict itself with standard dropout applied as noise (Gao et al., ). $\mathcal{L}_{sop}$ is the sequence ordering prediction loss, which emphasizes inter-sentence conherence (Lan et al., ). We include details of these losses in Appendix C.1.

The multi-objective loss $\mathcal{L}$ is optimized on the raw text dataset $\mathcal{D}$ for continued pretraining of a Transformer encoder. The learnt encoder can then be finetuned on downstream NLP datasets. In existing works (Gao et al., ), the tradeoff parameters $\lambda$ for different pretraining objectives require manual hyperparameter tuning, which is time-consuming and often leads to suboptimal results. We address this challenge by applying TapWeight for the continued pretraining of a Transformer encoder, enabling the automatic determination of the importance for each objective.

**Table 3: Results of different methods in GLUE benchmark. All methods are applied to a RoBERTa-base model. The best results are shown in bold.**
| Method | MNLI | QNLI | QQP | RTE | SST | MRPC | CoLA | STSB | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset Size | 392,702 | 104,743 | 363,871 | 2,490 | 67,349 | 3,668 | 8,551 | 5,749 |  |
| Finetuning | 86.7 | 92.8 | 90.3 | 77.8 | 94.8 | 89.3 | 61.6 | 91.2 | 85.6 |
| SimCSE | 85.6 | 90.1 | 90.7 | 74.6 | 91.1 | 89.2 | 59.7 | 91.0 | 83.6 |
| TAPT | 85.2 | 91.3 | 90.2 | 78.2 | 93.7 | 90.1 | 61.5 | 90.9 | 85.1 |
| TapWeight (ours) | 86.8 | 92.5 | 91.1 | 80.7 | 94.9 | 90.2 | 62.3 | 91.2 | 86.2 |

#### 4.2.2 Experimental Settings

We perform continued training of a pretrained RoBERTa-base model on a raw-text dataset $\mathcal{D}$ consisting of 1 million sentences from Wikipedia (Gao et al., ). For downstream evaluation, we use the GLUE benchmark, which comprises 8 natural language understanding tasks, including sentiment analysis, semantic similarity prediction, and grammaticality classification (Wang et al., ). Following standard practices, we use the original GLUE development set as the test set in our experiments, and randomly split the original training set into a training set and validation set with a ratio of 8:1. We use Matthew’s Correlation for the CoLA dataset, Pearson/Spearman Correlation for the STS-B dataset, and accuracy for all other datasets. Our baseline methods are all based on a RoBERTa-base model, including direct finetuning, TAPT based continued pretraining and SimCSE based continued pretraining. When applying TapWeight on the RoBERTa-base encoder, we set the unrolling step in the MLO framework to 1. We use an Adam optimizer with a step learning rate scheduler across all three optimization levels. All experiments are conducted on NVIDIA A100 GPUs. More detailed descriptions are provided in the Appendix for the datasets (C.2), baselines (C.3), and hyperparameter settings (C.4).

#### 4.2.3 Results

Table 3 presents the results of various methods on 8 natural language understanding tasks from the GLUE benchmark. TapWeight consistently outperforms all baseline methods across all 8 datasets, showcasing the effectiveness of our method. On average, our method achieved a score of 86.2, while finetuning a RoBERTa model without continued pretraining only got 85.6. The superior performance on both molecule property prediction and natural language understanding highlights the generalizability of our method across multiple data modalities and downstream tasks. Moreover, our method surpasses the SimCSE method on all 8 tasks, showcasing the effectiveness of reweighting pretraining objectives, as SimCSE uses a fixed ratio between MLM and CL losses during continued pretraining. Additionally, TapWeight outperforms the RoBERTa+TAPT approach, demonstrating that our strategy of leveraging downstream datasets by reweighting pretraining objectives is more effective than simply pretraining the model with unlabeled downstream data, as TAPT does.

### 4.3 Ablation Studies

**Table 4: Ablation Studies. Results of molecular property classification using our method and baseline methods, in terms of AUROC. The best results are shown in bold.**
| Method | BACE | BBBP | ClinTox | Sider | Tox21 | ToxCast | HIV | MUV | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| CP w/o Reweighting | 78.8 | 66.1 | 77.4 | 60.3 | 74.6 | 62.7 | 76.9 | 71.6 | 71.1 |
| TapWeight w/o MLO | 83.0 | 68.5 | 79.5 | 63.5 | 76.3 | 65.9 | 77.2 | 77.3 | 73.9 |
| TapWeight | 83.1 | 71.2 | 81.3 | 64.5 | 77.0 | 66.1 | 78.4 | 80.5 | 75.3 |

In this section, we perform ablation studies to evaluate the effectiveness of individual components within our framework. All experiments are conducted on the classification tasks in the molecular property prediction benchmark.

##### Pretraining Objective Reweighting

We validate the effectiveness of our pretraining objective reweighting strategy by comparing our method to continued pretraining with a fixed importance for each objective. As shown in Table 4, our method outperforms this baseline (CP w/o Reweighting) across all datasets, demonstrating the advantage of dynamically reweighting pretraining objectives in the continued pretraining process.

##### Multi-level Optimization

We validate the effectiveness of the multi-level (tri-level) optimization (MLO) framework by reducing our method to a bi-level optimization (BLO) (Xie, ) based method. Specifically, we merge the first and second level of problems from the TapWeight framework to form the lower-level problem in the new BLO baseline, where the model is optimized jointly using both the unsupervised pretraining loss on the unlabeled continued pretraining dataset $\mathcal{D}_{pt}$ and the finetuning loss on the training split of the downstream dataset $\mathcal{D}_{tr}$. In the upper-level problem, the importance for each pretraining objective is learned using the validation split of the downstream dataset. Formally, we define the following BLO problem:

$$ $\displaystyle\min_{\lambda}\mathcal{L}_{val}(\theta^{*}(\lambda),\mathcal{D}_{ val})$ (9) $\displaystyle s.t.\quad\theta^{*}(\lambda)$ $\displaystyle=\operatorname*{arg\,min}_{\theta}\mathcal{L}_{pt}(\theta,\lambda ,\mathcal{D}_{pt})+\gamma\mathcal{L}_{tr}(\theta,\mathcal{D}_{tr})$ $$

However, optimizing these two types of losses in the lower level requires extensive tuning of the tradeoff parameters $\gamma$, and often leads to competition between losses which results in performance decrease. As shown in Table 4, our MLO based reweighting method outperforms the BLO based approach across all datasets, highlighting the advantage of formulating multiple optimization problems. Nevertheless, BLO method still outperforms the baseline continued pretraining methods with fixed tradeoff parameters, indicating the necessity of using reweighting strategies.

Figure: Figure 2: Evolution of the tradeoff parameter $\lambda$ over the training steps of TapWeight on the following downstream datasets: (a) Esol, (b) Lipo, (c) Freesolv, (d) Tox21, (e) Toxcast, and (f) Clintox.
Refer to caption: x2.png

### 4.4 Qualititive Analysis

In this section, we present the evolution trend of the pretraining objective weights along the training trajectory using our method. As shown in Figure 2, we plot the value of $\lambda$ for 3 regression tasks (Esol, Lipo, Freesolv) and 3 classification tasks (Tox21, Toxcast, Clintox) with respect to the global training step. Our observations reveal that different downstream datasets require varying importance for each pretraining objective. For example, the JPP pretraining objective, $\mathcal{L}_{jpp}$, plays an key role in Lipo and Toxcast datasets, whereas the MG3 pretraining objective, $\mathcal{L}_{mg3}$, is more critical for Esol, Freesolv and Tox21 datasets. The diverse requirements of pretraining objectives across downstream datasets emphasize the need for a reweighting method like TapWeight, providing a clear explanation for why our method outperforms baseline approaches.

### 4.5 Computation Cost

**Table 5: Training cost of baseline methods and our method TapWeight.**
| Dataset | FT | CP+FT | TapWeight |
| --- | --- | --- | --- |
| MUV | $\times 1$ | $\times 2.18$ | $\times 3.29$ |
| Qm9 | $\times 1$ | $\times 2.76$ | $\times 3.93$ |
| QQP | $\times 1$ | $\times 2.54$ | $\times 3.76$ |

In this section, we compare the training time of our method with baseline methods on the QQP, MUV and Qm9 datasets, as shown in Table 5. We use finetuning (FT) and continued pretraining with a fixed tradeoff ratio (CP+FT) as baselines, normalizing the time cost of FT as 1. While TapWeight results in an increase in training time compared to FT and CP+FT, its substantial improvement across multiple downstream tasks generally justifies the additional cost. However, in real-world applications where training time is a critical factor, TapWeight may not be the ideal choice, representing a limitation of our approach.

## 5 Conclusion and Future Work

In this paper, we propose a task-adaptive continued pretraining method that dynamically reweights each pretraining objective within a multi-level optimization framework. Unlike previous approaches that use a fixed ratio between pretraining objectives, our method adjusts the importance of each objective based on feedback from downstream datasets. Experiments in both molecule property prediction and natural language understanding validate the effectiveness and generalizability of our method.

Given the success of TapWeight, several promising future research directions emerge. For instance, large multimodal pretrained models have recently gained popularity (Liu et al., ; Zhu et al., ). The combination of multiple modalities introduces a greater number of potential continued pretraining objectives, presenting necessities of applying TapWeight in this context. Additionally, exploring objective reweighting strategies for general pretraining, rather than task-adaptive pretraining (TAP), is another promising direction. Unlike TAP, pretraining objective reweighting in general domain presents greater challenges for algorithm efficiency, as general pretraining typically incurs significantly higher computational costs.

## 6 Reproducibility Statement

We provide the code of TapWeight at [https://anonymous.4open.science/r/TapWeight-9A2E](https://anonymous.4open.science/r/TapWeight-9A2E). In the code repo, we provide instructions on how to reproduce experimental results for both molecule property prediction and natural language understanding. Furthermore, we include detailed experimental settings of molecule property prediction in Section 4.1.2, with more information on selection of hyperparameters in Appendix B.4. For natural language understanding, we include detailed experimental settings in Section 4.2.2, with more information on selection of hyperparameters in Appendix C.4.

## Appendix A Optimization Algorithm

In this section, we give an example to briefly illustrate how to use Implicit Function Theorem (IFT) to compute best-response Jacobian matrices. Take $\pdv{\theta^{*}}{\lambda}$ term in Equation 6 as an example: although $\theta^{*}$ is an implicit function of $\lambda$, the exact value of $\theta^{*}(\lambda)$ given a value of $\lambda$ is usually approximated with gradient descent algorithms. As there is no analytical solution of $\theta^{*}(\lambda)$, it is difficult to directly compute the gradient $\pdv{\theta^{*}}{\lambda}$. To tackle this challenge, we compute this gradient using IFT following previous literature (Lorraine et al., ):

$$ $\displaystyle{\pdv{\theta^{*}}{\lambda}}={-\color[rgb]{1,0,0}[\nabla^{2} \mathcal{L}_{pt}(\theta)]^{-1}}\times{\color[rgb]{0.5,0.5,0}\pdv{\mathcal{L}_{ pt}}{\theta,\lambda^{T}}}$ (10) $$

The olive term, a second-order mixed partial derivative matrix, can be directly computed using automatic differentiation. Nevertheless, directly computing the red term, which is the invert of a Hessian matrix $\nabla^{2}\mathcal{L}_{pt}(\theta)$, is computational expensive due to its $O(n^{3})$ complexity. Various methods have been proposed to approximate the inverted Hessian matrix, including Neumann series (Lorraine et al., ), conjugate gradients (Rajeswaran et al., ) and finite difference (Zhang et al., ). In TapWeight, we select finite difference as the approximation method, thus enabling efficient computation of best-response Jacobian matrices.

## Appendix B Molecule Property prediction

### B.1 Pretraining Objectives

We use 3 types of pretraining objectives for continued pretraining of an Imagemol model (Zeng et al., ) to enhance its performance on molecule property prediction tasks.

##### Multi-Granularity Clustering

In this pretraining objective, we first perform K-means clustering to the unlabeled training dataset of molecules using their chemical structural fingerprint. After clustering, each molecule is assigned with a pseudo-label, and the molecular encoder model is pretrained by predicting this label. Formally,

$$ $\displaystyle\mathcal{L}_{mg1}=\sum_{i=1}^{n}\mathcal{L}(C^{100}(f_{\theta}(x_ {i})),y_{i}^{100})$ (11) $\displaystyle\mathcal{L}_{mg2}=\sum_{i=1}^{n}\mathcal{L}(C^{1,000}(f_{\theta}( x_{i})),y_{i}^{1,000})$ (12) $\displaystyle\mathcal{L}_{mg3}=\sum_{i=1}^{n}\mathcal{L}(C^{10,000}(f_{\theta} (x_{i})),y_{i}^{10,000})$ (13) $$

where $f_{\theta}$ is the molecular encoder, and $C$ are task-specific fully-connected neural networks for clustering label prediction.

##### Mask-based Contrastive Learning

In this pretraining objective, we use a 16 × 16 square area to randomly mask a molecular image $x$ to generate the masked image $\hat{x}$. We then perform constrastive learning on the image pair $(x,\hat{x})$ by minimizing the distance between representations of both images to promote consistency. Formally,

$$ $\displaystyle\mathcal{L}_{mcl}=\sum_{i=1}^{n}||f_{\theta}(x_{i}),f_{\theta}( \hat{x_{i}})||_{2}$ (14) $$

where $||\cdot||$ denotes the Euclidean distance between two molecular representation generated from the encoder.

##### Jigsaw Puzzle Prediction

In this pretraining objective, we introduce 100 types of different permutations with number 1 to 100, denoted as $y^{jig}$. We also assign a label of 0 for original molecular image without any We apply the permutation to molecular images $x$ to get permuted ones $\hat{x}$. The encoder $f_{\theta}$ is pretrained by predicting the permutation label. Formally,

$$ $\displaystyle\mathcal{L}_{jpp}=\sum_{i=1}^{n}\mathcal{L}(C(f_{\theta}(\hat{x_{ i}})),y_{i}^{jig})$ (15) $$

where $C$ is a task-specific fully-connected neural network for permutation label prediction.

### B.2 Datasets

We use the datasets from MoleculeNet benchmark for molecule property prediction Wu et al. ().

##### Quantum Mechanics

Qm7 and Qm9 are both molecular datasets for regression task on quantum mechanics properties of molecules. Qm7 dataset collects electronic properties of molecules determined using ab-initio density functional theory (DFT). Qm9 dataset collects geometric, energetic, electronic and thermodynamic properties of DFT-modelled small molecules.

##### Physical Chemistry

Esol, FreeSolv and Lipophilicity (Lipo) are all datasets for regression task on physical chemistry properties of molecules. ESOL dataset collects water solubility data for common organic small molecules. FreeSolv dataset collects experimental and calculated hydration free energy of small molecules in water. Lipo dataset collects experimental results of octanol/water distribution coefficient.

##### Biophysics

Bace, HIV and MUV are all datasets for classification tasks on biophysics properties of molecules. BACE dataset collects binary label of molecular binding results for a set of inhibitors of human $\beta$-secretase 1 (BACE-1). HIV dataset collects experimentally measured abilities of a molecule to inhibit HIV replication. MUV is a subset of PubChem BioAssay by applying a refined nearest neighbor analysis, designed for validation of virtual screening techniques.

##### Physiology

BBBP, Clintox, Sider, Toxcast and Tox21 are all datasets for classification tasks on physiology properties of molecules. BBBP dataset contains binary labels of blood-brain barrier penetration (permeability) ability for molecules. ClinTox dataset consists of qualitative data of drug molecules approved by the FDA and those that have failed clinical trials for toxicity reasons. Sider is a database of marketed drugs and adverse drug reactions (ADR), grouped into 27 system organ classes. ToxCast dataset contains toxicology data for a large library of compounds based on in vitro high-throughput screening, including experiments on over 600 tasks. Tox21 dataset collects qualitative toxicity measurements of molecules on 12 biological targets, including nuclear receptors and stress response pathways.

### B.3 Baseline Methods

##### Attribute Masking

Attribute masking (AttrMask) based pretraining captures domain knowledge by learning the regularities of the
node/edge attributes distributed over graph structure (Hu et al., ). Inspired by BERT (Devlin et al., ), it pretrains a graph neural network (GNN) by first masking node/edge attributes and then letting GNNs predict those attributes based on neighboring structure.

##### Context Prediction

Context Prediction uses subgraphs to predict their surrounding graph structures (Hu et al., ). It pretrains a GNN so that it maps nodes appearing in similar structural contexts to nearby embeddings. Specifically, the method first encodes the context into a fixed vector using an auxiliary GNN, and then trains the GNN encoder with negative sampling.

##### GraphMVP

The Graph Multi-View Pre-training (GraphMVP) framework applies self-supervised learning (SSL) by utilizing the correspondence and consistency between 2D topological structures and 3D geometric views (Liu et al., ). It introduces a novel contrastive learning loss, using the 2D and 3D representations of the same molecule as positive pairs.

##### Imagemol

ImageMol is an unsupervised pretraining deep learning framework pretrained on 10 million unlabelled drug-like, bioactive molecules, to predict molecular targets of candidate compounds (Zeng et al., ). The ImageMol framework is designed to pretrain chemical representations from unlabelled molecular images on the basis of local and global structural characteristics of molecules from pixels.

### B.4 Hyperparameter Settings

##### Classification

We set the global learning steps to be 30,000 for MUV dataset, 20,000 for HIV dataset, 10,000 for Tox21 and Toxcast datasets, and 3,000 for all other datasets. We set the batch size in level I to be 1024, and that in level II and level III to be 64 for all datasets. We set the learning rate to be 0.02 in level I, 0.05 in level II, and that in level III to be 200 for all datasets. We set the $\gamma$ value in Equation 3 to be 0.001.

##### Regression

We set the global learning steps to be 30,000 for qm9 dataset and 10,000 for all other datasets. The batch size and $\gamma$ are the same as those in classification tasks. We set the learning rate to be 0.02 in level I, 0.001 in level II, and 1 in level III for Lipo, Esol and FreeSolv datasets. We set the learning rate to be 0.01 in level I, 0.0001 in level II, and 0.1 in level III for Qm7 and Qm9 datasets.

## Appendix C Natural Language Understanding

### C.1 Pretraining Objectives

We use 3 types of losses for continued pretraining of an RoBERTa model Liu et al. () to enhance its performance on natural language understanding tasks.

##### Mask Language Modeling

This pretraining objective randomly mask some percentage of the input tokens, and then predict those masked tokens using embedding generated from the pretrained model (Devlin et al., ). In BERT and RoBERTa, 15% of the tokens are masked in the pretraining stage.

##### Constrastive Learning

This pretraining objective applies dropout noise to the encoder $f_{\theta}$ when taking in a sentence $x$ to get a negative sample of encoding $h^{\prime}=f_{\theta}(x)$ (Gao et al., ). We use $h$ to denote those positive encodings without dropout noise. The encoder is then trained by minimizing a constrastive learning loss:

$$ $\displaystyle\mathcal{L}_{cl}=-\sum_{i=1}^{n}\log(\frac{e^{sim(h_{i},h_{i}^{ \prime})}}{\sum_{j=1}^{n}e^{sim(h_{i},h_{j}^{\prime})}})$ (16) $$

where $sim$ is a similarity measure between two encodings.

##### Sentence Order Prediction

This pretraining objective uses two consecutive segments from the same document as positive examples. It generates negative examples using the same two consecutive segments but with their order swapped (Lan et al., ). The model is pretrained by predicting the label of these two types of examples.

### C.2 Datasets

We use 8 datasets from GLUE benchmark in natural language understanding tasks (Wang et al., ).

##### Single Sentence Tasks

The Corpus of Linguistic Acceptability (CoLA) contains English acceptability judgments sourced from books and journal articles on linguistic theory. The Stanford Sentiment Treebank (SST-2) features sentences from movie reviews annotated by humans for sentiment analysis.

##### Similarity and Paraphrase Tasks

The Microsoft Research Paraphrase Corpus (MRPC) is a dataset of sentence pairs extracted from online news sources, annotated by humans for semantic equivalence. The Quora Question Pairs (QQP) dataset includes question pairs from the Quora website, where the task is to determine if the questions are semantically equivalent. The Semantic Textual Similarity Benchmark (STS-B) contains sentence pairs from news headlines, video and image captions, and natural language inference datasets, with the task of predicting a human-annotated similarity score.

##### Inference Tasks

The Multi-Genre Natural Language Inference Corpus (MNLI) is a crowdsourced dataset of sentence pairs annotated for textual entailment, where the task is to predict the relationship between a premise and a hypothesis. Question-answering Natural Language Inference (QNLI) involves question-paragraph pairs, with the task of determining whether the paragraph contains the answer to the question. The Recognizing Textual Entailment (RTE) datasets consist of sentence pairs from news and Wikipedia, where the task is to predict the entailment between two sentences.

### C.3 Baseline Methods

##### RoBERTa

The Robustly Optimized BERT Pretraining (RoBERTa) paper (Liu et al., ) thoroughly evaluates the impact of key hyperparameters and training data size in BERT. RoBERTa uses the same architecture as BERT but is pretrained with an optimized strategy, leading to significant improvements in performance across various downstream tasks. The main differences between RoBERTa and BERT are: (1) training for a longer duration with larger batches and more data; (2) removing the next sentence prediction objective; (3) training on longer sequences; and (4) dynamically adjusting the masking patterns applied to the training data.

##### SimCSE

The Simple Contrastive Learning of Sentence Embeddings (SimCSE) framework includes both unsupervised and supervised approaches. In the unsupervised approach, SimCSE takes an input sentence and predicts the same sentence using a contrastive objective, where standard dropout serves as the noise. In the supervised approach, it integrates annotated pairs from natural language inference datasets into the contrastive framework, using human-labeled ”entailment” pairs as positive examples and ”contradiction” pairs as hard negatives.

### C.4 Hyperparameter Settings

We set the global learning steps to 20,000 for the QQP and MNLI datasets, and 10,000 for all other datasets. The batch size for level I is set to 512, while for levels II and III, it is set to 32 across all datasets. The learning rate for levels I and II is 2e-5, and for level III, it is set to 1 for all datasets. We set the $\gamma$ value in Equation 3 to be 0.005.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: see weakness section
soundness: 3
strengths: * The proposed three-level approach is flexible and demonstrates transparent performance improvements.
* The qualitative analysis effectively shows how different downstream tasks require different weights for pretraining objectives, justifying the need for adaptive weighting.
summary: This paper proposes TapWeight, which automatically determines optimal loss weights for different pretraining objectives through a three-level optimization approach.
weaknesses: Major weakness: 
* The evaluation on GLUE benchmark, which has been saturated in recent years, would benefit from including more challenging datasets. E.g., Crossfit (160 downstream tasks) https://arxiv.org/pdf/2104.08835
* The effectiveness of the pretraining objectives on more powerful generative models or LLMs remains unexplored. For example 1.  gpt3.5 result on more challenging glue-x: https://aclanthology.org/2023.findings-acl.806.pdf 2. three-stage LLM method for task adaptive pretraining: https://arxiv.org/abs/2402.05140
* Missing reference and comparison with other recent MLM models (they have higher performance on the GLUE benchmark): e.g., Pre-training Language Model as a Multi-perspective Course Learner https://arxiv.org/pdf/2305.03981

Other weakness:
* The selection criteria for multi-objective pretraining tasks in Level I needs justification, particularly why certain tasks were chosen over alternatives like token detection or swap detection.
* line 452-457: The effects of pretraining task to downstream tasks need a more in-depth investigation. Particularly, according to the trade-off score result, how will these findings better inform us on designing new or matching existing pretraining tasks with downstream tasks?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - Can the authors include experiment results of more types of tasks other than NLU and molecular prediction tasks? E.g., legal domain, financial domain, etc. Otherwise, I don't think this method has wide applications.
soundness: 3
strengths: - The proposed TapWeight method can effectively increase model performance on targeted downstream tasks that are different from the pre-training domain.
- The authors also conduct a few ablation studies to show the importance of reweighting and multi-level optimization.
- The code is released for reproducibility.
summary: This paper proposes TapWeight, an approach to adapt weights for different pre-training objectives according to downstream tasks. The authors propose a three-level optimization, where the first level is the traditional pre-training level, the second level is the downstream fine-tuning level, and the third-level is downstream validation level to finally provide signals for weight updates. In the experiments, the authors demonstrate that TapWeight can outperform baselines on molecular tasks, and outperform the pre-train fine-tuning pipeline on NLU tasks. They also conduct a few ablation studies to show the importance of reweighting and multi-level optimization.
weaknesses: - The proposed multi-level optimization is not new, and can actually be seen as meta-learning.
- The experiment results are not well presented. For example, table 2 is about "results" of molecular prediction, but the authors do not explain the metrics in the table, and do not use up-arrow or down-arrow to show whether the numbers should be bigger or smaller.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Could the authors provide more information about how the TapWeight update work? What is the granuarly? How many tokens of CP vs. SFT is seen per-step? Is it equal? Is my understanding correct that CP+SFT is jointly optimized in TapWeight? 

If above is true, is it true then that the loss weights adaptively update during training? Have you tried taking the final weights and doing another run with the final weights, fixed?

Is extra SFT applied after TapWeight?

How many A100s are used in each phase of training, and for how long? (More information about Table 5).
soundness: 3
strengths: 1. Finding effective ways of optimizing the choice of pretraining objective to optimize downstream performance is an important problem that would be of interest to the greater community. The authors provide a novel connection to MLO in this work, and the formal solving of this problem is valuable and welcome given the amount of empirical hillclimbing typically done in this space. 

2. The paper shows convincing results that their formal optimization of the whole MLO objective works from a quality perspective, and yields results that improves over the baselines they consider. Given the complexity of solving such an MLO problem, these results are promising and this technique may be of interest to the community. 

3. The presentation of the work is nice. It is well written, structured, and easy to follow, free of any grammatical mistakes. There are some aspects about the method that could be clearer, but overall this does not detract from the presentation of the paper.

4. Evaluation of the method on both molecular biology and language demonstrate well the generality of the method. Especially as many (5) pretraining objectives are considered for the molecular image model. 

5. Reproducibility is excellent. The authors have already open sourced the code, and provide all information needed to reproduce their experiment.
summary: This paper introduced a method called TapWeight, to automatically tune the loss weights of multiple pretraining objectives during task-adapted continued pretraining. In task adapted continued pretraining, a new mixture of pretraining objectives is introduced, targeted to improve a specific suite of downstream tasks. Downstream task performance is measured by applying SFT to the new pretraining checkpoint, and evaluating on the validation or test split. The main problem the paper aims to solve is how to balance the pretraining objectives in the continued pretraining mixture in order to maximize downstream task suite performance.

TapWeight tackles this problem by formally framing this as a three tiered multi-level optimization (MLO) problem, where Tier 1: Continued Pretraining given some loss weights, Tier 2: SFT, Tier 3: Optimizing pretraining loss weights to maximize validation performance. Note that this paper only studies the work as applied to Encoder-only models, and SFT finetuning (until convergence.) The SFT phase is slightly altered in the training setup to deduce dependence with other tiers (using the continued pretrain weights as regularization instead of initialization), and the MLO optimization employs the Betty library to solve (with some tricks to approximate the inverted Hessian, as noted in the Appendix).

The approach is evaluated on a popular, standard, molecular property prediction model, Imagemol, as well as on NLP tasks (GLUE), via RoBERTa. In both domains, the method outperforms compared baselines. The authors also provide ablation result for each component of their method.
weaknesses: 1.) Soundess: Most of the experiments in this work compare against baselines that do not apply task-adaptive continued pretraining. It is only in Table 4 that the authors show this ablation result for the molecular results (but not the NLP results.) In Table 3, it is unclear whether the method is gaining because of continued pretraining in general or because of TapWeight. TAPT represents a continued pretraining comparison, but from a different data distribution. 

This work would benefit from having more baselines that represent basic ways of hillclimbing the pretraining mixture. Some examples:
- All pretraining objectives weighted evenly in the mixture.
- Leave-one-out or one-at-a-time ablations for the pretraining mixture. 
- Even weighting of only the objectives that maximize certain tasks in the eval suite + any other grid-search like approach.
- This dataset ablation can take place over smaller models (to use less compute), before being applied to the final model size.

In Table 5, TapWeight is 1.5x the cost of normal training, so it is possible that TapWeight may be a way to save compute over traditional hillclimbing, but having these explicit compute & quality comparisons against more baselines would largely strengthen the paper. 

2.) Contribution: The contribution of this work is solid for the targeted application: continued pretraining (CP) of encoder-only models to maximize SFT tasks. However, there there is also a implicit assumption that the continued pretraining computation budget is small. In Table 5, the training budget of CP is only a little bit more expensive than that of SFT. It is unclear how true this assumption is for most modern models. It is unclear also how TapWeight would perform or handle much longer training horizons, or where CP>>SFT in terms of computational budget and they need to see a different budget of tokens.

Also note: Table 5 also only reports time rather than FLOPs, and its not possible for the reader to understand whether each phase occurs on the same amount of hardware.

The paper is missing a crucial baseline that might help the reader understand the wider applicability of the technique: training CP+SFT (Table 5) for the same time / FLOPs as TapWeight, to observe if the gains from TapWeight outweight the computation cost. (e.g. If you make the cheaper alternatives use as much compute, would they do better? Is TapWeight yield training compute efficiency gains?)

Note however that the authors are relatively upfront about limitations here. 

3.) Presentation: The description of the method relies on an understanding of MLO, which is OK, but the paper would be clearer if more descriptions were provided, especially as it applies to CP. For example: it seems like MLO does CP and SFT in a single phase, as implied by Table 5, and the optimization equations. However, Figure 1 somewhat suggests that CP is done then an entire SFT run is done, before updating the loss weights. Exactly how many tokens of CP vs. SFT is seen per step is also unclear (or is it assumed they see the same number of tokens, 1:1 for each step?)

There's also some work that may be relevant, but missing from related works:
- UL2: https://arxiv.org/pdf/2205.05131
- U-PaLM: https://arxiv.org/abs/2210.11399 (CP version)
A lot of manual ablation was done for the construction of UL2 and UL2R in U-PaLM. These would serve as very natural candidates for a more learned continual pretraining approach. 

4.) Soundness: The authors make the claim Line 88 “Moreover, TapWeight is applicable to any pretrained model with multiple pretraining objectives”. This is likely true in theory, but the authors lack empirical evidence for this given their limited setting (e.g. how about decoder-only models?). 

TapWeight as described also is only defined over loss weights of multiple pretraining objectives, but this is not well defined for pretraining objectives that need to see different amounts of different pretraining objectives, like UL2 mentioned above. Please correct me if I am wrong here.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Could you provide more details on the convergence rate of TapWeight, especially compared to traditional pretraining methods with fixed weights? How does the computational overhead impact practical training times on large-scale datasets?

2. I am curious to know whether there are any patterns observed in the reweighting behavior across tasks, such as prioritizing certain objectives over others based on task type. Were there some pre-training tasks that were given negligible weights? 

3. If possible, can you provide an analysis using larger models, and generative tasks, instead of classification? Additionally, a framework would be considered stable if an ablation could be shown across model families and varying scales. 

4. To obtain the robustness of the approach, a random task which has no relation with the downstream task can be augmented with the set of other pre-training tasks, just to see of TapWeight assigns a very low weightage to the random task. 

5. **A suggestion, please keep the tables in the page where the corresponding description is present, improves the readability of the paper.
soundness: 3
strengths: 1. The proposed framework automates the weighting of pretraining objectives, by optimizing the importance of each pretraining objective in a dynamic, data-driven manner.

2. The authors showcase experiments across two domains - molecular property prediction, and NLU

3. The authors provide an open-source implementation for their approach, adding transparency and reproducibility.
summary: The paper introduces a new task-adaptive pretraining (TAP) framework, called TapWeight, designed to improve downstream model performance by automatically optimizing the importance of each pretraining objective. Unlike traditional TAP methods that rely on manually set tradeoff parameters between multiple pretraining objectives, TapWeight dynamically learns these parameters by solving a multi-level optimization (MLO) problem. The authors report empirical results from the domain of molecular property prediction and natural language understanding, demonstrating its performance across 13 molecular and 8 language datasets compared to baseline TAP methods.
weaknesses: 1. The tasks chosen for evaluation are primarily classification. Given the current state-of-the-art LLMs, and VLMs, and the corresponding benchmarks, no analysis in that regard has been shown. For example, some maths problem solving tasks such as Math QA, MetaMath etc. can be used to constitute the pre-training tasks, and evaluated on a task such as GSM8K. 

2. The number of baselines used for comparison is limited. For example, the paper "Don’t Stop Pretraining? Make Prompt-based Fine-tuning Powerful Learner", is a strong baseline candidate for the proposed method, as the paper builds further upon the TAP method. 

3. Increased computational complexity limits the scalability of the method to larger models. 

4. The paper does not discuss or provide an empirical analysis of how quickly the model converges compared to fixed-weight pretraining methods.

Meta-review: This paper introduces TapWeight, a task-adaptive pretraining framework that automatically optimizes the importance of various pretraining objectives based on downstream feedback through a multi-level optimization problem. TapWeight aims to improve the performance of models on specific downstream tasks by making pretraining more relevant to the target domain. 

However, as pointed out by the reviewers, there are several concerns about the methodology and experimental setup: The use of the GLUE benchmark, which is considered to be saturated, limits the ability to truly gauge the improvements made by TapWeight; there is a lack of comparison with more current and challenging datasets and with recent generative models or LLMs which might have provided a clearer picture of TapWeight's advancements; the selection criteria for multi-objective pretraining tasks are not sufficiently justified, leaving questions about the rationale behind choosing certain tasks over others; and the experimental results presentation lacks clarity in some aspects, such as metrics explanation and the significance of numerical outcomes.

Even though the authors tried to address these concerns during the rebuttal, the reviewers were not fully satisfied, particularly with the response to the novelty and the generalizability of the approach.

Decision: Reject