# 【論文筆記】Stitchable Neural Networks

論文連結：https://arxiv.org/abs/2302.06586
發表於 CVPR 2023

## Overview

本篇研究提出 Stitchable Neural Networks (SN-Net)，是一個可擴展的深度學習架構。具體而言，選擇屬於同一個 model family 但不同大小的 pretrained models，例如：DeiT-Ti/S/B，彼此之間用一些簡單的 stitching layers 連接在一起成為一個網路，經過少量的 fine-tuning steps 之後，這個網路即為 SN-Net。只要在 runtime 選擇不同的 stitching layers，就可以像插值一般產生不同 complexity 和 performance 的子網路。這些子網路的規模大小都不同，因此可以應對各種資源限制。


這樣的做法優勢在於，可以有效利用現存的 pretrained model，而不需要從頭訓練特定大小的模型。和過往的 scalable deep learning frameworks 的差異在於 SN-Net 縫合多個模型，並且也可以產生多個不同 scale 的子網路，是一個 many-to-many 的架構。


## Methods

Model stitching 的概念在過往的研究就有被提出來，研究發現同樣架構但使用不同 random seeds 進行訓練的模型可以被縫合在一起，也有研究指出不同架構的模型也可以縫合。根據 model stitching 的概念，作者詳細規劃了 SN-Net 的構建方法。

### What to stitch

選擇要縫合在一起的模型需要考慮的問題包含：
- 不同模型架構學到的 representation 通常會有比較大的差別，縫合得到的結果可能不會很好
- 不同 dataset 學到的 representations 會有差別，因此選擇的模型應該要在同一個 pretrained domain 上

因此實驗主要選擇 pretrained 在同一份 dataset 上，架構相似但 scale 不同的模型，例如：DeiT-Ti/S/B。

### How to stitch

因為縫合的目的不是要提升表現，而是要將某個 network 中特定 layer 的 feature map 進行轉換，因此這個轉換應該要越簡單越好。作者決定讓 stitching layer 就只是單純的 1x1 convolution，透過這個 convolution layer 把 feature map 從原先的 feature space 轉換到另一個 network 的 feature space。


### Where to stitch

給定兩個 models，總共有兩種方式可以縫合，分別是 Fast-to-slow 以及 Slow-to-fast。Fast-to-slow 指的是讓 stitching layer 從比較小且比較快的模型傳遞到比較大且比較慢的模型，Slow-to-fast 則是相反。實驗發現 Fast-to-slow 表現比較好。

另外在縫合的模型上，作者提出了 nearest stitching 這個限制，限制縫合的模型只能是相鄰的兩個複雜度的，例如 Ti/S 或 S/B，這是因為他們發現當縫合的模型跨度太大的時候，會導致表現不好。


### Way to stitch
因為相鄰的兩個 layers 通常會共享相似的 representation，他們提出使用 sliding window 的方式框住縫合的 blocks。左邊的 paired stitching 適用在兩個模型深度相同的情況，同一個 windows 下的 blocks 共享同一個 stitching layers，可以控制的參數有 window size $k$ 和 stride $s$。右邊的 unpaired stitching 適用在兩個模型深度不一樣的時候，有時候小的模型比較淺，大的模型比較深，於是讓小的模型的每一個 block 都縫合到大模型中的數個 blocks。

> 這裡 paper 沒有解釋如果共享同一個 stitching layer 的 blocks 有不同 hidden dimension 該如何處理，看 code 也看不出來，查了一下發現他們主要用來實驗的 DeiT 模型，每個 layers 的 hidden dimension 都一樣（例如 DeiT-Ti 都是 192），可能是因為這樣才沒有特別解釋？

### Training
給定數個 pretrained models，利用上面的方法可以縫合成一個 joint network，目標就是希望可以訓練這個 network，讓它可以符合多種的資源限制。訓練的演算法非常簡單，給定 pretrained models 並決定好縫合的位置後，在訓練的每一次迭代都隨機抽取一個 stitch，並計算 loss 更新參數。

## Experiments

下面的實驗分別縫合 DeiT-based models 以及 Swin-based models，並且在 ImageNet 上訓練 50 個 epochs，每一個點都代表選擇某一個 stitching layer 得到的子網路結果。可以看到 Stitchable Neural Network 能夠符合多種的資源和準確度需求，隨著縫合更多比較大的模型的 blocks，表現也會越好。

不同的子網路是可以在 runtime 的時候隨時切換的，也就代表說可以依照當下的需求隨時調整模型的 scale，這個是很多模型無法做到的。


為了要比較 SN-Net 和 train from scratch 的差異，這裡用 DeiT-based model 作為對象，選擇了幾種不同的網路結構，比較特定網路結構從頭訓練在 ImageNet 上 300 epochs 和 SN-Net 當中選取出來的子網路兩者的表現差異。

可以看到如果要單獨訓練並得到這麼多不同種類的模型，需要很大量的計算資源，但 SN-Net 只需要用到現有的 pretrained model 並訓練 50 epochs 就可以有差不多的表現，並且整個 SN-Net 只有 118.4M 的參數。


## Ablation Studies

Ablation studies 也是以 DeiT-based SN-Net 為基礎進行實驗。下圖比較不同種 stitching 的方向所得到的結果，發現 Fast-to-slow 的縫合方式可以有比較好的表現，作者解釋現今的模型設計原則是隨著模型深度加深，逐步加寬 hidden dimension，但如果以 Slow-to-fast 方向縫合模型，等同於在深度變深的時候縮減了 hidden dimension，所以才造成表現不佳。


下圖比較 nearest stitching 的影響。如同前面所提到的，nearest stitching 限制縫合的模型只能是相鄰的兩個複雜度的，如果縫合兩個規模差距太大的模型，例如 Ti 和 B，則表現會比較不好。


下圖比較微調整個 SN-Net 和只微調 stitching layers 的表現。會嘗試只微調 stitching layers 的原因是因為我們用來縫合的模型其實都是已經完整 pretrained 好的，因此作者好奇能不能在訓練的時候只更新 stitching layers 的參數。結果顯示微調整個模型可以有更好的效果。