人脸修复论文写作学习——OSDFace: One-Step Diffusion Model for Face Restoration

论文标题：OSDFace: One-Step Diffusion Model for Face Restoration

作者：Jingkai Wang1, Jue Gong1, Lin Zhang1, Zheng Chen1，Xing Liu2, Hong Gu2, Yutong Liu1†, Yulun Zhang1†, Xiaokang Yang1

单位：1 Shanghai Jiao Tong University, 2 vivo Mobile Communication Co., Ltd

摘要：扩散模型在人脸修复任务中已展现出卓越的性能。然而，其多步推理过程计算开销较大，限制了其在实际场景中的应用。此外，现有方法往往难以生成和谐、逼真且与主体身份一致的人脸图像。在本研究中，我们提出了 OSDFace，一种新颖的单步扩散模型用于人脸修复。具体而言，我们设计了一个视觉表征嵌入器（Visual Representation Embedder, VRE），以更好地捕获先验信息并理解输入人脸。在 VRE 中，低质量人脸首先由视觉分词器（visual tokenizer）处理，随后通过向量量化字典（vector-quantized dictionary）嵌入，生成视觉提示（visual prompts）。此外，我们引入了一种基于人脸识别的面部身份损失（facial identity loss），以进一步确保身份一致性。我们还采用生成对抗网络（GAN）作为引导模型，促使修复后的人脸分布与真实图像对齐。实验结果表明，OSDFace在视觉质量和定量指标上均超越了当前最先进的（SOTA）方法，能够生成高保真、自然且身份一致性高的人脸图像。代码和模型将发布于： https://github.com/jkwang28/OSDFace 。

笔者主要将依据该篇论文来学习目前在人脸修复这个方向上的一些论文的写作方式，当然重心还是在于把这篇论文的要点进行分析。

1 Introduction

“Face restoration aims to restore high-quality (HQ) face images from degraded low-quality (LQ) inputs caused by complex degradation processes. Blur, noise, downsampling, and JPEG compression are common degradation types that can significantly affect the quality of face images, leading to an ill-posed problem. ”

人脸修复问题简单来说就是对HQ的退化路径求逆的过程，现在的大部分模型实际是在研究盲人脸修复问题，即退化因素未知的情况，我们一般认为非盲任务是下游任务，因为在训练时为了拟合复杂的退化路径，会将多种退化过程进行复合来生成HQ对应的LQ输入，以CodeFormer为例，它的LQ生成方法如下：

\[I_{LQ} = [(( I_{HQ} ⊗ k_σ )_{↓s} + n)_{JPEG_q}]_{↑s} \]

其中k表示未知模糊核，n表示随机噪声，s表示上/下采样尺度，⊗表示叉乘。

特别地，在超分辨任务中会使用多阶段退化的方式来做进一步处理，这种方法能更好地拟合现实LR图片退化路径多源、复合的特征，提升模型的泛化能力。这一思路在Restoreformer++中被提出作为一个故事点，叫做拓展退化模型（EDM，extending degrading model），这个模型主要添加了两个要素，一个是对人脸做空间平移来缓解由人脸对齐而带来的误差，另一个是额外引入了雾化效果（haze）和非均匀退化（uneven degradation），然后通过mask将两张图片合成为一张，如下：

\[I_{LQ} = [[(α(I_{HQ}⊗k_σ) + (1 − α)I_{haze})_{↓r} +n_δ]_{JPEG_q}]_{↑r} \]

“Recent advancements focus on leveraging CNN-based methods and Transformer-based methods to restore HQ face images, and have shown remarkable performance. However, due to the severe degradation of face images, generative models with strong prior knowledge exhibit irreplaceable advantages. ”

Diffusion的方法所展现的潜力是CNN和Transformer这类方法所难以达到的，这一点是得到广泛认可的，另外还有一种叫做流方法的生成模型，可以通过可逆的确定性路径做到唯一确定的复原结果，避免GANs的模式崩溃或Diffusion的随机性过强问题。但这个方法的难点在于如何拟合这个确定性路径，还存在一定的困难需要进一步解决。
CNN-based methods and Transformer-based methods作者列出了DAEFR、Restoreformer++、PLTrans和CodeFormer，其中PLTrans算是集百家module之长所弄出的一个Latent Transformer方法的杂烩版本，实际的架构创新方面不多，但效果提升了不少。（有意思的是实验部分提供的是x16、x32、x64的下采样退化实验对比数据，code目前还未给出，我们无从得知这是否是一种效率与效果balance后的结果）

“Generative adversarial networks (GANs) and diffusion models have gained considerable attention for generating high-fidelity images with perceptual details. Many efforts, such as DifFace and DiffBIR, have also been applied to restoring HQ face images. However, GANs are challenging to train and may encounter issues such as mode collapse and training instability. Conversely, diffusion models produce promising results but require multiple forward passes during inference, leading to increased computational cost and longer inference time.
To speed up the diffusion inference, one-step diffusion (OSD) models have emerged as a promising research topic. Benefiting from the recently developed generative diffusion models, especially large-scale pretrained text-to-image (T2I) models, OSD models enjoy the powerful restoration capability and fast inference speed. It is anticipated that OSD models could be highly competitive in face restoration practical applications.”

介绍了本文的主要动机，GAN模型最大的弊端在于模式崩塌和训练不稳定性，而diffusion模型的问题在于推理速度慢，计算开销大。作者认为，OSD模型是一个解决inference速度问题的好方法，大型的T2I模型已经受益于此，将其应用到脸部修复上应当具有很大的竞争力。

“As shown in Fig.1, most diffusion-based face restoration methods can effectively recover basic facial features such as eyes, mouth, and contour. However, details, such as realistic-looking hair and complex backgrounds, often remain unharmonious. The primary issue lies in the insufficient incorporation of face priors. Some diffusion-based methods rely solely on the diffusion model without considering face priors, resulting in unrealistic restorations. Other methods attempt to use face priors, but either limit the generative capabilities or cause information reduction. On the other hand, existing general OSD image restoration models, such as OSEDiff, are not specifically designed for facial features. OSEDiff employs image-to-tag prior and works well for natural images. However, it meets challenges when applied to faces, as humans are highly sensitive to facial features, and even minor inconsistencies or unrealistic details are easily noticeable.”

“Thus, it is worthwhile to investigate how to design a prior capturer and apply it to OSD model. One key issue is that designing priors to capture rich facial features is challenging. Vector-quantized (VQ) prior-based methods, such as CodeFormer and DAEFR, effectively leverage codebooks trained in HQ data as priors for face restoration. These models perform well in generating facial features. However, the face images often lack fine details and display unreal backgrounds because they are generated by the codebook items directly. Furthermore, incorporating priors into diffusion models for face restoration is also a critical focus. Numerous efforts have been undertaken to incorporate face priors. For example, PGDiff employs a face prior by pretraining a VQ-based restorer and using it as a target to guide the diffusion model. Due to limitations imposed by the target model, such a configuration inhibits the generative capabilities. In addition, CLRFace passes denoised latent vectors to a pretrained codebook and a subsequent decoder. Relying on direct image generation with a fixed-size codebook limits the diversity of results.”

第一段介绍了diffusion相关方法的一些不足，1.BFRffusion、DifFace、DiffBIR对面部先验信息的利用不足；2.CLR-Face、PGDiff试图使用面部先验，但要么限制了生成能力，要么导致信息丢失；3.现有的通用 OSD 图像修复模型，如OSEDiff，并非专为面部特征设计。（OSEDiff采用图像到标签（image-to-tag）的先验，在自然图像上效果良好。然而，将其应用于人脸时会遇到挑战，因为人类对面部特征高度敏感，即使是微小的不一致或不真实的细节也极易被察觉。）

第二段主要介绍了diffusion在VQ这个方向上的一些研究背景，VQ类方法往往因为被codebook限制了细节表达与真实背景的生成能力。将VQ先验引入diffusion已经有了一些工作，但这些工作也都受限于VQ框架而在生成能力上缺乏了多样性。
比较有意思的是，笔者认为VQ类方法的卓越之处就在于latent code这种压缩采集信息的方法，而codebook之所以会限制真实背景的生成能力，确实是由于codebook本身无法覆盖实际需要的丰富语义信息，为了解决码本大小限制生成内容的问题，Adacode学习了一个码本生成网络，根据输入动态生成一个临时的专属码本来解决不同图片对于有限码本的需求。尽管因为码本的训练困难，实际利用率低，显存消耗大，但笔者倾向于这不是codebook限制VQ类方法生成能力的原因，如果存在一种方法能够高效利用codebook，并且在长度和维度上展现突破，想必效果是不错的。这里有两个想法：
一个是改变codebook的更新方式。 在VQVAE中，原本的VQ是通过反向传播中的直通估计（Straight-Through Estimator，STE）来解决argmin无法传播梯度的问题的，如果直接让重建损失的梯度同时更新编码器和码本，会出现梯度方向冲突：编码器希望 \(z_e\) ​更接近某个 \(e_k\) 以减小量化误差；码本也希望 \(e_k​\) 更接近 \(z_e\) ；但如果两者同时被同一梯度更新，可能导致震荡或训练崩溃。因此，codebook与encoder的更新方式是通过一个stop-gradient的机制所建立的 reconstruction_loss + codebook_loss + commitment_loss 来更新的。我们不难看到对于codebook来说其只更新了argmin所选取的几个code，这在训练中很容易出现大部分码本空间没有利用的问题，这也是被人诟病的“码本崩溃”。在后续VQVAE2和VQGAN中，一些作者考虑到了码本利用的问题，尝试用指数移动平均EMA来解决unused code的问题，尽管需要多维护计数器和滑动均值两个统计量，但实际也只是缓解了这个问题。Adacode则是将大码本替换为一个码本生成网络中去更新以作为权衡，这样大码本训练不稳定的问题也就得到了解决。最近的VQGAN-LC与FVQ也认为码本崩溃的一个关键原因就是在于稀疏码本梯度 (Sparse Codebook Gradients) ，因此作者想到去使用一个线性投影或一个lightweight的ViT去全局关注码本来进行更新。笔者觉得这种思路是很有意思的，可以管中窥豹，看看codebook的上限究竟能否超过diffsion。
另一个是改变codebook的使用方式，大部分VQ类方法在使用codebook的时候都是argmin的直接选取，在上一点我们也有提到一些方法通过网络来更新整个码本，但其没有改变最终选取code的时候仍然是单个最近邻选取的问题，也就是说这仍然维持了VQ本身discrete latent space的特性。这个方法的一个核心问题在于其仍然采用了STE直接复制梯度的方式，那么编码器（Encoder）的更新优化是不精确的。同时，FVQ的论文作者提出了VQ的另一个问题：一步滞后更新 (One-step-behind Update) ，即码本的更新依赖于上一步的编码器输出，这种时间上的错位导致训练过程不稳定，就像一个总是看着后视镜开车的新手。那么adacode的解决方法其实是一个可行解，用网络去从码本中得到一个可行解，将离散的code转换到连续的空间上。当然，比Adacode还要简单的一个想法是去训练一个回归器；进一步地，可以将这个codebook用于Transformer的交叉注意力中进行embedding的输出。这些都可以是人脸修复方向对VR改进的一个简易解。
但值得注意的是，100%码本使用率是否有意义呢？这约等于将codebook改造成了一个全连通没有死神经元的神经网络，更具体的讨论这里不再展开，之后会单开一篇文章去做讨论。另外一个关于codebook最新的见解是，2D的潜在空间对于图像生成任务来看似乎是没有必要的，这种对应图像位置的约束只会造成码本利用率的降低，TiTok与FVQ都用实验说明了这一点。

“To alleviate the limitations, we propose a novel method, OSDFace, for face restoration. Firstly, OSDFace is an OSD model that leverages the powerful image restoration capabilities of diffusion models, as shown in Fig. 2, while offering a fast inference speed (about 0.1s for a 512×512 image). Secondly, our OSDFace integrates a visual representation embedder (VRE) to extract rich facial information directly from LQ inputs. Our proposed VRE consists of two components: a visual tokenizer and a VQ embedder. The visual tokenizer contains a VAE encoder and a VQ dictionary matching function. Trained on LQ data, it could efficiently capture information from LQ faces. In the same LQ semantic space, the VQ embedder performs a dictionary lookup of feature categories to obtain the visual prompt in O(1) time. Thirdly, to better keep a consistent identity, we incorporate a facial identity loss derived from face recognition into our training process. This approach significantly reduces the distance between the generated faces and the ground truth in the compact deep feature space. Besides, we employ a GAN discriminator as a guidance model, which encourages the distribution alignment between the generated faces and the ground truth. The alignment enables the model suitable for faces in complex environments.”

介绍了OSDFace模型的主要框架与效果，主要是通过一个由视觉分词器和VQ嵌入器构成的VRE与OSD组成，视觉分词器通过传统的VQVAE的Encoder与VQ方法捕获LQ的面部信息，VQ嵌入器则提供了一个特征分类的词典查询功能来得到视觉prompt。最后为了保持身份一致性，还从人脸识别模型中借用了一个脸部身份的loss去参与训练过程，作者认为这种方法显著缩小了在紧凑的深度特征空间中生成人脸与真实人脸之间的距离。同时，一个GAN判别器被用于指导模型来帮助这两种人脸的对齐。

"Our contributions can be summarized as follows.
• We propose OSDFace, a novel and effective one-step diffusion model for face restoration. This is the first attempt to utilize one-step diffusion for restoring faces.
• We design the visual representation embedder (VRE). Using LQ dictionary, VRE captures rich prior from LQ images for a deeper understanding of visual content.
• We customize a comprehensive method for realistic face alignment, incorporating facial identity loss for identity consistency and GAN loss for distribution alignment.
• Our OSDFace achieves significant SOTA face restoration performance, excelling in visual quality and quantitative metrics with reduced computational costs."

贡献总结如下：
提出了 OSDFace，一种新颖且有效的用于人脸修复的单步扩散模型。这是首次尝试利用单步扩散来修复人脸。
设计了视觉表征嵌入器（VRE）。利用 LQ 字典，VRE 能从 LQ 图像中捕获丰富的先验信息，以更深入地理解视觉内容。
定制了一套综合方法来实现逼真的人脸对齐，结合了用于身份一致性的面部身份损失和用于分布对齐的 GAN 损失。
OSDFace 取得了显著的 SOTA 人脸修复性能，在视觉质量和定量指标上均表现出色，同时降低了计算成本。

2 Introduction

2.1. Face Restoration

“Face restoration aims to recover high-quality (HQ) faces from low-quality (LQ) inputs affected by complex degradations. A key challenge for researchers is how to grasp face prior efficiently and effectively. Statistical priors are widely used in traditional image processing techniques. Recently, with the development of deep learning, researchers have begun to explore more advanced learning-based methods.
Geometric priors, while useful, are limited by the degree of degradation in low-quality images, restricting their ability to recover high-quality details. Reference priors aim to supplement missing information using the corresponding high-quality reference images. However, obtaining references that match the target data is often impractical. Generative priors from pretrained models are widely used for face restoration. These priors are typically incorporated into encoder-decoder networks through iterative latent optimization or direct latent encoding to enhance fidelity. This approach has demonstrated stronger restoration potential compared to earlier methods. Recently, diffusion priors have gained attention with advancements in diffusion models. The degraded images are generated from intermediate outputs at each diffusion iteration. In parallel, VQ prior-based methods have shown great promise, achieving impressive results in face restoration. They typically reconstruct images by matching features extracted from LQ inputs to items in a codebook, and then directly output the restored images. However, the limited capacity of VQ priors often leads to blurry backgrounds and imprecise details. Therefore, how to leverage face prior information more effectively remains a crucial problem.”

2.2. Diffusion Models
“Diffusion models, known for powerful generative capacity, transform random noise into structured data through iterative denoising processes. Moreover, diffusion models have achieved promising results in image restoration recently, showing strong performance in image-to-image tasks using efficient guidance strategies. However, some models possess highly complex structures. Most of them require numerous iterative steps, limiting their popularization for real-world applications.
Reducing the inference steps of diffusion models is critical to accelerating generation speeds and reducing computational costs. Yet, further reduction often results in a significant performance drop, making it crucial to balance inference speed with model capability. Most one-step diffusion (OSD) methods employ distillation to learn from a teacher model, ensuring the quality of the generated images. D3SR achieves promising results with OSD, free from the limitations of model distillation. However, all models are designed primarily for natural image restoration and lack strong generalization for face-specific tasks. It is essential to incorporate prior knowledge specific to human faces to enhance the applicability of OSD models in face restoration. Consequently, the models could generate more realistic and high-fidelity face images.”

**2.1节，作者将面部修复的叙述与引用重心放在了先验上，作者认为通用的三种先验（几何先验，参照先验和生成先验）中都是通过阶段性潜空间优化或是直接的潜变量编码来融入encoder-decoder网络中的，而扩散先验作为扩散模型发展而生的一种新先验，退化图像可以认为是扩散过程的中间产物，VQ先验因其有限的表达能力导致模糊背景与不准确的细节表达，如何去利用脸部先验的信息仍然是一个关键问题；2.2节，扩散模型的性能十分强大，但一些模型的结构过于复杂且迭代步数较多，限制了其可能的广泛应用。OSD主要是通过蒸馏保留了教师模型的能力，尽管有大量的模型已经在自然图像修复任务中得到应用，但鲜有在人脸任务上专用的模型。

这里重新介绍一下人脸修复领域的工作线：

基于几何先验的深度修复方法：利用人脸的几何和空间分布信息（如关键点、热力图、解析图）。代表工作如FSRNet（端到端估计几何先验）、Super-FAN（联合进行超分和关键点定位）。
基于参考先验的深度修复方法：利用额外的高质量参考图像（如同一身份的正面高清图）或从高质量图像中构建的人脸组件字典来指导修复。代表工作如GFRNet、GWAInet（使用引导图像）、DFDNet（使用组件字典）。
基于生成先验的深度修复方法：利用预训练人脸GAN（如StyleGAN）的潜在空间作为丰富的人脸先验。代表工作如PULSE（优化潜在码）、mGANprior（多潜在码）、GFP-GAN、GPEN（将预训练GAN作为解码器，平衡真实感和保真度）。
非先验的深度修复方法：不依赖额外先验，直接学习从低质量到高质量的映射。早期工作如BCCNN、CBN。后续方法引入了多路径结构、注意力机制、对抗学习、强骨干网络（如ViT, Diffusion）等技术来提升性能。
**

3 Method

“We aim to extract as many features as possible from low-quality (LQ) source images. These features are used as prompts to guide the diffusion model in generating realistic faces that closely resemble the individuals in the original images. To accomplish this, we propose an innovative one-step diffusion model focused on face restoration. This model enhances the realism of the generated faces while preserving the characteristics of LQ images.”

这里作者认为，要想恢复出来的HQ图像足够逼真，首先要尽可能地从LQ中提取特征，然后用这些特征作为prompt引导扩散模型来生成接近原始图像的身份信息。因此很自然地，正文部分从两方面去叙述这个故事。

3.1. One-Step Diffusion (OSD) Model

Latent diffusion models are structured around both forward and reverse operations. During the forward diffusion phase, Gaussian noise with a variance \(\beta_t∈(0, 1)\) is incrementally added to the latent vector z at each timestep, resulting in \(z_t=\sqrt{\hat{\alpha}_t}z + \sqrt{(1 − \hat{\alpha}t)} \epsilon\), where \(\epsilon ∼ N(0, I)\). Here, αt is defined as \(1−\beta_t\), and \(\alpha_t\) denotes the cumulative product of αs up to timestep t: \(\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s\).

In the reverse phase, the clean latent vector ẑ0 can be estimated directly using the predicted noise \(\hat{\epsilon}\):

\[\hat{z}_0 = \frac{z_t - \sqrt{1 - \bar{\alpha}_t} \, \hat{\varepsilon}}{\sqrt{\bar{\alpha}_t}}. \]

The predicted noise \(\hat{\epsilon}\) could be formally expressed as \(\epsilon_\theta(z_t; p, t)\), where p denotes prompt embedding.

As illustrated in Fig. 3, we first employ the encoder \(E_\theta\) to map the low-quality (LQ) image \(I_L\) into the latent space, yielding \(Z_L = E_\theta (I_L)\). Next, we perform one denoising step to obtain the predicted noise \(\hat{\epsilon}\) , allowing us to compute the predicted high-quality (HQ) latent vector \(\hat{z}_H\) :

\[\hat{z}_H = \frac{z_L - \sqrt{1 - \bar{\alpha}_{T_L}} \, \varepsilon_\theta(z_L; p, T_L)}{\sqrt{\bar{\alpha}_{T_L}}}. \]

where \(\varepsilon_\theta\) denotes the denoising network parameterized by \(\theta\), and \(T_L\) is the diffusion timestep.

Unlike one-step T2I diffusion models, the UNet input in OSDFace is not entirely Gaussian noise. Therefore, we predefine a parameter \(T_L \in [0, T]\) and feed it into the UNet, where \(T\) represents the total diffusion timesteps (in Stable Diffusion, \(T\) is set to 1,000). Finally, we use the decoder \(D_\theta\) to reconstruct the HQ face \(\hat{I}_H\) from \(\hat{z}_H\):

\[\hat{I}_H = D_\theta(\hat{z}_H). \]

Denoting the entire generator as \(G\), the overall computation procedure can be represented as:

\[\hat{I}_H = G_\theta(I_L; p). \]

图三的框架图清晰地给出了OSDFace的整体思路，即VQ部分用来做重建，Diffusion部分用来做生成，那怎么用VQ的信息指导Diffusion呢，这就是prompt的工作了。VQ部分参考的是DAEFR的工作，Diffusion参考的是OSD的工作，这样一个人脸方面的应用工作就诞生了。

3.2. Stage 1. Visual Representation Embedder

Previous studies have emphasized the critical importance of understanding input images and integrating face prior into models. To incorporate priors, we first train a feature extraction module, i.e., visual representation embedder (VRE), in stage 1 for LQ inputs. In stage 2, we utilize the pretrained VRE to guide the diffusion model, allowing it to leverage prior knowledge from the input image more effectively.

We design a variational autoencoder (VAE) with a vector-quantized (VQ) dictionary for effective feature extraction. The architecture consists of the VRE and a VAE decoder utilized exclusively during the first training stage. Through self-reconstruction training, the VRE builds a dictionary of image categories as priors, enabling it to act as a multi-class embedder using the VQ dictionary.

Given a low-quality (LQ) face \(I_L \in \mathbb{R}^{H \times W \times 3}\), the encoder \(E_L\) processes the image, obtaining a set of feature vector \(\{z_{L,k} \in \mathbb{R}^d\}_{k=1}^K\). Then, each feature \(z_L\) is transformed through a matching mechanism, generating a token \(q \in \mathbb{N}\) that corresponds to an item in the learnable low-quality VQ dictionary \(C_L = \{c_q \in \mathbb{R}^d\}_{q=1}^N\).

Visual Tokenizer. Tokenization is the process of dividing the input face into smaller units, i.e., tokens. In our framework, the visual tokenizer maps the input face to categories in the low-quality VQ dictionary:

\[M \circ E_L : \mathbb{R}^{H \times W \times 3} \to \mathbb{N}^K. \]

For an LQ face denoted as \(I_L\), we apply the VAE encoder \(E_L\) followed by the matching function \(M\) to obtain the token set \(Q_L = M(E_L(I_L))\). The token set \(Q_L\) represents the predicted categories corresponding to LQ face.

VQ Embedder. The VQ embedder processes the token set by retrieving the corresponding items from the VQ dictionary. Specifically, for each image token \(q\) at position \(k\), the embedder uses \(q\) as an index to access the dictionary item:

\(\begin{equation} z_k = \text{dict}(q), \quad q \in Q_L, \tag{6} \end{equation}\)

where \(z_k\) represents the \(k\)-th element of the visual prompt.

Unlike using image-to-tag models to generate tags as textual prompts, VRE directly tokenizes each face and converts the tokens into embeddings. This approach eliminates the information loss during the image-tag-embedding process. As illustrated in Fig. 4, VRE shows a strong recovery capacity. To study the working mechanism of VRE, we visualize the weight matrix of the final self-attention layer and the latent vector output by the VRE encoder. They are labeled as “att” and “enc” in the figure captions, respectively. The visualizations clearly indicate that the VRE focuses on both facial and nonfacial features.

In summary, we can embed the LQ face \(I_L\) using VRE to obtain the visual prompt \(p_L = \text{VRE}(I_L)\).

We construct two VQ dictionaries corresponding to the HQ and LQ image categories and train VQVAE using vector quantization with self-reconstruction. Following VQGAN, we use GAN loss to promote extracting real-world features over noise. However, the LQ encoder sometimes focuses on meaningless categories due to strong degradation. Thus, we employ an alignment training strategy inspired by the CLIP model to align the categories between LQ and HQ faces. By enhancing the diagonal correlation within VQ dictionaries, we could guide the LQ encoder to align its attention with the HQ encoder.

Training Objectives. We divide the loss functions into two parts. The first is a series of losses designed specifically for training the VQ dictionary, and the second is for association. Following previous works, we incorporate four types of losses during training. Three of them are image restoration losses, i.e., absolute differences \(\mathcal{L}_1\), perceptual function \(\mathcal{L}_{\text{per}}\) , and discriminative loss \(\mathcal{L}_{\text{dis}}\). Since the VQ is non-differentiable, we adopt a classic quantization loss \(\mathcal{L}_{\text{VQ}}\) additionally.

\[\begin{equation} \begin{split} \mathcal{L}_1 &= \|I_h - I_h^{\text{rec}}\|_1, \\ \mathcal{L}_{\text{per}} &= \|\Phi(I_h) - \Phi(I_h^{\text{rec}})\|_2^2, \\ \mathcal{L}_{\text{dis}} &= \mathbb{E}[\log D(I_h) + \log(1 - D(I_h^{\text{rec}}))], \\ \mathcal{L}_{\text{VQ}} &= \|\operatorname{sg}(\hat{Z}_h) - Z_h^c\|_2^2 + \beta \|\hat{Z}_h - \operatorname{sg}(Z_h^c)\|_2^2. \end{split} \tag{7} \end{equation}\]

Here, \(\Phi(\cdot)\) represents the feature extractor of VGG19, \(D\) is a patch-based discriminator, and \(\operatorname{sg}(\cdot)\) denotes the stop-gradient operation. \(\beta\) is set to 0.25.

The combined loss \(\mathcal{L}_{\text{VQ}}\) is then defined as

\[\begin{equation} \mathcal{L}_{\text{VQ}} = \mathcal{L}_1 + \lambda_{\text{per}} \cdot \mathcal{L}_{\text{per}} + \lambda_{\text{dis}} \cdot \mathcal{L}_{\text{dis}} + \mathcal{L}_{\text{VQ}}. \tag{8} \end{equation}\]

where we set \(\lambda_{\text{per}} = 1.0\) and \(\lambda_{\text{dis}} = 0.8\) in this setup.

Furthermore, we employ the cross-entropy loss introduced by DAEFR to enhance the correlation between the HQ and LQ features. The similarity matrix \(M_{\text{assoc}} \in \mathbb{R}^{K \times K}\) is constructed from the encoded feature \(z_H\) and \(z_L\). We define the similarity score along the HQ axis of \(M_{\text{assoc}}\) as \(p_{i,j}^h\), and along the LQ axis as \(p_{i,j}^l\). The cross-entropy losses are then defined as

\[\begin{align} \begin{split} \mathcal{L}_{\text{CHE}} &= -\frac{1}{K} \sum_{i=1}^{K} \sum_{j=1}^{K} y_{i,j} \log(p_{i,j}^h), \\ \mathcal{L}_{\text{CLE}} &= -\frac{1}{K} \sum_{i=1}^{K} \sum_{j=1}^{K} y_{i,j} \log(p_{i,j}^l), \end{split} \tag{9} \end{align}\]

where \(y_{i,j}\) represents the ground truth label. The final target for the feature association component is then defined as

\[\begin{equation} \mathcal{L}_{\text{assoc}} = \frac{\mathcal{L}_{\text{CHE}} + \mathcal{L}_{\text{CLE}}}{2}. \tag{10} \end{equation}\]

The overall loss function for training the VRE at this stage can be expressed as

\[\begin{equation} \mathcal{L}_{\text{total}} = \mathcal{L}_{\text{VQ}} + \lambda_{\text{assoc}} \cdot \mathcal{L}_{\text{assoc}}. \tag{11} \end{equation}\]

where \(\lambda_{\text{assoc}}\) is set to 0 for the initial epochs, and subsequently adjusted to 1 for the remaining epochs.

3.3. Stage 2. Realistic Face Alignment

Face restoration presents unique challenges and higher demands than general image restoration, as humans are exceptionally familiar with and sensitive to faces. Even minor inconsistencies or unnatural details in a generated face image are immediately noticeable (see Fig. 1), often resulting in an unrealistic appearance overall.

To ensure highly realistic face restoration, we design a series of losses aimed at guiding the model to generate globally harmonious and locally coherent face images. These carefully crafted losses play a key role in achieving high fidelity and visual authenticity. The overall loss function for the generator is defined as

\[\begin{equation} \begin{aligned} \mathcal{L}_{\text{gen}} ={}& \lambda_{\text{dis}} \cdot \mathcal{L}_G(z_H, \hat{z}_H) + \lambda_{\text{ID}} \cdot \mathcal{L}_{\text{ID}}(I_H, \hat{I}_H) \\ &+ \lambda_{\text{per}} \cdot \mathcal{L}_{\text{EA-DISTS}}(I_H, \hat{I}_H) + \text{MSE}(I_H, \hat{I}_H). \end{aligned} \tag{12} \end{equation}\]

Facial Identity Loss. A key challenge in face restoration is how to match each LQ face to a real, unique face portrait. However, during reconstruction, the model often struggles to render ambiguous features, leading to disharmonious facial attributes in the generated images. Inspired by advancements in face recognition tasks, we introduce a facial identity loss, specifically designed for face restoration. This loss leverages a pretrained face recognition model to measure feature similarity between the generated face and the target HQ image. The facial identity loss helps to enhance both the overall harmony of the generated face and the precise alignment of facial features.

We utilize a pretrained ArcFace model to encode both the generated and HQ faces into identity embeddings (IDs). The cosine similarity between these IDs is then computed and used as a loss to guide model training. Denote the facial feature extraction model as \(F\). The facial identity loss \(\mathcal{L}_{\text{ID}}\) can be defined as follows:

\(\begin{equation} \mathcal{L}_{\text{ID}} = 1 - \cos(F(I_H), F(\hat{I}_H)). \tag{13} \end{equation}\)

Perceptual Loss. Although LPIPS is widely used for perceptual evaluation, it can sometimes introduce visible artifacts, especially in diffusion models. To alleviate this issue, we use DISTS to better capture texture details. During the reconstruction of LQ faces, DISTS preserves texture details more effectively while maintaining perceptual similarity, particularly in areas like hair and skin.

In addition, the accurate generation of edge details is essential for achieving perceived sharpness in face restoration. Thus, following the practice in D3SR, we introduce an edge-aware DISTS (EA-DISTS) as our perceptual loss:

\(\begin{equation} \mathcal{L}_{\text{EA-DISTS}}(\hat{I}_H, I_H) = \mathcal{L}_{\text{DISTS}}(\hat{I}_H, I_H) + \mathcal{L}_{\text{DISTS}}(S(\hat{I}_H), S(I_H)). \tag{14} \end{equation}\)

where \(S(\cdot)\) is the Sobel operator.

GAN Loss. OSD models face challenges in generating stable images, due to their limited computational capacity. Previous studies typically employ distillation techniques to transfer knowledge from multi-step diffusion models. However, these methods are restricted by the performance limitations of the teacher models.

As an alternative, we use a discriminative network to enhance the realism of generated faces. This approach provides greater flexibility and improves computational efficiency. The adversarial loss, used to update both the generator \(G_\theta\) and discriminator \(D_\theta\), is defined as

\(\begin{align} \mathcal{L}_G &= -\mathbb{E}_t[\log D_\theta(F(\hat{z}_H, t))], \tag{15} \\ \mathcal{L}_D &= -\mathbb{E}_t[\log(1 - D_\theta(F(\hat{z}_H, t)))] - \mathbb{E}_t[\log D_\theta(F(z_H, t))], \tag{16} \end{align}\)

where \(z_H\) is the latent vector of HQ face, and \(F(\cdot, t)\) denotes the forward diffusion process at timestep \(t \in [0, T]\).

身份对齐一直是人脸修复的一个重要问题，在本文中，作者主要通过引入一个预训练的脸部识别模型作为一种衡量监督放入loss中，又引入了ArcFace模型作为身份监督加入loss。另外，人脸修复一贯通过LPIPS引入的“感知loss”被作者认为存在一定的视觉伪影问题，特别是对于扩散模型来说。这一点其实是很有意思的提出，但可惜作者并没有进一步对此进行解释。而是通过添加另外一个DIST loss来缓解这个问题。
总的来看，OSDFace更像是为了OSD在人脸修复领域应用所包的一盘饺子，其并没有提出任何比较新颖的理论观点，而是给出了OSD一个可能的工程化应用路径，从这一点来看，他站在了扩散模型应用的前沿，但简单是否有效，还需要看实验部分。

4. Experiments

4.1. Experimental Settings

Training Datasets. Our model is trained on FFHQ and its retouched version, containing 70,000 different high-quality face images. Images are resized to 512×512 pixels. Synthetic training data is generated using a dual-stage degradation model, with parameters following WaveFace. This dual-stage degradation process closely aligns the synthetic data with real-world degradation scenarios, handling both mild and severe degradations.

Testing Datasets. We evaluate our method on the synthetic dataset CelebA-Test from DAEFR and three widely used real-world datasets: Wider-Test, LFW-Test, and WebPhoto-Test, following the settings in previous literature. These datasets exhibit diverse and complex degradations. For example, the degradations in LFW-Test are relatively slight and regular, whereas those in Wider-Test are stronger.

Metrics.

For the Synthetic Dataset, we employ LPIPS and DISTS as reference-based perceptual quality measures, together with MUSIQ and NIQE as no-reference image quality measures. To evaluate the distribution similarity between the ground truth and restored faces, we calculate FID with both FFHQ and CelebA-Test HQ, showing the ability to recover “real” faces and HQ images. Following previous works, we use the embedding angle of ArcFace, namely “Deg.”, and landmark distance, namely “LMD”, as additional identity and fidelity metrics.

For the Real-world Datasets, we employ non-reference metrics, including FID (with FFHQ as reference), as well as CLIPIQA, MANIQA, MUSIQ, and NIQE. We utilize the evaluation codes provided by py-iqa and VQFR for all metrics.

Implementation Details.

In the first stage, we train the VRE using the Adam optimizer with a learning rate of \(1.44×10^{−4}\) and a batch size of 32. In our implementation, the input face image has dimensions of 512×512, and the quantized feature map has dimensions of 256×512 after flattening. The VQ dictionaries contain N=1,024 code items, each with a channel size of 512. The HQ and LQ dictionaries are trained for 50 and 10 epochs, respectively. The association parameter \(λ_{assoc}\) is initially set to 0, followed by an additional 10 epochs with \(λ_{assoc}=1\). The VRE is trained on 8 NVIDIA A800 GPUs.

In the second stage, our OSDFace is trained by the AdamW optimizer, with learning rate \(1×10^{−4}\) and batch size 2 for both the generator and discriminator. The LoRA rank in the UNet is set to 16. The SD 2.1-base serves as the pretrained OSD model, with the VRE used to construct the prompt embeddings. The setting of the discriminator is followed by D3SR, and the generator and discriminator are trained alternately. Training is performed for 150K iterations on 2 NVIDIA A6000 GPUs.

Compared State-of-the-Art (SOTA) Methods. We compare OSDFace with several SOTA methods, including RestoreFormer++, VQFR, CodeFormer, DAEFR, PGDiff, DifFace and DiffBIR. We also provide the current SOTA OSD model OSEDiff, and retrain it on FFHQ, namely OSEDiff*.

4.2. Performance on Synthetic Data

Quantitative Results. The results in Tab.1 demonstrate that OSDFace outperforms other methods in Deg., indicating better alignment with the ID features of HQ images. It also achieves higher scores on LPIPS and DISTS, highlighting superior perceptual coherence and detail preservation. Performance in FID(HQ) shows that OSDFace effectively recovers the distribution of the original HQ dataset.

Qualitative Results. Visual comparisons in Fig.5 illustrate that our method closely resembles real images. OSDFace enhances realism in elements such as earrings, eyelashes, eyebrows, hair strands, and subtle skin textures. Our approach restores hair most naturally, avoiding additional streaks or other regular textures, and maintains consistent hair color without random changes.

4.3. Performance on Real-world Data
Quantitative Results. In Tab. 2, we can see that OSDFace performs significantly better than another OSD method across all metrics. OSDFace outperforms non-diffusion models in most metrics. For multi-step diffusion models, we generally surpass PGDiff and DifFace. For DiffBIR, we maintain a small gap and even exceed it in certain metrics.

Qualitative Results. Figure 7 visualizes representative images in three real-world datasets. Existing face restoration methods often emphasize facial features while overlooking the precise reconstruction of facial accessories like glasses, hair, hats, headwear, and background details. Some methods, especially DiffBIR, are over-smoothing and lack fine facial textures. OSDFace naturally and clearly restores multiple faces within a single image, showing its effectiveness for unaligned faces. OSDFace ensures high-fidelity restoration with balanced and consistent visual quality throughout the image. We also test its zero-shot ability on cartoon images in Fig. 6, using Anime Face Dataset.

4.4. Complexity Analysis.
Table 3 provides a detailed comparison of model complexity, including number of steps, inference time, parameter number, and multiply-accumulate operations (MACs) during inference. All evaluations are carried out on an NVIDIA A6000 GPU for consistency. First, by relying on one-step diffusion process, OSDFace significantly outperforms multi-step diffusion models in both inference time and MACs, demonstrating remarkable efficiency. Furthermore, OSDFace maintains an advantage over other OSD model. We do not incorporate text encoders (e.g., the DAPE used by OSEDiff), nor do we introduce additional modules that increase Stable Diffusion model complexity (e.g., the ControlNet used by DiffBIR). Instead, we utilize a more efficient image encoder, VRE. Consequently, OSDFace achieves the lowest MACs count and fastest inference speed among all evaluated diffusion-based methods. Compared to the other OSD model, OSEDiff, OSDFace demonstrates a 25% reduction in parameters and a 23% improvement in inference speed.

4.5. Ablation Studies
Visual Representation Embedder (VRE). We evaluated the impact of various prompt embedding generation methods on the training of OSDFace, including VRE, learnable prompt embeddings, and Degradation-Aware Prompt Extractor (DAPE). As shown in Tab. 4, the experimental results indicate that our method using VRE outperforms others across most metrics. As shown in Fig. 4, using VRE significantly enhances the facial details and the coherence of non-facial areas in the generated faces.

Furthermore, we trained the VRE using two approaches: with and without the HQ-LQ feature association loss, further validating the effectiveness of our method. As shown in the last two rows of Tab. 4, training with association loss achieves superior performance. This demonstrates that our VRE can extract valuable information from LQ images to guide face restoration effectively. Notably, with the inclusion of the feature association loss, the VRE gains the ability to identify and reconstruct certain HQ-like categories from LQ inputs, enhancing face restoration quality.

Realistic Face Alignment Loss. For the series of losses used in training the OSDFace, such as facial identity loss, EA-DISTS, and adversarial loss, we conducted extensive ablation studies. Additionally, we introduce the EA-LPIPS loss, defined similarly to EA-DISTS in Sec. 3.3. This loss first processes the image with the Sobel operator, followed by calculating the LPIPS loss. As shown in Tab. 5, the losses we utilized significantly contribute to model performance, with most metrics reaching the best scores.

结论：指标提升不算小，但和同类型的扩散模型相比来说，它的提升空间大概是不在于此的。从视觉效果上来看，OSDFace的头发修复效果是相当明显的，并相比较于OSEDiff来说，人物身份的一致性有了一定的保障性，并且推理速度与VQGAN类方法几乎保持了一致，也算是针对扩散模型在人脸修复这一方面应用的不小提升了。

代码介绍

https://github.com/jkwang28/OSDFace
我的数据库里什么也没有哦~（training code no update yet）