引入多狀態記憶機制的迭代軟閾值學習算法

羅美露 余 磊 張海劍

(武漢大學電子信息學院, 湖北武漢 430072)

1 引言

稀疏恢復問題作為壓縮感知中研究和討論的重要問題之一,被廣泛應用于信號處理[1]、醫學成像[2]以及雷達系統[3- 4]等跨學科領域中。理論上,由受噪聲干擾的低維觀測向量重構原始稀疏信號,通常是一個欠定和病態的問題,目前主要有三類算法對其進行優化求解,它們分別是:貪婪算法[5- 6]、凸松弛算法[7- 8]和迭代閾值算法[9-12]。貪婪算法主要包括各類匹配追蹤算法,通過多次迭代從過完備原子庫里尋找稀疏向量的支撐集,并使用最小二乘估計重構信號。這類算法通常需要的測量數據多,計算量大,收斂速度也較慢。凸松弛算法主要通過凸優化方法解決凸或擬凸優化問題, 從而恢復出稀疏信號。相對于貪婪算法,該類算法大大減少了重構信號所需的觀測次數,得到的解是滿足條件的最稀疏解。迭代閾值算法通過閾值確定迭代中的保留項,進一步更新待恢復信號。文獻[11]中提出的迭代軟閾值算法(Iterative Soft-Thresholding Algorithm,ISTA)通過一個軟閾值操作更新信號,復雜度小且結構簡單,受到廣泛的研究與應用。但ISTA收斂速度較慢。文獻[12]基于ISTA的框架提出一種快速迭代軟閾值算法(Fast Iterative Soft-Thresholding Algorithm,FISTA),該算法選擇前兩次迭代點的線性組合作為下一次迭代時近似函數的起始點,進一步提高收斂速率,但收斂速度依舊有限。

近年來,許多國內外學者開始將深度學習技術應用到稀疏恢復求解中。文獻[13]介紹了一種迭代軟閾值學習算法(Learned Iterative Soft-Thresholding Algorithm,LISTA),將ISTA展開為遞歸前饋神經網絡,通過學習的方式優化更新字典矩陣,可更快地趨向最優解。文獻[14]提出一種快速迭代軟閾值學習算法(Learned Fast Iterative Soft-Thresholding Algorithm,LFISTA),構建了與LISTA類似的遞歸神經網絡,進一步提高了算法的收斂速度。文獻[15-16]基于LISTA分別引入了不同的約束條件以達到加速收斂的目的。文獻[15]利用LISTA結構中兩個可學習矩陣的耦合性,并采取了多種閾值選取方案,在簡化網絡結構的同時加快了稀疏恢復的速度。文獻[16]選取了一個可解析的自由權值矩陣,優化并固定該解析矩陣,通過學習調整權值步長和閾值參數即可加速收斂。從上述研究內容可知,深度學習技術的引入可以有效地提高ISTA的收斂速度,在此基礎上,通過增加約束條件可進一步加速收斂。

實際上,稀疏恢復優化問題屬于約束最小化問題的一種情況,一階迭代固定步長算法[17-19]一般適用于解決該類問題。文獻[20]利用李雅普諾夫函數證明了一階迭代固定步長算法可線性收斂。受到該思想啟發,本文提出一種引入多狀態記憶機制的改進LISTA(Learned Iterative Soft-Thresholding Algorithm with Multi-state Memory,LISTA-MM)算法。該算法設置了一個狀態連接度數,連接多個先前迭代點,學習迭代點的權重矩陣,從而靈活地利用稀疏信息,確保了稀疏估計過程中信息被正確傳遞并充分利用,能夠有效提高收斂速度,并且優化稀疏恢復性能。此外,本文還將LISTA-MM的卷積形式應用在圖像超分辨率中,訓練了一個圖像超分辨率網絡模型,并將該網絡的重構結果和其他幾種圖像超分辨率網絡的結果進行比較。實驗結果表明,基于LISTA-MM的網絡在圖像質量評估指標上優于其他幾種網絡,并且重構圖像的可視化結果受偽影影響小,具有更接近原始圖像的清晰規則紋理。

本文貢獻如下:(1)本文提出一種引入多狀態記憶機制的改進迭代軟閾值學習算法,設置狀態連接度數,靈活利用多個迭代狀態點信息,提高了算法的收斂速度;(2)通過仿真實驗驗證了本文算法在保證稀疏恢復精度的同時有效提高了收斂速度;(3)將LISTA-MM擴展為卷積形式,訓練了一個圖像超分辨率網絡模型,探索其在圖像超分辨率中的應用,實驗結果表明該網絡比其他幾種圖像超分辨率網絡具有更好的性能。

本文安排如下:第2節介紹稀疏恢復的求解問題并且詳細敘述本文提出的改進迭代軟閾值學習算法LISTA-MM;第3節介紹了LISTA-MM的卷積形式及其在圖像超分辨率中的應用;第4節展示了本文改進算法與幾種迭代軟閾值學習算法的對比結果,以及基于LISTA-MM的網絡與其他幾種圖像超分辨率網絡的重構結果對比;第5節為本文的結論。

2 引入多狀態記憶機制的迭代軟閾值學習算法

2.1 稀疏恢復問題的求解

在稀疏恢復問題中,令待恢復信號為z∈RN,假設其非零元素不超過K個(K-稀疏性),D∈RM×N為測量矩陣或過完備字典矩陣(M?N),則需要從低維觀測向量中恢復未知的稀疏信號:

y=Dz+n

(1)

其中,y∈RM是觀測向量,n∈RM為噪聲向量。上述問題是一個欠定且病態的問題,不能通過簡單的線性運算直接求解,常見求解方式是把噪聲條件下的稀疏恢復問題表述為下列無約束最小化問題:

(2)

其中||z||1表示稀疏正則項,λ是正則系數。上述問題是一個經典的LASSO問題,ISTA是求解該問題的典型一階算法,該算法利用第k次迭代時近似函數取得最小值的點zk作為下一次迭代的起始點,每一步的迭代僅涉及矩陣和向量的乘積及軟閾值的計算。

ISTA算法被認為是梯度算法的一種自然延伸,首先考慮連續可導的無約束最小化問題:

min{g(z):z∈RN}

(3)

用梯度下降法求解的迭代形式為:

zk+1=zk-tk?g(zk)

(4)

這里tk>0是迭代步長,梯度下降法可以表示成g在點zk處的近端正則化,其等價形式表示為:

(5)

(6)

忽略式(6)中的常數項可以得到:

(7)

由于l1范數是可分的,式(6)最終可寫為

zk+1=hλtk(zk-tk?g(zk))

(8)

(9)

其中zk表示第k次迭代的稀疏估計值,L通常作為DTD的最大特征值,T表示轉置運算符。

zk+1=hθ(Szk+Wey),k=1,…,K

(10)

其中兩個字典矩陣We、S和閾值參數θ可通過最小化重構誤差的方式來學習優化。

2.2 引入多狀態記憶機制的迭代軟閾值學習算法

值得注意的是,公式(2)中的稀疏恢復問題的目標函數是凸函數,該問題屬于約束最小化問題的一種情況:

(11)

其中f:RN→R表示包括l-光滑函數和u-強凸函數在內的一系列問題[20],一階迭代固定步長算法[17-19]一般適用于解決上述問題。類似于加速近端梯度法的思想,一階迭代固定步長算法利用先前V次迭代中近似函數取得最小值的點,將這些點通過步長γ加權表示為新點yk,作為下一次迭代的起始點,使得下一步的梯度更新方向與當前的梯度方向不要相差太大,能夠有效加速迭代點的收斂,其具體求解形式為:

(12)

其中α,βj,γj是(固定)步長,V(≥0)被稱為度數?；谝浑A迭代固定步長算法的求解思想,本文提出多狀態記憶機制并應用到LISTA框架中,其具體數學表達式為:

zk+1=hθ(Sqk+Wey)

(13)

其中權系數γi和矩陣S是相互獨立的可學習參數,VLISTA-MM被定義為模型的狀態連接度數。本文將每一次更新的迭代點稱為一個迭代狀態,狀態連接度數決定每次迭代需要記憶的先前狀態個數。多狀態記憶機制是一種信息記憶與傳遞機制,針對LISTA每次迭代只依賴于前一迭代狀態對算法收斂速度的限制,通過狀態連接度數設置需要記憶的狀態個數,在模型迭代更新時,靈活地記憶并傳遞多個迭代狀態的信息,并利用權系數γ加權得到新狀態點qk,將其作為下一次迭代的起始點,確保了迭代過程中不會出現過大的偏差,加快了迭代點向期望值收斂的速度。

將上式軟閾值函數中的qk用狀態加權和代入,并合并權系數γi和矩陣S,式(13)可重新表示為:

(14)

式中Wi=γiS是迭代點zi的可學習權值矩陣。與LISTA相同,通過最小化重構誤差優化更新矩陣We和Wi。圖1展示了LISTA-MM的網絡部分展開結構。為了方便觀察,圖中只保留了與迭代點zk+1有關的連接?？梢钥闯?先前的迭代狀態通過權值矩陣被選擇和傳遞到當前迭代中,靈活地組合這些迭代狀態的信息,從而減小了迭代過程中的估計偏差帶來的影響,提高稀疏估計的效率和正確性。

圖1 LISTA-MM網絡的部分展開結構Fig.1 The partially unfolded network structure of LISTA-MM

3 LISAT-MM在圖像超分辨率的應用

圖像超分辨率的目的是由一張低分辨率圖像ILR重構出對應的高分辨率圖像ISR。從稀疏的角度來看,低分辨圖像(高分辨圖像)可以被表示為過完備字典Dl(Dh)中少量原子的線性組合,這些原子被稱為線性稀疏系數gl(gh)∈RN。當字典Dl和Dh被正確定義時,假設每個低分辨圖像/高分辨圖像對共享相同的稀疏編碼,圖像超分辨率可以表示為形如公式(2)的稀疏編碼問題,本節將探索卷積形式的LISTA-MM在圖像超分辨率中的應用。

利用卷積運算符代替矩陣乘法,公式(14)可以被重新表述為卷積結構形式,其類似于文獻[21]提出的一個稀疏自動編碼器模型:

(15)

為了簡化計算和提高效率,這里令所有迭代層都共享同一個卷積算子S,分離出相互獨立的權重系數分別用w0,w1,w2,…,wV表示。在非負稀疏編碼的假設下,非負軟閾值函數和ReLU激活函數被認為是等效函數。因此,將hθ(x)替代成ReLU(x-θ),卷積LISTA-MM模塊可以采用深度學習技術訓練:

(16)

此時?表示卷積算子,y∈Rb×c×h×w表示通道數為c、尺寸為h×w的b幅圖像,We∈Rm×c×s1×s1和S∈Rm×c×s1×s1是兩個互不相關的可訓練變量,s1表示卷積濾波器的卷積核大小。卷積LISTA-MM的展開結構如圖2所示,這里以第k+1次迭代的展開為例。

圖2 卷積LISTA-MM的部分展開結構Fig.2 The partially unfolded architecture of the convolutional LISTA-MM

圖像超分辨率網絡的整體構成如圖3所示,主要分為特征提取、LISTA-MM和重構三大模塊。特征提取模塊和重構模塊均由卷積層和ReLU激活算子組成,分別提取圖像的低維特征和高維特征。中間LISTA-MM模塊通過遞歸方式學習更新特征圖的稀疏編碼。該網絡的初始輸入是插值后的低分辨率圖像ILR,經過網絡重構出期望的高分辨率圖像ISR。

圖3 圖像超分辨率網絡的整體結構Fig.3 The architecture of the image super-resolution network

4 實驗結果及分析

為了驗證本文所提算法的有效性,本文將選取仿真數據和自然場景圖像數據集[22-23]分別進行實驗。實驗分為兩個部分:首先,本文所提LISTA-MM算法與LISTA及其衍生的幾種迭代軟閾值學習算法進行比較,驗證多狀態記憶機制在提高稀疏恢復收斂速度方面的有效性。其次,本文將卷積LISTA-MM應用到圖像超分辨率網絡中,并與現有的一些圖像超分辨率網絡進行比較。

4.1 仿真信號稀疏恢復4.1.1 數據集及實驗細節

(17)

其中z是作為參照的稀疏信號,[||·||]表示能量的期望值。本節實驗采用歸一化均方誤差(Normalized Mean Square Error,NMSE)作為主要評價指標,以便更直觀地反映稀疏恢復算法的性能:

NMSE=10log10

(18)

為了驗證本文所提LISTA-MM在仿真數據上的有效性,實驗將公式(14)中的狀態連接度數VLISTA-MM分別設置為1～4,訓練不同的LISTA-MM模型(分別表示為LISTA-MM1,LISTA-MM2,LISTA-MM3,LISTA-MM4),討論了多狀態記憶機制對收斂速度的影響,并且將其與LISTA和文獻[15-16]提出的稀疏恢復算法(本文分別稱之為LISTA-CP[15],LISTA-CPSS[15],TiLISTA[16])進行比較。實驗環境為PyTorch和MATLAB R2017b,模型建立在PyTorch上,采取Adam優化器并將初始學習率設為0.02。

4.1.2 仿真信號稀疏恢復結果

為了探究多狀態記憶機制對稀疏恢復收斂速度的影響,本節實驗首先預訓練LISTA模型,并將預訓練參數固定,從頭訓練LISTA-MM模型,只更新公式(14)中的權重矩陣Wi。圖4展示了不同噪聲條件下LISTA-MM與LISTA算法的稀疏恢復性能結果,其中x軸表示迭代次數,y軸表示NMSE,以分貝(dB)為單位。從圖4可以看出,在不同SNR水平下,與LISTA相比,LISTA-MM達到同等大小的NMSE值所需的迭代次數更少,從而驗證其在加快收斂速度方面的有效性,且隨著SNR的增大,這種優勢更為明顯。另外,從圖4 (d)可以看出,當稀疏恢復結果開始趨于穩定時,LISTA-MM的恢復誤差比LISTA低3～10 dB。

圖4 不同信噪比水平下,LISTA-MM和LISTA的稀疏恢復性能結果(LISTA-MM模型的數字后綴表示狀態連接度數)Fig.4 Sparse recovery performance results of LISTA-MM and LISTA at different SNR levels (the digital suffix of LISTA-MM model represents the state connection degree)

為了進一步探究LISTA-MM的性能表現,本文進行無預訓練模型實驗,分別從頭訓練LISTA模型和LISTA-MM模型,并與LISTA-CP、LISTA-CPSS、TiLISTA進行比較。圖5展示了不同噪聲條件下上述算法的稀疏恢復性能結果。從圖中可知,當SNR=15 dB,LISTA-MM的收斂速度明顯優于其他四種算法。當SNR≥20 dB時,LISTA-MM和LISTA-CP的NMSE結果較為接近,收斂速度相當,且性能明顯優于其他三種算法。

圖5 不同信噪比水平下,LISTA-MM和LISTA,LISTA-CP,LISTA-CPSS,TiLISTA的稀疏恢復性能結果(LISTA-MM模型的數字后綴表示狀態連接度數)Fig.5 Sparse recovery performance results of LISTA-MM and LISTA,LISTA-CP,LISTA-CPSS,TiLISTA at different SNR levels (the digital suffix of LISTA-MM model represents the state connection degree)

此外,本文還展示了LISTA和LISTA-MM算法重構信號的可視化對比結果,如圖6所示,其中x軸表示序列長度,y軸表示信號幅度。從圖中可以看出,LISTA-MM算法的重構信號基本恢復了信號的稀疏信息,并且相較于LISTA,其信號幅度更接近原始稀疏信號,重構誤差更低。綜上所述,實驗結果表明,多狀態記憶機制的引入不僅可以加快算法的收斂速度,還能提高信號的稀疏恢復精度。

4.2 圖像超分辨率重建結果

參照現有的圖像超分辨率網絡,本節實驗使用與文獻[28]相同的訓練集。該訓練集共291張自然場景圖像,分別來自文獻[22]和Berkeley分割數據集[23]。由于訓練圖像具有局限性,訓練前本文采用數據增強方式,使用MATLAB R2017a對訓練集進行不同尺度(×2,×3,×4)的下采樣,并訓練了一個圖像超分辨率網絡模型。實驗使用的測試集是Set5[24],Set14[25],BSD100[23],Urban100[26],并將圖像由RGB顏色空間變換到YCbCr顏色空間,在Y通道上計算峰值信噪比(Peak Signal-to-Noise Ratio,PSNR)和結構相似性(Structural Similarity,SSIM)作為實驗評價指標,其中PSNR值越大表示重構圖像失真越少,SSIM值越大表示兩幅圖像相似度越高,重構圖像質量越好。

本節實驗采用卷積LISTA-MM(狀態連接度數VLISTA-MM設置為4)搭建圖像超分辨率網絡,將該網絡的重建圖像和實驗評估指標與現有的一些圖像超分辨率網絡進行比較。本節實驗中網絡迭代層數設置為25,實驗環境為PyTorch,在NVIDIA Titan Xp gpu下進行訓練,使用Adam優化器和圖像超分辨率網絡訓練常用的l1損失函數,初始學習率設置為4×10-4,每訓練10個階段減半,共訓練了35個階段。

實驗對LISTA-MM與雙三次插值算法,以及文獻[27-31]中的圖像超分辨率網絡進行比較(本文分別依次稱為Bicubic、SRCNN[27]、VDSR[28]、DRCN[29]、DRRN[30]和MemNet[31]),實驗評估指標在MATLAB中計算。表1總結了上述網絡在標準數據集上的實驗定量結果(在尺度因子分別為2,3,4時,數據集Set5,Set14,BSD100和Urban10的PSNR/SSIMs平均值),最佳性能用加粗數字表示。從表中可以看出,基于LISTA-MM的圖像超分辨率網絡在四個數據集上的評價指標均優于其他六種網絡。此外,圖7展示了不同網絡下圖像超分辨率的可視化結果。從圖中可以看出,其他六種網絡容易被明顯的偽影影響,重建圖像出現不正確的紋理和過于平滑的邊緣。相較其他六種網絡,LISTA-MM的重構圖像更清晰,細節紋理信息更豐富準確。實驗結果表明,基于LISTA-MM的網絡在評價指標和可視化效果上均優于其他網絡,證明了卷積LISTA-MM在圖像超分辨率中的優勢。

圖6 LISTA和LISTA-MM算法的重構信號可視化結果(SNR=25 dB,迭代層數等于20)Fig.6 Visualization of recovered signal of LISTA and LISTA-MM (SNR=25 dB,iterative layer equals to 20)

表1 不同網絡下標準數據集測試結果(最佳性能用加粗數字表示)

圖7 圖像可視化結果比較Fig.7 Comparison of image visualization results

5 結論

本文基于LISTA框架引入一種多狀態記憶機制,設置狀態連接度數,使多個先前迭代點通過可學習權值矩陣組合連接到當前迭代,從而確保了信息利用的正確性與高效性,達到加速算法收斂的目的。實驗結果表明,在同等誤差范圍內,LISTA-MM所需的迭代次數少于LISTA及其衍生算法,從而驗證LISTA-MM在加快收斂速度方面的有效性。此外,本文擴展LISTA-MM的卷積形式并將其應用于圖像超分辨率中,相應的實驗結果展示出該網絡在重構高分辨率圖像中的優勢。