研究者表示，他們將邊緣訓練看作一個優化問題，從而發現了在給定內存預算下實現最小能耗的最優調度。??

目前，智能手機和嵌入式平臺等邊緣設備上已經廣泛部署深度學習模型來進行推理。其中，訓練仍然主要是在具有 GPU 等高通量加速器的大型云服務器上完成。集中式云訓練模型需要將照片和按鍵等敏感數據從邊緣設備傳輸到云端，從而犧牲了用戶隱私并導致了額外的數據移動成本。

圖注：推特 @Shishir Patil?

因此，為了使用戶在不犧牲隱私的情況下個性化他們的模型，聯邦學習等基于設備的訓練方法不需要將數據整合到云端，也能執行本地訓練更新。這些方法已被部署在谷歌 Gboard 鍵盤上以個性化鍵盤建議，也被 iPhones 手機用來提升自動語音識別。同時，當前基于設備的訓練方法不支持訓練現代架構和大模型。在邊緣設備上訓練更大的模型不可行，主要是有限的設備內存無法存儲反向傳播激活。ResNet-50 的單次訓練迭代所需的內存是推理的 200 多倍。

以往工作提出的策略包括分頁到輔助內存和重新實現，以減少云端訓練的內存占用。但是，這些方法會顯著增加整體能耗。與分頁方法相關的數據傳輸通常需要比重計算數據更多的能量。隨著內存預算的縮減，重新實現會以 O(n^2 ) 的速度增加能耗。

在 UC 伯克利最近的一篇論文中，幾位研究者表明分頁和重新實現是高度互補的。通過對簡單操作重新實現，同時將復雜操作的結果分頁到閃存或 SD 卡等輔助存儲器上，他們能夠以最小的能耗擴展有效的內存容量。并且，通過這兩種方法的結合，研究者還證明了在移動級邊緣設備上訓練 BERT 等模型是可能的。通過將邊緣訓練看作一個優化問題，他們發現了在給定內存預算下實現最小能耗的最優調度。

• 論文地址：https://arxiv.org/pdf/2207.07697.pdf
• 項目主頁：https://poet.cs.berkeley.edu/
• GitHub 地址：https://github.com/shishirpatil/poet

研究者提出了 POET（Private Optimal Energy Training），這是一種在內存受限邊緣設備上對現代神經網絡進行能量最優訓練的算法，其架構如下圖 1 所示。鑒于為反向傳播緩存所有激活張量的成本極高，POET 對激活進行優化分頁和重新實現，因而可以將內存消耗最高減少兩倍。他們將邊緣訓練問題重新表述為整數線性程規劃（ILP），發現可以通過求解器在 10 分鐘內將其求解到最優。

圖注：POET 在邊緣設備上對 SOTA 機器學習模型的訓練進行優化。

對于部署在真實世界邊緣設備上的模型，當邊緣設備出現空閑并可以計算周期時就會進行訓練，例如谷歌 Gboard 會在手機充電時安排模型更新。因此，POET 也包含了嚴格的訓練限制。給定內存限制和訓練 epoch 的數量，POET 生成的解決方案也能滿足給定的訓練截止期限。此外，研究者還利用 POET 開發了一個全面的成本模型，并證明它在數學上是保值的（即不做近似），適用于現有的開箱即用架構。

論文一作 Shishir Patil 在演示視頻中表示，POET 算法可以在智能手機等商用邊緣設備上訓練任何需要極大內存的 SOTA 模型。他們也成為了首個展示在智能手機和 ARM Cortex-M 設備上訓練 BERT 和 ResNet 等 SOTA 機器學習模型的研究團隊。

集成分頁和重新實現

重新實現和分頁是降低大型 SOTA ML 模型內存消耗的兩種技術。在重新實現中，一旦不再需要激活張量就會被刪除，最常見的是在前向傳播期間。從而釋放了寶貴的內存，可用于存儲后續層的激活。當再次需要刪除的張量時，該方法會根據譜系的規定從其他相關的激活中重新計算。而分頁，也稱為 offloading，是一種減少內存的補充技術。在分頁中，不是立即需要的激活張量從主存儲器調出到二級存儲器，例如閃存或 SD 卡。當再次需要張量時，將其分頁。

圖 2 顯示了一個八層神經網絡的執行時間表。沿著 X 軸，每個單元對應神經網絡的每一層（共 8 層 L8）。Y 軸表示一個 epoch 內的邏輯時間步長。圖中占用的單元（用顏色填充）表示在相應的時間步執行的操作（前向 / 后向傳播計算、重新實現或分頁）。

例如，我們可以看到 L1 的激活是在第一個時間步 (T1) 計算的。在 T2 和 T3 時刻，分別計算 L2 和 L3 的激活量。假設層 L2 和 L3 恰好是內存密集型但計算成本較低的運算，例如非線性 (tanH、ReLU 等)，那么重新實現就成為了最佳選擇。我們可以刪除激活({T3, L2}， {T4, L3}) 來釋放內存，當后向傳播過程中需要這些激活時，可以再重新實現它們({T14, L3}， {T16, L2})。

假設 L5 和 L6 層是計算密集型運算，例如卷積、密集矩陣乘法等。對于此類運算，重新實現將導致運行時間和能量的增加，并且這種方式是次優的。對于這些層，最好將激活張量分頁到輔助存儲（{T6，L5}，{T7，L6}），并在需要時分頁到（{T10，L6}，{T11，L5 }）。

分頁的一個主要優點是，根據內存總線的占用情況，可以進行 pipelin 處理，以隱藏延遲。這是因為現代系統具有 DMA（直接內存訪問）特性，它可以在計算引擎并行運行時將激活張量從輔助存儲移動到主內存。例如，在時間步 T7，可以同時將 L6 調出并計算 L7。但是，重新實現是計算密集型的，不能并行化，這導致運行時間增加。例如，我們必須將時間步 T14 用于重新計算 L3，從而延遲其余反向傳播執行。

POET

該研究提出了 POET，這是一種用于深度神經網絡的圖形級編譯器，它重寫了大型模型的訓練 DAG，以適應邊緣設備的內存限制，同時保持高能效。

POET 是硬件感知的，它首先跟蹤前向和后向傳播的執行以及相關的內存分配請求、運行時間以及每次操作的內存和能源消耗。對于給定的硬件，每個工作負載的這種細粒度分析只發生一次，具有自動化、便宜等特性，并且為 POET 提供了最準確的成本模型。

POET 然后生成可以有效求解的混合整數線性規劃 (MILP)。POET 優化器搜索有效的重新實現和分頁調度，以最大限度地減少受內存限制的端到端能源消耗。然后使用得到的調度生成一個新的 DAG，在邊緣設備上執行。

雖然 MILP 是在商用硬件上解決的，但發送到邊緣設備的調度表只有幾百字節，因此內存效率很高。

對于計算成本低但內存密集型的操作，重新實現是最有效的。然而，分頁最適合于計算密集型操作，在這種操作中，重新實現將導致大量的能量開銷。POET 在一個集成搜索空間中共同考慮重新實現和分頁。

本文方法可擴展到復雜、現實的架構中，POET 優化器算法如下。

該研究在優化問題中引入了一個新的目標函數，以最小化計算、page-in 和 page-out 的綜合能耗，分頁和重新實現能耗結合的新目標函數為：

其中Φ_compute、Φ_pagein 和Φ_pageout 分別表示每個節點在計算、page-in 和 page-out 時所消耗的能量。

POET 根據圖的哪些節點 (k) 進行了重新實現，以及在每個時間步長 (t) 將哪些節點 page-in 或 page-out 來輸出 DAG 調度。

 實驗結果?

在對 POET 的評估中，研究者試圖回答三個關鍵問題。首先，POET 在不同的模型和平臺上能夠減少多少能耗？其次，POET 如何從混合分頁和重新實現策略中獲益？最后，POET 如何適應不同的運行時預算？

研究者在下表 2 中列出四種不同的硬件設備，分別為 ARM Cortex M0 MKR1000、ARM Cortex M4F nrf52840、A72 Raspberry Pi 4B + 和 Nvidia Jetson TX2。POET 是完全硬件感知的，依賴于細粒度的分析。

下圖 3 顯示了單次訓練 epoch 的能耗，每列分別對應不同的硬件平臺。研究者發現，POET 在所有平臺上生成節能耗最優的調度（Y 軸），同時減少峰值內存消耗（X 軸）并符合時間預算。

在下圖 5 中，研究者在 A72 上訓練 ResNet-18 時對 POET 和 Capuchin 進行了基準測試。隨著 RAM 預算的減少，Capuchin 比具有完整內存的基線多了 73% 到 141% 的能耗。相比之下，POET 產生的能耗不到 1%。這種趨勢適用于測試的所有架構和平臺。

表 3 中，該研究在 Nvidia 的 Jetson TX2 上訓練 ResNet-18 時對 POET 和 POFO 進行了基準測試。研究發現 POET 找到了一個集成的重新實現和分頁調度，可將峰值內存消耗降低 8.3%，并將吞吐量提高 13%。這展示了 POET 的 MILP 求解器的優勢，它能夠在更大的搜索空間上進行優化。雖然 POFO 僅支持線性模型，但 POET 可以推廣到非線性模型，如圖 3 所示。

圖 4 強調了 POET 在不同時間約束下采用集成策略的好處。對于每個運行時，下圖繪制了總能耗圖。? ?