Self-Supervised Learning based on Heat Equation

TL;DR

This paper introduces QB-Heat, a simple heat equation-inspired self-supervised learning method that effectively pre-trains CNNs and lightweight models for image classification and object detection, outperforming existing methods in certain settings.

Contribution

It extends the heat equation into high-dimensional feature space and simplifies it for masked image modeling, enabling effective self-supervised pre-training without complex architectures.

Findings

QB-Heat performs on par with MoCo-v2 in linear probing on ImageNet.

QB-Heat outperforms MoCo-v2 in non-linear probing with a transformer.

QB-Heat surpasses supervised pre-training in object detection AP when using a frozen backbone.

Abstract

This paper presents a new perspective of self-supervised learning based on extending heat equation into high dimensional feature space. In particular, we remove time dependence by steady-state condition, and extend the remaining 2D Laplacian from x--y isotropic to linear correlated. Furthermore, we simplify it by splitting x and y axes as two first-order linear differential equations. Such simplification explicitly models the spatial invariance along horizontal and vertical directions separately, supporting prediction across image blocks. This introduces a very simple masked image modeling (MIM) method, named QB-Heat. QB-Heat leaves a single block with size of quarter image unmasked and extrapolates other three masked quarters linearly. It brings MIM to CNNs without bells and whistles, and even works well for pre-training light-weight networks that are suitable for both image classification and object detection without fine-tuning. Compared with MoCo-v2 on pre-training a Mobile-Former with 5.8M parameters and 285M FLOPs, QB-Heat is on par in linear probing on ImageNet, but clearly outperforms in non-linear probing that adds a transformer block before linear classifier (65.6% vs. 52.9%). When transferring to object detection with frozen backbone, QB-Heat outperforms MoCo-v2 and supervised pre-training on ImageNet by 7.9 and 4.5 AP respectively. This work provides an insightful hypothesis on the invariance within visual representation over different shapes and textures: the linear relationship between horizontal and vertical derivatives. The code will be publicly released.