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

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

• 網(wǎng)站：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

背景

設(shè)備上的訓(xùn)練（On-device Training）允許預(yù)訓(xùn)練的模型在部署后適應(yīng)新環(huán)境。通過在移動(dòng)端進(jìn)行本地訓(xùn)練和適應(yīng)，模型可以不斷改進(jìn)其結(jié)果并為用戶定制模型。例如，微調(diào)語言模型讓其能從輸入歷史中學(xué)習(xí)；調(diào)整視覺模型使得智能相機(jī)能夠不斷識(shí)別新的物體。通過讓訓(xùn)練更接近終端而不是云端，我們能有效在提升模型質(zhì)量的同時(shí)保護(hù)用戶隱私，尤其是在處理醫(yī)療數(shù)據(jù)、輸入歷史記錄這類隱私信息時(shí)。

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

方法與結(jié)果

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

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

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

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

結(jié)論

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