該研究提出了第一個在單片機上實現(xiàn)訓練的解決方案，并且系統(tǒng)協(xié)同設計（System-Algorithm Co-design）大大減少了訓練所需內存。

說到神經網絡訓練，大家的第一印象都是 GPU + 服務器 + 云平臺。傳統(tǒng)的訓練由于其巨大的內存開銷，往往是云端進行訓練而邊緣平臺僅負責推理。然而，這樣的設計使得 AI 模型很難適應新的數(shù)據(jù)：畢竟現(xiàn)實世界是一個動態(tài)的，變化的，發(fā)展的場景，一次訓練怎么能覆蓋所有場景呢？

為了使得模型能夠不斷的適應新數(shù)據(jù)，我們能否在邊緣進行訓練（on-device training），使設備不斷的自我學習？在這項工作中，我們僅用了不到 256KB 內存就實現(xiàn)了設備上的訓練，開銷不到 PyTorch 的 1/1000，同時在視覺喚醒詞任務上 (VWW) 達到了云端訓練的準確率。該項技術使得模型能夠適應新傳感器數(shù)據(jù)。用戶在享受定制的服務的同時而無需將數(shù)據(jù)上傳到云端，從而保護隱私。

網站：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）允許預訓練的模型在部署后適應新環(huán)境。通過在移動端進行本地訓練和適應，模型可以不斷改進其結果并為用戶定制模型。例如，微調語言模型讓其能從輸入歷史中學習；調整視覺模型使得智能相機能夠不斷識別新的物體。通過讓訓練更接近終端而不是云端，我們能有效在提升模型質量的同時保護用戶隱私，尤其是在處理醫(yī)療數(shù)據(jù)、輸入歷史記錄這類隱私信息時。

然而，在小型的 IoT 設備進行訓練與云訓練有著本質的區(qū)別，非常具有挑戰(zhàn)性，首先， AIoT 設備（MCU）的 SRAM 大小通常有限（256KB）。這種級別的內存做推理都十分勉強，更不用說訓練了。再者，現(xiàn)有的低成本高效轉移學習算法，例如只訓練最后一層分類器 （last FC），只進行學習 bias 項，往往準確率都不盡如人意，無法用于實踐，更不用說現(xiàn)有的深度學習框架無法將這些算法的理論數(shù)字轉化為實測的節(jié)省。最后，現(xiàn)代深度訓練框架（PyTorch，TensorFlow）通常是為云服務器設計的，即便把 batch-size 設置為 1，訓練小模型 (MobileNetV2-w0.35) 也需要大量的內存占用。因此，我們需要協(xié)同設計算法和系統(tǒng)，以實現(xiàn)智能終端設備上的訓練。

方法與結果

我們發(fā)現(xiàn)設備上訓練有兩個獨特的挑戰(zhàn)：（1）模型在邊緣設備上是量化的。一個真正的量化圖（如下圖所示）由于低精度的張量和缺乏批量歸一化層而難以優(yōu)化；（2）小型硬件的有限硬件資源（內存和計算）不允許完全反向傳播，其內存用量很容易超過微控制器的 SRAM 的限制（一個數(shù)量級以上），但如果只更新最后一層，最后的精度又難免差強人意。

為了應對優(yōu)化的困難，我們提出了 Quantization-Aware Scaling (QAS) 來自動縮放不同位精度的張量的梯度（如下左圖所示）。QAS 在不需要額外超參數(shù)的同時，可以自動匹配梯度和參數(shù) scale 并穩(wěn)定訓練。在 8 個數(shù)據(jù)集上，QAS 均可以達到與浮點訓練一致的性能（如下右圖）。

為了減少反向傳播所需要的內存占用，我們提出了 Sparse Update，以跳過不太重要的層和子張的梯度計算。我們開發(fā)了一種基于貢獻分析的自動方法來尋找最佳更新方案。對比以往的 bias-only, last-k layers update, 我們搜索到的 sparse update 方案擁有 4.5 倍到 7.5 倍的內存節(jié)省，在 8 個下游數(shù)據(jù)集上的平均精度甚至更高。

為了將算法中的理論減少轉換為實際數(shù)值，我們設計了 Tiny Training Engine（TTE）：它將自動微分的工作轉到編譯時，并使用 codegen 來減少運行時開銷。它還支持 graph pruning 和 reordering，以實現(xiàn)真正的節(jié)省與加速。與 Full Update 相比，Sparse Update 有效地減少了 7-9 倍的峰值內存，并且可以通過 reorder 進一步提升至 20-21 倍的總內存節(jié)省。相比于 TF-Lite，TTE 里經過優(yōu)化的內核和 sparse update 使整體訓練速度提高了 23-25 倍。

結論

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

審核編輯 ：李倩