計(jì)算需求降為1%！ 清華大學(xué)首次提出「二值化光譜重建算法」，代碼全開(kāi)源｜NeurIPS 2023
相比于常規(guī)的三通道 RGB 圖像，高光譜圖像包含幾十上百個(gè)波段，從而捕獲了關(guān)于成像場(chǎng)景更豐富的信息。也正因?yàn)檫@一重要特性，高光譜圖像被廣泛地應(yīng)用于醫(yī)療，地形勘探，農(nóng)業(yè)等領(lǐng)域。
如圖1所示，在醫(yī)院進(jìn)行檢查時(shí)，如果只看常規(guī)的RGB圖像可能很難診斷病因，但是如果采用高光譜圖像捕獲并在特定波長(zhǎng)下渲染的話，就可以看清楚各類血管，骨骼結(jié)構(gòu)等，從而輔助醫(yī)生診斷。同樣的原理也可應(yīng)用在遙感地形勘探和農(nóng)業(yè)病蟲害檢測(cè)。
圖1 高光譜圖像的應(yīng)用
然而高光譜圖像并不容易獲取，傳統(tǒng)的成像設(shè)備采用光譜儀對(duì)成像場(chǎng)景進(jìn)行逐波段的掃描，費(fèi)時(shí)費(fèi)力，難以捕捉運(yùn)動(dòng)場(chǎng)景。
近些年，科學(xué)家們專門設(shè)計(jì)了單曝光壓縮成像（Snapshot Compressive Imaging，SCI）系統(tǒng)來(lái)解決這一問(wèn)題。
其光路結(jié)構(gòu)如圖2所示。該系統(tǒng)首先通過(guò)一個(gè)編碼孔徑掩膜對(duì)成像場(chǎng)景的各光譜通道進(jìn)行調(diào)制，然后通過(guò)一個(gè)三棱鏡進(jìn)行色散后在相機(jī)上捕獲到一個(gè)二維的快照估計(jì)圖（compressive measurement）。
通過(guò)這個(gè)光路系統(tǒng)，我們便可將三維的光譜立方塊壓縮成一個(gè)二維的圖像。而光譜壓縮重建的任務(wù)便是從這個(gè)二維的壓縮估計(jì)圖上恢復(fù)出三維的高光譜數(shù)據(jù)。
圖2 單曝光壓縮成像系統(tǒng)
當(dāng)前的主流方法是采用的是全精度模型如 CNN 或者 Transformer 來(lái)學(xué)一個(gè)從壓縮估計(jì)圖到三維光譜立方塊的映射。幾個(gè)比較經(jīng)典算法有我們之前的工作 MST，MST++，CST，DAUHST，HDNet 等（這些方法全部都開(kāi)源在我們的工具箱 MST 當(dāng)中）。
這類方法雖然取得了很好的重建效果，但卻難以部署到移動(dòng)端設(shè)備（如智能手機(jī)、相機(jī)、無(wú)人機(jī)等）上，因?yàn)橐苿?dòng)端設(shè)備的內(nèi)存空間，計(jì)算資源和電力均有限，無(wú)法運(yùn)轉(zhuǎn)全精度模型。
另一方面，全精度模型的一些計(jì)算單元如深度展開(kāi)算子和多頭自注意力機(jī)制等相對(duì)復(fù)雜，移動(dòng)端設(shè)備無(wú)法支持。
為了推動(dòng)光譜重建算法的實(shí)際應(yīng)用，清華大學(xué)等機(jī)構(gòu)的研究人員提出了一個(gè)全新的工作，其主要貢獻(xiàn)為：
圖片
論文鏈接：https://arxiv.org/abs/2305.10299
代碼鏈接：https://github.com/caiyuanhao1998/BiSCI
代碼鏈接：https://github.com/caiyuanhao1998/MST
（1）提出了光譜壓縮重建領(lǐng)域內(nèi)首個(gè)基于二值神經(jīng)網(wǎng)絡(luò)（Binarized Neural Network，BNN）的算法 — 二值化光譜重分布網(wǎng)絡(luò)（Binarized Spectral-Redistribution Network，BiSRNet）；
圖3 BiSRNet 與當(dāng)前最先進(jìn)的二值網(wǎng)絡(luò)的性能對(duì)比
（2）設(shè)計(jì)了一個(gè)新的二值化卷積單元 — 二值化光譜重分布卷積（Binarized Spectral-Redistribution Convolution，BiSR-Conv）。該卷積單元可以調(diào)整光譜表征的強(qiáng)度和分布，同時(shí)在反向傳播中更好地逼近二值化符號(hào)（Sign）函數(shù)，從而讓求得的梯度更加準(zhǔn)確；
（3）制作了四個(gè)二值化卷積模塊來(lái)解決特征圖形變過(guò)程中的維度不匹配問(wèn)題，從而讓全精度信息能流通整個(gè)模型的每一層卷積單元以彌補(bǔ)二值卷積造成的信息損失；
（4）我們的 BiSRNet 顯著地超越了當(dāng)前最先進(jìn)的二值化算法，甚至取得了能與全精度CNN比肩的效果，然而我們的BiSRNet算法卻只需要極低的存儲(chǔ)空間（約 0.06 %）和計(jì)算代價(jià)（約 1%）。
目前所有的代碼，預(yù)訓(xùn)練模型和測(cè)試結(jié)果均已開(kāi)源在一個(gè)二值化光譜壓縮重建工具包 BiSCI 內(nèi)，該工具包支持八類最主要的二值網(wǎng)絡(luò)，歡迎大家來(lái)使用。
同時(shí)，我們還將 BiSRNet 嵌入到了我們之前開(kāi)發(fā)的光譜重建工具箱 MST 當(dāng)中。目前 MST 工具包已支持超過(guò) 12 類深度學(xué)習(xí)算法，并包含各種配套的可視化函數(shù)，歡迎大家來(lái)對(duì)比。
方法介紹
基礎(chǔ)模型
一般來(lái)說(shuō)，用于二值化的全精度模型應(yīng)該是比較輕量的且它的計(jì)算單元可以再移動(dòng)端設(shè)備上運(yùn)行。然而，現(xiàn)存的CNN 或 Transformer 模型均不滿足這一要求。為此，我們重新設(shè)計(jì)了一個(gè)簡(jiǎn)單，輕量，易于部署的基礎(chǔ)模型（Base Model）。
受到之前工作 MST，MST++，CST，DAUHST 的啟發(fā)，我們?cè)O(shè)計(jì)的基礎(chǔ)模型也采用一個(gè) U 形結(jié)構(gòu)，如圖4所示。
圖4 基礎(chǔ)模型的網(wǎng)絡(luò)結(jié)構(gòu)
在這個(gè)基礎(chǔ)模型中，采用的所有計(jì)算單元都可以被移動(dòng)端設(shè)備支持，同時(shí)也不涉及計(jì)算復(fù)雜度高的操作。
二值化光譜重分布卷積單元
二值化光譜重分布卷積的細(xì)節(jié)如圖 4（c）所示。將輸入的全精度激活信號(hào)記為Xf，可以注意到，受到特定成像波長(zhǎng)的限制，高光譜信號(hào)沿著光譜維度有著不同的強(qiáng)度與分布。
為了適應(yīng)光譜圖像的這一特性，我們提議在二值化激活之前，按通道對(duì)高光譜圖像表示進(jìn)行重分布：
圖片
其中的Xr表示重分布后的激活，k和b表示可學(xué)習(xí)參數(shù)，然后Xr經(jīng)過(guò)一個(gè)符號(hào)函數(shù)后被量化到1位的激活Xb：
圖片
如圖 5 (b) 和 (c) 所示，由于符號(hào)函數(shù)是不可導(dǎo)的，之前的方法大都采用一個(gè)分段線性函數(shù) Clip(x) 或者二次函數(shù) Quad(x) 來(lái)在反向傳播中擬合符號(hào)函數(shù)。
圖片
圖 5 符號(hào)函數(shù)與各類逼近函數(shù)的對(duì)比圖
Clip(x) 與 Quad(x) 的具體表達(dá)式如下：
圖片
然而分段線性函數(shù)知識(shí)一個(gè)粗略的估計(jì)，它與符號(hào)函數(shù)之間依舊有著很大的誤差。
圖 5 中的陰影部分面積就表示這個(gè)誤差的大小，Clip(x) 的誤差是1。此外，一旦激活或者權(quán)重的值落在了 [-1, 1] 之外，他們就不會(huì)再被更新。
盡管分段二次函數(shù)是一個(gè)更精確的逼近（誤差為2/3），上述的兩個(gè)問(wèn)題依舊存在。
為此，我們?cè)O(shè)計(jì)了一個(gè)可縮放的雙曲正切函數(shù)來(lái)在反向傳播中擬合符號(hào)函數(shù)：
圖片
其中α>0 是一個(gè)可學(xué)習(xí)的參數(shù)，自適應(yīng)地調(diào)整 Tanh(αx) 與 Sign(x) 之間的距離。
我們證明當(dāng) α→+∞ 時(shí)，Tanh(αx)→Sign(x)如下：
圖片
我們仿照之前的二值化算法，定義：
圖片
然后可以得到：
圖片
然后計(jì)算圖 5 (d) 中的陰影面積：
圖片
由此可見(jiàn)，當(dāng)α足夠大時(shí)，Tanh(αx)可以任意地逼近 Sign(x)，同時(shí)，我們的 Tanh(x) 既不是分段函數(shù)，也不會(huì)因?yàn)?x 落在 [-1, 1] 之外而無(wú)法更新。
更進(jìn)一步，我們比較三個(gè)逼近函數(shù)的一階導(dǎo)數(shù)，如圖 5 下邊一行所示， ?Clip(x)/?x與?Quad(x)?x的取值范圍均為 [0,1]，這與 ?Sign(x)?x∈[0,+∞) 有著本質(zhì)的區(qū)別。
而我們的 ?Tanh(αx)?x 可以通過(guò)改變參數(shù)α來(lái)調(diào)整它的取值范圍 (0, α) 和形狀，從而更加靈活地逼近 ?Sign(x)?x
在二值化卷積層當(dāng)中，32位的權(quán)重Wf也被二值化到一位的權(quán)重Wb：
圖片
其中的E表示取均值。由此，計(jì)算繁重的全精度卷積操作便可由純邏輯異或非運(yùn)算（XNOR）與位計(jì)數(shù)（bit-count）操作替換如下：
圖片
因?yàn)槿刃畔f的取值范圍與二值卷積的輸出Yb差別較大，直接引入殘差連接將兩者相加容易導(dǎo)致 Yb的信息被淹沒(méi)掉，所以先采用一個(gè) RPReLU 激活函數(shù)來(lái)重塑Yb的取值范圍：
圖片
然后再引入殘差鏈接：
圖片
通過(guò)我們的設(shè)計(jì)，全精度信息流就不會(huì)被二值卷積給阻斷，從而能夠順暢地流過(guò)我們?cè)O(shè)計(jì)的 BiSR-Conv 單元，如圖 4 (c) 中的紅色箭頭所示。
圖片
圖6 本文卷積二值卷積塊與普通二值卷積塊在處理特征圖維度變化時(shí)的對(duì)比
全精度信息流在二值化算法中非常重要，因?yàn)樗谝欢ǔ潭壬蠌浹a(bǔ)了量化導(dǎo)致的信息損失。然而在特征圖的上下采樣過(guò)程中，由于維度發(fā)生了變化，難以直接引入殘差鏈接以補(bǔ)充全精度信息，如圖 6 中每個(gè)子圖的左半部分所示。
藍(lán)色箭頭表示二值信號(hào)，紅色箭頭表示全精度信息流。普通的二值卷積模塊會(huì)在特征度維度改變的同時(shí)阻斷全精度信息流的傳遞。然而我們的二值卷積模塊通過(guò)采用分割合并的技巧，讓全精度信息流不被阻斷，從而在整個(gè)二值模型中流通。
實(shí)驗(yàn)結(jié)果
量化指標(biāo)
表1 BiSRNet 與 SOTA BNN，傳統(tǒng)方法，全精度CNN方法的量化指標(biāo)對(duì)比
表 1 展示了我們的 BiSRNet 與 SOTA BNN，傳統(tǒng)方法，全精度 CNN 方法的量化指標(biāo)對(duì)比。可以看到，我們的 BiSRNet 顯著超越了當(dāng)前最先進(jìn)的 BNN 算法 BTM 2.55 dB。
同時(shí)超越了所有的 64 位傳統(tǒng)算法，增幅超過(guò)了4.49dB，值得注意的是，我們的 BiSRNet 取得了能與全精度 CNN 方法比肩的性能，比如超過(guò)了 λ?Net 1.23 dB，然而卻只占用了 0.06 % 的存儲(chǔ)空間和 1 % 計(jì)算代價(jià)。
視覺(jué)對(duì)比
圖7 BiSRNet 與其他 BNN 方法在仿真數(shù)據(jù)及上的視覺(jué)對(duì)比
圖8 BiSRNet 與其他 BNN 方法在真實(shí)數(shù)據(jù)集上的視覺(jué)對(duì)比
參考資料：
