Abstract

介紹權重選擇，一種通過從預訓練模型的較大模型中選擇權重子集來初始化較小模型的方法。這使得知識從預訓練的權重轉移到更小的模型。
它還可以與知識蒸餾一起使用。

權重選擇提供了一種在資源受限的環境中利用預訓練模型力量的新方法，希望能夠成為大模型時代訓練小模型的有用工具。

OscarXZQ/weight-selection (github.com)

1 Introduction

神經網絡權重的初始化對于其優化至關重要，正確的初始化有助于模型收斂并防止梯度消失等問題。
本文引入一種權重初始化的方法解決該問題，使用大型預訓練模型來訓練小型模型。

引入了權重選擇(weight selection)，選擇預訓練大型模型的權重子集來初始化較小的模型。
使用權重選擇來初始化小模型非常簡單，與從頭開始訓練相比，不會增加額外的計算成本。

3 Weight Selection

給定預訓練模型，目標是為同一模型系列中較小尺寸的模型獲得有效的權重初始化。借用知識蒸餾中的術語，將預訓練模型稱為教師，將我們旨在初始化的模型稱為學生。

3.1 Approach

模型可以通過增加嵌入維度或每個塊中的通道數量來加寬，并通過堆疊更多層來加深。
它還使我們能夠按照三個步驟執行權重選擇：層選擇、組件映射和元素選擇。

圖2：權重選擇(Weight Selection)，為了初始化預訓練模型的較小變體，從預訓練模型的相應組件中統一選擇參數。

3.1.1 Layer Selection

第一步是從教師模型中選層。對于學生的每一層，教師模型中被選中的層都是用于初始化的數據來源。對于各個同性架構和分層架構，層選擇的過程略有不同。

各向同性架構是指神經網絡的每一層在整個模型中都表現出一致且統一的分層設計。

分層結構的特點是多尺度表示和嵌入維度的層次結構。分層結構通常具有不同規模和嵌入維度的階段，如VGG這樣的經典卷積網絡逐漸減少空間維度，同時增加通道維度，捕獲多尺度特征。

對于各向同性架構，我們選擇教師的前$N$層，$N$為學生模型的層數，記作前$N$層選擇(first-$N$ selection)。
處理分層結構時，在每一個獨立的階段單獨應用前$N$層選擇。

另一種方法是統一層選擇(uniform layer selection)，即選擇教師中均勻間隔的層。

3.1.2 Component mapping

第二步，在學生和教師之間映射組件。上一步獲得了從教師到學生的層的映射，將任務簡化為用一個教師層初始化一個學生層。
得益于現代神經網絡設計采用的模塊化方法，同一系列模型中的層具有一組相同的組件，只是寬度不同。
因此，匹配相應組件的過程是自然的一對一映射。

組件(component)：在人工智能領域中，神經網絡的component通常指的是網絡的組成部分，也可以理解為神經網絡的各個層（layers）或模塊（modules）。神經網絡由許多不同類型的層組成，每一層都有自己特定的功能和參數。這些層可以分為輸入層、隱藏層和輸出層，其中隱藏層可以進一步細分為卷積層、池化層、全連接層等。

3.1.3 Element selection

建立組件映射后，下一步是用教師的較大對應組件來初始化學生的較小組件，使用uniform selection，從教師tensor中均勻選擇間隔的元素，具體方法在下一部分介紹。

3.2 Methods for element selection

本節制定元素選擇并介紹不同的選擇標準。

考慮一個學生的權重tensor $W_s$，想找到用于初始化的教師權重tensor $W_t$。如果$W_t$具有shape：$t_1,t_2,?,t_n$，那么具有相同組件類型的$W_s$，在選擇上也要跨過這$n$個維度。
目標是選擇$W_t$的子集來初始化$W_s$。

下面討論幾種可能的元素選擇方法，并在4.3節中對其表現進行比較。只要保持了一致性（如具有一致性的隨機選擇段落中所述），權重選擇就可以達到類似性能的水平。推薦實踐中使用uniform selection作為權重選擇的默認值。

一致性(consistancy)：通常指的是在模型訓練、推斷或應用過程中，確保系統行為和輸出的穩定性和一致性。

1. 數據一致性：確保數據在不同階段（采集、預處理、訓練、測試）的一致性，相同的數據在不同情況下應有相同的表現和輸出。
2. 模型一致性：確保模型在不同環境或輸入下的行為一致，在不同硬件、軟件或數據集上訓練的模型應該產生相似的結果。
3. 結果一致性：確保模型的輸出在不同運行中保持一致，給定相同的輸入或條件，模型應該產生相同的輸出或行為。

3.2.1 Uniform selection (default)

對于$W_t$的第$i$維，從$t_i$中均勻間隔地選擇$s_i$。
比如，想使用$4\times 6$的線性層$W_t$來初始化$2\times 3$的線性層$W_s$，沿著第一個維度選擇$1,3$數據（均勻從$1,2,3,4$中選擇2個），沿著第二個維度選擇$1,3,5$切片（均勻從$1,2,3,4,5,6$中選擇3個）。

3.2.2 Consecutive selection

對于$W_t$的第$i$維，從$t_i$中選出連續的$s_i$條數據。
對于具有分組組件的架構，連續選擇會選擇一些整個組，同時忽略對比。
對于沒有這種分組組件的架構，連續選擇相當于均勻選擇。

3.2.3 Random selection (with consistency)

對于所有權重tensor，對于$W_t$的第$i$維，選擇$s_i$大小的$t_i$的隨機子集。
通過驗證實驗得到，一致性（為所有權重矩陣選擇相同的指數）是權重選擇達到最佳性能的關鍵。
保持一致性的動機源自于殘余連接的存在——添加到教師模型中的神經元應該將其操作保留在學生當中。

此外，保持一致性可以在元素選擇過程中保留完整的神經元，因為只選擇了一致的位置。
均勻選擇和連續選擇本質上保持了一致性，它們都是具有一致性的隨機選擇的特殊實例。

3.2.4 Random selection (without consistency)

對于所有權重tensor，對于$W_t$的第$i$維，選擇$s_i$大小的$t_i$的隨機子集。

該方法不需要為每一個權重tensor選擇相同的索引。設計這個方法用來測試保持一致性的重要程度。

4 Experiments

表3：與經典初始化方法的比較，使用Cifar-100數據集測試多種初始化方法的準確度，具有一致性的權重選擇方法優于經典的初始化方法。

5 Analysis

5.1 Reduction in training time

weight selection可以顯著減少訓練時間。直接測量了使用不同數量的epochs，使用weight selection訓練ViT-T省去的時間，將結果展示在圖4(a)。有weight selection后，只需要1/3的epoch數量即可打到傳統初始化方法的表現。

圖4：更快地訓練。和隨機初始化比較，在Cifar-100上，ViT-T只需要1/3的epoch就可以達到其相同的表現。和在ImageNet-1K上預訓練+微調對比，weight selection可以在預訓練的第60個epoch達到其表現，省區6.12倍的時間。

5.2 Comparison with transfer learning

通過實驗來找到預訓練所需的訓練預算，以匹配權重選擇的準確性。在本實驗中，我們在ImageNet-1K上以不同的epoch數訓練ViT-T，然后在Cifar-100上微調300個 epoch。如圖 4b 所示，在 ImageNet-1K 上進行 60 輪預訓練才能在Cifar-100上達到相同的性能。在此設置下，與通過預訓練達到相同性能相比，權重選擇速度提高了6.12倍，且無需訪問用于預訓練的數據集。