基于塊金字塔記憶模塊的無監督異常檢測

鄢 寧，李岳陽，羅海馳

（1.江南大學 人工智能與計算機學院，江蘇 無錫 214122；2.江南大學 物聯網工程學院，江蘇 無錫 214122）

0 概述

異常檢測主要是判斷輸入數據是否與訓練數據“類似”，主要是針對少數、不可預測或不確定的罕見的事件。在現實世界中，由于異常種類繁多，出現的概率較低，且形式多種多樣，因此難以收集到大量的、所有種類的異常樣本。但是人們可以很容易地獲取到大量的正常樣本，這將導致正負樣本的比例極其不均衡，另外手工標注工作也非常耗時。因此，無監督異常檢測［1-3］的方法至關重要，在訓練時，通過使用正常樣本進行訓練從而得到正常樣本的數據分布；在推斷時，具有與正常樣本不同分布的數據樣本被認為是異常樣本。

在無監督異常檢測方法中，基于重建的方法在訓練過程中只需要正常樣本，不需要額外的訓練數據和預訓練模型［4］，因此被廣泛應用。如基于自編碼器［5-7］或生成對抗網絡［8-10］的重建方法已經得到了廣泛的研究，這種方法是基于正常樣本比異常樣本能夠通過潛在特征空間更好地被重建這一假設［11-12］，使重建網絡在正常圖像上進行訓練后能夠充分地學習到正常圖像的特征，從而在測試時能夠很好地重建正常圖像。但是，對于異常圖像來說，由于在訓練階段沒有學習異常特征，因此在測試時其重建能力較差，需通過計算輸入圖像和重建圖像之間的誤差來進行異常檢測。但是這一假設并不總是成立，由于卷積神經網絡具有強大的泛化能力，有時異常特征也能夠被重建出來，因此無法通過比較異常圖像和其重建圖像的差異性來判斷是否存在異常。

為了解決上述問題，文獻［13-15］通過在自編碼器中增加記憶模塊，并在訓練階段用于存儲正常樣本的典型模式，在測試階段通過聚合內存中與輸入相關的正常項來減少重建異常。但是，這種方式容易忽略圖像細節，導致輸出圖像模糊，且原始的記憶模塊［13］是以逐像素的方式進行存儲和讀取的，由于像素點的模式變化是有限的，異常像素點和正常像素點可能會具有相同的模式，那么異常像素點就很容易被重建，最終導致重建異常，因此該方法性能提升有限。

文獻［16］提出記憶模塊使用合適的塊進行存儲，異常塊與正常塊很可能會具有不同塊模式，從而對異常塊進行較差的重建。為了得到更好的檢測結果，選擇合適的塊的尺寸就顯得尤為重要。如果塊的尺寸過大，那么模型重建的圖像會丟失很多細節信息，從而導致重建模糊；如果塊的尺寸過小，異常就會被重建，從而導致檢測精度降低。此外，對于不同種類的數據集，由于其特點不同，合適的塊的尺寸也不同。對于某個種類的數據集，其需要檢測的異常的面積大小差距較大，對于面積較大的異常，如果選用較小尺寸的塊存儲和讀取，則容易重建異常；對于面積較小的異常，如果選用較大尺寸的塊存儲和讀取，則容易重建模糊。為此，需要進行大量的實驗以得到合適塊的尺寸。

基于上述問題，本文提出一種基于改進的塊記憶模塊的無監督異常檢測方法。在原有的塊記憶模塊［16］中加入多個尺度的塊記憶模塊，各尺度塊記憶模塊通過讀取、聚合得到多特征圖融合的輸出特征圖，最大限度地保留不同尺度下樣本的特征信息，以更好地重建數據，減少重建異常的情況。通過增加skip connection［17］結構，將編碼器輸出的多個尺度的特征圖使用塊金字塔記憶模塊得到多個增強特征圖，將skip connection 送入解碼器進行重建，避免直接將異常特征送入解碼器進行重建，以在提高重建圖像的清晰度的同時減少重建異常的情況。最終采用SSIM［18］損失，通過亮度、對比度和結構3 個維度來增強圖像重建效果，并且作為異常分數組成部分，增強原始圖像和重建圖像的差異性。

1 相關工作

1.1 基于重建的無監督異常檢測

基于重建的方法通常使用自編碼器［19-21］進行圖像重建。自編碼器由編碼器和解碼器組成，編碼器將輸入圖像映射到低維特征空間，解碼器試圖從低維特征空間中恢復數據，獲得重建圖像。文獻［8，22］引入對抗學習判別器，使用對抗訓練來提高圖像重建質量。除了使用基于圖像的重建損失和對抗損失來訓練自編碼器外，最小化原始圖像和重建圖像在潛在特征空間之間的距離也有助于學習正常樣本的數據分布。在推理階段，除計算原始圖像和重建圖像的差異分數外，原始圖像和重建圖像潛在特征空間距離的分數在異常評判中也發揮重要的作用。

1.2 記憶模塊

由于卷積神經網絡強大的泛化能力，且有些異常與正常特征具有相同的局部模式，異常也能夠被重建。為此，文獻［13］提出在自編碼器中增加記憶模塊，將編碼器的輸出不直接傳入解碼器，而是作為查詢來檢索內存中相關的項，然后將這些項聚合后送入解碼器。由于記憶模塊存儲的是正常數據的典型模式，因此在測試階段，從內存中獲取正常的數據信息用于重建，來增強異常的重建誤差。

記憶模塊如圖1 所示，其中內存M?RN×C包含N個維度為C的實值向量，mi表示內存M的第i個內存項，mi?RC,?i?{1,2,…,N}。內存M是 可尋址內存，尋址方案是根據查詢與內存項的相似性來計算注意權值w，w?R1×C。查詢z關于內存M中的內存項mi的權值wi的計算公式如下：

圖1 記憶模塊網絡結構Fig.1 Memory module network structure

然后通過使用尋址權值w對內存M中的內存項進行線性組合得到?RC，計算公式如下：

但是一些異常仍然能夠被重建，為了減少這種情況，并且避免受不重要信息的影響，對w進行稀疏化操作，計算公式如下：

其中：λ為收縮閾值（在實驗中將其設置為1/N）；?是一個非常小的正參量；max(?,0)為ReLU 激活函數。然后進行歸一化操作，得到通過稀疏化操作，使用少量且更相關的內存項來聚合得 到。

2 本文方法

2.1 整體網絡結構

提出網絡結構由重建子網絡G和判別子網絡D組 成，網絡結 構如圖2 所 示。采 用UNet［17］作為基 礎重建子網絡，其由編碼器GE和解碼器GD組成。給定輸入圖像x?Rh×w×c，其中，h為高，w為寬，c為通道維數。編碼器GE將x映射到低維特征空間，解碼器GD試圖從投影的低維空間中恢復數據得到重建圖像?Rh×w×c，其中，h為高，w為寬，c為通道維數。為了減少圖像重建模糊和重建異常的情況，在編碼器GE和解碼器GD之間增加skip connection 和塊金字塔記憶模塊得到改進后的重建子網絡G，將編碼器輸出的特征圖作為查詢檢索內存中相關的內存項，聚合后通過skip connection 傳入解碼器用于重建。

圖2 整體網絡結構Fig.2 Overall network structure

判別網絡D的網絡結構是按照DCGAN［23］的判別網絡結構設計的，引入判別網絡的目的主要有：1）使用對抗學習思想提升重建精度，在訓練階段，判別網絡D用來預測給定輸入的類標簽，對輸入圖像x和重建圖像進行分類，將輸入圖像x判別為真，重建圖像判別為假；2）判別網絡D除作為分類器外，還作為特征提取器，提取輸入圖像x和重建圖像的潛在特征表示，再計算兩者之間的差異性，訓練階段通過最小化這一差異性來進一步提升重建效果；在測試階段，將其作為異常評價指標的組成部分提升檢測精度。

2.2 塊金字塔記憶模塊

對于輸入圖像x，編碼器GE輸出一組特征，記為{F1,F2,…,FL}，其中，L表示編碼器GE共輸出L層特征圖。本文將每層特征圖作為其對應的塊金字塔記憶模塊的輸入，對其進行不同尺度的存儲和讀取，以減少異常重建。

對于編碼器GE輸出的第l層特征圖其中，Hl為高，Wl為寬，Cl為通道維數。對其使用塊金字塔記憶模塊進行存儲和讀取如圖3所示。

圖3 塊金字塔記憶模塊網絡結構Fig.3 Block pyramid memory module network structure

對于第i個分支，將其均分為個塊，其 中，rHi、rWi和rCi分別為 沿高度、寬度和 通道維度的劃分率。每個塊的高度為hli=rHi×Hl，寬度為wli=rWi×Wl，通道長度為cli=rCi×Cl。對于每一個塊，將其展平后得到一個向量作為查詢，檢索內存（表示內存Mli包含Nli個維度為Pli的實值向量，Pli=hli×wli×cli）中相關的內存項，再經過聚合得到的特征向量，最后變換為與Zl相同的維度大小的特征圖

本文通過增加多個尺度的塊記憶模塊來減少重建模糊和重建異常的情況。分支的劃分率設置如下：

2.3 損失函數

本文方法只使用正常圖像進行訓練，在測試階段同時使用正常圖像和異常圖像。因為異常圖像沒有參與訓練，所以在測試時異常圖像與其重建圖像在圖像空間和潛在的特征空間之間的差異性較大。本文結合3 個損失函數來訓練模型，以保證模型可以更好地重建正常圖像。

1）第1 項為重建損失，其中包括L1 損失和基于塊的SSIM［18］損失。

L1 損失確保模型能夠對輸入圖像x進行精確的重建，通過把輸入圖像x與重建圖像（即G(x)）的絕對差值的總和最小化來優化G：

基于塊的SSIM 損失的計算公式如下：

其中：h和w為圖像x（或者）的高和寬；NP為x（或者）的像素總個數；SSIM(x,)i,j表示輸入圖像x與重建圖像中以(i,j)為中心的兩個塊的SSIM 值。

因此，重建損失為：

其中：λ為SSIM 損失的權重參數。

2）第2 項損失為潛在特征空間損失，為確保輸入圖像x和重建圖像有盡可能相似的潛在特征表示，使用判別網絡D最后一個卷積層的輸出來提取輸入圖像x和重建圖像的潛在特征表示，令z=f(x)和=f()，采用L2 損失：

3）第3 項為對抗損失［24］，使用判別網絡進行對抗訓練，采用對抗損失來使生成網絡G盡可能逼真地重建正常圖像，判別網絡D對輸入圖像x與重建圖像進行分類，采用二分類的交叉熵損失來進行優化：

因此，總的訓練損失為上述3 個損失的加權總和：

其中：α、β和γ是調節單個損失對總體目標函數影響的權重參數。

2.4 異常分數

對于測試圖像I，其異常分數A(I)為：

其中：a為權重參數。

根據式（10）計算出測試集中每個測試樣本的異常分數，將測試集中所有樣本的異常分數表示為向量A，并采用文獻［9］方式將其縮放到[0,1]內，因此測試樣本I的最終異常得分為：

3 實驗

3.1 數據集和實驗設置

本文使用工業異常檢測數據集MVTec AD［25］，該數據集專為無監督工業視覺異常檢測而開發，是目前該領域的重要基準，包括用來訓練的3 629 張正常圖像和用于測試的1 725 張圖像（包括正常圖像和異常圖像），由15 種類別的對象組成，包括10 種物體類和5 種紋理類，每類圖像包括一個訓練集和一個測試集。圖像的分辨率在700～1 024 像素范圍內，該數據集的挑戰在于有些異常圖像與正常圖像的差別較小，且異常面積的大小差別較大。在實驗時，本文將圖像大小調整為256×256 像素。

本文方法是使用PyTorch 1.8.1 和CUDA 11.1 實現的，所有實驗均在RTX 3090 GPU 上運行。使用Adam［26］優化器，設 置lr=0.000 1，β1=0.500，β2=0.999，weight-decay=0，Epoch 為800，batch_size=8。所有對比實驗和消融實驗均在同一環境下完成。

3.2 評價指標

本文實驗采用AUC 來評價模型的性能，當AUC分數越高時，模型性能越好。AUC 表示隨機抽取一對樣本（一個正樣本和一個負樣本），然后用訓練得到的分類器來對這兩個樣本進行預測，該正樣本預測值大于該負樣本預測值的概率為：

其中：M表示正樣本個數；N表示負樣本個數；rank(i)表示第i個樣本的序號（M個正樣本按照預測值從小到大進行排序）。

3.3 主流方法和本文方法對比

為驗證本文方法的有效性，將主流基于重建的無監督異常檢測方法GANomaly［8］、Skip-GANomaly［9］、MemAE［13］和DAAD+［16］與本文方法在MVTec AD 數據集上進行對比，實驗結果如表1 所示，其中加粗數字為最優值。

表1 主流方法模型與本文改進方法的AUC 對比Table 1 Comparison of AUC between mainstream method models and improved method in this paper

分析表1 的數據可知，本文方法明顯優于主流的基于重建的無監督異常檢測方法。GANomaly 存在的重建模糊問題和Skip-GANomaly 存在的重建異常問題導致了兩者性能均不佳。MemAE 方法以逐像素點的方式存儲正常樣本的典型模式，在進行異常檢測時，輸入特征圖的每個像素點都通過加權聚合相似的內存項得到。由于異常像素點與正常像素點可能具有相同的模式，異常很容易被重建，因此性能提升有限。而本文方法是以記憶塊的方式進行存儲和讀取的，異常塊與正常塊往往具有不同的模式，因此可以通過減少異常重建的情況來提高檢測精度。此外，MemAE 方法只對最后一層輸出特征圖應用記憶模塊，導致重建模糊的問題，本文方法通過skip connection 結構減少了重建模糊的問題，進而提升了檢測性能。DAAD+方法以統一尺寸的記憶塊存儲和讀取，但是異常的大小不一，難以選擇合適的記憶塊尺寸，若尺寸過大，則導致重建模糊問題，若尺寸過小，則異常塊很可能與正常塊有相同的模式，從而導致重建異常問題。而本文方法通過使用多個不同尺度的記憶塊來適應不同大小的異常，既可以提升重建清晰度，又可以減少重建異常的情況，因此本文方法檢測精度高于DAAD+方法，在15 個類別中，該方法在11 個類別上取得了最高的AUC 分數，并達到最高的平均AUC，比DAAD+方法高出1.5%。

3.4 消融實驗

通過在MVTec AD 數據集上做消融實驗，以驗證所提出的各個模塊的有效性，實驗結果如表2 所示。將GANomaly 生成網絡替換為UNet 并刪去skip connection 為基準，對加入的模塊進行消融實驗。

表2 不同算法對模型性能的影響Table 2 The impact of different algorithms on model performance

分析表2 的數據可知，當對基準網絡直接加入skip connection 后，平均AUC 下降了4.3%，因為skip connection 雖然降低了網絡的重建難度，增強重建清晰度，但是對于異常圖像，異常特征從編碼器通過skip connection 直接傳入解碼器用于重建，導致輸入圖像與重建圖像差別不大，難以判斷是否存在異常，不同算法的檢測結果對比見圖4。如圖4（b）和圖4（c）所示，圖4（b）重建圖像較為模糊，圖4（c）重建圖像與原始圖像極為相似（其中，圖4（b）、圖4（c）和圖4（d）的第1 列表示預測圖像，第2 列表示預測圖像和其對應的原始圖像的殘差圖像）。加入塊金字塔記憶模塊后，平均AUC 提升了14.2%，因為通過塊金字塔記憶模塊存儲正常圖像的典型特征，過濾掉大部分異常特征，避免了異常特征通過skip connection 傳入解碼器用于重建，減少了重建異常的情況，增大了異常圖像和其重建圖像之間的差異性，同時還提升了重建清晰度，如圖4（d）所示，重建圖像清晰，明顯改善重建異常問題。加入SSIM 損失，平均AUC 提升了0.4%，提升了網絡的重建效果。

圖4 不同算法的檢測結果對比Fig.4 Comparison of detection results of different algorithms

4 結束語

本文提出一種基于塊金字塔記憶模塊的無監督異常檢測方法。針對重建模糊和重建異常等問題，設計塊金字塔記憶模塊用來存儲不同尺度的正樣本信息，在進行異常檢測時，該模塊將異常特征轉化為正常特征，抑制異常重建，從而通過重建誤差來檢測異常，并且該模塊能夠聚合與輸入特征相似的不同尺度的記憶項，更好地抑制較大面積差的異常重建，提升細膩度重建能力。利用skip connection 結構將經過塊金字塔記憶模塊增強后的不同尺度的特征圖送入解碼器能夠提高重建的清晰度，并采用SSIM 損失函數增強原始圖像和重建圖像的差異性。在MVTec AD 數據集上的實驗結果表明，本文方法相較于其他主流方法，在平均ROC 指標上的得分有明顯提升，減少了異常重建情況和提升了重建清晰度。下一步將針對塊金字塔記憶模塊多個尺度的增強特征圖，探索更加有效的融合方式來提升模型的檢測精度，同時優化模型結構，加快推理速度。