Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Masking meets Supervision: A Strong Learning Alliance

Byeongho Heo, Taekyung Kim, Sangdoo Yun, Dongyoon Han

TL;DR

MaskSub tackles the instability of strong masking augmentations in supervised learning by introducing a two-branch framework: a main-branch trained with standard supervision ($r_{mask}=0$) and a masked sub-branch ($r_{mask}=r$) guided by a relaxed self-distillation loss that averages the main and sub-branch outputs. This approach stabilizes optimization under masks and accelerates convergence while delivering consistent accuracy gains across ViT-based models, hierarchical architectures like ResNet and Swin, and even cross-domain tasks including BERT and CLIP finetuning. The authors further generalize the idea to DropSub and PathSub, showing benefits across pretraining, finetuning, transfer learning, and cross-domain applications (GLUE and CLIP), often with modest compute overhead. Overall, MaskSub provides a practical, broadly applicable augmentation mechanism that enhances supervised learning by leveraging masking in a stable, distillation-inspired two-branch framework.

Abstract

Pre-training with random masked inputs has emerged as a novel trend in self-supervised training. However, supervised learning still 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-branch (MaskSub). MaskSub consists of the main-branch and sub-branch, the latter being a part of the former. The main-branch undergoes conventional training recipes, while the sub-branch merits intensive masking augmentations, during training. MaskSub tackles 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 faster than in standard training, which suggests our method stabilizes the training process. We further validate MaskSub across diverse training scenarios and models, including DeiT-III training, MAE finetuning, CLIP finetuning, BERT training, and hierarchical architectures (ResNet and Swin Transformer). Our results show that MaskSub consistently achieves impressive 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

Masking meets Supervision: A Strong Learning Alliance

TL;DR

MaskSub tackles the instability of strong masking augmentations in supervised learning by introducing a two-branch framework: a main-branch trained with standard supervision () and a masked sub-branch () guided by a relaxed self-distillation loss that averages the main and sub-branch outputs. This approach stabilizes optimization under masks and accelerates convergence while delivering consistent accuracy gains across ViT-based models, hierarchical architectures like ResNet and Swin, and even cross-domain tasks including BERT and CLIP finetuning. The authors further generalize the idea to DropSub and PathSub, showing benefits across pretraining, finetuning, transfer learning, and cross-domain applications (GLUE and CLIP), often with modest compute overhead. Overall, MaskSub provides a practical, broadly applicable augmentation mechanism that enhances supervised learning by leveraging masking in a stable, distillation-inspired two-branch framework.

Abstract

Pre-training with random masked inputs has emerged as a novel trend in self-supervised training. However, supervised learning still 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-branch (MaskSub). MaskSub consists of the main-branch and sub-branch, the latter being a part of the former. The main-branch undergoes conventional training recipes, while the sub-branch merits intensive masking augmentations, during training. MaskSub tackles 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 faster than in standard training, which suggests our method stabilizes the training process. We further validate MaskSub across diverse training scenarios and models, including DeiT-III training, MAE finetuning, CLIP finetuning, BERT training, and hierarchical architectures (ResNet and Swin Transformer). Our results show that MaskSub consistently achieves impressive 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
Paper Structure (24 sections, 3 equations, 4 figures, 28 tables, 1 algorithm)

This paper contains 24 sections, 3 equations, 4 figures, 28 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Masked Sub-branch (MaskSub)
  5. Analysis
  6. Expand to drop regularizations
  7. Experiments
  8. Training from scratch (pretraining)
  9. Finetuning
  10. Hierarchical architecture
  11. Self-distillation
  12. Training budget
  13. Transfer learning
  14. Beyond vision domain
  15. Extending to drop regularizations
  16. ...and 9 more sections

Figures (4)

  • Figure 1: Overview of Masked Sub-branch (MaskSub). (a) standard supervised training; (b) masking augmentation training. The masking is applied to the main model, which degrades performance; (c) our MaskSub training, which separates the masking from the main model using the sub-branch and relaxes loss with self-distillation. MaskSub substantially improves the state-of-the-art training recipes touvron2022deit3wightman2021resnet.
  • Figure 2: MaskSub training analysis. We use 50%-random masking to compare three training settings: standard \ref{['eq:cross_entropy']}, masking \ref{['eq:reg_cross_entropy']}, and MaskSub. We visualize (a) validation accuracy; (b) train loss without masking; (c) train loss with masking; (d) gradient norm.
  • Figure : MaskSub in PyTorch-style pseudo-code
  • Figure A.1: MaskSub training analysis for training budget. The figures show training analysis experiments in Fig. 2 with the training budget (GPU days) as x-axis. Even considering its additional training budgets, MaskSub effectively improves convergence and accuracy.