Masking Augmentation for Supervised Learning

2023-06-20

AI Summary

In supervised learning, strong random masking augmentation often induces training instability. To address this, we propose Masked Sub-model (MaskSub): a dual-model framework where the main model undergoes standard supervised training, while a dedicated sub-model handles masked inputs; a class-wise self-distillation with relaxed loss is introduced to mitigate optimization oscillations. This is the first method to stably enable strong masking augmentation in purely supervised settings—without requiring contrastive objectives or reconstruction targets. Our core innovations are the synergistic dual-model architecture and the relaxed distillation loss, which jointly balance augmentation strength and training stability. Extensive experiments across diverse architectures—including DeiT-III, MAE, CLIP, ResNet, and Swin—and multiple training paradigms demonstrate consistent performance gains, accelerated loss convergence, and superiority over state-of-the-art baselines. The code is publicly available.

Abstract

Pre-training using random masking has emerged as a novel trend in training techniques. However, supervised learning faces a challenge in adopting masking augmentations, primarily due to unstable training. In this paper, we propose a novel way to involve masking augmentations dubbed Masked Sub-model (MaskSub). MaskSub consists of the main-model and sub-model; while the former enjoys conventional training recipes, the latter leverages the benefit of strong masking augmentations in training. MaskSub addresses the challenge by mitigating adverse effects through a relaxed loss function similar to a self-distillation loss. Our analysis shows that MaskSub improves performance, with the training loss converging even faster than regular training, which suggests our method facilitates training. We further validate MaskSub across diverse training recipes and models, including DeiT-III, MAE fine-tuning, CLIP fine-tuning, ResNet, and Swin Transformer. Our results show that MaskSub consistently provides significant performance gains across all the cases. MaskSub provides a practical and effective solution for introducing additional regularization under various training recipes. Code available at https://github.com/naver-ai/augsub

Problem

• Adopting masking augmentations in supervised learning
• Mitigating unstable training with masking augmentations
• Improving performance across diverse models and scenarios

Innovation

• Uses main-branch and sub-branch architecture
• Applies masking augmentations to sub-branch
• Employs relaxed self-distillation loss function