[Paper Review] Weighted Training for Cross-Task Learning

Shuxiao Chen, Koby Crammer+3 more|arXiv (Cornell University)|May 28, 2021

Topic Modeling34 references6 citations

TL;DR

This paper proposes Target-Aware Weighted Training (TAWT), a method for cross-task learning that adaptively assigns sample weights by minimizing a representation-based task distance between source and target tasks. TAWT improves BERT's performance on low-resource NLP sequence tagging tasks by an average of 3.1% absolute improvement, with theoretical guarantees and minimal hyperparameter tuning.

ABSTRACT

In this paper, we introduce Target-Aware Weighted Training (TAWT), a weighted training algorithm for cross-task learning based on minimizing a representation-based task distance between the source and target tasks. We show that TAWT is easy to implement, is computationally efficient, requires little hyperparameter tuning, and enjoys non-asymptotic learning-theoretic guarantees. The effectiveness of TAWT is corroborated through extensive experiments with BERT on four sequence tagging tasks in natural language processing (NLP), including part-of-speech (PoS) tagging, chunking, predicate detection, and named entity recognition (NER). As a byproduct, the proposed representation-based task distance allows one to reason in a theoretically principled way about several critical aspects of cross-task learning, such as the choice of the source data and the impact of fine-tuning.

Motivation & Objective

• To address the lack of theoretical understanding in cross-task learning, especially in deep learning settings.
• To develop a practical, efficient, and theoretically grounded method for improving sample efficiency in cross-task learning.
• To enable principled reasoning about source data value and the necessity of fine-tuning in cross-task transfer.
• To overcome the strong assumption of shared representations across tasks in existing transfer learning frameworks.
• To provide a lightweight, trainable weighting scheme that enhances pre-training and joint training without requiring extensive hyperparameter tuning.

Proposed method

• TAWT minimizes a representation-based task distance between source and target tasks to guide adaptive sample weighting.
• The method computes task distance based on the discrepancy between optimal representations of source and target tasks, not task-specific functions.
• It assigns trainable, adaptive weights to source samples or tasks to improve alignment with the target task's representation.
• The algorithm is compatible with existing paradigms such as pre-training and joint training in multi-task learning.
• Theoretical analysis provides a non-asymptotic generalization bound that includes a vanishing term dependent on the minimized task distance.
• The approach avoids strong assumptions like a single shared representation across all tasks, enabling more flexible and interpretable transfer.

Experimental results

Research questions

• RQ1 How can we theoretically justify and improve the efficiency of cross-task learning in deep neural networks?
• RQ2 What is a principled way to quantify the value of source data for a given target task?
• RQ3 When is fine-tuning necessary in cross-task learning, and how can we detect it via representation discrepancy?
• RQ4 Can we design a lightweight, effective, and theoretically grounded weighting scheme for source tasks or samples?
• RQ5 How does minimizing representation-based task distance lead to improved generalization and performance in low-resource settings?

Key findings

• TAWT achieves an average absolute improvement of 3.1% on four NLP sequence tagging tasks—PoS tagging, chunking, predicate detection, and NER—under low-resource target data settings.
• The representation-based task distance provides an interpretable measure of source data quality, reflecting the discrepancy between optimal source and target representations.
• The method demonstrates that fine-tuning is necessary when the representation-based task distance is non-zero, as source representations do not converge to the optimal target representation.
• TAWT requires minimal hyperparameter tuning and is computationally efficient, making it practical for real-world deployment.
• Theoretical analysis confirms a non-asymptotic generalization bound where the task distance term becomes negligible due to re-weighting, supporting improved generalization.
• Empirical results show consistent performance gains across diverse NLP tasks using BERT, even with limited target data, validating the method’s robustness and effectiveness.