TapWeight: Reweighting Pretraining Objectives for Task-Adaptive Pretraining
Authors: Ruiyi Zhang
Sai Ashish Somayajula
Pengtao Xie
Venue: Accepted by TMLR ()

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 processing tasks, significantly surpassing baseline methods. Experimental results validate the effectiveness and generalizability of TapWeight.

Published in Transactions on Machine Learning Research (06/2025)
TapWeight: Reweighting Pretraining Objectives for Task- Adaptive Pretraining
Ruiyi Zhang UC San Diego
rua0488rusd edu
Sai Ashish Somayajula UC San Diego
ssomayajBuussd edu
Pengtao Xie
plaic frcssd edu
UC San Diego
Reviewed on OpenReview: https: // openreview. net/ forum? DCCu2CEVFS
Abstract
Large-scale general domain pretraining followed by downstreaa--seeiifc finetuning has be- come a predominant paradigm in machine learning. However, discrepancies between the pretraining and target domains can still lead to performance degradation certain cases, un- derscoring the need for task-adaptive continued pretraining (TAP). TAP methods typically involve continued pretraining on lask-specific unlabeled datasets or introducing additional unsupervised learning objectives 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 oulcomes and higher computational costs. In this paper, we propose Weight, task-adaptive pretrain- ing framework which automatically determines the optimal importance of each pretraining objective based on downstream feedback. WapWhight reweights each pretraining objective by solving a multi-level optimization problem. We applied TeiWeight to both molecular property prediction and natural language processing tasks, significantly surpassing baseline methods. Experimental results validate the effectiveness and generalizability of Teight Our code is available at https://github. com/ruz044/Tapieighh
1 Introduction
Foundation models pretrained on large-scale general domain corpora have achieved state-of-the-art perfor- mance across a wide range of tasks (He et al., 2021a Devlin et al.) 2019 Brown et al. 2020). These models, which capture general knowledge for specific modalities such as text or images through unsupervised learn- ing, are typically adapted to downstream tasks via finetuning. However, when there is 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. 2020 Chen et al. 2023 Xie et al. 2024 To address this challenge, downstream task-adaptive continued pretraining, 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. (2020) conducts task-adaptive pretraining by performing unsupervised learning on the unlabeled data of the downstream task. [Wu et al.) (2021) introduces an additional perturbation masking objective during continued pretraining a BERT model (Devlin et al. 2019) enhancing its performance dialogue understanding tasks.

---

Published in Transactions on Machine Learning Research (06/2025)
Trainable Importance Pretraining Objectives
Update Imporiance of Pretraining Objectives
30%
50%
20%
Contrastive
Sentence Order Prediction
Downstream Task
Downstream
Masked Language
Learing
[Training Split]
(Validation Split)
Inset
ganteeee satull uerd
Modeling
you
Grareaticll
Model
Pretrained Encoder
Finetuned Encoder
Level
Level Finetime Encoder
Level Update Importance OR Downstream Validation Split
Continued Pretraining Encoder
with Pretraining Objective Importance Fixed
Downstream Training Split
Figure 1: An Overview of Tap Weight. In the first level, the model undergoes multi-objective pretraining with fixed tradeoff ratios between objectives. In the second level, the pretrained model is finetured the training split of the downstream dataset. In the third level, the finetuned model is evaluated on the validation split the downstream dataset to compute a loss, and the trainable tradcoff parameters fixed in the first level are learned by minimizing this validation loss.
Among these, many existing task-adaptive pretraining methods consist of multiple pretraining objectives (Wu et al. 2021 Gao et al. 2021 Cui et al. 2023), making it challenging to determine the relative importance of each objective. Some TAP methods assign equal weight to each pretraining objective (Lee et al.) 2020:
Wu et al. 2021), disregarding their varying impact on downstream performance. For instance, Gao et al. (2021) shows that pretraining BERT with a contrastive learning (CL) objective results better downstream performance on semantic textual similarity (STS) datascts 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. 2021 which often results in suboptimal performance and increased computational costs. This issue becomes particularly severe when the mumber of pretraining objectives is large, such as with the task-adaptive pretraining of popular molecular model Imagemol, which involves 5 distinct pretraining objectives (Zeng et al. 2022). Raghu et al. (2021) propose a multi-level framework to learn the importance of pretraining objectives, with locus on supervised pretraining. However, their approach faces challenges when scaling to large-scale unsupervised pretraining due computational cost.
To address the aforementioned challenges, we propose 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 downstream task, achieves the best downstream task performance. Our approach involves a three-level optimization framework to learn these parameters. In the first level, we perform adsk-adaptive pretrain- ing using initial tradeoff parameters, denoted as X. These parameters are kept fixed in this level and will be updated in subsequent levels. The resulting pretrained model thus a function of X. In the second level, the pretrained model from the first level is finetuned on the training split of the downstream dataset. Consequently, the finetuned model becomes implicit function of the tradeoff parameters. In the third level, the finetuned model is evaluated on the validation split of the downstream dataset, and the tradcoff parameters X are optimized by minimizing the validation loss. This end-to-end process allows the opti- mization problems in three levels to dynamically influence one another, forming an integrated framework that optimizes task-adaptive pretraining process and enhances downstream task performance. Morcover, TapWeight is broadly applicable to pretrained models with multiple pretraining objectives across various data modalities and downstream task types, demonstrating superior generalizability compared to existing
2

---

Published in Transactions on Machine Learning Research (06/2025)
task-adaptive pretraining methods (Nishida et al. 2021 Cui et al. 2023 Figure illustrates the complete
framework of Weigtt
We apply Weigt for task-adaptive pretraining of molecule representation model, Imagemol (Zeng et al. 2022), and language models including RoBERTa (Liu et al., 2019c) and DeBERTa (He et al. 2021b) Evaluating its performance across 13 molecular property prediction datasets and natural language pro- cessing tasks, WeiWt significantly outperforms baseline methods across different tasks and model sizes. The superior performance of WepWeight highlights its effectiveness and generalizability. Our contribution can be summarized follows:
We propose Wepheight, an approach that automatically scarches for the tradcoff parameters across multiple pretraining objectives and performs reweighted task-adaptive pretraining. TapWeight is formulated within a multi-level optimization (MI.)) framework. We employ an efficient gradient descent algorithm to solve the MLO problem, obtaining the optimal tradeoff parameters for multiple pretraining objectives.
We apply TeiWtight for lask-adaptive pretraining of a molecule representation model and a language model. Extensive experiments 13 downstream datasets in molecular property prediction and 11 datasets in natural language processing 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 specific domain, domain- adaptive pretraining (DAP) and task-adaptive pretraining (TAP) have been introduced (Gururangan et al. 2020). 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. 2020 Similarly, LegalBERT continues to pretrain BERT model on legal documents to improve performance on legal NLP tasks (Chakiidis al. 2020), while SeiBERT leverages a large multi-domain corpus of scientific publications for further pretraining, enhancing its effectiveness scientific NLP tasks (Beltagy et al., 2019). 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. 2023
U-PaLM (Tay et al., 2023) performs continued pretraining on PaLM [Chowdhery et al. 2024) model with 540 billion parameters using UL2's mixture-of-deenoser pretraining objective (Tay et al. 2022 achicving performance improvement on many few-shot lasks, such as MMLU and GSMSK.
Although DAP significanlly 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. 2020). 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 al., 2021). SimCSE introduces additional constrastive learning loss in addition to the original masked language modelling loss to further pretrain RoBERTa model, specifically enhancing its capability on standard semantic textual similarity tasks (Gao et al. 2021]. PCP combines the idea of instruction tuning with conventional continued pre- training, consistently improving the performance of state-of-the-ar! prompt-based finetuning approaches on 21 benchmarks (Shi & Lipani 2023). While existing TAP methods are effective, they are typically tailored to specific downstream tasks or data modalities (Wu al., [021 Cui et al., (2023). In contrast, Tap is applicable to pretrained models across diverse modalities and tasks with multiple pretraining objectives, underscoring its broad generalizability.
3

---

Published in Transactions on Machine Learning Research (06/2025)
2.2 Multi-level Optimization
Many machine learning tasks can be formulated as multi-level optimization (MLO) problems, such as neural
architecture search [iu et al. 2019a, Chen et al. 2019 Xu et al. 2020), meta learning (Finn et al. 2017
Rajeswaran et al. 2019 Zhang et al. 2024h, and hyperaarameter optimization (Lorraine et al. 2020 Lorraine
& Duvenaud 2018 Mackay et al. 2019). MLO problems consist of multiple levels of optimization problems that are mutually dependent, making it challenging for common automatic differentiation algorithms
handle them. To tackle this challenge, multiple algorithms (Lorraine et al., 2020: Liu et al.
2019a; Rajeswaran
problems. et al. 2019) and librarics [Choe ct al. 2023c ab have been proposed to efficiently compute gradients in ML.
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, Meta Wei learns an explicit weighting function for each data point to maximize the performance on a small amount of unbiased meta-data (Shu et al.) 2019). DoGE optimizes weights for each data domain using small proxy model to guide the pretraining of larger models (Fan ct al. 2024 MetaWeighting learns tradeoff parameters for each task in multi-task learning to minimize general izaatin loss (Mao et. al. 2022). Raghu et. al. (2021) proposes a multi-level optimization framework to learn the importance of supervised pretraining objectives based on feedback from downstream performance. Their method simulates the finctuning process by unrolling a few finctuning steps within one optimization level. Our method also falls within this category, with a specific focus on reweighting pretraining objectives for downstream task-adaptive continued pretraining. While our method shares three-level optimization framework with Raghu et al. (2021) for identifying the importance of pretraining objectives, we adopt proximal regularization to simulate the finetuning process, reducing computational cost and enabling scalable unsupervised pretraining on large models.
2.3 Multi-task Learning
Multi-task learning (MTL) enables models to learn multiple tasks simultancously, promoting knowledge shar- ing and transfer while mitigating task conflict. Existing MTL methods can be broadly categorized into two groups. The first group focuses primarily on architecture design for parameter sharing. For example, |Ruder (207)) introduces two typical approaches for parameter sharing: hard sharing and soft sharing. Misra et al.
2016 proposes using trainable linear mappings to dynamically select different combinations of activation maps for different tasks. Rosenbaum et al. (2018) introduces a trainable router nelwork to ileratively select functional blocks for different tasks. Liu et al. (2019b) employs soft attention modules to extract task-specific features from shared representations. Similarly, y., Yang et al. (2020) also utilizes a trainable router network, as in Rosenbaum et al. (2018), but with soft modularization that assigns probability weights to each connection between blocks.
The second group of MTL methods is based on optimization and gradient operations. For instance, Sener & Koltun (2018) formulates MTL as multi-objective optimization problem, defining the overall objective as finding Pareto-optimal solution. Chen et al. (2018) proposes a gradient normalization algorithm that dynamically adjusts the gradient magnitudes of each task. Yu et al. (2020) introduces gradient surgery, which projects a task's gradient onto the normal plane of another task's gradient if task confficts exist. Mao et al. b applies meta-learning to search for optimal task weights. Achituve et al. (2024) proposes the first Bayesian formulation for gradient aggregation in MTL and develops a new optimization algorithm basod on posterior estimation.
Our method, TapWeight, shares more similarity with the second category as it also operates gradients by reweighting each pretraining objective. However, TapWeight differs critically from all the MTL methods mentioned above, which generally assume that training and lesting tasks are identical and do not consider a pretraining finetuning scheme. In contrast, Tap Weight specifically addresses the continued pretraining (CP) problem, where the MTL tasks during CP differ from those during finctuning. In this setting, optimizing solely for performance on pretraining tasks does not guarantee improved finetuning outcomes, making direct application of existing MTI. methods inappropriate. To address this challenge, TapWeight introduces
4

---

Published in Transactions on Machine Learning Research (06/2025)
novel multi-level optimization framework that searches for pretraining task weights to maximize finetuning performance, evfectively overcoming the limitations of previous MTL approaches in the CP context.
3
Method
3.1
Overview
Given n continued pretraining objectives T1, T., and their corresponding training losses L1, L2 we formulate the multi-objective continued pretraining loss Lye 8s:
(0, E--E>xAi p) i=1
(1)
where Da the unsupervised pretraining dataset, A denotes the pretraining model parameters, and X is the tradeoff parameter for each pretraining objective. We denote the target downstream task as Dst and split it into Dir Deat and Dts which are training, validation and test splits respectively.
In our framework, TapWeight, (A.,...,An) so that the pretrained model achieves the highest performance on the test split Dr, after being fine-tuned on the downstream dataset Dtr. To achieve this, our method consists of three levels of opti- mization problems. In the first level, we perform continued pretraining of the model, with tradeoff weights tentatively fixed. In the second level, we conduct finetuning the pretrained model on the training split of the downstrcam dataset. In the third level, we compute a loss by applying the finctuned model on the validation split of the downstream dataset, and optimize the tradeoff parameters by minimizing this loss. We next formally define these three levels under multi-level optimization framework.
aim to automatically search for the optimal tradeoff weights
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 f to minimize the multi-objective pretraining loss Lpt unlabeled dataset Dpr
6" (A) = arg min Lee X, Dpt)
(2)
Since the optimal solution 0* to this problem depends on the value of the tradeoff parameter, it is implicit function of X, denoted 8" (2)..
Level II In the second level, aim to finetune the pretrained model with optimal parameters 0* ob- tained from previous level on the downstream dataset. Raghu et al. (2021) formulates the optimization problem using the same set of parameters at the lower level and differentiates through the entire gradient update trajectory, which imposes relatively high computational and memory burdens. contrast, optimiz- ing distinct sets of parameters at different levels enables the use of implicit differentiation methods, which significantly reduces computational costs, detailed Section 3.3 Therefore, we create a model with new parameters that are different from those in the pretrained model, but with regularization loss R between and 0 to encourage them to be close, inspired by [Rajeswaran el. al. (2019). This proximal constraint casts strong dependence between wt and 0", closely resembling the real finetuning process. Formally, the OP in this level is to optimize by minimizing the weighted summation of finetuning loss C., and the proximal regularization loss R:
(0) arg min Cor (wi, Drr) + RR((,,00
(3)
5

---

Published in Transactions on Machine Learning Research (06/2025)
where Der is the training split of the downstream dataset, and y is tradeoff hyperparameter to balance the finetuming loss and regularization loss. In practice we select the mean squared error (MSE) loss
the regularization loss R, and we further elaborate the reason for this choice in Appendix E The optimal
solution of in this level a function of 8r due to the loss term R, which is in turn a function of X, denoted as w" (A)).
Level III In the third level, we aim search for the optimal tradeoff parameters A between pretraining objectives. Formally, the OP in this level is to optimize X to minimize the validation loss Leal:
min (w (A), Deal)
(4)
where Dont is the validation split of the downstream dataset. In practice, we reparameterize each tradeoff weight A via softmax over unconstrained variables B:
ex)
A T2
ビ (xp
We optimize 3 rather than directly. This both guarantees tradeoff weights to be non-negative (A, 2 0) and imposes a constraint on the sum (Ei X 1) automatically.
Multi-level Optimization Framework In this way, we formulate three-level optimization problem with OPs in different levels mutually dependent on each other:
min Ceat (A), Deal)
(5)
8.1. (0) = arg min Ler (w: Der) + RR(,, (A))
(A) arg min Lps (0, X, Dpt)
By solving this multi-level optimization problem, are able to reweight each continued pretraining objective based beek from validation performance downstream tasks. In practice, both and in Equation 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 Lval in the third level with respect to the tradeoff parameter X. This full derivative dbaut can be computed with the following equation using chain rule:
dA
d.Caal
OLvat
Ow
80* 0X
(6)
d)
Own
a
In the right hand side of Equation 6 the green term, a partial derivative vector, can be directly computed with popular automatic differentiation libraries, such Pytorch (Paszke et al. 2019). 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 ct al. 2020 Zhang et al.) 2021), 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 0 efficiently, thereby obtaining the gradient of Leal with respect to X. We then optimize the tradcoff parameter with gradient descent. The complete algorithm is implemented using the Betty library (Choe al. 2023c bj. We present the complete optimization algorithm of Weightght in Algorithm

---

Published in Transactions on Machine Learning Research (06/2025)
Algorithm TTapWeighi-ophimizatonn
Input: Unsupervised pretraining dataset Dpc, training datasct Der, validation datasct Doat, total global optimization steps M
1: for 1...., .M do
2: Sample mini- batches: Xpt ~ Dpc, Xer ~ Drr, Xvol ~ Doat
Compute gradient VoLi x; Xpe), where Li Lye X, Xpr) (Equation
Update: 0 t 0 01 VoL1
dL2
5:
Compute hyper-gradient
where L2 = Lir(w, Xer) + ((,,0 (A)) (Equation
du
dL2
6: Update: w - w - 02
dw
dL3
7: Compute hyper-gradient
where L3 = Lnaalw" (A), X (Equation
dx
dL3
Update: A 03
dd
end for
Output: Optimal weights 0" w* X
4 Experiments
4.1
Molecular Property Prediction
In this section, use WapWeight 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 D (ar)isisn containing millions of molecules, we define a multi-objective continued pretraining loss inspired by Imagemol (Zeng et al. 2022):
[(x) => figi (x) + >2Lmmgz(1) + da AaLmga(ar) + App (x) + A (m
(7)
where x represents a molecular image, and = {ssis are tradeoff parameters. Lmgl, Lmg2, and Lmg3 are MACCS key (Durant et al. 2002) clustering- based classification losses with different number clusters. Cipp is a jigsaw puzzle prediction loss, where the model solves a jigsaw puzzle on the same molecular image. Lmat is mask-based contrastive learning loss, which generates constrastive pairs by masking molccular images. Details of these pretraining objectives can be found in Appendix C.I
The multi-objective loss C is optimized on the complete unlabeled dataset D to train a molecular image encodor. The lcarnt encoder can be further finetuned on downstream datasets for various molecular tasks. Existing approaches typically set the tradeoff parameters equally across different pretraining objectives, overlooking the varying contributions of each objective to specific downstream tasks (Zeng et al. 2022). We address this challenge by applying TepWeight 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 D, consisting of million molecules from PubChem (Kim et al. (023). For downstream tasks, we employ the MoleculeNet bench- mark, which includes 8 classification datascts focused on predicting biophysical and physiological properties essential for drug discovery (Wu al.] 2017 We generate the training, validation and test split of these
7

---

Published in Transactions on Machine Learning Research (06/2025)
Method
BАCE
BBBP
ClinTox
Sider
Tox21
ToxCast
HIV
MUV
Avg.
Dataset Size
1,513
2,089
1,478
1,427
7,831
6,575
41,127
98,087
AttrMask
77.2
70.2
68.6
60.4
74.2
62.5
74.3
73.9
70.2
PoddextPred
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
Finetuning (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
Table 1: Results of molecular property prediction on 8 classification tasks in MoleculeNet benchmark, terms of AUROC. Higher values better for all results, and the best results are shown in bold.
downstream datasets by applying scaffold splitting ! with an 8:1:1 ratio. We use AUROC the evaluation metric for all classification datascts, MAE for Qm7 and Qm9 datascts, 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 predic- tion Hu et al. (2020), and GraphMVP (Liu et al.) 2022). The pretrained molecular image encoder is basod on a ResNet model, with the final classification layer removed (He et al. 2015). More detailed descriptions are provided the Appendix for the datasets
C.2 baselines [C.)). and hyperparameter settings (C.4).
4.1.3 Results
Table show the results of various methods across molecular property classification tasks from MoleculeNet benchnark. Our method outperforms all baseline methods on all 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 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 downstteemm-guide 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. Tt worth losses,
each lask. Experimental
mentioning that TAPT (Gururangan al., 2020 is not applicable this task, as clustering- such Lmg3, 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, WepWeight does not face such limitations, demonstrating its generalizability.
4.2 Natural Language Processing
Method Dataset Size
Freesolv
Esol
Lipo
Qm7
Qm9
In this section, we validate the effec- liveness of TapWeight for continued pre- training of a masked language model (MLM) with its application to natural language processing tasks.
642
1,128
4,200
6,830
133,685
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
Finetuning (Tmagemol)
3.04
1.11
0.76
141.0
4.52
Weight (ours)
1.91
1.06
0.76
126.0
4.28
4.2.1 Prelimimary
Table 2: Results of molecular property prediction regression
Given a large-scale raw-text dataset (athiiin consisting of millions of sen- tences, we define the following continued pretraining loss:
tasks in MoleculeNet benchmark. Lower values are better for all results, and the best results are shown in bold.
Scaffold splitting partitions molecules hased their core structural frameworks that structurally similar com- pounds do not appear across training and test sets, thereby providing a more rigorous evaluation of model generalization.
8

---

Published in Transactions on Machine Learning Research (06/2025)
((z) A1 X2Ca(()) A3Coop(z)
(8)
where & a sentence, and À {Arigigs are tradeoff parameters. Cmlm represents the masked language model loss, which involves randomly masking tokens in the input sentences and predicting these masked tokens (Devlin al., 2019). Let denotes the contrastive learning loss, where an input sentence is used to predict itself with standard dropout applied noise (Gao et al.] 2021 Soop is the soquence ordering prediction loss, which emphasizes inler-sentence conherence (Lan al. 2020 We include details of these
losses in Appendix D.1
The multi-objective loss is optimized on the raw text dataset 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. 2021); the tradeoff parameters for different pretraining objectives require manual hyperparameter tuning, which is time-consuming and often leads suboptimal results. We address this challenge by applying Weight for the contimued pretraining of a Transformer encoder, enabling the automatic determination of the importance for each objective.
4.2.2 Experimental Settings
We perform continued pretraining of masked language models on a raw-lext dataset. D consisting of million sentences from Wikipedia (Gao et al. 2021). For downstream evaluation, we use RCT (Dennoncourt & Lee 2017). AGNews (Zhang ct al. 2015) and IMDB (Maas al., 2011) datasets, which are widely used for evaluation of TAP methods (Gururangan
al. 2020 Shi & Lipani 2023). We also the GLUE benchmark for evaluation, which comprises natural language understanding tasks, including sentiment analysis, semantic similarity prediction, and grammaticality classification (Wang et al. 2019). We use Matthew's Correlation for the CoLA dataset, Pearson/Spearran Correlation for the STS-B dataset, and accuracy for all other datasets. Our baseline methods are based on RoBERTa (Liu et al. 2019c) and DeBERTa (He et al. 2021b) models, including direct finetuning, TAPT based continued pretraining, PCP based continued pretraining (Shi & Lipani] 2023), and SimCSE based continued pretraining. More detailed descriptions are provided in the Appendix for the datasets (.2),, baselines D.3), and hyperparameter scttings D.4]
4.2.3 Results
Method
RCT
AGNews
IMDB
3 reports the performance of various methods on three datascts: RCT, AGNews, and IMDB, evaluated using, RoBERTa-base (125M parameters), RoBERTa-large (355M pa- rameters) and DeBERTa-xxarge (750M param- eters). The RCT dataset involves classifying sentences in biomedical texts based on their functional roles, AGNews focuses on topic clas- sification of news articles, and IMDB dataset for sentiment analysis of movie reviews. The results show that Weight consistently outperforms baseline methods across all tasks and all 3 model sizes. These findings validate the robustness of Weight across diverse do- mains (biomedical, news, and reviews), while also highlighting its scalability across different. model sizes. Moreover, our method surpasses the SimCSE method on all 3 tasks, showcasing the effectiveness of reweighting pretraining ob- jectives, SimCSE uses fixed ratio between
Table
Dataset Size
78,387
127,600
50,000
Finetuning (Rs)
86.3
93.2
94.5
SinCSE (R.)
85.9
93.0
94.1
TAPT (Rs)
86.4
93.5
94.7
Tap Weight (Rs)
86.7
93.8
95.1
Finetuning (Rt)
86.9
94.0
95.2
SimCSE (Re)
86.5
93.8
95.0
TAPT (R))
86.9
94.2
95.1
TapWeight (Rt)
87.4
94.8
95.5
Finetuning (Dat)
87.5
94.7
95.4
SimCSE (Da)
87.0
93.9
94.7
TAPT (Daa)
87.7
94.1
95.4
Weight (Dx:)
87.9
95.0
96.1
Table 3: Results of RoBERTa-base (Ro), RoBERTh-large (Rt) and DeBERRaaxxarge (D₂z) RCT, AGNews and IMDB datascts in terms of accuracy. Higher values are bet- ter for all results, and the best results are shown in bold.
9

---

Published in Transactions on Machine Learning Research (06/2025)
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
PCP
86.5
91.5
90.6
80.1
93.9
89.8
61.2
91.2
85.6
Weight (ours)
86.8
92.5
91.1
80.7
94.9
90.2
62.3
91.2
86.2
Table 4: Results of different methods in GLUE benchmark. All methods are applied to RoBERTa-base model. Higher values are better for all results, and the best results are shown in bold.
MLM and CL losses during continued pretrain- ing. Additionally, Tap Weight outperforms the RoBERIa+TAPT approach, demonstrating that our strategy of leveraging downstream datasets by reweighting pretraining objectives more effective than simply pre- training the model with unlabeled downstream data, as TAPT does.
Table 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 tasks other than the 3 datasets mentioned above. On average, our method achieved a score of 86.2, while finetuning RoBERTa model without continued pretraining only got 85.6. Furthermore, Weighight outperforms the RoBERTa+PCP approach, further underscoring its effectiveness and advantage compared to popular TAP methods. To sum up, 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.
We further applicd TaiWhight to more recent pretrained LM, DeBERTaa-v-larggs (He et al. 2023), to validate the gen- eralizability of TapWeight. Specifically, we compared Tap- Weight to direct finetuning of DeBERTa-v:-largg on both the CoLA dataset and the MRPC dataset from the GLUE benchmark, with the results shown Table 5 We ob- serve that TeiWttght still clearly outperforms the baseline method on both datasets, highlighting its generalizability to more recent pretrained LMs.
Dataset
CoLA
MRPC
Finetuming 73.0
91.5
Tap Weight (ours)
74.1
92.1
Table 5: Results of TaiWh and Finetuning using DeBERTa-v3-latte on CoLA and MRPC datasets.
4.3 Ablation Studies
In this section, we perform ablation studies to evaluate the effectiveness of individual components within our framework. All experiments are conducted the classification tasks in the molecular property prediction benchmark.
Pretraining Objective Reweighting We validate the effectiveness of our pretraining objective reweight- ing stratogy by comparing our mothod to contimued pretraining with a fixed importance for cach objective. As shown Table 6 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 bi-level optimization (BLO) Xic 2023) based method. Specifically, we merge the first and second level of problems from the TapWeight framework to form the lower-level prob- Jem in the new BLO baseline, where the model is optimized jointly using both the unsupervised pretraining loss on the unlabcled continued pretraining dataset pp and the finetuning loss on the training split of the downstream dataset Drr In the upper-level problem, the importance for each pretraining objective is learned using the validation split of the downstream dataset. Formally, define the following BT.0 problem:
10

---

Published in Transactions on Machine Learning Research (06/2025)
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
Weight
83.1
71.2
81.3
64.5
77.0
66.1
78.4
80.5
75.3
Results of molecular property classification using our method and baseline methods, in terms of AUROC. Higher values are better for all results, and the best results are shown in bold.
Table 6: Ablation Studies.
min Lval (A), Deat)
(9)
s.1. (*(A) = arg min Lyt (0, X, Doe) + 7.rr Der)
However, optimizing these two types of losses in the lower level requires extensive tuning of the tradeoff pa- rameters y, and often leads to competition between losses which results in performance decrease. As shown in Table [ our MT. based reweighting method outperforms the BT. 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 paramcters, indicating the ne- cessity of using reweighting strategies.
CP with Learned Final Weights We conducted additional experiments on three MoleculeNet datasets (Bace, Clintox, and HIV) to evaluate the performance of contimued pretraining (CP) using the tradcoff weights identified by WeiWeight (CP-final-weights)) with results shown in Table 7 On average, this base- line achieved score of 80.6, which is comparable to TapWeight's s score of 80.9, demonstrating the effec- tiveness of the tradeoff weights found by Weight.. However, this approach incurs additional computa- tional cost due to an extra round of CP, without yield- ing significant performance gains. Therefore, it is more practical directly use the model weights obtained from Tap Weight without conducting an additional CP stage.
Dataset
BАCE
Clintox
HIV
-eeihhs--eiights Tap Weight (ours)
82.3
81.1
78.5
83.1
81.3
78.4
Table 7: Performance of continued pretraining us- ing the tradcoff weights previously identified by TapWeight. (CP-final-weights) on three MoleculeNet datasets.
4.4 Qualititive Analysis
In this section, we present the evolution trend of the pretraining objective weights along the training tra- jectory using our method. As shown Figure 2 we plot the value of for 3 regression lasks (Esol, Lipo, Freesolv) and 3 classification tasks (Tox21, Toxcast, Clintox) with respect to the global training step. Our observations reveal that different downstream datascts require varying importance for cach pretraining ob- jective. For example, the JPP pretraining objective, Cipp. plays key role in Lipo and Toxcast datasets, whereas the MG3 pretraining objective, Lmg3, is more critical for Esol, Freesoly and Tox21 datasets. The di- verse requirements of pretraining objectives across downstream datasets emphasize the need for a reweighting method like Tap Weight, providing a clear explanation for why our method outperforms baseline approaches. Furthermore, similar downstream tasks exhibit some degree of similarity the weights assigned to pre- training tasks. For instance, the Esol and Frecsoly datasets, both focused on predicting physical chemistry properties of molecules, assign large weights to the MG3 pretraining objective. In contrast, the ToxCast and ClinTox datasets, which involve predicting molecular toxicity, assign smaller weights to the MG3 objective.
4.5 Computation Cost
11

---

Published in Transactions on Machine Learning Research (06/2025)
(a)
Esol
(b)
Lipo
(c)
Freesolv
0.35
0.24
0.22
0.25
0.25
0.2
0.18
0.15
0.15
0.16
Step
Step
1k
2k
3k
4k
5k
1k
2k
3k
4k
5k
2k
3k
4k
5k
(d)
Tox21
(e)
Toxcast
(f)
Clintox
0.22
0.24
0.21
0.22
0.24
0.22
0.2
0.2
0.18
0.18
0.16
0.18
Step
0.14
Step
0.14
Step
1k
2k
3k
4k
5k
1k
JPP MCL MG1 MG2
MG3
Figure Evolution of the tradeoff parameter X over the training steps of TapWeight on the following downstream datasets: (a) Esol, (b) Lipo, (c) Freesolv, (d) Tox21, (e) Toxcast, and (f) Clintox.
Datasct
FT
CP+FT TAPT PCP
Tap Weight
QQP
x1
x2.54
x1.57
X2.26
x3.76
Table 8: Training cost of baseline methods and our method TapWeight on QQP dataset.
In this section, we compare the training time (wall time) of our method with baseline methods on the QQP, MUV and Qm9 datasots, as shown in Table g and Table g We use finetuning (FT), continued pretraining with a fixed tradeoff ratio (CP+FT), TAPT and PCP as baselines, normalizing the time cost of FT as While TapWeight results an increase in training time compared bascline methods, its substantial improvement across multiple downstream tasks generally justifies the additional cost. However, in real-world applications where training time is critical factor, TapWeight may not be the ideal choice, representing limitation of our approach.
Dataset
FT
CP+FT TapWeight
MUV
x1
x2.18
x3.29
Qm9
x1
x2.76
x3.93
Table Training cost of baseline methods and our method Weigh on MUV and Qm9 datasets.
5 Conclusion and Future Work
In this paper, propose task-adaptive continued pretraining method that dynamically reweights each pretraining objective within a multi-level optimization framework. Experiments in both molccule property prediction and natural language processing validate the eflectiveness and generalizability of our method. Given the success of Weight,h several promising future research directions emerge. For instance, large multimodal pretrained models have recently gained popularity (Liu et al. 2023 Zhu al. 2024). The combination of multiple modalities introduces a greater number of potential continued pretraining objectives, presenting necessities of applying Weight in this context.
12

---

Published in Transactions on Machine Learning Research (06/2025)
6 Acknowledgement
This research supported by NSF IIS2405974, NSF IIS2339216, NIH R35GM157217, and NIH R21GM154171.
References
Idan Achituve, Idit Diamant, Arnon Netzer, Gal Chechik, and Ethan Fetaya. Bayesian uncertainty for gradient aggregation in multi-task learning. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 117-134. PMLR, 2024. URL https: httt::s//procceednggs www.sssss....r.///ee....e..e..ee ppress/v235// wtw.//221i/cccaecee achi aww.a/244iuuve2aa...mm tuvee4a. html.
Iz Beltagy, Kyle Lo, and Arman Cohan. SciBERT: A pretrained language model for scientific text. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIP), pp. 3615-3620, Hong Kong. China, November 2019. Association for Computational Linguistics. doi: 10. 1..1865/v11D199-11 1371. URL https://aallanthologgy org wwwp:/orggg///-0 1371
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Lan- guage models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Cur- ran Associates, Inc., 2020. URL https::/proceeedngs.neoom wwtpseeeecpepppppperrrree..eee cc/paper faaaer_ffleese files/ paalee//apprr/020///ea pdf
Tlas Chalkidis, Manos Fergadiotis, Prodromos Malakasiotis, Nikolaos Aletras, and Ion Androutsopoulos. LEGAL-BERT: The muppets straight out of law school. In Trevor Cohn, Yulan He, and Yang Liu (eds.), Findings of the Association for Compulational Linguisties: EMNLP 2020, pp- 2898 2904, Online, 261.
November 2020. Association for Computational Linguistics. doi: /1.18653/v1/2020.imm URL
org/2020 findings- emnlp. 261
Xin Chen, Lingxi Xie, Jun Wu, and Qi Tian. Progressive differentiable architecture search: Bridging the depth gap between search and evaluation. In Proceedings of the IEFF International Conference on Com- puter Vision, pp. 1294 1303, 2019.
Zeming Chen, Alejandro Hern' andez Cano, Angelika Romanou, Antoine Bonnct, Kyle Matoba, Francesco Salvi, Matteo Pagliardini, Simin Fan, Andreas Kopf, Amirkeivan Mohtashami, Alexandre Sallinen, Alireza Sakhaeirad, Vinitra Swamy, Igor Krawczuk, Deniz Bayazit, Axel Marmet, Syrielle Montariol, Mary-Anne Hartley, Martin Jaggi, and Antoine Bosselut. Meditron 70b: Scaling medical pretraining for large language models. Arxie, abs/ 0bs/2331.11601 2023. URL https://appl www.nticscholaa.... org/Corpusull Cooppopsa::055552222333333 265456229
Zhao Chen, Vijay Badrinarayanan, Chen-Yu Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normal- ization for adaptive loss balancing in deep multitask networks. In Proceedings of the 35th International Conference Machine Learning (ICML), 2018. URL https://arxiv wtgg/bss////g
Sang Keun Choe, Sanket Vaibhav Mchta, Hwijeen Ahn, Willie Neiswanger, Pengtao Xie, Emma Strubell, and Eric Xing. Making scalable meta learning practical. In shirth- Conference on Neural In Information Processing Systems, 2023a. URL https://openrevieer
Sang Keun Choc, Sanket Vaibhav Mchta, Hwijcen Ahn, Willie Neiswanger, Pengtao Xie, Emma Strubell, and Eric Xing. Making scalable meta learning practical. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b. URL http:://ggaaeeeeeer..e/e//e/aada
13

---

Published in Transactions on Machine Learning Research (06/2025)
Sang Keun Choe, Willie Neiswanger, Pengtao Xie, and Eric Xing. Betty: An automatic differentiation library for multilevel optimization. In The Eleventh International Conference on Learning Representations, 2023c.
URL https://openreevie..co http://wwww.eete/otiii//i/./
Aakanksha Chowdhery, Sharan Narang. Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sashank Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prab- hakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Is- ard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thamumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Heler-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: scaling language modeling with pathways. J. Mach. Learn. Res., 24(1), March 2024. ISSN 1532-4435.
Yibo Cui, Liang Xie, Yakun Zhang, Meishan Zhang, Ye Yan, and Erwei Yin. Grounded entity-landmark adaptive pre-training for vision-and-language navigation. 2023 IEEE/CVI CCF International Conference on Computer Vision (ICCV), pp. 12009- 12009-12019, 2023. URL /ttp::tp:///ap..... httpps:/aap..emmatticchh..... semantccccool.rr wwwpihhollrrrrr/ooppueeeee oxg/CorpusID:
261101130
Franck Dernoncourt and Ji Young Lee. Pubmed 200k ret: a dataset for sequential sentence classification in medical abstracts. In International Joint Conference on Natural Language Processing, 2017. URL
/
https://api.. wewanti
wowppoopu//oouulliic ID: 151184
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pro-training of deep bidircc- tional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171-4186, Min- ncapolis, Minnesota, June 2019. Association for Computational Linguisties. doi: 00.18653/v1/N11 1423. URL https://aclanthoology ott.0/19190////1111111333
Joseph L. Durant, Burton A. Leland, Douglas R. Henry, and James G. Nourse. Reoplimization of mdl keys for use in drug discovery. Journal of Chemical Information and Computer Sciences, 42(6):1273 1280, 2002. ISSN 0095-2338. doi: 10.1021/ci01012:: URL http:////ooi..rgg1.0100/////////
Simin Fan, Matteo Pagliardini, and Martin Jaggi. DOGE: Domain reweighting with generalization estima- tion. In Forty-first International Conference on Machine Learning, 2024. URL https::/openreviee.
net/forum?id::rftnn
Chelsca Finn, Picter Abbeel, and Sorgey Levine. Model-agnostic meta-lcarning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1126 1135. PMLR, 06 11 Aug 2017. URL https::/pprcceeedng.... mlr wttppp/e//////o//////// htt.://dennga.a......rree//////m httl
Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE: Simple contrastive learning of sentence embeddings. In Empirical Methods in Natural Language Processing (EMNLP), 2021.
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A. Smith. Don't stop pretraining: Adapt language models to domains and tasks. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeling of the Association for Computational Linguistics, pp. 8342 8360, Online, July 2020. Association for Computational Linguistics. doi: 20006633///2222cclmm.mm 740.. URL https::/ecllaathooogg) wrg.e/20rg220000 acl- mell-mann.77 740)
Kaiming Hc, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016
IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2015. URL https:
/api. www.seettnncccocr..rooo.o.o wtt.://ccoaa.a..oooooopououaaaaa
14

---

Published in Transactions on Machine Learning Research (06/2025)
Kaiming He, Xinlci Chen, Saining Xie, Yanghao Li, Piotr Doll'an ar, and Ross B. Girshick. Masked autoencoders are scalable vision learners. 2022 TEEE/CV Conference on Computer Vision and Pattern Recognation (CVPR), pp. 15979 15988, 2021a. URL htw.s///op...mmtticccccccc...... wrg/CorpusID.-C.C...o... 243985998
Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decodiin-enhanced bert with disentangled attention. In International Conference Learning Representations, 2021b. URL https:
//openreview.c //openteeieev.ne//foorutrr
Pengcheng He, Jianfeng Gao, and Weizhu Chen. Debertav3: Improving deberta using electra-style pre- training with gradient-disentangled embedding sharing. In International Conference Learning Repre- sentations, 2023. URL http::///ppaarrreeeeeeeeerrrrrrs
XXLxHHA
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategics for pre-training graph neural networks. In International Conference on Learning Representa- tions, 2020. URL |nttps://openrevies: http://wwweeetttooorrdddddaaaa
Sunghwan Kim, Jian Chen, Tao Cheng, Asta Gindulyte, Jian He, Sherry He, Qiang Li, Bradley A Shoemaker, Paul A Thiessen, Bo Yu, Ludmila Zaslavsky, Jian Zhang, and Evan E Bolton. Pubchem 2023 update. Nucleic Acids Research, (1(DD))D1373-D1388 D13800 Jan 2023. doi: 10.1093/nar/gkacce
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020. URL htapss///opparevvieenettffooooo
Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jacwoo Kang. Biobert: pre-trained biomedical language representation model for biomedical lext mining, Bioinfor- matics, (6(4):1234-124444 2020.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARIS: Differentiable architecture search. In Inter- national Conference on Learning Representations, 2019a. URL http:s////ppareeeeee.nnettrrtrts
S1eYHoC5FX
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jac Lec. Visual instruction tuning, 2023.
Shengchao Liu, Hanchen Wang, Weiyang Liu, Joan Lasenby, Hongyu Guo, and Jian Tang. Pre-training molecular graph representation with 3d geometry. In International Conference on Learning Representa- tions, 2022. URL httpp:://ppenrevierr http:/enrrreiieeeetfffoooeeeeeee pidi xQUekkaaa
Shikum Liu, Edward Johns, and Andrew J. Davison. End-to-end multi-task learning with attention. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition (CVPR), 2019b. URL
https://arxiv. httpe////arii..ogg/aaa///aa/ wwwp/r////////oooooeeeeee//////o 10704
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019c. URL
http://arxiv.orrg 20tg/abbs1990// 11692 cite arxiv:190..11692.
Jonathan Lorraine and David Duvenaud. Stochastic hyperparameter optimization through hypernetworks. In Advances in Neural Information Processing Systems (NeurIPS) Meta-learning Workshop, February 26 2018.
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In Silvia Chiappa and Roberto Calandra (ods.), Proceedings of the Twenty Third In- ternational Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 1540-1552. PMLR, 26-28 Aug 2020. URL v108/lorraine2a..1 html
mlr. .trrs//
Andrew L. Maas, Raymond E. Daly, Peter Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Dekang Lin, Yuji Matsumoto, and Rada Mihalcea (eds.), Proceedings of the 49th Annoual Meeting of the Association for Compulational Linguisties: Human
15

---

Published in Transactions on Machine Learning Research (06/2025)
Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL luttps://aclanthholog httpps///ocaabbbooooggg////11555
org/P11-1015
Matthew Mackay, Paul Vicol, Jonathan Lorraine, David Duvenaud, and Roger Grosse. Self-tuning net- works: Bilevel optimization of hyperparameters using structured best-response functions. In International Conference on Learning Representations, 2019. URL
Yuren Mao, Zekai Wang, Weiwei Liu, Xuemin Lin, and Pengtao Xie. MetaWeighting: Learning to weight tasks in multi-task learning. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Findings of the Association for Computational Linguistics: ACL 2022, pp. 3436 3448, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: /2118663/v122222222 http://////222...nnnnsnee/eee acl.12, URL http:///aclanttoooy.ooor////////
findings-acl.277
Ishan Misra, Abhinav Shrivastava, Abhinav Gupla, and Martial Hebert. Cross-stitch networks for multi-task learming. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recog- nition (CVPR), 2016. URL https://openaceese Cross-Stitch. Netthhksetwwsss..cooooooo for CVPR_ Laiaa2216_______ 1001__apper... pdf
ctteevv..coc/oooteeee cvpp__cvp______ 2016/papers/Misra.
Kosuke Nishida, Kyosuke Nishida, and Sen Yoshida. Task-adaptive pre-training of language models with word embedding regularization. In Findings of the Association for Computational Linguistics: ACL 2021, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: an imperative style, high-performance deep learning library. Curran Associates Inc., Red Hook, NY, USA, 2019.
Aniruddh Raghu, Jonathan Lorraine, Simon Korublith, Matthew McDermott, and David Duvenaud. Meta- learning to improve pre-training. In Proceedings of the 35th International Conference on Neural In- formation Processing Systems, NIPS '21, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 9781713845393.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In H. Wallach, H. Larochelle, A. Beygelzimer, F. BAlché-Buc, E. Fox, and R. Garnett (eds.), Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedingss neurips cwtuuips.ccppapeee.. wwtppcppppaaaieeeeeeee ptl.e///pppr///0////i.. 07260000160224433344222229022222 pdf
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer, Routing networks: Adaptive selection of non- linear functions for multi-task learning. In International Conference on Learning Represendations (TCI.R), 2018. URL httptttts///iiii www.oogg////.....ooooo..o... /rgg/bss//11.....////g//// 01239
Sebastian Ruder. An overview of multi-task learning in deep neural networks.
ar preprint
arXiv: 05098, 2017. URL https:://rxiv..ort http://rxxi..oogab//////// abs /ttbs///10.... 05098.
Ozan Sener and Vladlen Koltun. Multi-task learning as hnaliioobjeediieeeoptiiiiiiaiiii In Advances in Newral Information Processing Systems (NeurIPS), 2018. URL httpp:://aad...org/aaas////////s 04650
Zhe Zhengxiang Shi and Aldo Lipani. Don't pretraining? make prompt-based fine-tuning powerful learner. In seirthh Conference on Neural Information Processing Systems, 2023. URL https:
Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Meta-weight-net:
Learning an explicit mapping for sample weighting. In NeurIPS, 2019.
Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier García, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Denny Zhou, Neil Houlsby, and Donald Metzler. U12: Unifying language learning paradigms. In International Conference on Learning Representations, 2022. URL atttps///apt http://pppp.smtaasccccooo//oo/oe wwwanticscholar.c... hrtproaaar...rrroooopppii///222
16

---

Published in Transactions on Machine Learning Research (06/2025)
Yi Tay, Jason Wei, Hyung Chung, Vinh Tran, David So, Siamak Shakeri, Xavier Garcia, Steven Zheng, Jin- feng Rao, Aakanksha Chowdhery, Denny Zhou, Donald Metzler, Slav Petrov, Neil Houlsby, Quoe Le, and Mostafa Dehghani. Transcending scaling laws with 0.1% extra compute. In Houda Bouamor, Juan Pino, and Kalika Bali (ods.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 1471-1486, Singapore, December 2023. Association for Computational Linguistics. doi: emnlp-main.91
10.18653/v1 /200/3 aar..mmmmmlmmane URL https: attps://aallanthooggy www.aanttooooooooooooaeoooo
Alex Wang, Amanprect Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2019. URL htttps://openrreveer hwt.://wwwwe//ooommmieroroaaaaa.
Han Wu, Kun Xu, Linfeng Song, Lifeng Jin, Haisong Zhang, and Linqi Song. Domain-adaptive pretraining methods for dialogue understanding. In Anmual Meeting of the Association for Computational Linguistics, 2021. URL https: www.nticscoollaa.e...o.o wwwpeehollarrrro//ooeoo :
Zhenqin Wu, Bharath Ramsundar, Evan N. Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S. Pappu, Karl Leswing, and Vijay S. Pande. Moleculenet: benchmark for molecular machine learning. Chemical Science, 9:513 530, 2017. URL hattps://ayy https::/app..semtttiiooo wetppppssmaatticchll.o.. oegeeaaannssscooooroooo///rr/rpp 217680306
Pengtao Xie. Improving bi-level optimization based methods with inspiration from humans' classroom study techniques. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.); Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 38137 38137-38186. PMLR, 23 29 Jul 2023. URL https: fttps://proceedingg mlr. wwtpp//e/e///iiee///eep html
Yong Xie, Karan Aggarwal, and Aitzaz Ahmad. Efficient continual pre-training for building domain specific large language models. In Annual Meeting of the Association for Computational Linguistics, 2024. URL
Yuhui Xu, Lingxi Xie, Xiaopeng Zhang, Xin Chen, Guo-Jun Qi, Qi Tian, and Hongkai Xiong. {PC}- {darts]: Partial channel connections for memory-offiiint architecture search. In International Conference on Learning Representations, 2020. htttps:///openrrevieee htt.a//as.////.gaeerrreeee.eeeea
Ruihan Yang, Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft mod-
ularization. In Advances in Neural Information Processing Systems (NewrIPS), 2020. URL https:
//arxiv. Slarxiv.org/all 1166111.166.111111
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. In Advances in Neural Information Processing Systems (NeurTPS), 2020. URL
wttp::////...ooo.....o 06782
Xiangxiang Zeng, Hongxin Xiang, Linhui Yu, Jianmin Wang, Kenli Li, Ruth Nussinov, and Feixiong Cheng. Accurale prediction of molecular properties and drug largets using self-supervised image representation learning framework. Nature Machine Intelligence, 4(11):1004- 1016, 2022.
Miao Zhang, Steven W. Su, Shirui Pan, Xiaojun Chang, Ehsan M Abbasnejad, and Reza Haffari. idarts: Differentiable architecture search with stochastic implicit gradients. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 12557-12566. PMIR, 18-24 Jul 2021. URL https://procediing. mlr. press/v139//ee 2at00//haagg2..10. html
Ruiyi Zhang, Rushi Qiang, Sai Ashish Somayajula, and Pengtao Xic. AutoLoRA: Automatically tuning matrix ranks in low-rank adaptation based on meta learning. In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 5048-5000 Mexico City, Mexico, June 2024. Association for Computational Linguistics. doi: 10.18653/v1/2000 naacl-long 282. URL httph:tt//sae http::://aclammhoooggoooo/////// na00024.naaclloooo 282
17

---

Published in Transactions on Machine Learning Research (06/2025)
Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classi- fication. In Newral Information. Processing Systems, 2015. URL https://api wwwanticccholar..rrro org/ CorpusID:36818:
Deyao Zhu, Jun Chen, Xiaoqian Shen, Xiang Li, and Mohamed Elhoseiny. MmnGGPT-4: Enhancing vision- language understanding with advanced large language modcls. In The Twelfth International Conference on Learning Represenlalions, 2024. URL
18

---

Published in Transactions on Machine Learning Research (06/2025)
A Optimization Algorithm
In this soction, we give example to briefty illustrate how to use Implicit Function Theorem (IFT) to compute best-response Jacobian matrices. Take 86 term in Equation as example: although 0r is an implicit function of X, the exact value of (*) given value of À is usually approximated with gradient descent algorithms. As there is no analytical solution of 0* (A), it is difficult to directly compute the gradient 00"
Rx To tackle this challenge, we compute this gradient using IFT following previous literature (Lorraine
et al. 2020
00*
8 Cpt 0x - IV Lma 00 OXT
(10)
The green term, second-order mixed partial derivative matrix, can be directly computed using automatic differentiation. Nevertheless, directly computing the red term, which is the invert of Hessian matrix VrLrr( is computational expensive due to its O(n") complexity. Various methods have been proposed to approximate the inverted Hessian matrix, including Neumann series (Lorraine et al.) 2020), conjugate gradients (Rajeswaran et al. 2019) and finite difference (Zhang et al. 2021). In Tap we select finite difference as the approximation method, thus enabling efficient computation of best-response Jacobian matrices.
B Complexity Analysis of TapWeight
In this section, we discuss the time complexity of Weighight. In summary, our optimization algorithm scales linearly with the number of model parameters, owing to the approximation of the best-response Jacobian matrix-vector multiplication, inspired by Lorraine et al. (2220) and Zhang et. al. (2021). We outline our analysis below:
Assume the number of optimizable parameters in the model is n. Directly computing the full derivative in Equation @ would incur a complexity of Onn), because the best-response Jacobian matrix requires computing an inverse Hessian-vector product (iHVP), and directly inverting a Hessian itself is O(n)). To address this, we leverage the Neumann series approximation to reduce the complexity to O(n²), using the following formula inspired by paper Lorraine et al. (2020):
M
v (V(L((w) Rs v. (I-W ((--er
1=0
where M is hyperparameter in practice. However, if directly compute the Hessian above, the result- ing quadratic complexity is still computationally intensive for models with billions of parameters. Inspired by paper Zhang et al. (2021). we further reduce the complexity by leveraging the Hessian-vecte product (HVP) and finite differences. HVP has a complexity of O(n) as implemented in modern automatic differ- entiation (AD) libraries (e.g., torch. autograd. functionnn hvp in Pytorch (Paszke et al., 2019)), which is more efficient than directly computing the Hessian. A detailed algorithm for computing the best-response Jacobian-vector product can be found in Appendix B ('Practical Implementation of Hypergradient") of Zhang et al. (2021), where only a limited number of first-order gradient compulations are required. Since gradient computation is O(n) in modern AD libraries, the total complexity of multiplying the first two terms in Equation 0 is also Onn). The resulting vector is then multiplied by subsequent term, which can also be resolved similarly in ((n) time. Overall, computing the right-hand side (RHS) of Equation has a complexity of (n).
Other components of our optimization and parameter updates-ail operate in O(n) time. Therefore, the total complexity of the TapWeight optimization algorithm grows linearly with the model size n.
loss computation, gradient backpropagation,
19

---

Published in Transactions on Machine Learning Research (06/2025)
C Molecule Property prediction
C.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.
2022
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,
Multi-Granulee
Lmg1
Σ
(11)
Emg2
1.000) 2,0000))
(12)
2m1
Lig3
C(C110000(fo 2000)))),0000) 10,000)
(13)
where fo is the molecular encoder, and C' are task-specific fully-connected neural networks for clustering label prediction. In the optimization problem the first level of TapWeight, C' is optimized as well as A. However, C is not relevant to the optimization problems in the other levels and is discarded after TapWeight training, did not explicitly include it in Equation
Mask-based Contrastive Learning In this pretraining objective, randomly mask molecular image to generate the masked image f. We then perform constrastive learning on the image pair (x, by minimizing the distance between representations of both images to promote consistency. Formally,
16 16 square area to
Liel
(14)
where ||||| 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 to 100, denoted as وو We also assign a label of for original molecular image without any We apply the permutation to molecular images to get permuted ones å. The encoder fo is pretrained by predicting the permutation label. Formally,
Jip Σ v')
(15)
where C' is task-specific fully-connected neural network for permutation label prediction.
C.2 Datasets
We the datasets from MoleculeNet benchmark for molecule property prediction Wu al. 2017
20

---

Published in Transactions on Machine Learning Research (06/2025)
Quantum Mechanics Qm? and Qm9 are both molecular datasets for regression task quantum mechan- ics 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 physical chemistry properties of molecules. ESSOT, 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/war water distribution coefficient.
Biophysics Bace, HIV and MUV are all datasets for classification tasks biophysics properties of molecules. BACE dataset collects binary label of molecular binding results for set of inhibitors of hu- man 8-secretase 1 (BACE-I). 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 classífication tasks on physiology properties of molecules. BBBP dataset contains binary labels of blood-brain barrier penetration (permeabil, ity) 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 database of marketed drugs and adverse drug reactions (ADR), grouped into system organ classes. ToxCast dataset contains toxicology data for large library of compounds based on in vitro high-throughhut screening, including experiments over 600 tasks. Tox21 dataset collects qualitative toxicity measurements of molecules on 12 biological targets, including nuclear receptors and stress response pathways.
C.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. 2020). Inspired by BERT (Devlin et al. 2019). it pretrains 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 strue- tures (Hu et al. 2020). 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 fixed vector using an auxiliary GNN, and then trains the GNN encoder with nogative sampling.
The Graph Multi- 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. 2022). It introduces a novel contrastive learning loss, using the 2D and 3D representations of the same molecule as positive pairs.
GraphMVP
Imagemol ImageMol is unsupervised pretraining deep learning framework pretrained 10 million unlabelled drug-like, bioactive molecules, to predict molecular targets of candidate compounds (Zeng et al. 2022). The ImageMol framework is designed to pretrain chemical representations from unlabelled molecular images on the basis of local and global structural charactoristics of molocules from pixels.
C.4 Hyperparameter Settings
We set the number of clusters for the loss terms Emgl, Emg2, and Lmg3 to 100, 1,000, and 10,000, respectively. During the continued pretraining, we set the unrolling step in the MLO framework to be We use the SGD optimizer with a step learning rate scheduler across all three optimization levels. All experiments are conducted on NVIDIA A100 GPU.
21

---

Published in Transactions on Machine Learning Research (06/2025)
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 to be 1024, and that in level Il 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 level III to be 200 for all datasets. We set the y value in Equation to be 0.001.
Regression We set the global learning steps be 30,000 for qm9 dataset and 10,000 for all other datasets. The batch size and y the same 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.
D Natural Language Understanding
D.1 Pretraining Objectives
We use 3 types of losses for continued pretraining of an RoBERTa model Liu et al. (2019c) 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.
2019 In BERT and RoBERTa, 15% of the tokens are masked in the pretraining stage.
Constrastive Learning This pretraining objective applies dropout noise the encoder Jo when laking in sentence r to get negative sample of encoding h' (z) (Gao et al., 2021). We use to denote those positive encodings without dropout noise. The encoder is then trained by minimizing constrastive learning loss:
sim((h
Let =
Σ 1=1 ===1 K33
log(
(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. 2020). The model is pretrained by predicting the label of these two types of examples.
D.2 Datasets
We use 8 datasets from GLUE benchmark in natural language understanding tasks (aag et al. 2019 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.
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 dataset of sentence pairs extracted from online news sources, annolated 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 human-annotated similarity score.
22

---

Published in Transactions on Machine Learning Research (06/2025)
Inference Tasks The Multi-Genre Natural Language Inference Corpus (MNLI) is crowdsourced dataset of sentence pairs annotated for textual entailment, where the task is to predict the relationship wetn 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 to predict the entailment between two sentences.
D.3 Baseline Methods
RoBERTa
The Robustly Optimized BERT Pretraining (RoBERTa) paper [Lu et al. 2019c) 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 lasks. The main differences between RoBERTa and BERT are: (1) training for 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 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.
DeBERTa The Decoding-enhanced BERT with Disentangled Attention (DeBERTa) model (He et al. 2021b) introduces architectural improvements over BERT to enhance contexlualizer representations. De- BERTa differs from BERT primarily by incorporating (1) a disentangled attention mechanism that separately encodes content and positional information, instcading of adding both word (content) embedding and po- sition embedding together as that in BERT; and (2) an absolute position encoding correction to improve position-dependet generalization. These enhancements enable DeBERTa to achieve superior performance a wide range of natural language understanding tasks.
D.4 Hyperparameter Settings
When applying Weight the RoBERTa encoder, set the unrolling step in the MLO framework to
1. We use an Adam optimizer with step learning rate scheduler across all three optimization levels. All experiments are conducted on 1 NVIDIA A100 GPU. 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 is set to 512, while for levels II and III, it set to 32 across all datasets. The learning rate for levels and is 2e-5, and for level III, it is set for all datasets. We sct the y value Equation to be 0.005.
E Selection of Proximal Regularization Loss
In this soction, we discuss the selection of proximal regularization loss in the second level of TepWeight, as specified in Equation 3 In principle, both mean squared error (MSE) and Kullback Leibler divergence choose
(KLDiv) can be used the proximal regularization loss in Level II of TapWeight. However, MSE following existing works (Rajeswaran et al., 2019 Choe et al. 2023c), where MSE is consistently used as the proximal regularization loss when converting itcrative optimization in bi-level optimization problems inlo implicit optimization. In preliminary experiments, we observed that the MSE loss offers improved computational efficiency and numerical stability compared to the KLDiv loss, particularly in large-scale settings where both 0* and w comprise billions of parameters. This makes MSE a more practical choice for training at seale. To quantify the efficiency gap, we benchmarked the forward and backward computation times of MSE and KL.iv losses R(w, 0*). We use a RoBERTa-base backbone (approximately
23

---

Published in Transactions on Machine Learning Research (06/2025)
R
Forward
Backward
KLDiv MSE
0.385
0.484
0.101
0.108
Table 10: Average wall-clock times (in seconds) of forward and backward operation for MSE and KLDiv loss.
125M parameters) for w and 0", which is implemented in PyTorch on the same hardware (NVIDIA A100).
Average wall-clock times (in seconds) are reported in Table 10
As shown in Table T0 the MSE is approximately times faster than KLDiv both forward and backward passes, highlighting its computational efficiency. That said, it is possible that KLDiv has advantages in certain scenarios, which we aim to explore in future work.
24


Reviews:
[{"review_id": "KzsAfBWVDi", "reviewer": "Reviewer_XrvN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n  * The method is sound and the explanations are clear\n  * The method is evaluated on two different datasets/tasks\n  * The computational cost of the proposed method is discussed\n\n### Weaknesses\n\n  * Related work could be better discussed. In particular, it is not clear how the baselines differ from the proposed method in particular how they weigh the pretraining objectives; In addition, the task at hand seems highly related to multi-task learning which is only briefly addressed in the related work section (e.g. MetaWeighting). It  would be interesting to have a more in-depth discussion on the added value of the proposed method compared to those baselines\n\n  * Direct Finetuning seems to be a competitive baseline. While the introduction presented direct finetuning as a \"common approach which is often suboptimal\", the results for finetuning are only reported in one table (GLUE) where it seems to perform well (better than all baselines but 1). This combined with the fact that finetuning seems to be more efficient than the proposed approach in terms of training time (Table 6) raises the question of how practical the approach is compared to simple finetuning.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The article addresses the TAP (Task-Adaptive Pretraining) problem. The core challenges of TAP, like most multi-task problems, is to balance different pretraining objectives, in order to maximise the model's performance on the target downstream tasks.\n\nTAP optimises the objectives' weights end-to-end; more specifically The goal is to find the optimal pretraining loss weights $\\lambda$ that minimises the validation loss, for model weights obtained via two phases of training :\n  * **1)** pretraining with $\\lambda$-weighted objectives (in that phase, the target weights are fixed throughout pretraining)\n  * **2)** finetuning model weights on the downstream tasks's training set\n\nThe key issues of the proposed approach is to keep the whole optimisation process computationally feasible, avoiding costly gradient computations. This is addressed in two ways:\n  * The weights between the two pretraining phases are not the same, thus cutting the gradient flow\n  * Implicit Function Theorem (IFT) to approximate Jacobians\n\nThe method is evaluated on two tasks: NLP (GLUE) and molecule  representation", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* **Critical for acceptance:** Better reference and placement compared to the multi-task related literature\n\n* **Critical for acceptance** Table 6 includes latency results for Qm9 and MUV finetuning, however there is no finetuning results in their respective tables (1 and 2). It seems to method  finetuning is always an important baseline to report. In particular, on the GLUE benchmark we observe that finetuning is only **0.6%** worse than the proposed method but requires roughly 70% of the training time (Table 6, on QQP) \n\n* **Strengthen the work**: Clarify what category each baseline belongs to: i.e., do they all use uniform weights ? do they optimise the weights via grid search, or something else ? In particular, on the GLUE benchmark in Table 4, I find it peculiar that all baselines underperform finetuning, although the introduction highlights the issue that \"direct finetuning of the pretrained model often fails to deliver optimal results\". This would suggest that either the baselines are not adequate, or that direct finetuning actually performs better than expected.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ipRvmZ9GR7", "reviewer": "Reviewer_cC8s", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- The paper proposes an efficient optimization approach to update the pre-training objective weights $\\lambda$, even if the exact gradients cannot easily be directly computed.\n- The method is evaluated on a variety of tasks from 2 highly different domains (molecular properties and NLP), with generally positive results compared to other baselines and pre-training with equal weights.\n\n**Weaknesses**\n\n- The description of the optimization procedure was unclear to me. In particular, based on section 3, I initially thought that pre-training was run to convergence to find $\\theta^*(\\lambda)$, then task-specific training was completed to find $\\omega^*(\\theta^*(\\lambda))$, and then the pretraining weights were updated based on the full validation loss, with this approach repeated for a few iterations. However, based on figure 2, I now believe that batches from the 3 stages are interleaved, but I am not sure if that is correct.\n    - \"three end-to-end stages\" was confusing.\n- Some of the baseline models are arguably outdated. It would be nice to show if the proposed approach could also be useful for more recent open-weight large language models.\n- \"Scaffold splitting\" is introduced without a definition or citation. Some readers, especially those from a mainly NLP background, may not understand what it refers too. In general, acronyms should ideally be spelled out at least once.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Given a target downstream task, TapWeight attempts to find the optimal weight allocation across different pretraining objectives. In the first stage of a tri-level optimization problem, the weighted pre-training loss is minimized. In the second stage, using another set of model parameters, performance on the downstream task (training set) is optimized, with a regularization term that encourages these other parameters to remain close to those obtained by pretraining. In the third stage, to minimize the downstream task validation loss, the pre-training objective weights are updated using the chain rule and some Jacobian approximations. The approach is evaluated on molecular property and natural language tasks, using respectively 5 and 3 pre-training objectives.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Critical**\n- Clarify the optimization procedure (maybe with pseudo-code?).\n\n**Would strengthen**\n1.  Revise the paper to ensure it can be more easily understood by readers from different backgrounds.\n2. Did you evaluate downstream performance using a single pretraining objective at a time ($\\lambda_i=1, \\lambda_{j\\neq i}=0$)?\n    - Did you try using a few random $\\lambda$s?\n3. What are the results if you learn the final weights, then run continued pretraining and finetune the model (starting from the initial pretrained parameters)?\n4. Do you impose a constraint on the sum or norm of $\\lambda$? If not, would multiplying $\\lambda$ by a constant lead to the same results?\n5. Could you use the KL divergence to measure how close $\\theta^*$ and $\\omega$ are? If yes, then why did you decide to use the parameter mean squared error? Otherwise, why not?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Kqy1uG4LHe", "reviewer": "Reviewer_KqG3", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The proposed method is novel and achieves best performance among baselines included in the paper. \n2. Generality of the method across two quite different application domains (molecular property and NLP). \n3. Ablations show the proposed components of TagWeight are important in achieving good performance.\n4. Quality of paper presentation is high.\n\nWeaknesses:\n\n1. Proposed method involves differentiating through a multi-stage process. This can become very expensive computationally or becomes prohibitive for larger models. Section 4.5 attempts to compare (wall-?) time of different versions of TapWeight, but not with other baselines. Also 4.5 does not address scaling of algorithmic complexity with model size.\n2. Description of the algorithm is not completely clear. How many iterations of \\lambda are done? Is it every learning batch? Does this mean the method conducts an evaluation on the downstream task's validation task every gradient step?\n3. Minor clarifications in writing. See \"Requested Changes\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper focuses on the problem of determining optimal weights across multiple continued pre-training losses, where model performance is measured after further fine-tuning on downstream tasks. The a multi-stage approach is proposed, where the weights are initially fixed during continued pre-training, downstream fine-tuning, and evaluation on the validation set. The weights are then adjusted based on evaluation performance by differentiating through the multi-stage process. The proposed method is evaluated on two major domains: molecular property prediction and natural language processing. Results on both domains show the proposed method, TapWeight, improves over existing baseline methods for continued pre-training. The authors also conducted ablations of several aspects of TapWeight to determine their importance in improving performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. [Critical] Section 3.1 \"In our framework, TapWeight, we aim to automatically search for the optimal ... so that the pretrained model will achieve highest performance on Dval after finetuned on a downstream dataset Df t.\" This is the wrong objective. Surely you mean achieve highest **test** performance? Please confirm that the results reported in Section 4 are conducted on the test splits?\n2. [Critical] Add an Algorithm listing of TapWeight, clearly specifying when and how many times each stage is done, the number of $\\lambda$ iterations, when evaluations on downstream validation set is done, when gradients are calculated and with which algorithm, etc. Define what a \"global training step\" means. It should be clear enough that a reader can quickly grasp the computational complexity of TapWeight.\n3. [Strengthen] The paper write up starts to describe TapWeight in \"stages\" but switches to \"levels\" half way through the paper. Please be consistent throughout and be careful with your choice of terminology. \"Stages\" implies that many gradient updates are done during continued pre-train, and fine-tuning, then repeat, whereas \"levels\" implies all of the above are done per gradient update. \n4. [Critical] Specify computational complexity of the algorithm used in 3.3.\n5. [Strengthen] Add computational complexity comparison with baselines.\n6. [Strengthen] Discuss how TapWeight computational complexity scales with model size.\n7. [Strengthen] B.1 equations (1-3), is the model under training $f_{\\theta}$ or does it include the C's as well?", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a strategy to set task weights in pretrain/fine-tune setups by backpropagating through the pre-training, fine-tuning, and evaluation steps using the implicit function theorem. After the discussion period, all recommend or lean toward acceptance. The AE concurs with the reviewers that the paper meets TMLR's evaluation criteria.
Decision: Accept (poster)