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

TL;DR

This work tackles learning under class imbalance by introducing Iterative Model Weight Averaging (IMWA), an enhancement of Model Weight Averaging (MWA) that trains multiple models in parallel within episodes and averages their weights to form new initializations, iterating this process over $E$ episodes. IMWA preserves comparable computational cost to vanilla MWA while enabling benefits at each episode, and it can be combined with Exponential Moving Average (EMA) for additive gains. Empirical results across class-imbalanced image classification, semi-supervised variants, and semi-supervised object detection demonstrate that IMWA yields consistent improvements over vanilla MWA and often complements EMA, achieving state-of-the-art or near state-of-the-art performance on several benchmarks. The method is simple to integrate with existing approaches and robust across architectures and data regimes, indicating practical utility for imbalanced recognition tasks.

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.

Figures (6)

• Figure 1: Observations from applying MWA on image classification task. The averaged model is obtained by two trained individual ResNet-34 He_2016_CVPR. (a)Accuracy of the best individual model vs. average model on both class-balanced (i.e., CIFAR-10 and CIFAR-100) and class-imbalanced (i.e., CIFAR-10-LT and CIFAR-100-LT) datasets, where average model performs better on the class-imbalanced scenario. (b)Improvements of the average model versus each epoch, where performing model weight averaging in the early stage of training brings higher improvement. (c)L2 distance between two trained individual models is higher in the class-imbalanced datasets.
• Figure 2: The illustration of (a) vanilla Model Weight Averaging, and (b) our proposed Iterative Model Weight Averaging (IMWA). IMWA splits the whole training stage into several episodes. At each episode, multiple models are trained in parallel from a same initialization model weight, with same training iterations but different data orders, and then they are average into one model. Then, the weight of the average model will be treated as a new initialization for the next episodes, which forms a "iterative" manner.
• Figure 3: Comparison of our IMWA and other MWA approaches in terms of achieving performance improvements for (a) BCL zhu2022balanced and (b) LiVT xu2023learning on three CIIC benchmarks.
• Figure 4: Comparison of our IMWA and vanilla MWA approaches in terms of achieving performance improvements for DASO oh2022daso and ACR wei2023towards under consistent class distributions on CISSIC benchmarks, i.e., CIFAR10-LT and CIFAR100-LT.
• Figure 5: Comparison of our IMWA and vanilla MWA approaches in terms of achieving performance improvements for DASO oh2022daso and ACR wei2023towards under inconsistent class distributions on CISSIC benchmarks, i.e., CIFAR10-LT and STL-LT.