Should We Be Pre-training? An Argument for End-task Aware Training as an Alternative

TL;DR

This work argues that when a specific end-task is known in advance, end-task aware training (TARTAN) can outperform traditional end-task agnostic pre-training and continued pre-training. It introduces MT-TARTAN, a multi-tasking variant, and META-TARTAN, a meta-learning variant, to jointly optimize end-task and auxiliary objectives, achieving higher data efficiency and better generalization on three low-resource NLP tasks. The results show notable improvements over DAPT/TAPT, with META-TARTAN particularly effective at leveraging out-of-distribution auxiliary data and discovering robust task-weighting strategies. The study suggests a paradigm shift toward end-task aware pre-training for resource-constrained scenarios and outlines future directions for finer-grained weighting and broader application.

Abstract

In most settings of practical concern, machine learning practitioners know in advance what end-task they wish to boost with auxiliary tasks. However, widely used methods for leveraging auxiliary data like pre-training and its continued-pretraining variant are end-task agnostic: they rarely, if ever, exploit knowledge of the target task. We study replacing end-task agnostic continued training of pre-trained language models with end-task aware training of said models. We argue that for sufficiently important end-tasks, the benefits of leveraging auxiliary data in a task-aware fashion can justify forgoing the traditional approach of obtaining generic, end-task agnostic representations as with (continued) pre-training. On three different low-resource NLP tasks from two domains, we demonstrate that multi-tasking the end-task and auxiliary objectives results in significantly better downstream task performance than the widely-used task-agnostic continued pre-training paradigm of Gururangan et al. (2020). We next introduce an online meta-learning algorithm that learns a set of multi-task weights to better balance among our multiple auxiliary objectives, achieving further improvements on end-task performance and data efficiency.

Figures (4)

• Figure 1: Pre-training trains on auxiliary task $T_{\mathrm{aux}}$ before fine-tuning on primary task $T^*$. End-task aware training optimizes both $T_{\mathrm{aux}}$ and $T^*$ simultaneously and can find better minima since optimization is informed by the end-task.
• Figure 2: Compared to DAPT, TARTAN makes more efficient use of data. Large standard deviations are a result of the heterogeneity of the domain data used and the fact that our tasks are low-resource.
• Figure 3: Having a separate classification head for computing meta-gradients is important. Using the same head as when training up-weights the end-task and under-utilizes auxiliary tasks.
• Figure 4: The meta-learned task weightings show similar trajectories across different end-tasks.

Theorems & Definitions (1)

• proof