[Paper Review] Two-Stage Fine-Tuning: A Novel Strategy for Learning Class-Imbalanced Data

Taha ValizadehAslani, Yiwen Shi+6 more|arXiv (Cornell University)|Jul 22, 2022

Imbalanced Data Classification Techniques4 citations

TL;DR

This paper proposes a two-stage fine-tuning strategy to address class imbalance in text classification, where the final layer is first fine-tuned with class-balanced reweighting loss before standard fine-tuning. The method improves tail-class performance by protecting minority classes during initial representation learning, outperforming standard and reweighted fine-tuning across synthetic and real-world ADME semantic labeling datasets.

ABSTRACT

Classification on long-tailed distributed data is a challenging problem, which suffers from serious class-imbalance and hence poor performance on tail classes with only a few samples. Owing to this paucity of samples, learning on the tail classes is especially challenging for the fine-tuning when transferring a pretrained model to a downstream task. In this work, we present a simple modification of standard fine-tuning to cope with these challenges. Specifically, we propose a two-stage fine-tuning: we first fine-tune the final layer of the pretrained model with class-balanced reweighting loss, and then we perform the standard fine-tuning. Our modification has several benefits: (1) it leverages pretrained representations by only fine-tuning a small portion of the model parameters while keeping the rest untouched; (2) it allows the model to learn an initial representation of the specific task; and importantly (3) it protects the learning of tail classes from being at a disadvantage during the model updating. We conduct extensive experiments on synthetic datasets of both two-class and multi-class tasks of text classification as well as a real-world application to ADME (i.e., absorption, distribution, metabolism, and excretion) semantic labeling. The experimental results show that the proposed two-stage fine-tuning outperforms both fine-tuning with conventional loss and fine-tuning with a reweighting loss on the above datasets.

Motivation & Objective

• To address the challenge of poor performance on tail classes in long-tailed, class-imbalanced text classification datasets.
• To mitigate the disadvantage faced by underrepresented classes during standard fine-tuning of pretrained models.
• To develop a simple yet effective modification to standard fine-tuning that enhances learning on minority classes without retraining the entire model.
• To evaluate the method on both synthetic and real-world text classification benchmarks, including ADME semantic labeling.
• To demonstrate that initial fine-tuning with class-balanced loss improves downstream performance on long-tailed data distributions.

Proposed method

• First, the final classification layer of a pretrained model is fine-tuned using a class-balanced reweighting loss to prioritize underrepresented classes.
• The reweighting loss assigns higher weights to loss contributions from minority classes to counteract their scarcity.
• After the first stage, the entire model is fine-tuned using standard cross-entropy loss in a second stage.
• The rest of the model parameters remain frozen during the first stage, preserving pretrained representations.
• The method leverages the full model capacity in the second stage while ensuring initial learning focuses on class imbalance.
• The approach is applied to both multi-class and two-class text classification tasks with minimal architectural changes.

Experimental results

Research questions

• RQ1 Can a two-stage fine-tuning strategy improve performance on minority classes in long-tailed text classification datasets?
• RQ2 Does initial fine-tuning with class-balanced reweighting loss lead to better representation learning for underrepresented classes compared to standard fine-tuning?
• RQ3 How does the proposed method compare to standard fine-tuning and fine-tuning with reweighting loss in terms of tail-class performance?
• RQ4 Is the two-stage strategy effective across both synthetic and real-world text classification tasks with class imbalance?
• RQ5 Can the method achieve better performance without fine-tuning the entire model in the first stage?

Key findings

• The two-stage fine-tuning method achieved superior performance on tail classes compared to standard fine-tuning and fine-tuning with reweighting loss on synthetic two-class and multi-class text classification datasets.
• On the ADME semantic labeling task, the proposed method significantly improved F1-scores for minority classes, demonstrating its effectiveness in real-world applications.
• The first-stage reweighting loss effectively protected minority classes from being overshadowed during initial learning, leading to better generalization.
• The method maintained high performance on head classes while substantially improving performance on tail classes, indicating a balanced improvement across the class distribution.
• The ablation study confirmed that the two-stage process was essential, as skipping the first stage led to degraded performance on minority classes.