CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

機器之心發表於2022-04-08

深度學習促進人工智慧(AI)領域不斷髮展,實現了許多技術突破。與此同時,如何在有限硬體資源下挖掘模型潛能、提升部署模型的準確率成為了學界和業界的研究熱點。其中,知識蒸餾作為一種模型壓縮和增強的方法, 將泛化能力更強的「大網路模型」蘊含的知識「蒸餾」到「小網路模型」上,來提高小模型精度,廣泛地應用於 AI 領域的全監督、半監督、自監督、域遷移等各個方向。

近日, OPPO 研究院聯合上海交通大學將視角聚焦到知識蒸餾的正規化本身,提出了新的自蒸餾框架:DLB(Self-Distillation from Last Mini-Batch),模型無需額外的網路架構修改,對標籤噪聲具有魯棒性,並大幅節約訓練的空間複雜度。此外,在三個基準資料的實驗中,模型達到了 SOTA 效能。相關論文「Self-Distillation from the Last Mini-Batch for Consistency Regularization」已被 CVPR 2022 收錄。

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

論文連結:https://arxiv.org/pdf/2203.16172.pdf

DLB 自蒸餾框架

如何減少蒸餾計算複雜度?

知識蒸餾通常可以劃分為三類,即離線蒸餾、線上蒸餾和自蒸餾。其中,自蒸餾具有訓練輕量、知識遷移效率高的特點,最近受到更多研究者的重視。

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

圖 1:本文方法與其他自蒸餾方法的比較

但是傳統的自蒸餾,例如 Be Your Own Teacher,在模型訓練過程中需要對模型結構進行修改。除此以外,訓練成本高、計算冗餘且效率低下也是自蒸餾需要攻克的難題。

為了解決上述難題,讓模型更好地部署到手機等終端裝置中,OPPO 研究院和上海交通大學的研究員們提出了 DLB 自蒸餾框架。利用訓練時前後 Batch 預測結果的一致性,在無需對模型進行網路結構修改的前提下,就能降低訓練複雜度,增強模型泛化能力。

1. 本文的任務

提出更加輕量的自蒸餾方式,降低訓練的計算複雜度,提高模型準確率和泛化性。

2. 本文創新與貢獻

  • 提出 DLB,透過儲存與下個 Batch 部分樣本重疊的軟目標(soft targets)進行自蒸餾。節省計算記憶體,並且簡化訓練流程。

  • 讓訓練樣本的每次前向過程都與一次反向傳播過程相關聯,從而提升學習效率。

  • 實驗分析了 DLB 訓練方法的動態影響,發現其正則化效果來源於即時生效的平滑標籤帶來的訓練一致性,為自蒸餾的理論研究提供了實驗基礎。

DLB 自蒸餾框架訓練機制

DLB 在訓練階段的每個 iteration 中,目標網路扮演著「教師」和「學生」的雙重角色。其中,教師的作用是生成下一個 iteration 的軟目標進行正則化;學生的作用是從前一個 iteration 平滑的標籤中蒸餾,並最小化監督學習目標。

資料集定義為CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路,包含 n 個樣本的 Batch 定義為:CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路,以影像分類為例,圖片先進行資料增強,然後將其輸入神經網路,最佳化預測輸出與真值間的交叉熵損失:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

上式中 p_i 的表示式如下:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

θ為網路引數,K 表示分類類別數,τ表示溫度。

為了提高泛化能力,傳統的 vanilla 知識蒸餾透過額外最佳化的 KL 散度損失來遷移預先訓練好的 teacher 網路的知識,即:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

不同於以往採用預先訓練 teacher 模型的方式生成(P_i^τ ) ̃,DLB 採用訓練中前一個 Batch 蘊含的資訊生成(P_i^τ ) ̃,並將其作為正則化的即時平滑標籤。

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

圖 2:DLB 訓練方式示意圖

如圖 2 所示,資料樣本在第 t 次迭代時的定義為CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路,神經網路的引數為θ_t。

B_t 和 B_(t-1)是透過使用資料取樣器獲得,在前向過程後計算 L_CE。每個小批次的一半限制為與上一個 iteration 一致,其餘一半與下一個 iteration 一致。之後,前半部分小批次利用上一個 iteration 中生成的動態軟目標進行學習。即CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路由 t-1 次迭代的軟標籤(soften labels)CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路生成。因此,引入的正則化損失公式如下:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

儲存平滑標籤只需要很少的額外記憶體成本,所以帶來的額外計算成本很低。整體損失函式由下式表示:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

綜上,DLB 演算法整體的訓練的虛擬碼如下所示:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

實驗設定

研究員們採用三個影像分類基準資料集評估效能,包括 CIFAR-10、CIFAR-100TinyImageNet。實驗結果都達到了最佳效能,如下表所示:

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

具體而言,在平均錯誤率層面, DLB 在 CIAFR-100 改進幅度為 0.83% 至 2.50%,在 CIFAR-10 上為 0.37% 至 1.01%,在 TinyImageNet 上為 0.81% 至 3.17。值得一提的是,DLB 的表現明顯優於 Tf-KD 和 PS-KD,這證明了 DLB 在模型泛化提升上的優勢。

為了評估 DLB 與基於資料增強正則化方法的相容性,研究員在 CIFAR-10 和 CIFAR-100 上將 DLB 與 CutMix、CutOut 和 DDGSD 相結合。如下所示,實驗表明透過結合 DLB 和基於增強的正則化,可以實現額外的效能增益。

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

為了證明魯棒性,研究員在訓練之前向 CIFAR-100、CIFAR-10 隨機注入標籤噪聲,實驗結果如下圖所示,DLB 可以有效地抵抗標籤噪聲並提高整體效能。 

CVPR 2022 | 這個自蒸餾新框架新SOTA,降低了訓練成本,無需修改網路

結語

本文提出了一種基於自蒸餾思想的深度學習訓練策略,將自蒸餾思想融入到模型訓練過程中,對傳統知識蒸餾進行改進,無需額外預先訓練 teacher 的過程。透過在三個基準資料集上進行實驗,多維度論證了 DLB 訓練策略的有效性與普適性。

當下,深度學習網路模型結構複雜度不斷提升,使用有限硬體資源開發和部署 AI 模型成為新的科研問題。在本文中,研究員設計的 DLB 訓練策略,在某種程度上解決了「不增加模型複雜度,提高模型準確率」這一業界難題。

未來,OPPO 將持續研發該演算法,並將其賦能到相關業務場景中,不斷提升 OPPO 產品的使用者體驗。

招聘資訊:OPPO 研究院 智慧感知與互動研究部 計算機視覺演算法崗位(全職和實習生)熱招中!簡歷投遞郵箱:liyaqian@oppo.com

相關文章