TapWeight: Reweighting Pretraining Objectives for Task-Adaptive Pretraining

5. Patterns for Pretraining Objectives Given Downstream Tasks

We observe that similar downstream tasks tend to assign similar weights to specific pretraining objectives. For instance, the Esol and Freesolv datasets, which both involve predicting physical chemistry properties of molecules, assign large weights to the MG3 pretraining objective. Conversely, the Toxcast and Clintox datasets, which focus on predicting molecular toxicity, assign small weights to MG3. We have updated Section 4.4 of the manuscript to include these observations. As the interactions between pretraining objectives and downstream task types are inherently complex, we plan to further investigate this direction in future work.

6. Convergence of TapWeight Compared to General Continued Pretraining

All experiments in this paper, including those for TapWeight and fixed-weight continued pretraining (CP), have reached empirical convergence. Empirically, TapWeight requires slightly more time to converge compared to CP with fixed weights. For example, using the MUV dataset as a downstream task, TapWeight takes approximately 1.94 times longer to converge than CP. We plan to conduct a theoretical analysis of TapWeight's convergence rate in future work.

7. Table Locations in the Paper

We thank the reviewer for the helpful suggestion. In response, we have reorganized the tables in the revised manuscript as per the reviewers’ comments to improve readability and presentation clarity.

$1$ Don’t Stop Pretraining? Make Prompt-based Fine-tuning Powerful Learner, Zhengxiang Shi, Aldo Lipani, NeurIPS 2023.

We appreciate your constructive feedback very much, and provide our response to your questions as follows.

1. Experiments on Larger Models

We thank the reviewer for this insightful suggestion. In response, we conducted additional experiments by applying TapWeight to a larger language model, RoBERTa-large, and compared it with baseline methods on three additional datasets: RCT, AGNews, and IMDB. As shown in Table 4 of the revised manuscript, TapWeight consistently outperforms the baseline methods across all three datasets, demonstrating its robustness and scalability across different sizes of pretrained models.

2. Comparison with Additional Baselines

We thank the reviewer for the helpful comments. In response, we conducted additional experiments to compare our method with the more advanced baseline, PCP $1$ . As shown in Tables 3 of the revised manuscript, experimental results on the GLUE benchmark demonstrate that TapWeight outperforms the PCP baseline, highlighting its advantage over more recent and powerful continued pretraining methods.

3. Experiments on Additional Tasks

We thank the reviewer for this valuable suggestion. In the revised manuscript, we evaluated our method on three commonly used datasets for continued pretraining: RCT, AGNews, and IMDB. These datasets cover a variety of NLP tasks beyond the NLU tasks in the GLUE benchmark: RCT involves classifying sentences in biomedical texts based on their functional roles, AGNews focuses on topic classification of news articles, and IMDB addresses sentiment analysis of movie reviews. As shown in Table 4 of the revised manuscript, TapWeight consistently outperforms both the baseline without continued pretraining and the baseline continued pretraining methods, TAPT and SimCSE, across all three datasets. These additional experiments strengthen the evidence of TapWeight's effectiveness across diverse NLP tasks outside of natural language understanding in GLUE.

We would also like to clarify why our method was not applied to text generation tasks in this paper. Most state-of-the-art language models for text generation, such as the GPT and LLaMA series, predominantly rely on a single auto-regressive loss for next-token prediction during pretraining. As this objective is the standard choice for such models, we did not prioritize applying TapWeight to reweight multiple objectives in this context. In contrast, masked language models and foundational models in specialized domains such as biology and medicine often employ multiple pretraining objectives, creating a clear need for methods like TapWeight to automate the tradeoff parameter selection. Nonetheless, we recognize the potential of applying TapWeight to text generation models and plan to explore this direction in future work.

4. Experiments with a Random Task as Objective

We thank the reviewers for the insightful suggestion. Based on this, we conducted an experiment where an additional pretraining task was introduced: predicting labels randomly sampled from a Gaussian distribution. Since these labels have no dependency on the real input data, this task does not contribute meaningfully to the continued pretraining process and is expected to receive zero weight. We added this random task alongside the original five pretraining objectives for molecular property prediction in TapWeight. Experimental results demonstrate that the weight for this random task rapidly decreases to zero within 100 training iterations. This finding highlights the ability of TapWeight to identify and effectively minimize the influence of useless pretraining objectives.

5. Relevant Related Works

We thank the reviewer for the suggestion. We have revised the manuscript and added the suggested papers to the related work section to provide a more comprehensive literature review.

6. Comparison with CP+SFT Under the Same Budget

We conducted additional experiments on three datasets from MoleculeNet (Bace, Clintox, and HIV) to compare the performance of the CP w/o reweighting method under a similar training budget as TapWeight. The baseline achieved an average score of 78.0, which is lower than the 80.9 score achieved by TapWeight. These results demonstrate that simply increasing the training time of CP does not significantly improve downstream task performance. As shown in Table 1 and Table 5 of the revised manuscript, CP w/o reweighting often fails to outperform the simple finetuning baseline without CP. This is likely because the CP stage was conducted on an unlabeled text corpus in the general domain, which is not strongly related to the downstream tasks. Without feedback from downstream tasks, as implemented in TapWeight, performing CP on this general corpus does not positively impact some downstream tasks. This is especially true when the continued pretraining methods are applied on models like RoBERTa, which have already been pretrained on general domain text.

7. Comparison with CP Using TapWeight Searched Trade-Off Weights

We thank the reviewer for the insightful suggestion. To address this, we conducted additional experiments on three MoleculeNet datasets (Bace, Clintox, and HIV) to evaluate the performance of CP using tradeoff weights searched by TapWeight. On average, this baseline achieved a score of 80.6, which is comparable to the 80.9 score of TapWeight, demonstrating the effectiveness of the tradeoff weights identified by TapWeight. However, this approach requires additional computational cost due to an extra round of CP without yielding significant performance improvements. Therefore, it is more practical to directly use the model weights from TapWeight without conducting another CP stage.

8. Additional Details on Computational Resources

All experiments were conducted using a single A100 GPU. We have revised the experimental setup section in the manuscript to make this clear.

9. Revision of Line 88

We thank the reviewer for pointing this out. We have changed this line in the revised manuscript to ‘TapWeight is broadly applicable to pretrained models with multiple pretraining objectives across various data modalities and downstream task types’ to enhance its precision.

$1$ Don’t Stop Pretraining? Make Prompt-based Fine-tuning Powerful Learner, Zhengxiang Shi, Aldo Lipani, NeurIPS 2023.

We appreciate your constructive feedback very much, and provide our response to your questions as follows.

1. Comparison with Baseline Methods

We apologize for any confusion. As shown in Table 3, in the NLU experiments on the GLUE benchmark, the SimCSE baseline method uses the same continued pretraining data as TapWeight, with MLM and CL as the pretraining objectives. Results indicate that TapWeight outperforms SimCSE on the GLUE benchmark, demonstrating that TapWeight's superior performance on NLU tasks stems from its reweighting strategy rather than the continued pretraining process itself.

Additionally, we conducted further experiments to compare our method with a more advanced baseline, PCP $1$ , for continued pretraining. As shown in Table 3, experimental results on the GLUE benchmark demonstrate that TapWeight outperforms the PCP baseline, highlighting its advantages over more recent and powerful continued pretraining methods.

2. Comparison to Different Ways of Hill Climbing the Pretraining Mixture

We thank the reviewer for the valuable suggestion. Regarding the proposed baseline method where "all pretraining objectives are weighted evenly in the mixture," this corresponds to the CP w/o reweighting ablation setting in Table 5 of the revised manuscript. Results indicate that TapWeight outperforms this baseline, highlighting its effectiveness as a method for optimizing the pretraining mixture.

For the grid search-based baseline method, if each of the five pretraining objectives in the molecular property prediction task has four possible tradeoff weight options, this would result in approximately 10^3 times the time consumption of a single continued pretraining run. While this computational cost could potentially be reduced using a smaller proxy model, we were unable to complete this experiment in the response period. We plan to further explore this direction of optimizing the pretraining mixture in future work.

3. Scenario for Longer Training Horizons

We thank the reviewer for the insightful suggestion. For smaller downstream datasets such as RTE, Bace, or Sider, the training time for the continued pretraining (CP) stage is more than 5 times longer than that of the supervised fine-tuning (SFT) stage, due to the small size of these downstream datasets. In these scenarios, TapWeight consistently outperforms baseline methods, demonstrating its robustness across varying scales of training horizons. We plan to further investigate the robustness of TapWeight in scenarios where the CP training extends to durations as long as several weeks in future work.

4. Clarifications on Update Strategies of TapWeight Framework

Is CP+SFT jointly optimized in TapWeight?

CP (level I) and SFT (level II) are two distinct levels in the TapWeight framework, and their optimizations are performed iteratively within a multi-level optimization framework. Importantly, the models at these two levels have different parameters. This differs from a joint optimization of CP+SFT, where both stages are updated in the same optimization step, typically using the same set of parameters.

How many tokens in CP and SFT are seen per-step?

The number of tokens processed in CP and SFT is determined by the respective batch sizes of each level, as specified in the hyperparameter section (Appendices B.4 and C.4). For example, when applying TapWeight to the continued pretraining of a RoBERTa-base model with feedback from the MRPC dataset, the batch size is 512 for the CP stage (level I) and 32 for the SFT stage (level II). The number of tokens processed is therefore not equal, reflecting the specified hyperparameters.

Is it true that the loss weights adaptively update during training?

Yes, the loss weights adaptively update during training, as illustrated in Figure 2.

Is extra SFT applied after TapWeight?

Yes, after TapWeight continued pretraining, the model weights from the CP stage (level I) are used for supervised fine-tuning (SFT) to achieve optimal performance on the downstream task.

Can TapWeight be used for ‘pretraining objectives that need to see different amounts of different pretraining objectives, like UL2’?

In UL2, the amounts of different objectives are manually determined and serve a role similar to the tradeoff weights in TapWeight. However, UL2 samples these objectives during pretraining rather than optimizing them simultaneously, making TapWeight not directly applicable due to the non-differentiability of the sampling process. That said, applying TapWeight in such a scenario could be feasible with techniques like reparameterization, which we plan to explore further in the future.

Dear Reviewer,

We greatly appreciate the thoughtful feedback and suggestions you provided. We have tried our best to address your questions in our response. As the discussion period is nearing its conclusion, we kindly request your input on whether any additional questions remain or if our response has satisfactorily resolved your initial concerns. We welcome any further inquiries you may have.

Warm regards,

The Authors

We appreciate your constructive feedback very much, and provide our response to your questions as follows.

1. Novelty of the Proposed Method

While our method shares some similarities with meta-learning approaches, it is distinct in several key aspects. Gradient-based meta-learning typically operates under a bi-level optimization framework and trains models from scratch, aiming to learn model parameters that quickly adapt to specific domains with limited training examples. These methods are rarely applied in the context of continued pretraining.

In contrast, TapWeight is specifically designed for learning the tradeoff weights for multiple continued pretraining objectives. Its goal is to determine the optimal weights such that, after continued pretraining and fine-tuning, the model achieves the best performance on the validation split of the downstream dataset. Our approach introduces a novel three-level optimization framework, which fundamentally differs from the traditional bi-level meta-learning paradigm and highlights the unique contributions of our method.

2. Presentation of Experimental Results

We apologize for any confusion caused by the presentation of our results. In response to the reviewer's comments, we have revised the manuscript to clearly specify the metric used in each table and explicitly indicate in the captions whether a higher or lower value represents better performance.

3. Additional Experiments on Other Domains

We thank the reviewer for this valuable suggestion. In the revised manuscript, we evaluated our method on three commonly used text datasets from different domains for continued pretraining: RCT (biomedical), AGNews (news), and IMDB (movie reviews). As presented in Table 4 of the revised manuscript, TapWeight consistently outperforms both the baseline method without continued pretraining and the baseline continued pretraining methods, TAPT and SimCSE, across all three datasets. These additional experiments across diverse domains further validate the effectiveness and robustness of our method in various contexts.

Dear Reviewer,

We greatly appreciate the thoughtful feedback and suggestions you provided. We have tried our best to address your questions in our response. As the discussion period is about to end, we are wondering whether any additional questions remain or if our response has successfully resolved your initial concerns. We welcome any further inquiries you may have.

Warm regards,

The Authors

We appreciate your constructive feedback very much, and provide our response to your questions as follows.

1. Evaluation on Additional Datasets

We thank the reviewer for this valuable suggestion. In the revised manuscript, we have evaluated our method on three widely used datasets for continued pretraining: RCT, AGNews, and IMDB. As presented in Table 4 of the revised manuscript, TapWeight consistently outperforms both the baseline method without continued pretraining and the baseline continued pretraining methods, TAPT and SimCSE, across all three datasets. These additional experiments further validate the effectiveness of our approach.

We also thank the reviewers for the suggestions of using CrossFit dataset. While CrossFit is a comprehensive and challenging benchmark, it is primarily designed for evaluating few-shot learning methods with only 16 training examples per task. Although TapWeight could be applied to investigate robustness in low-resource scenarios, this is slightly beyond the scope of our current work, as TapWeight focuses on continued pretraining for general downstream tasks. Due to time constraints during the author response period, we were unable to conduct experiments on CrossFit. However, we plan to explore its potential on this benchmark in future work.

2. Comparison with Additional Baseline Methods and Larger LLMs

We thank the reviewers for the insightful suggestions. In response, we conducted additional experiments comparing our method to a more advanced baseline, PCP $1$ , for continued pretraining. As shown in Table 3 in the revised manuscript, results on the GLUE benchmark show that TapWeight outperforms PCP, demonstrating its superiority over recent and powerful continued pretraining methods.

Furthermore, we extended our experiments by applying TapWeight to a larger language model, RoBERTa-large, and comparing it against baseline methods, on RCT, AGNews, and IMDB datasets. As shown in Table 4 of the revised manuscript, the results indicate that TapWeight maintains its superior performance over baseline methods across all three datasets tested, highlighting its robustness across models of varying scales.

We would also like to clarify the goals of our paper and address any potential misunderstandings:
(1) We do not propose new pretraining objectives. Instead, our approach TapWeight leverages existing continued pretraining objectives, automatically learning tradeoff weights between these objectives based on feedback from downstream tasks.
(2) Our focus is on proposing a novel continued pretraining (CP) method and comparing it against other CP baselines such as TAPT and PCP under comparable computational budgets and model scales, as specified in those works. TapWeight demonstrates significant performance gains under these conditions. Due to computational constraints, we were unable to conduct continued pretraining experiments on models of the scale of GPT-3 or GPT-3.5, as referenced in the Tag-LLM and GLUE-X papers.

3. Selection of Pretraining Objectives

In our NLU experiments, we selected pretraining objectives that are widely used and generally recognized as effective, such as MLM (BERT), CL (SimCSE), and SOP (ALBERT). These objectives were chosen based on their demonstrated success in prior studies, particularly under continued pretraining settings. For instance, MLM and CL objectives have been validated as effective for continued pretraining by SimCSE. In contrast, objectives like token detection and swap detection have shown effectiveness in ELECTRA-style training, which involves both a generator and a discriminator. This setup differs from the continued pretraining framework employed in this paper. While these objectives were not prioritized in our previous experiments, we plan to incorporate them in future studies to assess their impact and offer more comprehensive insights.

4. Analysis of Tradeoff Score Results

We observe that downstream tasks with similar objectives tend to share similar weights for pretraining tasks to some extent. For instance, the Esol and Freesolv datasets, which both involve predicting physical chemistry properties of molecules, assign large weights to the MG3 pretraining objective. In contrast, the Toxcast and Clintox datasets, both focused on predicting molecular toxicity, assign smaller weights to the MG3 objective. We have updated Section 4.4 of the revised manuscript to incorporate these observations. As the interactions between pretraining objectives and downstream task types are complex, we plan to investigate this further in future work to provide deeper insights.

$1$ Don’t Stop Pretraining? Make Prompt-based Fine-tuning Powerful Learner, Zhengxiang Shi, Aldo Lipani, NeurIPS 2023.

Thank you again for your detailed and considerate review of our manuscript. Your acknowledgement of our efforts have been greatly motivating. Once more, we extend our gratitude for your valuable time and expertise.