Bi-MTDP：通過二值網(wǎng)絡(luò)加速多任務(wù)密集預(yù)測，又快又提點 | CVPR 2024

論文提出二值化多任務(wù)密集預(yù)測器 Bi-MTDP，通過二值神經(jīng)網(wǎng)絡(luò)（BNNs）顯著加速多任務(wù)密集預(yù)測模型，同時保持甚至提高模型性能。為了避免信息嚴重退化而導致二值化帶來性能下降，論文引入了深度信息瓶頸層，在前向傳播時強制要求下游任務(wù)表示滿足高斯分布；此外，還引入知識蒸餾機制來糾正反向傳播中信息流方向

來源：曉飛的算法工程筆記 公眾號

論文: Efficient Multitask Dense Predictor via Binarization

Introduction

? 在計算機視覺領(lǐng)域，有一個日益增長的趨勢，即以多任務(wù)學習的方式處理密集預(yù)測任務(wù)，例如語義分割、單眼深度估計和人體解析。由于多任務(wù)模型中通過交叉連接結(jié)構(gòu)進行信息補充機制的收益，一系列密集預(yù)測任務(wù)的整體性能得到了極大提升（見圖1）。然而，處理多個復雜像素級任務(wù)的當前最先進（SoTA）多任務(wù)密集預(yù)測模型具有相當大的計算需求，這種高計算要求限制了它們在資源受限環(huán)境中（如自動駕駛、機器人技術(shù)和虛擬現(xiàn)實等領(lǐng)域）的應(yīng)用。論文旨在優(yōu)化這些龐大的SoTA模型以適用于邊緣設(shè)備，在速度和性能之間取得平衡。

? 已經(jīng)不是研究探索神經(jīng)網(wǎng)絡(luò)壓縮策略，包括剪枝、網(wǎng)絡(luò)量化和知識蒸餾。值得注意的是，網(wǎng)絡(luò)二值化作為一種量化形式，將權(quán)重和激活最小化為 \pm 1 ，從而能夠用二值神經(jīng)網(wǎng)絡(luò)（BNNs）中更高效的XNOR-Bitcount操作替代全精度網(wǎng)絡(luò)中計算密集型的內(nèi)積運算。理論上，二值化可以將存儲成本降低32倍，并將推斷速度提高64倍，使得BNNs適用于邊緣設(shè)備。

? 雖然二值神經(jīng)網(wǎng)絡(luò)（BNNs）在圖像分類方面取得了令人印象深刻的成果，幾乎達到了全精度ResNet級別的準確性，但其應(yīng)用主要局限于小規(guī)模模型，無法用于其他計算密集型的計算機視覺任務(wù)。將BNNs擴展到更大的模型應(yīng)該是下一步的研究，但這種擴展受到過擬合和信息降級等問題的阻礙。在全精度模型中有效的技術(shù)，如標簽平滑、dropout和mixup，在BNNs上效果較差。此外，最先進的多任務(wù)密集預(yù)測任務(wù)通常需要深層復雜模型并配備多模態(tài)融合結(jié)構(gòu)，加劇了有效實施二值化所面臨的挑戰(zhàn)。

? 在應(yīng)用二值化進行多任務(wù)密集預(yù)測任務(wù)時，主要障礙是在深度模型中信息流的顯著降級，導致性能下降。為了解決這個問題，論文提出了一個Binary Multitask Dense Predictor（Bi-MTDP）基線，通過二值化模塊構(gòu)建了一個多任務(wù)密集預(yù)測器。基于對這個基線的徹底評估，論文得出結(jié)論：二值化操作破壞了多任務(wù)模型中的信息流，因此與其全精度對應(yīng)物相比，下游任務(wù)的表示不夠豐富。為解決這一問題，論文通過兩個方向上額外的信息流校準機制更新Bi-MTDP

  1. 在正向傳播中實施變分信息瓶頸機制，使嵌入遵循具有稀疏性質(zhì)的高斯分布以過濾掉與任務(wù)無關(guān)的因素。
  2. 在反向傳播中利用現(xiàn)有全精度模型通過基于特征的知識蒸餾來校準二值網(wǎng)絡(luò)梯度。

? Bi-MTDP的好處可以從兩個正交的角度進行分析：

  1. 從網(wǎng)絡(luò)二值化的角度來看，通過將二值化與多任務(wù)密集預(yù)測框架相結(jié)合，Bi-MTDP的成功證明了它可以有效地補充信息，從而提高了各個二值模型的性能。
  2. 從多任務(wù)密集預(yù)測任務(wù)的角度來看，加速這些繁瑣的模型有利于在其中設(shè)計更有效和高效的交叉模塊，如圖1所示。由于現(xiàn)有的密集預(yù)測模型在建模交叉模塊方面存在嚴重限制，因為它們過度依賴卷積操作，所以對于多任務(wù)密集預(yù)測來說，通過交叉機制學習涵蓋多任務(wù)上下文各個范圍的交互和推理是至關(guān)重要的。

Multitask Network Binarization

Binary Neural Network

? 首先，簡要回顧二值神經(jīng)網(wǎng)絡(luò)（BNNs）的一般思想，詳細闡述二值化的加速機制和信息流退化問題。定義一個具有 K 個層的全精度（FP）神經(jīng)網(wǎng)絡(luò)， f(\mathbf{x}) = (\mathbf{W}^{K}\times\sigma\times \mathbf{W}^{K-1} \cdots \sigma\times \mathbf{W}^{1})(\mathbf{x}) ，其中 \mathbf{x} 是輸入樣本， \mathbf{W}^{k}:\mathbb{R}^{d_{k-1}} \longmapsto \mathbb{R}^{d_{k}}(k=1,...,K) 代表連接第 (k-1) 層和第 k 層的權(quán)重矩陣， d_{k-1} 和 d_{k} 分別表示第 k 層網(wǎng)絡(luò)輸入和輸出的大小，函數(shù) \sigma(\cdot) 對特征圖執(zhí)行逐元素激活。

? BNNs在前向操作和反向梯度逼近方面與FP神經(jīng)網(wǎng)絡(luò)有所不同。具體而言，在前向傳播過程中，BNN保持FP潛在權(quán)重 \mathbf{W}_F 用于梯度更新，第 k 個權(quán)重矩陣 \mathbf{W}_F^k 被二值化為 \pm 1 ，通過二值化函數(shù) \textit{sign}(\cdot) 得到二值權(quán)重矩陣 \mathbf{W}_B^k ，即 \mathbf{W}_B^k = \textit{sign}(\mathbf{W}_F^k) 。然后第 k 層的中間激活圖（全精度）由 \mathbf{A}_F^{k} = \mathbf{W}_B^k \mathbf{A}_B^{k-1} 產(chǎn)生。相同的量化方法用于將全精度激活圖二值化為 \mathbf {A}_{ B }^{ k }=\operatorname { sign }\left( \mathbf { A } _ { F } ^ { k }\right) ，并且通過對這個過程進行 L 次迭代執(zhí)行整個二值化的前向傳遞，如圖2所示。對于BNNs,權(quán)重和激活都是1-bit，因此網(wǎng)絡(luò)在內(nèi)存成本方面加速了32倍。重要的是，在推斷階段，BNN加速了64倍, 因為FP網(wǎng)絡(luò)中的FP乘法被BNN中的Xnor-Bitcount所取代。

? 在反向傳播中，主要挑戰(zhàn)是普遍存在的sign函數(shù)在理論上是不可微的，因此通過梯度傳播極大地破壞了信息流。為了解決這個問題，利用直通估計器（STE）來數(shù)值近似整個BNN的導數(shù)，即

\begin{equation} \text{Backward:~~~} \frac{\partial\mathcal{L}}{\partial x}= \left\{ \begin{array}{ll} \frac{\partial\mathcal{L}}{\partial \textit{sgn}(x)} & \quad \vert x \vert \leq 1 \\ 0 & \quad \vert x \vert > 1. \end{array} \right. \label{eq:critic_new} \end{equation}

? 值得注意的是，在實踐中沒有采用上述的傳統(tǒng)逼近方法，而是利用普遍存在的Bi-Real和IR-Net逐漸逼近符號函數(shù)，這已被證明是更好的估計方法。盡管已經(jīng)有許多方法來消除二值化引起的信息流惡化，但由于權(quán)重、激活和梯度的嚴重精度損失，惡化仍然是不可避免的。因此，二值化破壞了復雜計算機視覺模型的性能。

Multitask Dense Predictor

? 在為密集預(yù)測任務(wù)應(yīng)用二值化技術(shù)后，二值化模型的性能是不可接受的，如圖1和表3中的二進制單任務(wù)結(jié)果所示。因為這些SoTA密集預(yù)測模型的架構(gòu)相對較重且較深（例如，HRNet-48或ResNet-101具有特定任務(wù)頭），而不是用于分類的模型（例如，ResNet-18具有全連接層作為分類頭）。此外，在二進制模型中通過反向傳播傳遞信息，特別是在深度模型中，效率極低。

? 密集預(yù)測任務(wù)可以相互補充信息，例如，法線和深度可以直接從彼此派生出來，這可以被建模為彼此的正則化。密集預(yù)測任務(wù)之間的相關(guān)性值得利用以提高模型的整體性能，例如在深度學習時代之前，一些開創(chuàng)性的工作利用帶有深度信息的RGB-D圖像來預(yù)測場景語義以改進預(yù)測質(zhì)量。在深度學習時代，最近引起關(guān)注的基于注意力機制的多任務(wù)學習方法中，明確和隱含地從其他任務(wù)中蒸餾信息作為補充信號，從而提高目標任務(wù)的性能。簡而言之，上述方法通過將現(xiàn)有骨干網(wǎng)絡(luò)與后續(xù)解碼過程結(jié)合起來實現(xiàn)初始任務(wù)預(yù)測，并如圖3(左)所示。

? 具體地，骨干網(wǎng)絡(luò)提取共享特征后通過一組特定任務(wù)頭進行處理，產(chǎn)生了一系列對于 T 個任務(wù)的初始預(yù)測，即 \{Y_i^k\}~(k = 1,\cdots, T) （骨干網(wǎng)絡(luò)和特定任務(wù)頭被稱為網(wǎng)絡(luò)的前端）。將 Y_i^t 轉(zhuǎn)換并二值化成1-bit特征圖形式，得到相應(yīng)場景的一組二值特征圖，即 \{F_{B,i}^t\}~(t = 1,\cdots,T) ，這比骨干網(wǎng)絡(luò)的共享二值特征更加關(guān)注任務(wù)。在最終進行任務(wù)預(yù)測之前，通過二值化注意力機制將這些來自特定任務(wù)的表示融合為多模態(tài)蒸餾信息。論文的方法也是可擴展到不同任務(wù)，尤其是可能有一些任務(wù)只在網(wǎng)絡(luò)前端（初始預(yù)測）中被使用。這些初始任務(wù)也被稱為輔助任務(wù)，因為它們充當代理以改善最終任務(wù)的性能，如圖3所示。

  • Multi-Modal Distillation (MMD) via Binarized Attention Mechanism

? 多模態(tài)蒸餾模塊是多任務(wù)密集預(yù)測模型的關(guān)鍵。具體而言，利用注意力機制來引導不同任務(wù)的不同模態(tài)生成的二值特征圖之間的信息傳遞。由于傳遞的信息流并不總是有幫助的，注意力可以作為一個門函數(shù)來控制信息的流動。換句話說，使網(wǎng)絡(luò)能夠自動學習集中或忽略來自其他二值特征的信息?；诙祷僮鳎梢酝ㄟ^二值化注意力公式化MMD。在傳遞信息給第 k 個任務(wù)時，首先獲得一個二值化注意力圖 \mathbf{A}_{B,i}^k ，即

\begin{equation} \mathbf{A}_{B,i}^{k} \longleftarrow \textit{bool}(\mathbf{W}_{B}^{k}\otimes\mathbf{F}_{B,i}^{k}) \end{equation}

? 其中， \mathbf{W}_{B}^{k} 是二值卷積層的參數(shù)， \mathbf{F}_{B,i}^{k} 是初始預(yù)測的二值特征圖， \otimes 表示卷積操作。然后，通過以下方式傳遞信息：

\begin{equation} \mathbf{F}_{B,i}^{o,k} \longleftarrow \textit{sign}\left[ \mathbf{F}_{B,i}^{k} + \sum_{t=1, t\neq k}^{T}\mathbf{A}_{B,i}^k\odot(\mathbf{W}_{B,t}\otimes \mathbf{F}_{B,i}^{t})\right] \end{equation}

? 其中， \odot 表示逐元素相乘。蒸餾過程如圖3左所示，然后在圖3右側(cè)的頭部中，輸出的二值特征圖 \mathbf{F}_{B,i}^{o,k} 被用于相應(yīng)的第 t 個任務(wù)。通過使用特定任務(wù)的蒸餾激活，網(wǎng)絡(luò)可以為每個任務(wù)保留更多信息，這對于BNNs尤其有益，因為惡化的信息流主要導致性能下降。

? 另一方面，多任務(wù)密集預(yù)測模型在性能方面受益于網(wǎng)絡(luò)二值化。盡管這些多任務(wù)模型已經(jīng)取得了令人期待的性能，但它們?nèi)匀皇艿交诰矸e的蒸餾模塊的限制，這些模塊在多尺度上被大量使用，模擬了相對局部感知領(lǐng)域中的關(guān)鍵空間和任務(wù)相關(guān)交互。從理論上講，在不同網(wǎng)絡(luò)節(jié)點中增加更多的蒸餾模塊可以提高模型性能，但由于計算限制，不能無限制地向現(xiàn)有模型中添加蒸餾模塊。幸運的是，通過二進制網(wǎng)絡(luò)節(jié)省的計算成本，可以在模型中實現(xiàn)額外的蒸餾模塊。

  • Binary Baseline for Multitask Dense Prediction, Bi-MTDP

? 為了在多任務(wù)學習框架下使用BNNs獲得密集預(yù)測，將MTI-Net二值化為二進制基線。具體來說，完整精度的MTI-Net中的主要模塊，包括骨干網(wǎng)絡(luò)、頭部和多模態(tài)蒸餾模塊，都被替換為二進制模塊（權(quán)重和激活都是1-bit），將這個基線稱為Bi-MTDP。

Information Flow Supplementation

? 盡管為多任務(wù)密集預(yù)測構(gòu)建了一個完全二值化的基線Bi-MTDP，并使用常見技術(shù)對管道進行訓練，但性能仍然是一個主要關(guān)注點。無關(guān)因素在前向傳播時過度擬合，而在反向傳播時優(yōu)化方向嚴重不匹配，導致信息嚴重退化。為了解決這些問題，論文進一步提出了Bi-MTDP的變體，Bi-MTDP-F。具體來說，在共享的二進制骨干網(wǎng)絡(luò)的輸出后引入了一個變分信息瓶頸（VIB）層，精確地強制特征提取器保留輸入數(shù)據(jù)的最小充分信息。此外，采用基于特征的知識蒸餾來引導優(yōu)化方向。

  • Variational Information Bottleneck for Filter-Out Nuisance Factors

? 通過共享的骨干網(wǎng)絡(luò)獲取輸入圖像的初始二進制表示，需要訓練一系列有針對性的頭部來對其進行分割，一個直接的策略是將這些表示送到后續(xù)的MMD模塊中。然而，二值化表示缺乏同質(zhì)化，導致了模型過擬合問題。因此，需要對二值化表示進行正則化，而正則化不應(yīng)該污染表示中的信息流。幸運的是，信息瓶頸（IB）原理直接涉及與最佳假設(shè)相關(guān)的壓縮，即數(shù)據(jù)不匹配和模型復雜性應(yīng)同時被最小化。

? 由于VIB能夠有效地捕捉相關(guān)部分并過濾掉輸入中的無關(guān)部分，論文在骨干網(wǎng)絡(luò)后設(shè)計了一個基于VIB的新穎層。具體來說，它明確要求特征提取器保留輸入數(shù)據(jù)的最小充分信息。換句話說，它可以幫助確保信息流靈活地學習針對性任務(wù)的清晰表示?；?code>VIB的分類目標函數(shù)可以被表述為信息損失項，如下所示：

\begin{equation} \mathcal{L}_{vib} = \textit{KL}\left[p(\mathbf{Z}\mid \mathbf{A_{B}}), r(\mathbf{Z}) \right], \end{equation}

? 其中， \mathbf{A_{B}} 是輸入的二進制骨干表示， \mathbf{B} 是潛在表示變量， p(\mathbf{Z}\mid \mathbf{A_{B}}) 是一個多變量高斯分布， r(\mathbf{Z}) 是一個標準正態(tài)分布。一般來說，后者是一個正則化項，控制過濾掉輸入信息的程度。

  • Feature-based Knowledge Distillation for Guiding the Direction of Information Flow

? 蒸餾是一種常見且重要的優(yōu)化方法，用于緩解量化模型在超低比特寬設(shè)置下性能下降的問題，可以靈活地應(yīng)用于任何架構(gòu)，利用全精度教師模型的知識。通常做法是逐層地從全精度教師向量量化對應(yīng)的激活進行蒸餾，即 \mathbf{F}_{B,l} 和 \mathbf{F}_{FP,l} （ l = 1, \cdots, L ，其中 L 表示網(wǎng)絡(luò)層數(shù)）。使用均方誤差（MSE）作為距離函數(shù)來衡量學生和教師之間相應(yīng)特征的差異。知識蒸餾損失可以寫成如下形式：

\begin{equation} \mathcal{L}_{kd} = \sum_{l=1}^{L}\textit{MSE}(\mathbf{F}_{B,l},\mathbf{F}_{FP,l}). \end{equation}

Counter-Intuitive Results of Bi-MTDP-A

? 直觀地說，在FP網(wǎng)絡(luò)上實施二值化必然會導致表示退化，因為sign函數(shù)的梯度無法完美估計。因此，二值化模型不可能優(yōu)于其全精度對應(yīng)模型。然而，Bi-MTDP的變體Bi-MTDP-C（即僅二值化多模態(tài)蒸餾的全精度骨干模型）的性能優(yōu)于其全FP版本。具體來說，僅對多模態(tài)蒸餾進行二值化處理就能同時將模型加速～39%，并將分割mIoU提高～4%。這一結(jié)果表明，論文的方法并不是在模型性能和效率之間進行天真的權(quán)衡，而是提升多任務(wù)密集預(yù)測器的有力工具。

? 這一令人興奮的成果甚至有點違反直覺，推測原因是：i) 在MMD上進行二值化可以過濾掉與任務(wù)無關(guān)的信息；ii) 網(wǎng)絡(luò)內(nèi)的信息流更加有效。為了驗證這一推測，從兩個方面進行了一系列實驗，即Bi-MTDP-C的表示能力和網(wǎng)絡(luò)內(nèi)信息流的補充。

  • Qualitative Study of Learned Features with Bi-MTDP

? 為了研究Bi-MTDP-C及其FP對應(yīng)模型的表示能力，使用t-SNE算法在2-D空間中可視化了二值化多模態(tài)蒸餾（MMD）模塊后面的特征圖，以及使用Grad-Cam算法確定網(wǎng)絡(luò)認為重要的區(qū)域，結(jié)果如圖5所示。顯然，二值化模型Bi-MTDP-C能夠通過二值化注意力模塊過濾掉無關(guān)信息（見圖5（a）），從而有助于學習更具有區(qū)分性的特征（見圖5（b）），導致更高的定量結(jié)果?？傮w而言，分割生成的空間特征圖更好。增強的代表能力可以促進更高的定量結(jié)果。

  • Analysis of Information Flow Supplementation within Network via Centered Kernel Alignment

? 分析神經(jīng)網(wǎng)絡(luò)層內(nèi)的分布信息流是具有挑戰(zhàn)性的，因為層的輸出分布在大量神經(jīng)元之間。中心化核對齊（CKA）可以解決這些挑戰(zhàn)，通過定量比較網(wǎng)絡(luò)內(nèi)部或跨網(wǎng)絡(luò)的激活。具體而言，對于一個由 m 個樣本輸入的網(wǎng)絡(luò)，CKA算法將 \mathbf{X} \in \mathbb{R}^{m\times p_1} 和 \mathbf{Y} \in \mathbb{R}^{m\times p_2} 作為輸入，它們是兩個層的輸出激活（分別具有 p_1 和 p_2 個神經(jīng)元）。令 \mathbf{K} \triangleq \mathbf{X}\mathbf{X}^{\top} 和 \mathbf{L} \triangleq \mathbf{Y}\mathbf{Y}^{\top} 表示兩個層的格拉姆矩陣，CKA計算如下：

\begin{equation} \mathrm{CKA}(\mathbf{K},\mathbf{L}) = \frac{\mathrm{HSIC}(\mathbf{K},\mathbf{L})}{\sqrt{\mathrm{HSIC}(\mathbf{K},\mathbf{K})\mathrm{HSIC}(\mathbf{L},\mathbf{L})}}, \end{equation}

? 在這里，HSIC是希爾伯特-施密特獨立準則。給定中心矩陣 \mathbf{H}=\mathbf{I}_n - \frac{1}{n}\mathbf{1}\mathbf{1}^{\top} 和中心化的格拉姆矩陣 \mathbf{K}^{\prime}=\mathbf{H}\mathbf{K}\mathbf{H} 以及 \mathbf{L}^{\prime}=\mathbf{H}\mathbf{L}\mathbf{H} ， \mathrm{HSIC}=\frac{\mathrm{vec}(\mathbf{K}^{\prime})\mathrm{vec}(\mathbf{L}^{\prime})}{(m-1)^2} ，表示這些中心化格拉姆矩陣之間的相似性。重要的是，CKA對于表示的正交變換（包括神經(jīng)元的排列）是不變的，而歸一化項確保了對各向同性縮放的不變性。這些特性使得能夠?qū)ι窠?jīng)網(wǎng)絡(luò)隱藏表示進行有意義的比較和分析。

? 因此，引入CKA來研究多任務(wù)密集預(yù)測模型中的信息流。在熱圖中，點越淺，兩個相應(yīng)層之間的相似度越高。兩個層的輸出表示之間的相似得分越高，意味著這兩個層共享更多信息。結(jié)果如圖5（c）所示，我們可以看到Bi-MTDP-C前層和后層之間的相似得分要比MTI-Net中的相似得分高得多。這表明Bi-MTDP-C能夠補充網(wǎng)絡(luò)內(nèi)部的信息流，從而提高模型性能。

Experiments

?
?
?

如果本文對你有幫助，麻煩點個贊或在看唄～
更多內(nèi)容請關(guān)注 微信公眾號【曉飛的算法工程筆記】

work-life balance.
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點，簡書系信息發(fā)布平臺，僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容