基于掩碼自動編碼器的圖像修復研究

收稿日期：2023-07-24

DOI：10.19850/j.cnki.2096-4706.2024.03.015

摘" 要：掩碼圖像建模（MIM）因為在視覺表示方面具有巨大潛力而備受關注。現有的使用簡單像素重構損失的MIM方法生成質量不高，輸出模糊，針對這個不足，提出了基于掩碼自動編碼器的圖像生成和自監督表示學習框架。對掩碼圖像建模研究的關鍵點是，該模型在輸入和輸出時使用VQGAN學習到的語義標記，并將其與掩碼相結合，增加對比損失函數和噪聲損失函數，以實現生成和表示學習的雙重目標。首先使用對比損失函數來塑造圖像樣本的嵌入空間，以促進有意義的表示學習。同時，利用噪聲損失函數來鼓勵模型重建圖像中的高頻成分，從而提高生成能力。這種綜合的方法使得掩碼自動編碼器成為一個強大且高效的模型，同時具備生成高質量圖像和學習有用的圖像表示的能力。

關鍵詞：掩碼；自動編碼器；VQGAN；修復效果

中圖分類號：TP391.4；TP181" 文獻標識碼：A" 文章編號：2096-4706（2024）03-0069-05

Research on Image Restoration Based on Mask Autoencoder

LUO Di， ZHANG Qian， BAI Wuer

（School of Data Science and Information Engineering， Guizhou Minzu University， Guiyang" 550025， China）

Abstract： Masked Image Modeling （MIM） has received significant attention due to its tremendous potential in visual representation. Existing MIM methods that use simple pixel to reconstruct loss suffer from generating low-quality image， blurry outputs. To address this shortcoming， a framework for image generation and self-supervised representation learning based on mask autoencoder is proposed. The key point of research on modeling masked images is that the model uses semantic labels learned by VQGAN in both input and output， and combines them with masks to add contrast loss functions and noise loss functions to achieve the dual goals of generation and representation learning. Firstly， use the contrast loss function to shape the embedding space of image samples to promote meaningful representation learning. At the same time， using the noise loss function to encourage the model to reconstruct high-frequency components in the image， thereby improving the generation ability. This comprehensive approach makes mask autoencoder a powerful and efficient model， while also possessing the ability to generate high-quality images and learn useful image representations.

Keywords： mask; autoencoder; VQGAN; restoration effect

0" 引" 言

近年來，基于掩碼圖像建模（MIM）的方法表明，將輸入圖像的部分區域進行掩碼，并通過預測這些掩碼區域來重構原始圖像，表示學習框架可以獲得高質量的表示。然而，現有的使用簡單像素重構損失的MIM方法生成質量不高，輸出模糊。例如，作為這類方法的代表，MAE的重構質量較差，細節和紋理丟失。許多其他MIM方法也存在類似問題，本論文的重點是彌合這一差距，提出了一個既能生成逼真圖像又能從圖像中提取高質量表示的框架。除了在預訓練過程中使用可變的掩碼比例外，與先前的MIM方法不同，本文模型的輸入和重構目標都是語義令牌，而不是像素。這種設計改進了生成和表示學習，克服了上述問題。對于生成，在令牌空間中操作不僅使模型能夠迭代地執行圖像生成任務，而且使模型能夠學習掩碼令牌的概率分布，而不是所有可能掩碼像素的平均值，從而產生多樣化的生成結果（見圖4）。對于表示學習，使用令牌作為輸入和輸出使網絡能夠在高語義級別上操作，同時不丟失低級細節，從而比現有的MIM方法具有顯著更高的線性探測性能。

對比學習和掩碼自編碼器（MAE）采用了非常不同的學習機制：前者訓練編碼器對語義保持數據變化具有不變性，而MAE學習空間統計相關性。此外，MAE方法在損失函數中獨立處理每個樣本，而對比方法則明確考慮批次中所有樣本之間的關系，通過減小或增大嵌入距離。鑒于此，本文假設這兩種方法是互補的，可以為給定的輸入提取不同的判別特征。如果這個假設成立，預計基于提取的特征將在各種下游任務中表現出改進的性能。這激發了本文對一個結合方法的探索。

此外，受擴散模型的進展啟發，本文引入了第三個損失，該損失基于在掩碼自編碼器重構過程中的噪聲預測。向未掩碼的輸入補充高斯噪聲，并訓練模型預測每個補丁中添加的噪聲。去噪鼓勵編碼器從輸入中提取高頻信息，而自編碼器的重構傾向于關注低頻信息。這個額外的損失有兩個目的：它提高了下游性能，并解決了MAE中的一個計算浪費問題，對FLOPS幾乎沒有影響：未掩碼補丁的重構被丟棄且未使用。

1" 相關工作

1.1" 使用Vision Transformers的掩碼圖像模型

Vision Transformer（ViT）作為一種基于Transformer架構的圖像分類模型，將圖像分割成一系列的圖像塊（或稱為圖像補丁），然后使用Transformer的自注意力機制對這些圖像塊進行建模。ViT的出現引發了對使用ViT作為主干網絡的強大自監督學習框架的研究興趣。一些先前針對ConvNet主干網絡開發的技術，例如DINO和MoCo-v3，在適當調整后也在ViT上展現出了競爭力。隨后，針對ViT的特定方法出現，其中一種方法是基于掩碼圖像建模，受到了自然語言處理中預訓練方法（例如BERT）的啟發。

基于掩碼圖像建模的方法在圖像中引入掩碼，并要求模型預測被掩碼隱藏的部分，從而促使模型學習對圖像的語義和視覺特征進行建模。

其中，MAE（masked auto encoder）是一種經典的基于掩碼自編碼器的方法，它將圖像的部分區域進行掩碼，并通過自編碼器的重構過程來預測被掩碼隱藏的區域。MAE方法在ViT上的應用表明，即使不需要將掩碼令牌通過編碼器傳遞，也可以通過這種方式在ViT上進行自監督預訓練，從而提高了效率。

綜上所述，自監督學習中的基于掩碼圖像建模的方法為擴展到未篩選數據集的大規模圖像任務提供了一種有效值得的解決方案。MAE展示了經典的掩碼自編碼方法可以用于在ViT上進行預訓練，而無須將掩碼令牌通過編碼器傳遞。這大大提高了效率；本文的方法類似地利用了這一點。

1.2" 噪聲損失

去噪自編碼器（Denoising Auto encoders， DAE）學習在給定噪聲輸入的情況下重構清晰數據。通過學習將低密度數據區域映射到高密度區域，DAE學習了數據流形的形狀。其中DAE學習之后的得分函數為s（x） = ?x log p（x）。這個關鍵觀察為生成擴散模型的重要進展奠定了基礎，這些模型利用得分函數的估計來生成樣本。然而，DAE在表示學習方面的最近成功尚未得到廣泛應用，除了一些例外情況。在本研究中，利用去噪自編碼器來消除掩碼圖像建模中重構未掩碼補丁但從未使用的MAE的低效問題。

1.3" 對比損失

在計算機視覺領域，對比學習是自監督學習的一種重要方法，它提供了一種無須監督的方式來提取通用特征。特別是通過強制執行對增強操作的不變性并利用負樣本，在球面上均勻分布嵌入，對比學習已經取得了最先進的性能。對比預訓練任務在概念上與基于掩碼的圖像模型（如MAE）有很大區別，后者學習了空間統計依賴關系。另一個區別是，自編碼器鼓勵在潛在表示中保留信息，而對比學習可能抑制特征。這使我們推測這兩種方法學習了不同的數據特征，因此可能是互補的學習機制。這激發了將對比學習和基于掩碼圖像建模相結合的動機，以開發一種獲得兩者優點的強化預訓練任務。

2" 預訓練

標記化：首先使用標記器將輸入圖像標記，為一系列語義標記。標記器使用與VQGAN模型，中的第一階段相同的設置。這個標記化步驟允許我們的模型對語義標記而不是原始像素進行操作，這對生成和表示學習都有好處，如圖1所示。

掩碼策略：為了進一步彌合生成建模和表征學習之間的差距，采用一種可變掩碼比的掩碼策略。首先從中心為0.55的截斷高斯分布中隨機采樣掩碼比mr，左截斷0.5，右截斷1。如果令牌輸入序列的長度為l，則隨機屏蔽掉令牌mr · l，并用一個可學習的掩碼令牌[M]代替。由于mr≥0.5，所以從這些掩碼令牌中隨機丟棄0.5l個令牌。丟棄大部分被掩碼的令牌顯著減少了預訓練時間和內存消耗，同時有助于生成和表示性能。

編解碼器設計：在掩碼和刪除輸入令牌之后，VIT編碼器將經過掩碼和丟棄的令牌序列作為輸入，并將其編碼到潛在特征空間中。在解碼之前，首先使用編碼器學習的類令牌特性[C]將編碼器的輸出填充到完整的輸入長度。類標記位置可以概括輸入圖像的全局特征。因此，不使用跨不同圖像共享的可學習掩蔽令牌，而是使用特定于每個圖像的[C]來填充編碼器輸出。與使用掩蔽令牌相比，該設計改進了生成和表示學習性能，然后解碼器獲取填充的特征來重建原始令牌。

重構訓練：假設" 代表從分詞器獲的潛在標記，其中N為標記序列長度， 為確定哪些令牌要被掩碼的對應二進制掩碼。訓練目標是重構來自未掩碼標記的掩碼標記。因此，在解碼器的輸出和基于真實標記的one-hot向量之間添加交叉熵損失。具體而言，將重構的標記與真實的one-hot向量比較，計算交叉熵損失。具體的公式如下：

其中，YM為Y中未被掩碼的標記子集，p（ yi |YM）為在未被掩碼的標記條件下，編碼器-解碼器網絡預測的概率。在MAE的方法中，只在掩碼標記上優化此損失（在所有標記上優化損失會降低生成和表示學習性能）。

對比損失：在這種方法中，編碼器返回的嵌入 ， ∈ RT×d通過沿第一維度進行簡單平均匯總，形成d維嵌入。然后，這些嵌入通過一個輕量的多層感知機（MLP）投影頭傳遞，將其映射到一個低維空間Rr（其中r＜d）。接著，這些嵌入被單位化以產生嵌入 ，i = 1，…，n。在每個批次項i中，從該批次中選擇其他2n-2個樣本，即 ，作為負樣本。然后，使用選定的正樣本" 以及選定的負樣本來計算InfoNCE損失。InfoNCE損失是一種常用的學習對數線性模型的損失函數，用于學習將正樣本與負樣本分開的嵌入。它的公式如下：

其中，τ為一個溫度參數，它平衡了正負樣本之間的差異。利用InfoNCE損失優化嵌入可以促進相似項在嵌入空間中更接近，不相似項在嵌入空間中更遠。這有助于學習更好的表示，并在下游任務中獲得更好的性能。

噪聲損失：去噪訓練目標的擴散建模和等價的基于分數的方法在自監督學習中取得的重要進展。首先對每個圖像添加獨立的各向同性高斯噪聲，使得" 變為 ，其中 ， 均勻采樣于區間[0，σmax]。這個帶噪聲的輸入被掩碼并傳遞給編碼器，當將編碼后的patch傳遞給解碼器時，向解碼器提供關于噪聲水平" 的信息，以幫助它將噪聲與真實圖像分離。這是受到去噪擴散方法的啟發，該方法將噪聲圖像和噪聲水平作為輸入傳遞給去噪模型。通過將" 作為位置編碼器在解碼器中進行處理來實現這一點。首先，生成" 的正弦嵌入，經過輕量級MLP處理得到可學習的嵌入 ，其維度與" 的潛在維度相匹配。最后將結果添加到每個嵌入標記上（包括缺失的標記[M]），以提供噪聲水平信息：，其中t = 1，…，T，然后將結果傳遞給解碼器生成 。其中去噪損失函數，僅在未被掩碼的像素上計算。具體的去噪損失函數為：

其中，?為 ，原始圖像" 和解碼器生成的重建圖像" 之間的差異乘以掩碼信息 。在這個公式中 ，t的取值為0或1，表示圖像的特定patch區域是否被掩碼。具體來說，對于差異圖像的第t個patch中的每個像素位置，乘法運算將根據對應的" 的取值，對差異進行逐元素的乘法操作。如果" 為1，表示該像素是未掩碼的，那么對應位置的差異值將保留。如果" 為0，表示該像素是掩碼的，那么對應位置的差異值將被置為零。這個公式的目的是計算僅在未掩碼的補丁上進行損失計算，將掩碼的區域排除在外，以避免在損失計算中考慮這些區域的影響。通過這種方式，可以專注于對可見部分的差異進行建模和優化，而忽略掩蔽部分的影響。需要注意的是，這個去噪損失非常輕量級，由于MLP引入的額外開銷非常小。噪聲補丁的重建不會增加任何額外成本，因為解碼器會生成所有補丁的重建結果，包括掩蔽和未掩蔽的補丁，即使Lrec只使用掩蔽補丁的重建結果。最后，在擴散建模文獻中經常觀察到，雖然訓練去噪模型估計噪聲與估計干凈輸入本身是等價的，但兩者之間存在很大的經驗差距，噪聲預測的效果更好。

總的損失函數：整個訓練目標由編碼器和解碼器優化三個損失組合：

Lall = λInfoceNCE LInfoceNCE + λreconstrctive Lreconstrctive + λdenoise Ldenoise

其中，0≤λInfoceNCE LInfoceNCE，λreconstrctive Lreconstrctive，λdenoise Ldenoise并且λInfoceNCE + λreconstrctive + λdenoise = 1權衡了這些目標，在實踐中，使用等式約束來消除一個變量，通過以下方式參數化權重：λreconstrctive = （1 - λInfoceNCE） · λ和λdenoise = （1 - λInfoceNCE） · （1 - λ），其中0≤λ≤1，這種參數化方式使得很容易控制重建損失Lreconstrctive、去噪損失Ldenoise與對比損失LInfoceNCE之間的相對權重。

3" 訓練結果分析

通過訓練得到的損失函數圖像如圖2所示，因為ImageNet100數據集較大，設備有限，所以在A100上訓練了80個epoch，學習率如圖3所示，將訓練的80個epoch的修復效果圖如圖4所示，為了和MAE的修復效果進行比較，其中使用的掩碼比也是0.75，可以明顯地看出MAE的修復效果有大面積的紋理缺失，和原圖像進行比較有很大的失真，通過本文提出的模型修復的圖像基本上保留了原圖像的紋理細節，但也存在個別的紋理細節上的差異，但總的來說，修復效果還是比較客觀的。

表1在ImageNet 256×256上與最先進的生成模型進行定量比較，以實現類無條件生成。參數的數量包括編碼器、解碼器和去令牌化器。

通過實驗結果可以看出，本文的模型可以很好地修復掩碼的圖像，在掩碼比例達到0.75的情況下，完全看不出原始圖像是什么，但通過模型可以很好地還原出來，雖然并沒有與原圖像做到完全的相同，但在整體上還是達到很好的效果，在圖像分類領域可以起到很好的作用。

圖2" 訓練損失函數圖像

圖3" 學習率函數圖像

圖4" 效果展示圖

表1" 生成模型進行定量比較表

方法 Model Acc. #params

MAE ViT-B 68.0 86 M

CAE [ ViT-B 70.4 86 M

CMAE ViT-B 73.9 86 M

Our ViT-B 74.7 110 M

4" 結" 論

通過訓練80個epoch的掩碼自動編碼器模型，在ImageNet100數據集上取得了良好的修復效果。與傳統的MAE（均方誤差）修復方法相比，本文提出的模型在修復圖像時能夠更好地保留原圖像的紋理細節，減少紋理缺失和失真。該掩碼圖像建模方法在視覺表示方面具有巨大潛力，是一個強大且高效的模型，既能夠生成高質量圖像，又能夠學習有用的圖像表示。然而，也需要注意，雖然修復效果整體上比較客觀，但在個別紋理細節上可能還存在差異。

參考文獻：

[1] KINGMA D P，SALIMANS T，POOLE B，et al. Variational Diffusion Models [J/OL].arXiv：2107.00630[cs.LG].[2023-05-26].https：//arxiv.org/abs/2107.00630v5.

[2] HOU X X，SHEN L L，SUN K，et al. Deep feature consistent variational autoen-coder [C]//2017 IEEE Winter Conference on Applications of Computer Vision （WACV）.Santa Rosa：IEEE，2017：1133–1141.

[3] HE K M，CHEN X L，XIE S N，et al. Masked Auto encoders Are Scalable Vision Learners [C]//2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition （CVPR）.New Orleans：IEEE，2022：15979-15988.

[4] DONG X Y，BAO J M，ZHANG T，et al. Peco：Perceptual Codebook for BERT Pre-Training of Vision Transformers [J/OL].arXiv：2111.12710 [cs.CV].[2023-06-26].https：//arxiv.org/abs/2111.12710.

[5] HENDRYCKS D，BASART S，MU N，et al. The many faces of robustness：A critical analysis of out-of-distribution generalization [C]//2021 IEEE/CVF International Conference on Computer Vision（ICCV）.Montreal：IEEE，2021：8320-8329.

[6] CHEN X K，DING M Y，WANG X D，et al. Context Auto encoder for Self-Supervised Representation Learning [J/OL].arXiv：2202.03026 [cs.CV].[2023-06-28].https：//arxiv.org/abs/2202.03026v2.

[7] VINCENT P，LAROCHELLE H，LAJOIE I，et al. Stacked Denoising Auto encoders：Learning Useful Representations in a Deep Networ kwith a Local Denoising Criterion [J/OL].Journal of machine learning research，2010：3371-3408.[2023-06-28].https：//www.jmlr.org/papers/volume11/vincent10a/vincent10a.pdf？ref=https：//githubhelp.com.

[8] JIA C，YANG Y F，XIA Y，et al. Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision [J/OL].arXiv：2102.05918 [cs.CV].[2023-06-28].https：//arxiv.org/abs/2102.05918.

[9] ZHAO L，ZHANG Z Z，CHEN T，et al. Improved transformer for high-resolutiongans [J/OL].[2023-06-28].https：//www.zhuanzhi.ai/paper/860c355383ec194784fd9edcb57db5c1.

[10] WEI C，FAN H Q，XIE S N，et al. Masked Feature Predic-tion for Self-Supervised Visual Pre-Training [EB/OL].[2023-06-28].https：//openaccess.thecvf.com/content/CVPR2022/papers/Wei_Masked_Feature_Prediction_for_Self-Supervised_Visual_Pre-Training_CVPR_2022_paper.pdf.

[11] ROMBACH R，BLATTMANN A，LORENZ D，et al. High-Resolution Image synthesis with Latent Diffusion Models [C]//2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition（CVPR）.New Orleans：IEEE，2022：10674–10685.

作者簡介：駱迪（1999—），女，仡佬族，貴州道真人，碩士在讀，主要研究方向：圖像處理；張乾（1984—），男，苗族，貴州貴定人，教授，博士，主要研究方向：機器學習、模式識別、計算機視覺；柏武貳（1995—），男，布依族，貴州貞豐人，碩士在讀，研究方向：圖像處理。