AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
本文論文一作是南開大學統計與資料科學學院研二碩士生楊雨辰，指導老師為南開大學統計與資料科學學院的徐君副教授。徐君老師團隊的研究重點是計算機視覺、生成式 AI 和高效機器學習，並在頂級會議和期刊上發表了多篇論文，谷歌學術引用超過 4700 次。
自從大型 Transformer 模型逐漸成為各個領域的統一架構，微調就成為了將預訓練大模型應用到下游任務的重要手段。然而，由於模型的尺寸日益增大，微調所需要的視訊記憶體也逐漸增加，如何高效地降低微調視訊記憶體就成了一個重要的問題。此前，微調 Transformer 模型時，為了節省視訊記憶體開銷，通常的做法是使用梯度檢查點（gradient checkpointing，也叫作啟用重算），以犧牲訓練速度為代價降低反向傳播（Backpropagation, BP）過程中的啟用視訊記憶體佔用。
最近，由南開大學統計與資料科學學院徐君老師團隊發表在 ICML 2024 上的論文《Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation》提出透過更改反向傳播（BP）過程，在不增加計算量的情況下，顯著減少峰值啟用視訊記憶體佔用。
論文：Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation
論文連結：https://arxiv.org/abs/2406.16282
專案連結：https://github.com/yyyyychen/LowMemoryBP
文章提出了兩種反向傳播改進策略，分別是 Approximate Backpropagation（Approx-BP）和 Memory-Sharing Backpropagation（MS-BP）。Approx-BP 和 MS-BP 分別代表了兩種提升反向傳播中記憶體效率的方案，可以將其統稱為 LowMemoryBP。無論是在理論還是實踐意義上，文章都對更高效的反向傳播訓練提供了開創性的指導。
在理論視訊記憶體分析中，LowMemoryBP 可以大幅降低來自啟用函式和標準化層的啟用視訊記憶體佔用，以 ViT 和 LLaMA 為例，可以對 ViT 微調降低 39.47% 的啟用視訊記憶體，可以對 LLaMA 微調降低 29.19% 的啟用視訊記憶體。
在實際實驗中，LowMemoryBP 可以有效地使包括 ViT, LLaMA, RoBERTa, BERT, Swin 在內的 Transformer 模型微調峰值視訊記憶體佔用降低 20%~30%，並且不會帶來訓練吞吐量和測試精度的損失。
Approx-BP
在傳統反向傳播訓練中，啟用函式梯度的反向回傳是嚴格對應其導函式的，對於 Transformer 模型中常用的 GELU 和 SiLU 函式，這意味著需要將輸入特徵張量完整地存入啟用視訊記憶體中。而本文的作者提出了一套反向傳播近似理論，即 Approx-BP 理論。在該理論的指導下，作者使用分段線性函式逼近啟用函式，並用分段線性函式的導數（階梯函式）替代 GELU/SiLU 梯度的反向回傳。這個方法匯出了兩個非對稱的記憶體高效啟用函式：ReGELU2 和 ReSiLU2。這類啟用函式由於使用 4 段階梯函式進行反向回傳，從而使得啟用儲存只需要使用 2bit 資料型別。
MS-BP
BP 網路每一層通常都會將輸入張量存入啟用視訊記憶體以用作反向傳播計算。作者指出如果可以將某一層的反向傳播改寫成依賴輸出的形式，那麼這一層和後一層就可以共享同一個啟用張量，從而降低啟用儲存的冗餘。
而文章指出 Transformer 模型中常用的 LayerNorm 和 RMSNorm，在將仿射引數合併到後一層的線性層之後，可以很好地符合 MS-BP 策略的要求。經過重新設計的 MS-LayerNorm 和 MS-RMSNorm 不再產生獨立的啟用視訊記憶體。
實驗結果
作者對計算機視覺和自然語言處理領域的若干個代表模型進行了微調實驗。其中，在 ViT，LLaMA 和 RoBERTa 的微調實驗中，文章提出的方法分別將峰值視訊記憶體佔用降低了 27%，29% 和 21%，並且沒有帶來訓練效果和訓練速度的損失。注意到，作為對比的 Mesa（一個 8-bit Activation Compressed Training 方法）使訓練速度降低了約 20%，而文章提出的 LowMemoryBP 方法則完全保持了訓練速度。
結論及意義
文章提出的兩種 BP 改進策略，Approx-BP 和 MS-BP，均在保持訓練效果和訓練速度的同時，實現了啟用視訊記憶體的顯著節省。這意味著從 BP 原理上進行最佳化是非常有前景的視訊記憶體節省方案。此外，文章提出的 Approx-BP 理論突破了傳統神經網路的最佳化框架，為使用非配對導數提供了理論可行性。其匯出的 ReGELU2 和 ReSiLU2 展現了這一做法的重要實踐價值。
歡迎大家閱讀論文或者程式碼去了解演算法的詳細細節，LowMemoryBP 專案的 github 倉庫上已經開源相關的模組。