[Paper Review] Attentional Biased Stochastic Gradient for Imbalanced Classification.

Qi Qi, Yi Xu+3 more|arXiv (Cornell University)|Dec 13, 2020

Imbalanced Data Classification Techniques69 references4 citations

TL;DR

This paper proposes Attentional Biased Stochastic Gradient (ABSGD), a theoretically grounded, efficient method for imbalanced deep learning that assigns individual example weights via an attention mechanism proportional to the exponential of scaled loss values. It achieves state-of-the-art performance with only one backward pass per iteration, outperforming class-level and meta-learning-based individual weighting methods by capturing intra-class diversity while ensuring convergence via distributionally robust optimization (DRO).

ABSTRACT

In this paper, we present a simple yet effective method (ABSGD) for addressing the data imbalance issue in deep learning. Our method is a simple modification to momentum SGD where we leverage an attentional mechanism to assign an individual importance weight to each gradient in the mini-batch. Unlike many existing heuristic-driven methods for tackling data imbalance, our method is grounded in {\it theoretically justified distributionally robust optimization (DRO)}, which is guaranteed to converge to a stationary point of an information-regularized DRO problem. The individual-level weight of a sampled data is systematically proportional to the exponential of a scaled loss value of the data, where the scaling factor is interpreted as the regularization parameter in the framework of information-regularized DRO. Compared with existing class-level weighting schemes, our method can capture the diversity between individual examples within each class. Compared with existing individual-level weighting methods using meta-learning that require three backward propagations for computing mini-batch stochastic gradients, our method is more efficient with only one backward propagation at each iteration as in standard deep learning methods. To balance between the learning of feature extraction layers and the learning of the classifier layer, we employ a two-stage method that uses SGD for pretraining followed by ABSGD for learning a robust classifier and finetuning lower layers. Our empirical studies on several benchmark datasets demonstrate the effectiveness of the proposed method.

Motivation & Objective

• Address data imbalance in deep learning by capturing intra-class example diversity beyond class-level weighting.
• Develop a method grounded in distributionally robust optimization (DRO) with theoretical convergence guarantees.
• Reduce computational cost compared to meta-learning-based individual weighting methods that require multiple backward passes.
• Balance feature learning and classifier training through a two-stage training strategy using SGD followed by ABSGD.
• Improve generalization on long-tailed and imbalanced benchmark datasets without architectural changes.

Proposed method

• Introduce an attentional mechanism that assigns individual gradient weights based on the exponential of scaled loss values per sample in a mini-batch.
• Formulate the method within the framework of information-regularized DRO, where the scaling factor acts as a regularization parameter.
• Ensure convergence to a stationary point of the regularized DRO problem through theoretical analysis.
• Use a single backward pass per iteration, maintaining efficiency comparable to standard SGD.
• Implement a two-stage training pipeline: first pretrain with standard SGD, then fine-tune with ABSGD to optimize the classifier and lower layers.
• Assign higher attention weights to examples with higher loss, focusing optimization on hard or underrepresented samples.

Experimental results

Research questions

• RQ1 Can an attention-based gradient weighting scheme improve generalization in imbalanced classification while maintaining computational efficiency?
• RQ2 Does the proposed method achieve better performance than class-level weighting schemes by capturing intra-class example diversity?
• RQ3 Can the method be theoretically justified via distributionally robust optimization (DRO) with convergence guarantees?
• RQ4 How does the method compare to meta-learning-based individual weighting methods in terms of training efficiency and accuracy?
• RQ5 Does the two-stage training strategy effectively balance feature learning and classifier optimization in imbalanced settings?

Key findings

• ABSGD achieves state-of-the-art performance on multiple benchmark datasets with long-tailed class distributions.
• The method outperforms class-level weighting schemes by capturing intra-class variation and focusing on hard examples.
• ABSGD maintains computational efficiency with only one backward pass per iteration, unlike meta-learning-based methods requiring three.
• Theoretical analysis confirms convergence to a stationary point of the information-regularized DRO problem.
• The two-stage training strategy improves robustness, especially in the classifier layer, by decoupling feature learning from classifier optimization.
• Empirical results show consistent accuracy gains across diverse imbalanced datasets, demonstrating generalization beyond specific architectures.