A Simple Recipe for Competitive Low-compute Self supervised Vision Models
This work addresses the problem of enabling practical self-supervised learning for low-compute vision models, which is incremental as it builds on existing distillation and SSL methods.
The paper tackles the performance drop of self-supervised learning in small vision architectures by proposing RoB, a simple distillation technique that repurposes joint-embedding methods to transfer knowledge from large teachers to small students, resulting in ViT-Tiny models that improve prior SSL state-of-the-art on ImageNet by 2.3% and match or exceed supervised distillation on five downstream tasks.
Self-supervised methods in vision have been mostly focused on large architectures as they seem to suffer from a significant performance drop for smaller architectures. In this paper, we propose a simple self-supervised distillation technique that can train high performance low-compute neural networks. Our main insight is that existing joint-embedding based SSL methods can be repurposed for knowledge distillation from a large self-supervised teacher to a small student model. Thus, we call our method Replace one Branch (RoB) as it simply replaces one branch of the joint-embedding training with a large teacher model. RoB is widely applicable to a number of architectures such as small ResNets, MobileNets and ViT, and pretrained models such as DINO, SwAV or iBOT. When pretraining on the ImageNet dataset, RoB yields models that compete with supervised knowledge distillation. When applied to MSN, RoB produces students with strong semi-supervised capabilities. Finally, our best ViT-Tiny models improve over prior SSL state-of-the-art on ImageNet by $2.3\%$ and are on par or better than a supervised distilled DeiT on five downstream transfer tasks (iNaturalist, CIFAR, Clevr/Count, Clevr/Dist and Places). We hope RoB enables practical self-supervision at smaller scale.