Paper Info.

Cao, Yun-Hao, Peiqin Sun, and Shuchang Zhou. "Three guidelines you should know for universally slimmable self-supervised learning." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

---

Abstract

• 우리는 다양한 device에 self-supervised model을 배포할 때 더 나은 accuracy-efficiency trade-offs를 달성하기 위해 universally slimmable self-supervised learning(US3L)을 제안한다
  self-supervised learning(SSL)을 universally slimmable networks에 직접 적용하면 training 과정에서 자주 collapse(붕괴 현상)이 발생하는 것을 관찰했다.
  이후, temporal consistent guidance가 univerally slimmable network에 SSL을 성공적으로 적용하는 데 중요한 요소임을 발견했으며,이를 위해 unified gradient 관점에서 temporal consistency를 보장하기 위한 loss design의 세 가지 guidelines을 제안한다.

• 또한, 우리는 training efficiency와 accuacy를 동시에 개선하기 위해 dynamic sampling and group regularization strategies를 제안한다.
  US3L 방법은 CNNs and vision transformer에서 실험적으로 검증되었으며,
  only once training and one copy of weights만으로도 recognition, object detction and instance segmentation 등 여러 benchmark에서 SOTA를 능가한다.

---

1. Introduction

• 지난 10년간 deep supervised learning은 큰 성공을 거뒀지만, 대규모 annotated training data에 크게 의존하는 것이 단점이다.
  SSL은 large-scale dataset에 annotating을 할 필요가 없다는 점에서 인기를 얻고 있다.
  contrastive learning의 등장 이후 SSL은 가속도가 붙었으며,
  최근 연구들에서는 supervised pretraining을 통한 downstream tasks에서 SSL이 유사하거나 더 나은 성능을 보이기도 했다.
  그러나 edge computing에 pretrained models을 배포하는 것은 해당 devices들의 memory, computation and storage capabilities의 제한으로 인해 여전히 어려운 문제이다.

• 이를 해결하기 위해 여러 model compression techniques들이 제안되었다.
  여기에는 lightweight architecture design, knowledge distillation, network pruning, quantization 등이 포함된다.
  supervised learning 상황에서 이러한 문제를 해결하기 위해 Slimmable Networks(S-Net)과 universally slimmable networks(US-Net)이라는 계열이 제안되었다.
  이 model들은 한 번의 학습만으로 다양한 widths를 자유롭게 switch할 수 있다.

• slimmable networks에 힘입어 다음과 같은 질문이 생긴다 :
  "Can we train a self-supervised model that can run at arbitrary width?"
  naive solution으로는 US-Net framework에서 supervised loss를 self-supervised loss로 대체하는 것이지만, 실험적으로 확인한 결과 이 방법은 직접 적용되지 않는다는 것을 알게 되었다.
  Table 1은 self-supervised 시나리오에서 현상이 매우 다르다는 것을 보여준다.
  인기 있는 SSL 방법인 SimSiam을 slimmable network에 적용한 후 model이 직접적으로 붕괴되었음을 확인할 수 있다.
  sub-networks에 대한 inplace distillation을 사용하면 model 붕괴는 막을 수 있지만, 여전히 SimSiam을 개별적으로 학습한 model의 결과와 S-Net+Distill 간에는 큰 차이가 있다.
  "그렇다면 SSL에서 상황이 어떻게 다른지, 그리고 성능을 어떻게 더 향상시킬 수 있을까?(즉, 격차를 어떻게 줄일 수 있을까?)"

• 이 논문에서는 차이를 설명하기 위한 통합적인 관점을 제시하고, 이러한 격차를 해결하기 위한 적절한 조치를 제안한다.
  unified gradient 관점에서 우리는 sub-networks에 대한 guidance가 iteration 사이에 consistent(일관성)을 유지해야 한다는 것이 핵심임을 발견했으며,
  SSL의 어떤 구성 요소가 temporal inconsistency 문제를 유발하는지와 US-Net이 왜 supervised learning에서 동작하는지 이유를 분석했다.
  이러한 이론적 분석을 바탕으로 우리는 US-Net training의 loss design에 대한 temporal consistency를 보장할 수 있는 세 가지 guidelines을 제안한다.
  이 중 하나라도 만족되면 US-Net은 supervised든 self-supervised든 잘 동작할 수 있다.
  또한, SSL의 특성과 US-Net의 한계를 고려하여, training overhead를 줄이고 accuracy를 향상시키기 위한 dynamic sampling과 group normalization을 제안한다.
  우리의 main contributions은 다음과 같다 :
  • 우리는 US-Net을 학습할 떄 supervised와 self-supervised learning 사이의 중요한 차이를 발견했다.
    이러한 관찰을 바탕으로 US-Net의 loss design에 대한 세 가지 guideline을 unified gradient 관점에서 분석하고 요약했다.
  • SSL에서 large data를 다루기 위해, accuracy를 희생하지 않고 training cost를 줄일 수 있는 dynamic sampling 전략을 제안
  • US-Net의 training 방식이 model capacity를 제한하는 방식을 분석하고, 서로 다른 channel에 서로 다른 freedom(자유도)을 부여하여 이를 해결하여 group regularization을 제안했다.
  • 우리의 방법은 CNN과 ViT 모두에서 효과를 입증했다.
    이 방법은 only once training과 a single model만을 요구하며, 각 model을 개별적으로 학습하는 결과를 초과할 수 있고(?),
    pretrained teacher로부터의 knowledge distillation과 비교할 만한 성능을 보여준다.

---

2. Related Works

Self-supervised Learning

• time-consuming and expensive data annotations을 피하기 위해,
  large-scale unlabeled images or videos로부터 visual representations을 학습하기 위한 여러 self-supervised methods가 제안되었다.
  SOTA SSL 방법의 주된 동력으로서, contrastive learning은 최근 몇 년간 representation learning의 성능을 크게 향상시켰다.
  contrastive learning은 유사한 sample은 가깝게 끌어당기고(pulling) 서로 다른 sample을 멀리 밀어내는 것(pushing)을 목표로 하는 discriminative approach이다.
  SimCLR [7]와 MoCo [16]는 모두 contrastive loss function인 InfoNCE를 사용하며, 이는 negative sample이 필요하다.
  BYOL [14]과 SimSiam [9]은 asymmetrical design을 사용하여 contrastive learning에서 negative sampling을 생략한다.

• self-supervised moel의 accuracy-efficiency trade-off를 개선하기 위해 여러 연구가 제안되어왔다.
  Fang et al. [13]은 lightweight models을 위한 self-supervised knowledge distillation(SEED)를 제안했다.
  하지만 서로 다른 widths(sparsities)를 가진 model을 각각 개별적으로 학습해야 하므로, 이는 상당한 computational cost and storage overhead를 발생시키며 large data volumes에 대해 지속적이지 않다.
  또한 pretrained teacher model이 필요하지만, 우리의 방법은 필요하지 않다.
  우리는 structured pruning에 중점을 두고, one-for-all training에 대한 loss-design의 통합적인 이론적 설명을 제공한다.

Slimmable Networks

• Slimmable networks, Universally slimmable networks, OFA 등은 모두 supervised learning paradigm에서 수행되는 반면,
  우리는 self-supervised 학습을 기반으로 한다.

---

3. Method

3.1. Preliminary

• 이 subsection에서는 두가지 representative SSL 방법인 SimSiam과 SlimCLR 뿐만 아니라
  (universally) slimmable networks도 설명할 것이다.

1) Self-supervised Losses

• input image $x_i$로부터 two randomly augmented views를 $x_{i,1} 과 x_{i,2}$로 정의.
  $f$는 backbone(예: ResNet)과 projection MLP head로 구성된 encoder network를 나타낸다.
  SimSiam은 한 image의 두 augmentations 간의 similartiy를 maximize한다.
  prediction MLP head $h$는 한 view의 output을 transform(변환)하여 다른 view에 맞춘다.
  $x_{i,1}$의 output vectors는 $z_{i,1} ≜ f(x_{i,1})$ 및 $p_{i,1} ≜ h(f(x_{i,1}))$로 표기한다.
  ($z_{i,2}$ 와 $p_{i,2}$도 유사하게 정의
  $x_{i,2}$의 output vectors는 $z_{i,2} ≜ f(x_{i,2})$ 및 $p_{i,2} ≜ h(f(x_{i,2}))$로 표기한다. )
  
  negative consine similarity는 $D(p, z) ≜ -\frac{p \cdot z}{\|p\|_2 \|z\|_2}$로 정의되며,
  이후 discussion을 간단하게 하기 위해 $z$와 $p$가 $L_2$-normalized 되었다고 가정한다.
  $SG(\cdot)$는 stop-gradient operation을 나타낸다.
  그러면 SimSiam에서의 손실 함수는 다음과 같다:
  
  SimCLR [7]과 MoCo [16]는 InfoNCE [30] loss을 사용하여 negative samples을 대조한다 :
  간단히 하기 위해 temperature parameter $\tau$를 생략했다.

2) Slimmable Networks

• slimmable networks는 다양한 scales로 실행할 수 있는 networks이다.
  training 동안, 각 iteration에서 the smallest, the largest and a few randomly sampled network만 loss를 계산하는 데 사용되며, 이를 sandwich rule이라고 한다.
  성능을 향상시키기 위해 inplace distillation이 도입되었으며, the largest network 내의 knowledge를 sub-networks에 distillation loss를 통해 전달한다.

3.2. The Proposed Method

• US3L method(Algorithm 1)의 구성요소를 설명하겠다.
  

3.2.1. Loss Design

• 우리 방법의 일반적인 framework는 Fig. 1a에 나타나 있다.
  여기서 loss function은 base loss(for base/largest network)와 distillation loss(for subnetworks)로 구성된다.
  기본적으로 우리는 targets을 생성하기 위해 momentum encoder, InfoNCE as the base loss, and MSE as the distillation loss를 사용한다.
  또한 teacher와 student 간의 capacity 차이로 인한 영향을 완화하기 위해 auxiliary distillation head를 사용해야 한다.
  전체 loss function은 다음과 같다 :
  
  $g(\cdot)$는 auxiliary distillation MLP head이며,
  $z^s$ 와 $z^m$ 은 각각 sub-network의 output과 momentum encoder의 output이다.
  Eq. (3)은 loss design의 유일한 option이 아니며, 우리가 제시한 guideline을 충족하는 한 잘 작동할 수 있다(Fig. 1b)
  (이에 대한 자세한 내용은 Sec. 3.3에서 논의)
  

3.2.2. Dynamic Sampling

• Yu et al. [32]가 training 과정에서 각 iteration마다 4개의 switch를 sampling했는데, 이는 SSL training에서 매우 time-consuming하다.
  따라서 우리는 training overhead를 줄이면서 성능을 향상시키기 위해 dynamic sampling strategy를 설계했다(Fig. 1c)
  

  1. good and consistent teacher가 sub-networks training에 필수적이기 때문에 [1], 초반에는 base network만 학습하면 충분하다.
  2. 둘째로, sub-networks의 학습은 gradual(점진적으로) 이루어져야 한다.
     구체적으로, the width of the smallest sub-network는 점차 줄여나가야 한다.
     
     위의 두 가지 sampling strategy를 결합하여 sampling 횟수를 4에서 3으로 줄이면서도 성능 저하 없이 성공적으로 학습을 수행할 수 있었다.
     (자세한 분석 및 결과는 appendix)

• 우리의 구현에서는 학습 과정을 두 stage로 나눈다.

  1. 첫 번째 stage에서는 the largest network만 학습한다. (즉, $s=1$).
  2. 두 번째 stage에서는 the largest, the smallest + a random width ($s=3$)를 sampling한다.
     그리고 the width of the smallest model은 점차 줄여간다.
     예를 들어, 두 번째 stage의 sampling width는 처음에는 $[0.75, 1.0]$이고, 이후 $[0.5, 1.0]$, 마지막으로 $[0.25, 1.0]$로 줄여간다.

3.2.3 Group Regularization

• 두 channel $k_1$과 $k_2$ ($k_1 < k_2$)가 있을 때, $k_2$가 US-Net에서 사용된다면 $k_1$도 반드시 사용되어야 한다.
  즉, 앞의 channels이 뒤의 channels들보다 더 자주 사용된다.
  따라서 US-Net의 학습에서는 sub-networks의 성능을 보장하기 위해 대다수의 weight가 앞의 channel에 집중된다.
  그러나 이러한 weight distribution은 base model의 capacity를 제한하여 성능에 영향을 미칠 수 있다.
  이 문제를 해결하기 위해, 우리는 뒤쪽 channels에 더 많은 freedom(즉, smaller regularization coefficients)를 부여하여 weight가 더 충분히 활용되도록 하는 group regularization을 제안한다. (Fig. 1d)
  전체 $K$개 channels을 $G$개의 group으로 나누고, 각 group은 $K_G = \lfloor K/G \rfloor$ 개의 channels을 포함한다.
  그 다음, 우리는 다음과 같이 정의한다 :
  여기서 $w$는 weight matrix를 나타내며, 이 논문에서는 $G=8$과 $\alpha=0.05$로 설정했다.
  $\alpha=0$일 때, group regularization은 standard $L_2$ regularization으로 변환된다.
  (또한 우리는 appendix에서 group regularization이 US-Net에 특화되어 있다는 것을 실험적으로 입증했다)

3.3. Three Guidelines for Loss Design

• US-Net training의 특별한 점은 sub-network training (즉, $L_{Distill}$)의 도입이며, 이에 따라 sub-networks의 training stability가 매우 중요하다.
  본 논문에서는 sub-networks에 대한 guidance가 temporal consistency를 가지는 것이 핵심이라고 발견했다.
  하나의 image에서 인접한 두 iteration에서 서로 다른 view를 가지며, 이는 model이 수렴되지 않아 불안정하기 때문에 다른 output을 생성하게 된다.
  우리는 same image의 다른 view에서 생성된 gradient가 iteration 간에도 서로 가깝기를 희망하며, 즉 변화에 강하고 sub-networks에 consistent guidance를 제공할 수 있기를 바란다.

• SSL의 맥락에서 (1)과 (2)를 distillation에 맞게 조정할 수 있으며,
  여기서 $z_s$와 $z_t$는 각각 sub-network와 base network의 output을 나타낸다.
  losses는 다음과 같다 :