Mugs: A Multi-Granular Self-Supervised Learning Framework

TL;DR

The paper tackles the limitation of single-granular representations in self-supervised learning by introducing Mugs, a framework that learns three complementary granular features via instance discrimination, local-group discrimination, and group discrimination. It combines a two-crop teacher-student setup with memory buffers and local-group transformers to build instance-, local-group-, and group-level representations, optimized through a joint objective with equal weighting. Empirical results on ImageNet-1K and downstream tasks show state-of-the-art linear probing accuracy (e.g., $82.1\%$ on ViT-L/16) and strong transfer performance in detection, segmentation, and video segmentation, demonstrating improved generality and transferability. The work highlights the significance of multi-granular features for versatile downstream performance and provides a scalable, parameter-efficient pathway to richer visual representations.

Abstract

In self-supervised learning, multi-granular features are heavily desired though rarely investigated, as different downstream tasks (e.g., general and fine-grained classification) often require different or multi-granular features, e.g.~fine- or coarse-grained one or their mixture. In this work, for the first time, we propose an effective MUlti-Granular Self-supervised learning (Mugs) framework to explicitly learn multi-granular visual features. Mugs has three complementary granular supervisions: 1) an instance discrimination supervision (IDS), 2) a novel local-group discrimination supervision (LGDS), and 3) a group discrimination supervision (GDS). IDS distinguishes different instances to learn instance-level fine-grained features. LGDS aggregates features of an image and its neighbors into a local-group feature, and pulls local-group features from different crops of the same image together and push them away for others. It provides complementary instance supervision to IDS via an extra alignment on local neighbors, and scatters different local-groups separately to increase discriminability. Accordingly, it helps learn high-level fine-grained features at a local-group level. Finally, to prevent similar local-groups from being scattered randomly or far away, GDS brings similar samples close and thus pulls similar local-groups together, capturing coarse-grained features at a (semantic) group level. Consequently, Mugs can capture three granular features that often enjoy higher generality on diverse downstream tasks over single-granular features, e.g.~instance-level fine-grained features in contrastive learning. By only pretraining on ImageNet-1K, Mugs sets new SoTA linear probing accuracy 82.1$\%$ on ImageNet-1K and improves previous SoTA by $1.1\%$. It also surpasses SoTAs on other tasks, e.g. transfer learning, detection and segmentation.

Figures (9)

• Figure 1: Attention visualization on "mugs" of ViT-B / 16 trained by our Mugs. See more examples in Sec. \ref{['resultsonImagenet']}.
• Figure 2: Comparison of linear probing accuracy on ImageNet-1K. By pretraining on ImageNet-1K, under different model sizes (see (a)) and pretraining epochs (see (b)), Mugs consistently improves previous SoTA (iBOT) by at least 0.8%.
• Figure 3: Overall framework of Mugs. (a) shows the overall framework. For each image, Mugs respectively performs two random augmentations and feeds two crops into backbones of student and teacher. Next, it adopts three granular supervisions/losses: 1) instance discrimination supervision, 2) local-group discrimination supervision, and 3) group discrimination supervision. Teacher is updated via exponential moving average of student. "sg" denotes stop-gradient. (b) shows the pipeline of local-group modules in both student and teacher. This module averages all patch tokens, and then finds top-$k$ neighbors from memory buffer. Next, it uses a transformer to aggregate the average and its $k$ neighbors to obtain a local-group feature (class token) and feeds it into a local-group head.
• Figure 4: T-SNE visualization of the learned feature by ViT-B/16. We show the fish classes in ImageNet-1K, i.e., the first six classes, including tench, goldfish, white shark, tiger shark, hammerhead, electric ray. See more examples in Appendix.
• Figure 4: Visualization of pretrained ViT-B/16 (a) and ViT-S/16 (b) & (c) by Mugs. See more examples in Appendix. Best viewed in $3\times$ sized color pdf.