CLR-Face: Conditional Latent Refinement for Blind Face Restoration Using Score-Based Diffusion Models

1 Introduction

Blind face restoration (BFR) is a challenging problem that aims to recover a high-quality facial image from a low-quality input. With the increasing availability of degraded images captured in various real-world scenarios, the demand for effective BFR techniques has increased significantly in recent years. The primary goal of BFR is to recover undegraded facial features while preserving the person's identity. A

complex combination of many factors, such as low resolution, blur, noise, compressions, etc., corrupts a facial image. Typically, traditional approaches [Baker and Kanade, 2000; Bourlai et al., 2011] rely on a degradation model and manually designed priors, which often lead to sub-optimal results and restricted ability to handle a variety of real-life images. Recently, generative priors [Wang et al., 2021; Yang et al., 2021] have shown remarkable improvement in restoration performance. Generative priors encapsulated in a welltrained high-quality face generator (e.g., StyleGAN [Karras et al., 2019]) are typically exploited in such approaches. The degraded image is first projected to the latent space to get a 'cleaner' latent vector either directly or after certain refinement operations. The decoder projects the latent vector back to the pixel space, producing a restored image. Instead of dealing with a continuous latent space, [Gu et al., 2022; Zhou et al., 2022] utilizes pretrained vector-quantized codebooks. The mechanism of vector quantization in a compressed latent space reduces the uncertainty of the task, making these methods robust to various degradations. Intuitively, these approaches exploit information from two sources when restoring an image. The degraded input image contains crucial information about the person's identity. On the other hand, the pretrained decoder is utilized for a high-quality generation. Thus, ultimately, the task boils down to predicting the corresponding clean latent given a degraded image/latent. Using a highly compressed latent space usually reduces the complexity of the prediction task and prioritizes high-quality generation but fails to maintain intricate facial details essential for identity preservation. It may also lead to repeated and less diverse texture patterns in the output. On the other hand, using a less compressed latent space can be potentially more expressive and flexible but is often extremely difficult to optimize when directly projecting the degraded image to the latent space or trying to predict the correct latent code as a classification problem. A lower compression often diminishes the advantages of latent-space-based generative models, leading to suboptimal degradation removal from the input. Similar behavior was recently discussed in [Gu et al., 2022], where the trade-off between fidelity and quality depending on the compression factor was analyzed. We argue that the performance of such face restoration approaches mainly depends on the nature and size of the latent space. It is non-trivial to balance the demands of face generation and restoration using existing latent refinement techniques. [Gu et al., 2022; Zhou et al., 2022] deploy a highly spatially-compressed latent space (downsampled by 32) and instead use skip-connections to improve the fidelity, which is often counter-productive for the output quality, as the information from the encoder is usually corrupted.

In this work, we design an alternative approach, where we first select a highly expressive latent space with lower compression and to handle its additional complexity, we introduce diffusion-based-prior to model it. We learn the clean latent space distribution by directly modeling the gradient of the log density of the data, known as the score function [Song and Ermon, 2019], using a neural network. Such gradient information can be utilized in reverse by stochastic sampling to generate diverse samples from the underlying distribution. Unlike unconditional generation as in [Song et al., 2020], for the current task, we model it as a conditional generation task, where given degraded latent, we perform multi-step refinement to produce the cleaner counterpart.

Our approach has several key advantages. 1. Similar to prior state-of-the-art (SOTA) blind face restoration works [Zhou et al., 2022], we also utilize a pre-trained VQVAE framework to produce high-quality face images given correct latent codes. 2. The significantly lower compression ration allows us to maintain higher fidelity, suitable for BFR task. 3. Compared to the single-step refinement strategy of prior arts, multi-step conditional refinement of the corrupted latent using score-based diffusion prior is far more effective in producing the clean embedding, enabling us to use relatively mild compression rates, achieving more faithful and sharper reconstructions. 4. Compared to pixel-space-based diffusion models used in [Saharia et al., 2022] (DDPM), our score-based latent-diffusion model has much better time complexity and requires significantly fewer refinement steps (≈ 20×).

The diffusion prior is extremely powerful for producing a clean latent. However, during multiple refinements, deviating from the trajectory of correct identity is plausible. To further constraint the refinement process, we deploy a separate Identity Recovery Network (IRN), which prioritizes recovering identity-specific intricacies given the input image using a face recognition loss. We utilize the output of the IRN to initialize the iterative refinement process. Next, at each step, we utilize a gradient-based score-update strategy using a pre-trained face recognition network to steer the refinement process along the identity recovered by the IRN. While this approach considerably diminishes the risk of identity distortion, the guidance signal from a face-recognition network may not align well with the other objective of producing visually pleasing results. Such a dilemma is observed in existing works as well [Wang et al., 2021], where usually a small amount of identity loss is introduced while prioritizing standard perceptual losses, which ultimately fails to recover the identity information adequately. Instead of blindly enforcing the identity constraint across the entire latent space, we introduce a learnable latent mask, where the gradient from the recognition network is harnessed to selectively update specific latent locations that contribute most significantly to identity information. This enables us to balance the restored identity and the perceptual quality derived from the diffusion prior.

To summarize, our main contributions are

1. We propose a diffusion-prior-based conditional latent refinement strategy for recovering the clean latent inside a pretrained VQVAE framework.
2. The strong modeling capability of such iterative refinement allows us to design a latent space with a lower spatialcompression ratio (f=4), improving the overall fidelity and sharpness of the output.
3. We further deploy an IRN to recover identity-specific features from the degraded input. The output of IRN, along with adaptive guidance using a learnable latent mask, steers the refinement process for better identity recovery without compromising the perceptual quality.
4. Extensive evaluation on multiple real and synthetic datasets demonstrates our proposed framework's superior perceptual and identity-preserving properties.

2 RelatedWorks

Various priors have been proposed, including geometric priors such as facial landmarks [Kim et al., 2019; Chen et al., 2018; Zhu et al., 2016] and parsing maps [Chen et al., 2021; Shen et al., 2018], as well as reference-based approaches [Li et al., 2020; Dogan et al., 2019]. However, these priors often struggle to estimate necessary information from corrupted images. Recent generative priors optimize the latent vector for GAN inversion techniques [Menon et al., 2020; Gu et al., 2020] or direct projection of the input image to the latent space [Richardson et al., 2021]. [Yang et al., 2021], and [Wang et al., 2021] exploited the generative prior inside an encoder-decoder, with structural details from the degraded input through skip connections. For severe degradations, such links lead to unwanted artifacts.

Diffusion and score-based models have shown improvements recently [Saharia et al., 2022; Whang et al., 2022; Choi et al., 2021; Yue and Loy, 2022; Chung et al., 2022]. Iterative refinement strategies has been adopted by [Saharia et al., 2022], [Whang et al., 2022] for super-resolution and motion deblurring tasks. [Choi et al., 2021] guided the reverse process with low-frequency information from a conditional image. However, such a conditioning strategy does not translate well for the BFR problem with significant degradation and may alter the identity. [Yue and Loy, 2022] uses an unconditional diffusion model and starts from an intermediate stage of the reverse diffusion process. But, as the underlying diffusion model is unconditional, the identity of the restored face changes considerably. If it is used for a smaller timespan, the visual quality and sharpness of the output suffer considerably. [Chung et al., 2022] addressed only non-blind super-resolution tasks, and its identity-preserving capability is yet to be tested for more difficult BFR scenarios. Our work mainly aims to balance identity preservation and perceptual quality for heavily degraded faces.

3 Method

Our framework has two main parts, a Vector-Quantized (VQ) autoencoder is our backbone that maps a degraded image to the latent space with a downsampling factor f = 4 and reverses this operation after refining the latent. A conditional

(a) Inference

Figure 1: An overview of our inference (left) and training framework (right).

diffusion model (denoising UNet) is used for the latent refinement task. We pre-train the quantized autoencoder through self-reconstruction to obtain a discrete codebook and the corresponding decoder for projecting the restored latent back to pixel space. In addition, we also deploy IRN to produce an initial estimate of the restored face, focusing on recovering identity-specific features, which are used to steer the iterative latent refinement process. An overview of our framework is shown in Fig. 1. In the following sections, we describe the details of different parts of our framework.

3.1 Learning Latent Space of Images via Vector-Quantized Codebook

We use x and z to represent image in pixel space and the corresponding latent representation, respectively. Our perceptual compression model utilizes a variant of the VectorQuantize codebook called VQGAN [Esser et al., 2021]. It learns a perceptual codebook through a combination of perceptual loss [Johnson et al., 2016] and adversarial training [Isola et al., 2017], preventing blurriness that can occur when relying solely on pixel-space L2 or L1 losses. The model consists of an encoder E, a decoder G, and a codebook Z = zkKk=1 ∈ RK×d containing a finite set of embedding vectors. Here, K represents the codebook size, and d denotes the dimension of the codes. Given an input image x ∈ R3×H×W , we project it to the latent space as z = E(x) ∈ Rd× Hf × Wf . Next, a spatial collection of image tokens zq is obtained using a spatial-wise quantizer Q(·), which maps each spatial feature zij into its closest codebook entry zk

The decoder reconstructs the image from the latent, giving ˜x = G(zq). The encoder E, the decoder G and the codebook Z are trained end-to-end via the following loss function

where sg(·) stands for the stop-gradient operator, ϕ(·) represents the VGG-based [Simonyan and Zisserman, 2014] feature extractor. The first three terms help reconstruct the input image, whereas the last two terms reduce the distance between codebook Z and input feature embeddings z. As we build upon the foundation of continuous space diffusion models, we absorb the quantization layer in the decoder. Given z, we use the conditional diffusion model to recover the correct latent embedding through iterative refinement. Next, quantization and the pre-trained high-quality codebook can handle any inaccuracies in the latent prediction for producing a clean, restored face. Compared to existing BFR works that use VQVAE [Gu et al., 2022; Zhou et al., 2022], our work has two critical differences: we use a milder compression (downsample f = 4 instead 32) and a larger codebook size (K = 8192 instead of 1024). Such an expressive yet highly complex latent space can be optimally modeled by iterative diffusionprior-based refinement. As a result, the restored image is of higher fidelity and realness than prior art.

3.2 Latent Refinement Using Diffusion Prior

Given the high dimensional latent with a lower compression ratio, existing methods encounter challenges in generating a cleaner latent through a single pass in the encoder or a classification-based refinement module, as elaborated in [Gu et al., 2022]. Instead, we aim to learn the probability distribution of the clean latent space (through gradients of log probability density functions) and map the degraded latent to the clean space in multiple refinement steps. Our approach is built upon the framework of [Song and Ermon, 2019], where stochastic differential equations drive the underlying diffusion process. We can represent this process using continuous time variables zt t=1 t=0 , where z0 represents the initial variable (i.e., clean latent representation) and zt represents its noisy version at time t. The diffusion process is defined by an Itˆo SDE [Rogers and Williams, 2000]

where f(·, ·) and g(·) are the drift and diffusion coefficients, respectively, and wt denotes the standard Brownian motion.

We denote the distribution of zt as p(zt). The SDE in Eq. 3 can be converted to a generative model by first sampling from z1 ∼ N(z1, 0, I) and then running the reverse-time SDE as

where ¯wt is a reverse-time standard Wiener process and dt is an infinitesimal negative time step. As shown in [Song et al., 2020], there exists a deterministic ordinary differential equation (ODE), whose trajectories share the same marginal probability densities as the SDE (Eqs. 3, 4)

In practice, blackbox ODE solvers can be used on Eq. 5 for faster and high quality sampling, assuming we know the score function. The score function (∇zt log p(zt)) is typically approximated using a neural network sθ(zt, t) [Vincent, 2011; Song and Ermon, 2019; Ho et al., 2020], trained to minimize the weighted Fisher's divergence with a positive weighting function λ as

Conditional Diffusion for BFR

Instead of unconditionally sampling from the underlying clean latent distribution using Eq. 5, for the task of BFR, our goal is to predict the clean counterpart, given a degraded latent zd of the input image xd . We are interested in p(z0|zd ), where z0 is the target clean latent distribution corresponding to the image xd and zd is the condition signal. Similar to Eq. 5, we can obtain a conditional reverse-time ODE as

We need to learn the conditional score function ∇zt log p(zt|zd ) in order to be able to sample from Eq. 7 using reverse-time diffusion. In [Song et al., 2020], ∇zt log p(zt|zd ) was estimated using

where ∇zt log p(zt) is learned using an unconditional model and ∇zt log p(zd |zt) is learned using a separate auxiliary network. But, modeling such a forward process is nontrivial for BFR, as there could be many possible zd given a zt. Instead, following the intuition of [Saharia et al., 2022; Tashiro et al., 2021], we can extend the formulation of Eq. 6 to a conditional case as

where we pass the condition zd directly to the denoising UNet, making it learn the conditional score. We empirically verify that it works well for BFR task. In [Song et al., 2020; Kadkhodaie and Simoncelli, 2020], the noise was added to both the target and the conditional observations, which we found to be harmful. We follow the simplified configuration of [Karras et al., 2022], where we set f(zt, t) = 0 and g(t) = √2t. Then, p(zt) can be expressed as p(z0) ∗ N(0, t2 I), where ∗ represents convolution operation and p(z0) is the underlying distribution of the clean or non-degraded images' latent space, which we aim to model. Once we know the score function (using a neural network), Eq. 7 can be re-written as

During the reverse process, we can proceed backward in time with a black box ODE solver, such as Heun solvers [Karras et al., 2022], to obtain the solution trajectory. The resulting ˆz0 can be considered an approximate embedding of the clean image corresponding to xd . ˆz0 is passed through the pre-trained decoder to produce the restored image. The decoder has the quantization layer, which handles minor errors in the continuous space prediction.

3.3 Identity Recovery Network (IRN)

Existing works typically introduce an identity-preserving loss to the same restoration network in addition to standard perceptual losses. But, often, minimizing both losses together hampers the perceptual quality and may introduce unwanted artifacts, especially if the coefficient for identity loss is high. Instead, to ease the learning process, we aim to disentangle these two objectives and use a separate IRN (fϕ) to focus on recovering the identity. We train fϕ, emphasizing the identity-preserving loss, with a small amount of L1 loss for better stability. The loss function for training fϕ is as follows:

where Dcos denotes the cosine distance between two feature vectors and farc denotes a pre-trained ArcFace model. We denote the second term in Eq. 11 as Larc. The IRN aims to generate a stable approximation xid that primarily recovers the identity information rather than focusing on visually pleasing sharp results. Such crucial identity features are also transferred to the latent embedding zid . To ensure that during reverse diffusion process, the output does not deviate from the already recovered identity in xid , we add a gradient term ∇ˆz0,t Larc to the estimated score at each step, where the Larc is calculated between the image obtained after projecting ˆz0,t to pixel-space using the pretrained decoder and xid , ˆz0,t is the denoised output at step t. It follows the intuition of classifier guidance in [Dhariwal and Nichol, 2021]. But, we have observed that simply adding the gradient for the whole latent again deteriorates the perceptual quality. We hypothesize that the latent embedding should have specific features that correlate to identity-specific details in the pixel space. We should ideally regularize the score function of only those features while keeping the rest untouched.

To verify this, we introduce a learnable latent mask Ml ∈ Rd× Hf × Wf , denoting the identity-specific latent features. We use a stack of a few convolutional layers followed by a sigmoid layer as fM, that takes the ˆzct as input and produces Ml. To train fM , we first update the score to obtain the corresponding ˆz0,t and project it back to pixel space. We use LPIPS, cosine-distance-based ID loss on the output image, and a sparsity constraint on Ml to ensure Ml selects only the subset of identity-specific features without harming the LPIPS score. As discussed in our ablation study, the observed improvement after introducing the latent mask validates our design choice. For all the cases, straight-through gradient estimator [Esser et al., 2021] is utilized to copy the gradients from the decoder to the continuous latent embedding to handle the non-differentiable quantization operation.

4 Experimental Results

4.1 Training Dataset

The FFHQ dataset [Karras et al., 2019] contains 70,000 highquality face images at 1024 × 1024 resolution. We resized the images to 512 × 512 for training. We synthesized degraded images on the FFHQ dataset using the degradation model proposed in [Yang et al., 2021; Wang et al., 2021; Wang et al., 2022; Chen et al., 2021]: xd = ((x ⊗ k) ↓s +nσ)q. Here, x, xd , k, nσ, s, and q are the clean face image, the corresponding degraded image, the blur kernel, the Gaussian noise with standard deviation σ, downscaling factor, and the JPEG-compression quality factor, respectively. We randomly and uniformly sampled σ, s, and q from [0,20], [1,32], and [30,90], respectively. We evaluated our approach on 3,000 synthetic images from CelebA-Test dataset [Karras et al., 2017]. The degraded images in this dataset were synthesized using the same degradation range as our training. Additionally, we tested our method on real-world datasets: WebPhoto-Test [Wang et al., 2021] (407 images), WIDER Face [Yang et al., 2016] (970 images), Celeb-Child [Wang et al., 2021] (180 images) and TURB. For the TURB dataset, we randomly sampled 139 images from the BRIAR [Cornett et al., 2023] and LRFID [Miller et al., 2019] datasets, which provided a more challenging scenario as our models were not trained on severe turbulence-affected images.

4.2 EvaluationMetrics

To assess the quality of the restored images quantitatively, we primarily rely on the Frechet Inception Distances (FID) [Heusel et al., 2017] and Learned Perceptual Image Patch Similarity (LPIPS) [Zhang et al., 2018] metrics. We also compute PSNR and SSIM for completeness, although they often fail to capture visual quality. To evaluate the face recognition performance of the restored images, we calculate the cosine similarity between the features of the restored image and the paired GT image (IDS). Higher cosine similarity val- Input GFPGAN

Input

Table 3: Quantitative comparisons of FID (↓) on real-world datasets in terms of FID. Bold and underline indicate the best and the second best performance, respectively.

ues indicate better recovery of the identity information. We use the same evaluation metrics and pre-trained models for LPIPS, FID, and ArcFace as used in the prior art [Wang et al., 2022; Wang et al., 2021]. We also adopt landmark distance (LMD) as the fidelity metric, following [Gu et al., 2022].

4.3 Comparisons with State-of-the-Art Methods

We compare with the following state-of-the-art (SOTA) methods: PSFRGAN [Chen et al., 2021], GPEN [Yang et al., 2021], GFPGAN [Wang et al., 2021], CodeFormer [Zhou et al., 2022], RestoreFormer [Wang et al., 2022], DifFace [Yue and Loy, 2022] and VQFR [Gu et al., 2022]. We use the official results and checkpoints provided by the authors.

Synthetic BFR: We reported restoration accuracy on the synthetic CelebA-Test dataset for BFR task in Table 1. Our approach achieves a much better balance between fidelity (IDS, LDM) and perceptual metrics (LPIPS, FID).

Extreme Upscaling: We tested the algorithms under extreme BFR conditions, where we applied a fixed downscaling factor of ×32 to 512×512 images, resulting in degraded images of size 16 × 16. To make the task more challenging, we added noise and blur to the images. The quantitative results can be found in Table 2. Despite the limited information in the input, our approach outperforms other algorithms in terms of perceptual quality and recognition accuracy.

Real-World Cases: Our approach outperforms existing GAN-based methods like GPEN and GFPGAN in terms of producing more realistic and faithful reconstructions on realworld degraded datasets. CodeForm, VQFR's outputs show visible artifacts or repetitive skin/hair texture due to highly compressed latent space. DifFace outputs are less sharp and may alter facial details as the underlying model is unconditional. In contrast, our approach achieves superior results with fewer artifacts, even for low to medium degradation levels. Real-world colorization examples are shown in Fig. 5, for which we finetune our network for colorization. We also report the accuracy for the downstream face-recognition task on the TURB dataset in Table 5. Our approach comfortably outperforms prior arts for recognition, as well.

5 Ablation Analysis

In Table 4, we analyze the effect of individual components of our approach on the perceptual quality and identitypreserving aspect. We use CelebA-Test with 512 × 512 images for our ablation. We empirically found that the IRN is IP

VQVAE-Rest GD (t=1000)

IRN

Ours GT

Figure 6: (Left): Reconstruction performance of VQVAE. (Right): Restoration performance of VQVAE, GD, IRN and our final model using latent space refinement.

model agnostic, and any SOTA restoration network results in a comparable accuracy boost. We provided detailed comparisons of possible options for IRN in the supplementary. We finally select SwinIR [Liang et al., 2021] that keeps a good balance between performance and accuracy. We design our VQ backbone with the following settings: f=4, codebook = 8192, latent = 3×128×128, feat = 128 (×1, ×2, ×3). For the diffusion model, we have used 40 time steps with heun sampler [Karras et al., 2022] and a UNet-based denoiser based on GD 1 . More implementation and network details, additional results are provided in the supplementary material.

Compression Factor: First, we compared VQVAE models with different latent sizes: f = 32 and f = 4. The f = 32 models, commonly used in previous works [Gu et al., 2022], resulted in poor reconstruction quality due to significant spatial compression to 16 × 16. In contrast, the f = 4 models achieved superior reconstruction quality, as shown in Fig. 6 (Left). We have verified that it is satisfactorily accurate in reconstructing degraded images as well (PSNR=37.9). Higher spatial compression led to the loss of fine facial details and overly smooth outputs. For example, in the first row, the eye color was changed, and in the second row, both the eye and hair patterns were altered in the f = 32 models.

Latent Refinement: We attempted to fine-tune the VQVAE model (with f = 4) directly for restoration, but it often failed to remove degradation effectively, as reported in [Gu et al., 2022]. To address this issue, we developed a modified version called VQVAE-Rest. This variant takes both the degraded image and the coarse estimate from the IRN as inputs, enabling it to focus on recovering residual details. However, even with this modification, the output of VQVAE-Rest still exhibited lower quality and noticeable artifacts (Figure 6 (Right), Column 2). These results show the need for appropriate latent-space refinement to achieve satisfactory restoration outcomes. We also compare with Guided-Diffusion (GD) for pixel-space-based diffusion (DDPM), which is run for 1000 timesteps in Figure 6 (Right). As can be observed, it is not only computationally intensive due to large no. of time steps, but also fails to reconstruct the image properly starting from pure noise. The outputs of the IRN shown in Figure 6 (Right), Column 4, may lack sharp details. However, it recovers crucial facial locations like eye, nose, etc.

ID Loss: Table 4 demonstrates the significance of the ArcFace [Deng et al., 2019] based identity-preserving loss for the IRN (IRNID vs. IRN), which prioritizes the recovery of identity-specific features, which are valuable for regularizing the diffusion process.

Adaptive Mask: While the diffusion model alone can achieve good perceptual quality, it may inadvertently alter certain identity-specific facial features. Adding IRNID -based gradient improves the IDS, but is detrimental for perceptual metrics. The final output could be slightly different from the output of IRN as we regularize only a subset of the latent in the diffusion stage. Our final model, combining the strengths of the Diff. Prior, IRNID and the learnable latent mask Ml achieves superior accuracy both in terms of perceptual metrics and fidelity. Inference time(s) for GFPGAN, CodeFormer, DifFace, Guided-Diffusion and ours are: 0.5, 0.15, 6,30, 0.61. Our design is still much faster than existing diffusion-based BFR methods, and can potentially be further accelerated by reducing the number of steps using knowledge-distillation techniques [Song et al., 2023]. We also have comparable or less parameters (16M) compared to VQFR (76.3M), DifFace (17.5), RestoreFormer (12M).

1 GD: https://github.com/openai/guided-diffusion

6 Conclusions

We propose a iterative latent-space refinement technique using diffusion prior for restoring severely degraded face images, achieving a better balance between restoration quality and identity recovery compared to existing methods. However, the performance of our approach is limited by the IRN's ability to preserve identity information. To improve efficiency, we plan to explore the development of a recognitionmodel-free reverse algorithm in future research.

Acknowledgements

This research is based upon work supported in part by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via [2022-21102100005]. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of ODNI, IARPA, or the U. S. Government. The US. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.