BIM: Block-Wise Self-Supervised Learning with Masked Image Modeling

TL;DR

BIM introduces a block-wise masked image modeling framework that reduces memory usage and enables efficient training of multiple neural network backbones simultaneously, making MIM more accessible for resource-constrained environments.

Contribution

The paper proposes a novel block-wise training method for MIM that decreases memory requirements and allows concurrent training of diverse DNN architectures.

Findings

Maintains high performance with reduced memory consumption.

Enables simultaneous training of multiple DNN backbones.

Reduces computational costs compared to traditional MIM training.

Abstract

Like masked language modeling (MLM) in natural language processing, masked image modeling (MIM) aims to extract valuable insights from image patches to enhance the feature extraction capabilities of the underlying deep neural network (DNN). Contrasted with other training paradigms like supervised learning and unsupervised contrastive learning, masked image modeling (MIM) pretraining typically demands significant computational resources in order to manage large training data batches (e.g., 4096). The significant memory and computation requirements pose a considerable challenge to its broad adoption. To mitigate this, we introduce a novel learning framework, termed~\textit{Block-Wise Masked Image Modeling} (BIM). This framework involves decomposing the MIM tasks into several sub-tasks with independent computation patterns, resulting in block-wise back-propagation operations instead of the traditional end-to-end approach. Our proposed BIM maintains superior performance compared to conventional MIM while greatly reducing peak memory consumption. Moreover, BIM naturally enables the concurrent training of numerous DNN backbones of varying depths. This leads to the creation of multiple trained DNN backbones, each tailored to different hardware platforms with distinct computing capabilities. This approach significantly reduces computational costs in comparison with training each DNN backbone individually. Our framework offers a promising solution for resource constrained training of MIM.