用少於256KB記憶體實現邊緣訓練，開銷不到PyTorch千分之一

說到神經網路訓練，大家的第一印像都是 GPU 伺服器 雲端平台。傳統的訓練由於其龐大的記憶體開銷，往往是雲端進行訓練而邊緣平台僅負責推理。然而，這樣的設計使得 AI 模型很難適應新的數據：畢竟現實世界是一個動態的，變化的，發展的場景，一次訓練怎麼能涵蓋所有場景？

為了使得模型能夠持續的適應新數據，我們能否在邊緣進行訓練（on-device training），使設備不斷的自我學習？在這項工作中，我們僅用了不到256KB 記憶體就實現了裝置上的訓練，開銷不到PyTorch 的1/1000，同時在視覺喚醒字詞任務上(VWW) 達到了雲端訓練的準確率。此項技術使得模型能夠適應新感測器資料。用戶在享受客製化的服務的同時而無需將資料上傳到雲端，從而保護隱私。

• 網址：https://tinytraining.mit.edu/
• 論文：https://arxiv.org/abs/2206.15472
• Demo: https://www.bilibili.com/ video/BV1qv4y1d7MV
• #程式碼: https://github.com/mit-han-lab/tiny-training

#背景

裝置上的訓練（On-device Training）允許預先訓練的模型在部署後適應新環境。透過在行動端進行本地訓練和適應，模型可以不斷改進其結果並為使用者自訂模型。例如，微調語言模型讓其能從輸入歷史中學習；調整視覺模型使得智慧相機能夠持續辨識新的物件。透過讓訓練更接近終端而不是雲端，我們能有效在提升模型品質的同時保護使用者隱私，尤其是在處理醫療資料、輸入歷史記錄這類隱私資訊時。

然而，在小型的 IoT 設備進行訓練與雲端訓練有著本質的區別，非常具有挑戰性，首先， AIoT 設備（MCU）的 SRAM 大小通常有限（256KB）。這種等級的記憶體做推理都十分勉強，更不用說訓練了。再者，現有的低成本高效轉移學習演算法，例如只訓練最後一層分類器（last FC），只進行學習bias 項，往往準確率都不盡如人意，無法用於實踐，更不用說現有的深度學習架構無法將這些演算法的理論數字轉換為實測的節省。最後，現代深度訓練框架（PyTorch，TensorFlow）通常是為雲端伺服器設計的，即使把 batch-size 設定為 1，訓練小模型 (MobileNetV2-w0.35) 也需要大量的記憶體佔用。因此，我們需要協同設計演算法和系統，以實現智慧終端設備上的訓練。

方法與結果

我們發現裝置上訓練有兩個獨特的挑戰：（1）模型在邊緣裝置上是量化的。一個真正的量化圖（如下圖所示）由於低精度的張量和缺乏批量歸一化層而難以優化；（2）小型硬體的有限硬體資源（內存和計算）不允許完全反向傳播，其記憶體用量很容易超過微控制器的SRAM 的限制（一個數量級以上），但如果只更新最後一層，最後的精確度又難免差強人意。

為了回應最佳化的困難，我們提出了Quantization-Aware Scaling (QAS) 來自動縮放不同位元精度的張量的梯度（如下左圖所示）。 QAS 在不需要額外超參數的同時，可以自動匹配梯度和參數 scale 並穩定訓練。在 8 個資料集上，QAS 均可以達到與浮點訓練一致的效能（如下右圖）。

為了減少反向傳播所需的記憶體佔用，我們提出了 Sparse Update，以跳過不太重要的層和子張的梯度計算。我們開發了一種基於貢獻分析的自動方法來尋找最佳更新方案。對比以往的 bias-only, last-k layers update, 我們搜尋到的 sparse update 方案擁有 4.5 倍到 7.5 倍的記憶體節省，在 8 個下游資料集上的平均精確度甚至更高。

為了將演算法中的理論減少轉換為實際數值，我們設計了Tiny Training Engine（TTE）：它將自動微分的工作轉到編譯時，並使用codegen 來減少執行時間開銷。它還支援 graph pruning 和 reordering，以實現真正的節省與加速。與 Full Update 相比，Sparse Update 有效地減少了 7-9 倍的峰值內存，並且可以透過 reorder 進一步提升至 20-21 倍的總內存節省。相較於 TF-Lite，TTE 裡經過最佳化的核心和 sparse update 使整體訓練速度提高了 23-25 倍。

結論

#本文中，我們提出了第一個在單晶片上實現訓練的解決方案（僅用256KB 記憶體和1MB 快閃記憶體）。我們的演算法系統協同設計（System-Algorithm Co-design）大大減少了訓練所需記憶體（1000 倍vs PyTorch）和訓練耗時（20 倍vs TF-Lite），並在下游任務上達到較高的準確率。 Tiny Training 可以賦能許多有趣的應用，例如手機可以根據用戶的郵件/ 輸入歷史來定制語言模型，智能相機可以不斷地識別新的面孔/ 物體，一些無法聯網的AI 場景也能持續學習（例如農業，海洋，工業流水線）。透過我們的工作，小型終端設備不僅可以進行推理，還可以進行訓練。在這過程中個人資料永遠不會上傳到雲端，因此沒有隱私風險，同時 AI 模型也可以不斷自我學習，以適應一個動態變化的世界！

以上是用少於256KB記憶體實現邊緣訓練，開銷不到PyTorch千分之一的詳細內容。更多資訊請關注PHP中文網其他相關文章！