Slimmable Networks for Contrastive Self-supervised Learning

TL;DR

SlimCLR introduces slimmable networks to self-supervised pretraining, enabling the extraction of multiple small models from one training run without large teacher models. The core challenge is interference among weight-sharing sub-networks, manifesting as gradient magnitude imbalance and gradient direction divergence; the paper mitigates this with slow start, online distillation, and loss reweighting, plus a switchable linear probe for linear evaluation. Empirical results on ImageNet with MoCov2 and MoCov3 show that sub-networks achieve substantial gains over isolated small models, while the full model may trade a bit of accuracy for overall efficiency. The work demonstrates practical benefits for deploying adaptable, computation-aware representations across devices and tasks, including transfer to COCO without extra teacher models. The approach provides a non-distillation, one-shot pathway to scalable, small-model self-supervised learning with strong transfer capabilities.

Abstract

Self-supervised learning makes significant progress in pre-training large models, but struggles with small models. Mainstream solutions to this problem rely mainly on knowledge distillation, which involves a two-stage procedure: first training a large teacher model and then distilling it to improve the generalization ability of smaller ones. In this work, we introduce another one-stage solution to obtain pre-trained small models without the need for extra teachers, namely, slimmable networks for contrastive self-supervised learning (SlimCLR). A slimmable network consists of a full network and several weight-sharing sub-networks, which can be pre-trained once to obtain various networks, including small ones with low computation costs. However, interference between weight-sharing networks leads to severe performance degradation in self-supervised cases, as evidenced by gradient magnitude imbalance and gradient direction divergence. The former indicates that a small proportion of parameters produce dominant gradients during backpropagation, while the main parameters may not be fully optimized. The latter shows that the gradient direction is disordered, and the optimization process is unstable. To address these issues, we introduce three techniques to make the main parameters produce dominant gradients and sub-networks have consistent outputs. These techniques include slow start training of sub-networks, online distillation, and loss re-weighting according to model sizes. Furthermore, theoretical results are presented to demonstrate that a single slimmable linear layer is sub-optimal during linear evaluation. Thus a switchable linear probe layer is applied during linear evaluation. We instantiate SlimCLR with typical contrastive learning frameworks and achieve better performance than previous arts with fewer parameters and FLOPs. The code is at https://github.com/mzhaoshuai/SlimCLR.

Figures (5)

• Figure 1: A slimmable ResNet-50 in supervised (left) and self-supervised (right) manners. $\text{ResNet-50}_{[1.0, 0.75, 0.5, 0.25]}$ means this slimmbale network can switch at width $[1.0, 0.75, 0.5, 0.25]$. The width ${0.25}$ represents that the number of channels is scaled by $0.25$ of the full network.
• Figure 2: The overall framework of SlimCLR. A slimmable network produces different outputs from weight-sharing networks with various widths $w_1,\ldots,w_n$, where $w_1$ is the width of the full model. $\theta$ are the network parameters and $\xi$ are an exponential moving average of $\theta$. $\mathtt{sg}$ means stop-gradient.
• Figure 3: Ratios of gradient norms: $\frac{\lVert \nabla_{\theta_{1.0}}\mathcal{L}_{\theta} \rVert_2}{\lVert\nabla_{\theta_{0.25}}\mathcal{L}_{\theta} \rVert_2}$ and $\frac{\lVert \nabla_{\theta_{1.0 \backslash 0.25}} \mathcal{L}_{\theta} \rVert_2}{ \lVert \nabla_{\theta_{0.25}} \mathcal{L}_{\theta} \rVert_2}$. The gradient norm of the network is calculated by averaging the layer-wise $\ell_2$ gradient norms. $\nabla_{\theta_{1.0 \backslash 0.25}} \mathcal{L}_{\theta}$ is the gradient of the final loss w.r.t. parameters $\theta_{1.0 \backslash 0.25} \in \Theta_{1.0} \backslash \Theta_{0.25}$, i.e., rest parameters of $\Theta_{1.0}$ besides $\Theta_{0.25}$.
• Figure 4: Gradient direction divergence. The principal gradient directions of the last linear layer are shown.
• Figure 5: Error surface and optimization trajectory. Best viewed in color with 300% zoom.