[Paper Review] AutoMix: Unveiling the Power of Mixup for Stronger Classifiers

Zicheng Liu, Siyuan Li+5 more|arXiv (Cornell University)|Mar 24, 2021

Advanced Neural Network Applications4 citations

TL;DR

AutoMix proposes an end-to-end, bi-level optimization framework that automatically learns adaptive mixup policies for improved image classification. By decoupling mixed sample generation and classification via a learnable Mix Block and a Momentum Pipeline, AutoMix achieves state-of-the-art accuracy across nine benchmarks with minimal computational overhead, outperforming prior hand-crafted and optimization-based mixup methods.

ABSTRACT

Data mixing augmentation have proved to be effective in improving the generalization ability of deep neural networks. While early methods mix samples by hand-crafted policies (e.g., linear interpolation), recent methods utilize saliency information to match the mixed samples and labels via complex offline optimization. However, there arises a trade-off between precise mixing policies and optimization complexity. To address this challenge, we propose a novel automatic mixup (AutoMix) framework, where the mixup policy is parameterized and serves the ultimate classification goal directly. Specifically, AutoMix reformulates the mixup classification into two sub-tasks (i.e., mixed sample generation and mixup classification) with corresponding sub-networks and solves them in a bi-level optimization framework. For the generation, a learnable lightweight mixup generator, Mix Block, is designed to generate mixed samples by modeling patch-wise relationships under the direct supervision of the corresponding mixed labels. To prevent the degradation and instability of bi-level optimization, we further introduce a momentum pipeline to train AutoMix in an end-to-end manner. Extensive experiments on nine image benchmarks prove the superiority of AutoMix compared with state-of-the-art in various classification scenarios and downstream tasks.

Motivation & Objective

• To address the trade-off between precise mixing policies and high optimization complexity in mixup data augmentation.
• To unify mixed sample generation and classification into a single, end-to-end trainable framework.
• To improve model generalization and stability by decoupling the training of the generation and classification sub-networks.
• To eliminate reliance on hand-crafted or saliency-guided mixing strategies in favor of a learnable, task-driven policy.
• To achieve superior performance across diverse datasets and network architectures without increasing inference cost.

Proposed method

• AutoMix reformulates mixup training into two sub-tasks: mixed sample generation and mixup classification, solved via a bi-level optimization framework.
• A learnable, lightweight module called Mix Block generates mixed samples by modeling patch-wise relationships using cross-attention on feature maps of paired inputs.
• The Mix Block is trained under direct supervision from the corresponding mixed labels to ensure alignment with the classification objective.
• A Momentum Pipeline is introduced to decouple and stabilize the bi-level optimization, preventing gradient entanglement and model collapse.
• The framework jointly optimizes the Mix Block and the classification head in an end-to-end manner, with hyper-parameters α (Beta distribution) and feature layer l for control.
• The method avoids expensive portfolio optimization by directly optimizing the mixup policy toward the classification loss.

Experimental results

Research questions

• RQ1 Can an end-to-end, learnable mixup policy outperform hand-crafted and saliency-guided mixup methods in terms of accuracy and efficiency?
• RQ2 How can mixed sample generation be effectively decoupled from classification training to improve stability and performance?
• RQ3 What is the impact of different feature layers and mixing ratio distributions on AutoMix's generalization ability?
• RQ4 Can the proposed Momentum Pipeline effectively stabilize the bi-level optimization of the generation and classification tasks?
• RQ5 Does AutoMix generalize across diverse datasets and network architectures, including vision transformers and ResNets?

Key findings

• AutoMix achieves state-of-the-art performance on nine image classification benchmarks, including CIFAR-10/100, Tiny-ImageNet, ImageNet-1k, CUB-200, FGVC-Aircraft, iNaturalist2017/2018, and Places205.
• On ImageNet-1k, AutoMix improves top-1 accuracy by up to 1.2% over CutMix and 0.8% over PuzzleMix using ResNet-18.
• With ResNet-50, AutoMix achieves 70.72% top-1 accuracy on Tiny-ImageNet, outperforming MixUp (63.86%) and CutMix (66.36%).
• The ablation study confirms that α=2 and feature layer l₃ yield the best trade-off between accuracy and complexity.
• AutoMix reduces training time by up to 70% compared to combinatorial optimization-based methods like PuzzleMix, while maintaining higher accuracy.
• The Momentum Pipeline significantly improves training stability, preventing the Mix Block from collapsing into trivial solutions.