字節豆包大模型團隊突破殘差連接侷限!預訓練收斂最快加速80%

機器之心發佈

機器之心編輯部

自從 ResNet 提出後,殘差連接已成爲深度學習模型的基礎組成部分。其主要作用是 —— 緩解梯度消失問題,使得網絡的訓練更加穩定。

但是,現有殘差連接變體在梯度消失和表示崩潰之間存在一種 “蹺蹺板式” 的權衡,無法同時解決。

爲此,字節豆包大模型 Foundation 團隊於近日提出超連接(Hyper-Connections),針對上述 “蹺蹺板式” 困境,實現了顯著提升。

該方法適用於大規模語言模型(LLMs)的預訓練,在面向 Dense 模型和 MoE 模型的實驗中,展示了顯著性能提升效果,使預訓練收斂速度最高可加速 80%。

研究團隊還發現,超連接在兩個小型的視覺任務中表現同樣優異,這表明,該方法在多個領域有廣泛的應用前景。

1. 超連接的核心思想

前文提及,殘差連接的兩種主要變體Pre-Norm和Post-Norm各自都有其侷限性,具體體現如下:

超連接的核心思路在於 —— 引入可學習的深度連接(Depth-connections)和寬度連接(Width-connections)。

從理論上,這使得模型不僅能夠動態調整不同層之間的連接強度,甚至能重新排列網絡層次結構,彌補了殘差連接在梯度消失和表示崩潰(Representation Collapse)之間的權衡困境。

深度連接與寬度連接

起初,該方法會將網絡輸入擴展爲 n 個隱向量(n 稱作 Expansion rate)。之後每一層的輸入都會是 n 個隱向量,超連接會對這些隱向量建立以下兩類連接:

靜態與動態超連接

超連接可以是靜態的,也可以是動態的。

其中,靜態超連接(Static Hyper-Connections, SHC)意味着連接權重在訓練結束後固定不變。而動態超連接(Dynamic Hyper-Connections, DHC)則對應連接權重可根據輸入動態調整。實驗表明,動態超連接效果更好。

3. 爲什麼使用超連接(Hyper-Connections)

研究團隊認爲,殘差連接的兩種變體,即前歸一化(Pre-Norm)和後歸一化(Post-Norm),可以被視爲不可訓練的超連接。

隨後,團隊引入了順序 - 並行二象性概念,展示了超連接如何動態優化層的排列以提升網絡性能。

順序 - 並行二象性

給定一系列神經網絡模塊,我們可以將它們順序排列或並行排列。作者認爲,超連接可以學習如何將這些層重新排列,形成順序和並行配置的混合。

在不失一般性的情況下,可以將擴展率設置爲 n=2。如果超連接以如下矩陣形式學習,神經網絡將被順序排列:

在這種情況下,深度連接退化爲殘差連接,如圖 (a) 所示。

當奇數層和偶數層的超連接矩陣分別定義爲以下形式時,神經網絡每兩層將被並行排列,類似於 Transformer 中的 parallel transformer block 的排列方式,如圖 (b) 所示。

因此,通過學習不同形式的超連接矩陣,網絡層的排列可以超越傳統的順序和並行配置,形成軟混合甚至動態排列。對於靜態超連接,網絡中的層排列在訓練後保持固定;而對於動態超連接,排列可以根據每個輸入動態調整。

4. 實驗結果

實驗主要集中在大規模語言模型的預訓練上,涵蓋了 Dense 模型和 MoE 模型。

實驗結果表明,使用超連接的模型顯著優於使用殘差連接的模型。

1B Dense 模型實驗

只要擴展率 > 1,效果就十分顯著,且訓練更穩定,消掉了訓練 loss 的 spikes。

7B Dense 模型實驗

團隊甚至 Scale 到了 7B 模型,效果也十分亮眼,同時可以看到有超連接的網絡訓練更穩定。

7B 候選激活 1.3B 的 MoE 模型實驗

可以看到,下游指標全漲,在 ARC-Challenge 上甚至漲了 6 個百分點。

綜上,研究團隊介紹了超連接(Hyper-Connections),它解決了殘差連接在梯度消失和表示崩潰之間的權衡問題。實驗結果表明,超連接在大規模語言模型的預訓練以及視覺任務中都表現出顯著的性能提升。

值得注意的是,超連接的引入幾乎不增加額外的計算開銷或參數量,團隊認爲,該成果具有廣泛的應用潛力,可以推廣到文音視圖模態的不同任務上,包括多模態理解、生成基座模型等。

5. 寫在最後

團隊關注底層問題,尤其在 LLMs 和多模態方面,期望實現更多突破。

更多團隊技術研究進展,可以通過「豆包大模型團隊」技術解讀欄目瞭解。