Towards Robust Blind Face Restoration with Codebook Lookup Transformer
Anonymous Author(s)
AffiliationAddress
Abstract
Blind face restoration is a highly ill-posed problem that often requires explicitauxiliary guidance to improve the mapping from a degraded input to a desiredrestored output.
In this paper, we demonstrate that the uncertainty and ambiguity ofthe mapping can be largely reduced by casting face restoration as a code predictiontask in a small, finite proxy feature space.
Under this paradigm, we propose a
Transformer-based prediction network, named CodeFormer , to exploit global contexts of the input for code prediction , enabling the discovery of a natural face thatclosely approximates the target high-quality image even when the input is severelydegraded.
To further enhance identity preservation, we propose a controllable feature transformation module to control information flow from the input image.
Sucha design allows a flexible trade-off between fidelity and quality so that one couldreduce the reliance on the input image in case of heavy degradation.
Equipped withthe proposed components, our CodeFormer suggests superior robustness againstdegradation, outperforming state of the arts in both quality and fidelity.
Extensiveanalysis is also conducted to verify the effectiveness of our method.
Introduction
Face images captured in the wild often suffer from various degradation, such as compression, blur, andnoise.
Restoring such images is highly ill-posed as the information loss induced by the degradationleads to infinite plausible high-quality (HQ) outputs given a low-quality (LQ) input.
The ill-posednessis further elevated in blind restoration, where the specific degradation is unknown.
Despite theprogress made with the emerge of deep learning, learning a LQ-HQ mapping without additionalguidance in the huge image space is still intractable, leading to the suboptimal restoration quality ofearlier approaches.
To improve output quality, auxiliary information that 1) reduces the uncertaintyof LQ-HQ mapping and 2) complements high-quality details is indispensable.
Various priors have been used to mitigate the ill-posedness of the problem, including geometricpriors [4, 5, 27, 40], reference priors [21–23] and generative priors [2, 34, 39].
While improvedtextures and details are observed, these approaches often suffer from limited expressiveness and highsensitivity to degradation .
Such priors provide insufficient guidance for restoration, and the networksessentially resort to the information in the input LQ images, which is usually highly corrupted.
Asa result, the output quality is deteriorated by the degradation of the input image, and the mappinguncertainty still exists.
Different from the aforementioned approaches, this work casts face restoration as a code predictiontask in a small discrete space, and exploits the learnable and discrete codebook prior to improve
E n c od e r 820…
[Figure]
n
[Figure]
r a t o r
+
[Figure]
E n c od e r
T r a n s f o r m e r 012…N -
[Figure]
t
[Figure]
r
[Figure]
feature expressiveness as well as the robustness against degradation.
Such an approach has beenextensively deployed in the task of image synthesis [3, 9, 31], but its adoption in restoration remainsnon-trivial.
As depicted in Fig.
1(b)(right), even when the encoder is trained with degradation, thefeatures extracted from the corrupted images are not well clustered, and hence nearest-neighborfeature matching is unrealiable in such case.
In this study, we propose a robust Transformer-based code prediction network, named CodeFormer ,to exploit global compositions and long-range dependencies for enhanced predictions.
Specifically,taking the input features, a Transformer module is employed to predict the code sequence as a discreteimage representation in the learned codebook space that perceptually approximates to the imagespace, as shown in Fig. 1(h).
Furthermore, we introduce a controllable feature transformation modulewith an adjustable coefficient to control the relative importance of the input.
Such a design allows aflexible trade-off between fidelity and quality , where one could reduce the reliance on the input imageto produce high-quality outputs in case of heavy degradation.
Our approach is different from existing works in the following aspects: (1) Most existing works [9, 31]adopting a vector-quantized codebook employs nearest-neighbor matching, which is infeasible whenthe input image is severely corrupted.
Tailored for restoration, we propose a Transformer-basedprediction module to incorporate global information.
Comparing the results depicted in Fig.
1(f)-(g),the proposed CodeFormer is able to recover more details than the nearest-neighbor matching, suchas the glasses and the eyebrows.
In contrast to generative priors [9, 34, 39] (Fig. 1(a)), thecombinations of codebook items form a discrete representation space with only finite cardinality.
Through mapping the input image to a much smaller proxy space, the uncertainty of the restoration isattenuated, promoting robustness against degradation.
With reduced uncertainty, the features possessgreater expressiveness, reducing the reliance on the input image and also improving the robustnessagainst degradation, as shown in Fig. 1(c)-(i).
(3) Unlike the hand-crafted reference priors [21–23],our dictionary ( i.e. , codebook) automatically learns optimal elements for face restoration, providingadditional flexibility and circumventing the laborious dictionary design.
Equipped with our components, our proposed CodeFormer demonstrates superior performance inexisting datasets and also our newly introduced WIDER-Test dataset that consists of 970 severely degraded faces collected from the WIDER-Face dataset [38].
In addition to face restoration, ourmethod also demonstrates its effectiveness on other challenging tasks such as face inpainting, wherelong-range clues from other regions are required.
Systematic studies and experiments are conductedto demonstrate the merits of our method over previous works.
Related Work
Blind Face Restoration.
Existing blind face restoration methods can be categorized into 1) geometricprior-based, 2) reference prior-based, and 3) generative prior-based methods.
Since face is highlystructured, geometric priors of faces are exploited for blind face restoration.
Some methods introducefacial landmarks [5], face parsing maps [4, 27, 37], facial component heatmaps [40] or 3D shapes [26,14] in their designs.
However, such prior information cannot be accurately acquired from degradedfaces.
Moreover, geometric priors are unable to provide rich details for high-quality face restoration.
Reference-based approaches have been proposed to circumvent the above limitations.
These methodsgenerally require the references to possess same identity as the input degraded face.
For example, Liet al.
[23] propose a guided face restoration network that consists of a warping subnetwork and areconstruction subnetwork, and a high-quality guided image of the same identity as input is used forbetter restoring the facial details.
However, such references are not always available.
To alleviate thisissue, DFDNet [21] pre-constructs dictionaries composed of high-quality facial component features.
However, the component-specific dictionary features are insufficient for high-quality face restoration,especially for regions out of the dictionary scope ( e.g. , skin, hair).
Moreover, the large size ( ∼ 3.2G)of the dictionary limits its efficiency.
The recent pre-trained GANs have provided a new insight for blind face restoration.
The generativepriors are obtained by projecting the input image into a latent feature space.
Different strategiessuch as iterative optimization [24] and multiple-code optimization [10] are proposed for effective
GAN inversion.
However, preserving high fidelity of the restored faces is challenging when onemaps the degraded faces into the low-dimensional latent space.
To relieve this issue, GLEAN [2],
GPEN [39], and GFPGAN [34] embed the generative prior into encoder-decoder network structures,with additional structural information as guidance.
Despite the improvement of fidelity, these methodsstill cannot bypass the intrinsic limitations of generative priors for faithful restoration.
Dictionary Learning.
Sparse dictionary learning in traditional methods aims to learn an overcomplete dictionary as a set of basis vectors such that input signals can be expressed as sparselinear combinations of tokens in the dictionary.
Sparse representation with learned dictionaries hasdemonstrated its superiority in image restoration tasks, such as super-resolution [11, 29, 30, 36]and denoising [8].
Their high-level insight of learning a dictionary to store HQ texture clues hasinspired deep learning-based image restoration [16] and synthesis methods [9, 31].
Jo and Kim [16]construct a look-up table (LUT) by transferring the network output values to a LUT, so that only asimple value retrieval is needed during inference.
However, storing HQ textures in the image domainusually requires a huge LUT, limiting its practicality.
Different from LUT, VQVAE [31] automaticallylearns a highly compressed codebook via a vector quantized autoencoder model.
VQGAN [9] furtheradopts the adversarial loss and perceptual loss to maintain perceptual quality at a high compressionrate, significantly reducing the codebook size without sacrificing its expressiveness.
Inspired by thelearnable discrete codebook, this paper investigates a discrete space for blind face restoration, and wetackle the challenges specific in restoration.
Methodology
The main focus of this work is to exploit a discrete representation space that reduces the uncertaintyof restoration mapping and complements high-quality details for the degraded inputs.
Since localtextures and details are lost and corrupted in low-quality inputs, we employ a Transformer modulethat considers global composition of natural faces for accurate code prediction, enabling high-qualityrestoration.
The overall framework is illustrated in Fig.
CNNDecoder ! !
CNNEncoder 012…N
Codebook " " !
CNNEncoder " # ! ! " ! "
"#$ " "#% " & !
$ $ !(#,%)
=argmin ! !
∈( $ )(#,%) −- *
Nearest-NeighborMatching 10 42 8 1 52 12 0 28 16 23 11 49 4 87 1(a) (b)
CNNDecoder ! ! !
T r a n s f o r m e r
CodePrediction 012…N
Codebook " # … … 108 #! $
CFT $ # $ ' %$ ' ./ ×" & ! "
# " $ #" $
ControllableFeatureTransformation (CFT)
Fixedmodules + + "
Multiplication Addition "+ .
̂.
We first incorporate the idea of quantized vector [9, 31] and pre-train an autoencoder throughreconstruction to obtain a discrete codebook and the corresponding decoder (Sec. 3.1).
The priorcaptured by the codebook and decoder is then used for restoration.
Based on this prior, we thenemploy a Transformer for accurate code prediction, through modeling the global face compositionfrom the low-quality inputs (Sec. 3.2).
In addition, a controllable feature transformation module isintroduced for further enhancing fidelity without compromising quality (Sec. 3.3).
The training ofour method is divided into three stages accordingly.
3.1 Codebook Learning (Stage I)
To reduce uncertainty of the LQ-HQ mapping and complement high-quality details for restoration,we employ a learnable and discrete codebook prior, which improves feature expressiveness as well asrobustness against degradation.
As shown in Fig.
2(a), the HQ face image I h ∈
R H
× W × 3 is first embeded as a compressed feature
Z h ∈ R m × n × d by an encoder E H .
Following VQVAE [31] and VQGAN [9], we replace each “pixel”in Z h with the nearest item in the codebook C = { c k ∈ R d } Nk =0 to obtain the quantized feature
Z c ∈ R m × n × d and the corresponding code entries sequence s ∈ { 0 , · · · , N − 1 } m × n :
[Equation]
The decoder D H then reconstructs the high-quality face image I rec given Z c .
Thus, a new latentdiscrete representation s in form of a m · n code token sequence that specify the respective codeindex in the learned codebook, i.e., Z ( i,j ) c
= c k when s ( i,j ) = k .
Training Objectives.
We pre-train a quantized autoencoder to learn a context-rich codebook usingthree image-level losses: L1 loss L 1 , perceptual loss [17, 41]
L per , and adversarial loss [9] L adv :
[Equation] [Equation]
where Φ denotes the feature extractor of VGG19 [28].
Since, image-level losses are underconstrainedwhen updating the codebook items, we also adopt the intermediate code-level loss [31, 9]
L featcode toreduce the distance between codebook C and input embeddings Z h :
[Equation]
where sg ( · ) stands for the stop-gradient operator and β = 0 .
25 is a weight trade-off for the updaterates of the encoder and codebook.
Since the quantization operation in Eq.
(1) is non-differentiable,we adopt straight-through gradient estimator [9, 31] to copy the gradients from the decoder to theencoder.
The complete objective of codebook learning L codebook is:
[Equation]
where λ adv is set to 0.8 in our experiments.
Codebook Settings.
Our encoder E H and decoder D H consist of 12 residual blocks and five resizelayers for downsampling and upsampling, respectively.
We obtain a large compression ratio ofr
= H/n = W/m = 32 .
As we will show later, this compression ratio leads to a great robustnessagainst degradation and a tractable computational cost.
Although more codebook items could easereconstruction, the redundant elements could cause ambiguity in subsequent code predictions.
Hence,we set the number of elements in the codebook to N = 1024 , which is sufficient for accurate facereconstruction.
Besides, the code dimension? d is set to 256.
3.2 Codebook Lookup Transformer Learning (Stage II)
Due to corruptions of textures in LQ faces, the nearest-neighbour (NN) matching in Eq.
(1) usuallyfails in finding accurate codes for face restoration.
As depicted in Fig. 1(b), LQ features with diversedegradation could deviate from the correct code and be grouped into nearby clusters, resultingin undesirable restoration results, as shown in Fig. 1(f).
To mitigate the problem, we employ a
Transformer to model global interrelations for better code prediction.
Built upon the learned model presented in Sec. 3.1, as shown in Fig.
2(b), we append a Transformer [33] module containing nine self-attention blocks to the encoder.
We fix the codebook C anddecoder D H and finetune the encoder E H for restoration.
The finetuned encoder is denoted as E L .
To obtain the LQ embeddings Z l ∈ R m × n × d through E L , we first unfold the embeddings into m · nvectors Z vl ∈ R ( m · n ) × d .
The s -th self-attention block of Transformer computes the as following:
[Equation]
where X
= Z vl .
The queries Q , keys K , and values V are obtained from X s through linearlayers.
We add a sinusoidal positional embedding [1, 6] P ∈ R ( m · n ) ×
d on the queries Q and thekeys K to increase its expressiveness of modeling sequential representation of codebook.
Finally,taking the encoding feature Z vl as an input, the Transformer predicts the m · n code sequenceˆ s ∈ { 0 , · · · , | N | − 1 } m · n in form of the probability of the N -way classification.
The predicted codesequence then retrieves the m · n respective code items from the learned codebook, forming thequantized feature ˆ Z c ∈ R m × n × d that produces a high-quality face image through the fixed decoder
D H .
Thanks to our large compression ratio ( i.e. , 32), our Transformer is effective and efficient inmodeling global relationships of LQ images.
Training Objectives.
We finetune the CNN encoder E L and Transformer blocks for restoration, whilethe codebook and decoder are kept fixed.
Instead of employing reconstruction loss and adversarialloss in the image-level, only code-level losses are required in this stage: 1) cross-entropy loss L tokencodefor code token prediction supervision, and 2) L2 loss L feat ′ code on code features to force LQ embedding
Z l to approach the HQ embedding Z c , which eases the difficulty of token prediction learning:
[Equation]
where the ground truth of latent code s i is obtained from the pre-trained autoencoder in Stage I(Sec. 3.1), thus the HQ embedding Z c is then retrieved according to the s i .
The final objective of
Transformer learning is:
[Equation]
where λ token is set to 0.5 in our method.
Note that the trained model in this stage has alreadydemonstrated robustness and effectiveness in face restoration, especially for severely degraded faces.
3.3 Controllable Feature Transformation (Stage III)
Despite our Stage II model has maintained a good image fidelity, we also investigate a flexibletradeoff between fidelity and quality for face restoration.
Thus, we propose the controllable featuretransformation (CFT) module to control information flow from the LQ encoder E L , which enhancesthe fidelity without compromising the restoration quality.
Specifically, as shown in Fig.
2, the LQ features F e are used to slightly tune the decoder features F d through feature transformation [35] withthe predicted affine parameters of α and β .
An adjustable coefficient w ∈ [0 , 1] is then used to controlthe relative importance of the input:
[Equation]
We adopt this module at multiple scales s ∈ { 32 , 64 , 128 , 256 } between encoder and decoder.
Sucha design allows a flexible trade-off between fidelity and quality, where one could use a small wto reduce the reliance on the input image, thus producing high-quality outputs in case of heavydegradation.
On the contrary, a large w will introduce more input information to enhance the fidelityin case of mild degradation.
Unless otherwise stated, we set the w = 0 by default in our experiments.
Training Objectives.
To train these controllable modules, we keep the code-level losses of L tf in Stage II, and also add image-level losses of L 1 , L per , and L adv , which are the same as that in Stage Iexcept that I rec is replaced by I res , the restoration output.
The complete loss for this stage is the sumof the above losses weighted with their original weight factors.
Experiments
4.1 Datasets
Training Dataset.
We train our models on the FFHQ dataset [19], which contains 70,000 high-quality(HQ) images, and all images are resized to 512 × 512 for training.
To form training pairs, we synthesize
LQ images I l from the HQ counterparts
I h with the following degradation model [21, 34, 39]:
[Equation]
[Equation]
) where the HQ image I h is first convolved with a Gaussian kernel k σ , followed by a downsampling ofscale r .
After that, additive Gaussian noise n δ is added to the images, and then JPEG compressionwith quality factor q is applied.
Finally, the LQ images are resize back to 512 × 512 .
We randomlysample α , r , δ , and q from [1 , 15] , [1 , 30] , [0 , 20] , and [30 , 90] , respectively.
Testing Datasets.
We evaluate our method on a synthetic dataset CelebA-Test and three real-worlddatasets: LFW-Test, WebPhoto-Test, and our proposed WIDER-Test.
CelebA-Test contains 3,000images selected from the CelebA-HQ dataset [18], and the LQ images are synthesized under thesame degradation range as our training settings.
The three real-world datasets respectively containthree different degrees of degradation, i.e. , mild (LFW-Test), medium (WebPhoto-Test), and heavy(WIDER-Test).
LFW-Test consists of the first image of each person whose name starting with A in
LFW dataset [15], containing 431 images.
WebPhoto-Test [34] consists of 407 low-quality facescollected from the Internet.
Our WIDER-Test comprises 970 severely degraded face images from the
WIDER Face dataset [38], providing a more challenging dataset to evaluate the generalizability ofblind face restoration methods.
4.2 Experimental Settings and Metrics
Settings.
We represent an input image of 512 × 512 × 3 as 16 × 16 code sequences.
For all stages oftraining, we use the Adam [20] optimizer with a batch size of 16.
We set the learning rate to 8 × 10 − 5for stages I and II, and adopt a smaller learning rate of 2 × 10 − 5 for stage III.
The three stages aretrained with 1.5M, 200K, and 20K iterations, respectively.
Our method is implemented with the
PyTorch framework and trained using four NVIDIA Tesla V100 GPUs.
Metrics.
For the evaluation on real-world datasets, we employ the widely-used non-referenceperceptual metrics: FID [13] and NIQE [25].
For the evaluation on CelebA-Test with ground truth,we adopt PSNR, SSIM, and LPIPS [41] as our metrics.
We also evaluate the identity preservationusing the cosine similarity of features from ArcFace network [7], denoted as IDS.
4.3 Comparisons with State-of-the-art Methods
We compare our approach with current state-of-the-art methods, including PULSE [24], DFDNet [21],
PSFRGAN [4], GLEAN [2], GFP-GAN [34], and GPEN [39].
We conduct extensive comparisons onsynthetic and real-world datasets.
More results are shown in the supplementary material.
[Figure]
[Figure]
Evaluation on Synthetic Dataset.
We first show the quantitative comparison on the CelebA-Testin Table 1.
In terms of the image quality metrics LPIPS, FID, and NIQE, CodeFormer achievesthe best scores than existing methods.
Besides, it also faithfully preserves the identities, reflectedby the highest IDS score and PSNR.
Additionally, we present the qualitative comparison in Fig.
The compared methods fail to produce pleasant restoration results.
For example, DFDNet [21],
PSFRGAN [4], GFP-GAN [34], and GPEN [39] introduce artifacts and GLEAN [2] produces oversmoothed results.
Moreover, all other methods are unable to preserve the identity.
Thanks to the effective code prediction in the discrete and finite space, CodeFormer not only produces high-qualityfaces but also preserves the identifies well, even when the input faces are highly degraded.
[Table] [Table]
Evaluation on Real-world Datasets.
As presented in Table 2, our CodeFormer achieves comparableperceptual quality with the compared methods on the real-world testing datasets with mild andmedium degradations and the best scores on the testing dataset with heavy degradations.
Although
PULSE [24] also obtains good perceptual scores, it cannot preserve the identity of input images, as the FID scores and visual results respectively suggested in Table 1 and Fig. 1(d).
From the visualcomparisons in Fig.
4, we see that our method generalizes well to real-world face restoration andproduces most visually pleasing results.
Notably, CodeFormer successfully preserves the identity andproduces natural results with rich details.
4.4 Ablation Studies
[Figure]
Effectiveness of Codebook Space.
We first investigate the effectiveness of the codebook space.
Asshown in Exp.
(a) of Table 3, removing the codebook ( i.e. , directly feeding the encoder features tothe decoder) results in worse LPIPS and IDS scores.
The results suggest that the discrete space ofcodebook is the key to guarantee the robustness of our model.
Superiority of Transformer-based Codebook Prediction.
To verify the superiority ofour Transformer-based codebook prediction, we compare it with two different solutions,
[Figure]
i.e. , nearest-neighbour (NN) matching and a CNNbased code prediction module, denoted as Code
Pred., ( i.e. , encoder E L followed by a Linear prediction layer).
As shown in Table 3, the comparison of Exp.
(b) and Exp.
(c) indicates that adoptinga code prediction module is more effective than optimizing the features for NN-matching.
However,the local nature of convolution operation of CNNsrestricts its modeling capability for long code sequence prediction.
In comparison to the CNNbased method, i.e. , Exp.
(c), our Transformerbased solution produces higher-fidelity results.
The superiority of our Transformer-based moduleis also demonstrated in Fig.
5, Fig. 6, and Fig. 8.
Importance of Fixed Decoder.
Unlike DFDNet [21], we adopt a compact codebook C ∈ R N × dwith N =1024 , d =256 .
We hypothesize that the efficacy of our method is brought by the codebookcoupled with the pre-trained decoder, instead of the codebook alone.
To verify our hypothesis, weconduct two studies: 1) fixing both the codebook and decoder, i.e. , Exp.
(g), and 2) fine-tuning thedecoder, i.e. , Exp.
(e), in the training process of stage II.
As shown in Table 3, fine-tuning decoderdeteriorates the performance.
This is because fine-tuning the decoder breaks the tight connectionsbetween the codebook and the decoder, resulting in suboptimal performance.
Therefore, we keep thedecoder fixed throughout training.
Flexibility of Controllable Feature Transformation Module.
Considering the diverse degradationsin real-world face images, we provide a controllable featrue transformation module (CFT) to allow aflexible trade-off between fidelity and quality.
As shown in Fig.
7, a smaller w tends to produce ahigh-quality output while a larger w improves the fidelity of output.
Such an appealing flexibilityimproves the adaptiveness of our method for different scenarios, but is rarely touched in previousworks.
In addition, as shown in Table 3, Exp.
(f), i.e. , setting the coefficient w to 1 increases theidentity score but decreases the quality score.
In this work, we prioritize quality, and hence the
[Figure]
=!
= 0. = 0.
= 0. = 0.
! =
[Figure]
coefficient w is set to 0 by default.
Note that one could adjust w without retraining to suit therequirement of the tasks.
4.5 Extension
[Figure]
The proposed Codeformer can be easily extended to face inpainting, and it shows great performanceeven in large mask ratios.
To build training pairs, we use a publicly available script [39] to randomlydraw irregular polyline masks for generating masked faces.
We compare our method with twostate-of-the-art face inpainting methods of CTSDG [12] and GPEN [39], as well as Nearest-Neighbor(NN) matching for code retrieval.
As shown in Fig.
8, CTSDG and GPEN show less effectiveness formasked faces with the large masks.
The NN matching based method reconstructs the face structures,but it also fails in restoring complete visual parts such as the glasses and the eyes.
In contrast,
CodeFormer generates high-quality faces without strokes and artifacts.
[Figure] [Figure]
Our method is built on a pre-trained codebook and decoder.
Thus, their capability may affect theperformance of our method.
For example, though the identity issue is significantly relieved by the
Transformer’s global dependency modeling, it still exists in some rare visual parts such as accessories,where the current codebook space cannot seamlessly replace the image space.
Using multiple scalesin the codebook space to explore more fine-grained visual representations may be a solution.
Conclusion
This paper addresses the fundamental challenges of generating high-quality faces when inputs areseverely corrupted.
Through learning a small discrete codebook space, we turn face restoration todiscrete code prediction, significantly reducing the uncertainty of restoration mapping and easing thelearning of restoration.
We explore global interaction from degraded inputs and model the code composition of face images via an expressive Transformer for accurate code prediction.
Benefiting fromthese designs, our method shows strong robustness against heavy degradation.
To enhance fidelitywithout compromising the restoration quality, we also propose a controllable feature transformationmodule that allows a flexible trade-off control between fidelity and quality.
Experimental resultssuggest the superiority and robustness of our method.