IMWA: Iterative Model Weight Averaging Benefits Class-Imbalanced Learning Tasks

TL;DR

This paper introduces IMWA, an iterative model weight averaging method that improves class-imbalanced learning by averaging models in multiple training episodes, outperforming vanilla MWA and complementing EMA strategies.

Contribution

The paper proposes a novel iterative MWA technique for class-imbalanced tasks, demonstrating its effectiveness and compatibility with existing EMA methods.

Findings

IMWA outperforms vanilla MWA in class-imbalanced learning.

Early-epoch averaging yields better performance.

IMWA enhances various class-imbalanced tasks, including image classification and object detection.

Abstract

Model Weight Averaging (MWA) is a technique that seeks to enhance model's performance by averaging the weights of multiple trained models. This paper first empirically finds that 1) the vanilla MWA can benefit the class-imbalanced learning, and 2) performing model averaging in the early epochs of training yields a greater performance improvement than doing that in later epochs. Inspired by these two observations, in this paper we propose a novel MWA technique for class-imbalanced learning tasks named Iterative Model Weight Averaging (IMWA). Specifically, IMWA divides the entire training stage into multiple episodes. Within each episode, multiple models are concurrently trained from the same initialized model weight, and subsequently averaged into a singular model. Then, the weight of this average model serves as a fresh initialization for the ensuing episode, thus establishing an iterative learning paradigm. Compared to vanilla MWA, IMWA achieves higher performance improvements with the same computational cost. Moreover, IMWA can further enhance the performance of those methods employing EMA strategy, demonstrating that IMWA and EMA can complement each other. Extensive experiments on various class-imbalanced learning tasks, i.e., class-imbalanced image classification, semi-supervised class-imbalanced image classification and semi-supervised object detection tasks showcase the effectiveness of our IMWA.