TapWeight: Reweighting Pretraining Objectives for Task-Adaptive Pretraining

TL;DR

TapWeight introduces an automated method to optimize the importance of multiple pretraining objectives for task-adaptive pretraining, improving performance on downstream tasks in NLP and molecular property prediction.

Contribution

It proposes a novel multi-level optimization framework that automatically reweights pretraining objectives based on downstream feedback, reducing manual tuning and computational costs.

Findings

Significantly outperforms baseline methods in molecular property prediction.

Achieves superior results in natural language understanding tasks.

Demonstrates robustness and generalizability across domains.

Abstract

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