ReF-LDM: A Latent Diffusion Model for Reference-based Face Image Restoration

Q4-1. The redundant use of the diffusion model will cause slower speed and more computation, limiting applications like video face restoration.

During inference, the CacheKV mechanism pre-computes and saves the keys and values (KVs) for each reference image with only one pass of diffusion network; then, the saved reference KVs are reused in each step of the multi-step diffusion denoising process. In a similar way, for video face restoration, we can pre-compute and save the CacheKV and reuse the same CacheKV for all frames of a video without redundant computation.

Q4-2. How is the time embedding selected when extracting CacheKV from reference images.

We choose the time embedding of timestep 0 (the same one used by the main diffusion U-net) when extracting CacheKV. Our motivation is that the reference images are clean, so we use the last and cleanest diffusion timestep $t=0$. We have also experimented with a learnable time embedding, while the results are similar to using that of timestep 0. So we stick to the original choice. The motivation of using timestep 0 is mentioned in paper L93-94, we will explain it more clearly in the revision.

Q4-3. The results with no reference image.

We show the outputs of ReF-LDM without reference image in PDF figure 5, which suggests that ReF-LDM can restore an LQ image to a reasonable HQ image without reference images. The reason may be twofold:

1. our denoising U-net has been pretrained with only LQ condition (paper L218-220)
2. we drop the reference condition during training with probability 0.1 (paper L419).

Q3-1. Explain how CacheKV works.

We include an illustration in the PDF figure 3 (which is a simplified version of the paper figure 2) and explain as follows.

During inference,

• We first extract the CacheKV from reference images. Specifically, for a reference image, we encode the image $x_\text{ref}$ to latent $z_\text{ref}$ with VAE-encoder. Then, we pass the latent $z_\text{ref}$ to the same U-net used by the main diffusion process, where we save the reference-image keys and values of each self-attention layer.
• In the main diffusion process, the saved reference-image keys and values (i.e. CacheKV) can be accessed for each denoising timestep $t=T,…,1$. Specifically, in each self-attention layer of the main U-net, we concatenate the reference-image keys (of the corresponding attention layer) with the main keys along the token dimension, and the same goes for the values. By doing so, the main U-net can utilize keys and values of the reference images while denoising.

We appreciate your question and will refine the description (paper L90-105) in the revision.

Q3-2. How does CacheKV use varying numbers of reference images when calculating QKV?

We separately pass each reference image through the VAE-encoder and the U-net. For each attention layer in the U-net, we append and save the keys/values from all the reference images along the token dimension, i.e., one reference ($N_\text{token}, d_\text{key}$) → multiple references ($N_\text{ref}* N_\text{token}, d _\text{key}$). Similarly, in the main diffusion process, we concatenate the reference keys/values with the main U-net keys/values along the token dimension. We appreciate your question and will explain this in a clearer way in the revision.

Q3-3. Why does CacheKV provide a more aligned feature space than ablation setting CrossAttention?

We include an illustration of CacheKV and CrossAttention in PDF figure 3.

• CacheKV uses the same U-net as the main diffusion process. So in the main diffusion process, each attention layer can access the reference-image keys and values of the corresponding attention layer (which is an aligned feature space).
• CrossAttention only encodes reference images with VAE-encoder and passes the same reference-image latent {$z_\text{ref}$} to all cross-attention layers. The VAE’s latent space is less aligned with the U-net of the main diffusion process.

We appreciate your question. We will include the illustrations and refine the explanation in the revision.

Q3-4. No comparison with GAN-based papers.

We have compared with some recent GAN-based BFR methods in Table 6, where DMDNet, DAEFR, and VQFR use discriminator and GAN loss; VQFR and CodeFormer use pretrained VQGAN. We didn't run some earlier GAN-based models on our benchmarks (ex. 2021 GFP-GAN, 2020 ASFNet), as the papers we have compared with (ex. 2024 DAEFR, 2022 CodeFormer, 2022 DMDNet) outperformed these earlier methods. Please let us know if the reviewer suggests that we compare with specific papers, we will try to provide additional comparison during the reviewer-author discussion period.

Q3-5. What will happen if we use reference images from another identity?

We included example results of using input reference images from two different people in PDF figure 4. ReF-LDM outputs reasonable high-quality images based on the LQ image with facial features affected by the reference identity (ex. wrinkles). We appreciate the reviewer’s interesting question, and we will include related results in the revision.

Q2-1. Different impacts of reference images in comparison to style transfer works.

The reference images affect model outputs differently because face image restoration and style transfer have different objectives. In our task, the model output is an HQ image, and the model inputs are reference images and an LQ image (which is synthesized by applying degradation to the HQ image). Therefore, our model is encouraged to follow the image layout and style of the input LQ image and to integrate the distinctive facial features (e.g., eye shape) of the input references by the LDM loss and timestep-scaled identity loss.

Q2-2. Loss function used to pretrain LDM with only LQ input.

We use only the typical LDM loss (without the proposed timestep-scaled identity loss) to pretrain LDM. We will state this clearly in the revision.

Q2-3. Why not use face ID embeddings of reference images as network condition?

We conducted an additional ablation experiment which uses pretrained ArcFace to encode input reference images. The example images (in PDF figure 2) and the lower face similarity metric (IDS↑ CacheKV 0.65 vs. ArcFace 0.24) suggest that it is more challenging for ReF-LDM to integrate the reference identity information with ArcFace than with CacheKV. This difficulty may arise because ArcFace embedding space is very different from that of the main diffusion U-net. In contrast, CacheKV uses the same diffusion U-net to provide reference-image keys and values with the corresponding attention layers, making it easier for the main diffusion U-net to integrate. Another intuitive explanation is that ArcFace is trained to distinguish different identities, thus, it is a good discriminator and provides effective supervision as the identity loss. However, ArcFace is not trained to generate a face image, so it may not contain enough appearance information for the restoration task.

Q2-4. Ablation study of timestep-scaled identity loss, apply loss after T steps.

In a training iteration of LDM-based models, we typically sample a $z_t$ and pass it though a denoising network once then compute the loss. To compute identity loss on a multi-step prediction, we need to pass an input $z_t$ through the denoising network recurrently for at most $T=1000$ passes. This will greatly increase the required computational time and GPU memory, so we didn’t experiment with such a setting.

Q2-5. The proposed FFHQ-Ref dataset contains no new images.

While FFHQ-Ref is indeed based on the images collected by FFHQ dataset, we believe its value lies in the boost of image quality and quantity over previous reference-based dataset (compared in paper Table 1). In addition, the careful curation and organization for the task of reference-based restoration (i.e., ID-based data splitting and FPS-sorted reference images) also enable fair and practical evaluation for future works in this field.

Q1-1. Limited improvement over no-reference methods in Table 6.

ReF-LDM achieves a significant boost in face similarity metric (IDS↑ 0.675 vs. 0.323 for the best no-reference method CodeFormer). As suggested by the reviewer, we evaluated an additional metric FID (with target distribution as FFHQ 70k images) and found that our ReF-LDM has better FID especially for severely degraded LQ images (FID↓ 37.4 vs. 49.1 for the best no-reference method DAEFR). We appreciate the reviewer’s suggestion and will include the FID evaluation in the revision.

We agree that we do not have the best LPIPS scores and that additional fusion modules (like CodeFormer’s CFT or VQFR’s dual decoders) may further improve LPIPS, while in this work we would like to focus on a simple network design to integrates an input LQ and multiple reference images into an LDM-based model.

Q1-2. More ablation studies of variants of timestep-scaled identity loss.

We conducted an additional experiment with variant setting $t<500$ (timestep-dependent scale = 1 when $t<500$ otherwise 0), which shares the same motivation as the default setting $\sqrt{\bar{\alpha_t}}$ to downscale the identity loss when a larger and noisier timestep t is sampled. The example images (in PDF figure 1) and the metrics in the table suggest that $t<500$ also works but the default $\sqrt{\bar{\alpha_t}}$ yields better identity similarity and quality. We appreciate the reviewer’s suggestion and will train and include more variants of timestep-dependent scales in the revision.

Q1-3. Illumination Changes.

Thank you for pointing this out. We checked the input reference images used in the examples where the illumination changed. In the case of Figure 8, the reference images have strong orange lightning and our output has a warmer illumination. Therefore, the cause may be the influence from the reference images. We will include case studies and discuss possible future work (such as illumination-preservation loss or different model architecture design) in the revision.

Q1-4. Our CacheKV vs. related work MasaCtrl.

MasaCtrl proposes a training-free method to exploit StableDiffusion for reference-based text-to-image synthesis. They also utilize the denoising U-net to extract keys and features (KVs) from a reference image (termed as source image in their paper). However, we are different in several aspects:

1. CacheKV passes the references through the U-net only once and reuses the saved KVs in each main diffusion denoising step without redundant computation, while MasaCtrl passes the reference through the U-net for every diffusion step.
2. We have an ablation experiment “spatial-concatenation” that is closer to MasaCtrl, as it also passes references to U-net for each step (paper L125-135). As shown in paper Table2, our CacheKV (without the redundant U-net passes) is more efficient, especially when we have multiple reference images.
3. Our model implicitly learns how to naturally integrate the references’ facial features into LQ’s image layout, while MasaCtrl explicitly uses object masks (computed with text condition’s cross-attention scores, ex. “dog”) to prevent the access of irrelevant reference features and messy results.

We appreciate the suggestion to discuss the related work MasaCtrl and will include it in the revision.

Q1-5. Our timestep-scaled identity loss vs. concurrent work PuLID (Arxiv 2024).

The concurrent work PuLID shares the same motivation with us, which is, to avoid using identity loss on a flawed one-step output from a larger and noisier timestep $t$. PuLID exploits a few-step distilled diffusion model and runs four denoising steps to obtain a better output for identity loss. They adopt off-the-shelf SDXL-lightning for their SDXl-based model. However, for a non-SD-based work (like our ReF-LDM), PuLID’s solution requires an additional few-step diffusion distillation training process. In contrast, we propose a simple and effective solution by applying a timestep-dependent scale to identity loss.