AE-NeRF: Audio Enhanced Neural Radiance Field for Few Shot
Talking Head Synthesis

Audio-driven talking head generation aims to animate a speaker according to input audios. Image based methods (Prajwal et al. 2020; Zakharov et al. 2019) utilize GANs (Goodfellow et al. 2020) and Auto-encoders (Kingma and Welling 2013) to generate talking faces with audio signals as conditional inputs. Model based methods (Chen et al. 2019; Thies et al. 2020; Das et al. 2020; Wang, Mallya, and Liu 2021; Zhou et al. 2020; Song et al. 2022) leverage structural information such as 2D landmarks or 3DMM parameters for better face modeling. For instance, (Chen et al. 2019) and (Thies et al. 2020) generate faces with predicted facial landmarks or 3DMM expression coefficients. These methods can quickly adapt to an unseen identity. However, the prediction error of the representations may lead to inferior image quality, and they usually require hundreds of videos for training. NeRF based methods (Guo et al. 2021; Liu et al. 2022; Shen et al. 2022) have brought a new trend of talking head synthesis. They perform optimization on the video clip of a single person, and can synthesize pose-controllable faces of any resolution with high fidelity. AD-NeRF (Guo et al. 2021) use two separated NeRFs to model the head and the torso part respectively. SSP-NeRF (Liu et al. 2022) performs rays re-sampling based on the loss magnitude of different semantic regions. Despite the above advantages, the generalization ability of NeRF based methods to new identities still needs to be improved, and they suffer from performance drop when the video clip is relatively short.

Neural Radiance Fields (NeRFs) (Mildenhall et al. 2020) combines MLPs with differentiable volume rendering and achieves photorealistic view synthesis results. Although impressive results are obtained, the original NeRF needs to be retrained for each new scene, which are both time consuming and computational expensive. Moreover, when only sparse views are available, because of the lack of the prior knowledge between scenes (Yu et al. 2021), the synthesis results can suffer from a large degradation in quality.

Few Shot Neural Rendering (Trevithick and Yang 2021; Yu et al. 2021; Chen et al. 2021; Gu et al. 2021; Xu et al. 2022) are proposed to alleviate these problems with the assistance of different kinds of priors such as 2D image features (Trevithick and Yang 2021; Yu et al. 2021; Wang et al. 2021), trainable latent codes (Jang and Agapito 2021; Gafni et al. 2021) and style inputs (Gu et al. 2021; Chan et al. 2022). Among them, the pixel level feature (Trevithick and Yang 2021; Yu et al. 2021; Wang et al. 2021) from randomly chose 2D reference images is the most commonly used prior to promote the NeRF’s rendering ability when only a few observations are available. When coming up with a new scene, the NeRF can perform quick generalization based on the reference image features from that scene. DFRF (Shen et al. 2022) directly uses the above reference scheme for few shot talking head generation. But it ignores the importance of audio features in talking head rendering.

The full pipeline of our AE-NeRF is shown in Fig. 2. Both audio features and visual features are extracted as the references. The proposed Audio Aware Aggregation module fuses these features with cross attention and yields a strong prior for fully leveraging the limited data. Then, the audio related and audio decoupled regions are modeled by our Audio-Aligned Face Generation strategy. Finally, the portrait of the speaker and semantic masks are synthesized through volume rendering.

The original NeRF encodes a static scene as a continuous volumetric radiance field $\mathcal{F}$, which is modeled by an MLP. It takes a 3D query point $\mathbf{x}$ and its view direction $\mathbf{d}$ as input, and outputs the corresponding density $\sigma$ and color $\mathbf{c}$: $F(\mathbf{x},\mathbf{d})=(\mathbf{c},\sigma)$. When applying NeRF to talking head, one will take the audio feature as an additional input, and the rendering process can be written as $\mathcal{F}(\mathbf{x},\mathbf{d},\mathbf{a})=(\mathbf{c},\sigma)$.

Specifically, given $L$ reference images, let $\mathbf{I}_{i}\in\mathbb{R}^{H_{i}\times W_{i}\times 3}$ and $\mathbf{P}_{i}\in\mathbb{R}^{3\times 4}$ denote the $i$-th image and camera projection matrix respectively ($i\in\{0,1,...,L-1\}$). A shallow convolutional network without downsampling is employed to extract dense features $\mathbf{F}_{i}\in\mathbb{R}^{H\times W\times D}$ from each image $\mathbf{I}_{i}$, where $H$, $W$ and $D$ are the height, width and the feature channel respectively. To facilitate the rendering of a 3D point $\mathbf{x}$ on target image, we first project $\mathbf{x}$ onto $i$-th reference image to obtain image features $\mathbf{f}^{ref}_{i}\in\mathbb{R}^{D}$. Then all extracted features $\mathbf{f}^{ref}_{1...L}$ are merged as a condition $\tilde{\mathbf{f}}$. So, an audio-driven NeRF model with reference scheme can be formulated as

$$\mathcal{F}({\mathbf{x}},\mathbf{d},\mathbf{a},\tilde{\mathbf{f}})=(\mathbf{c},\sigma).$$

(1)

In addition, we denote the projection coordinate of the 3D point $\mathbf{x}$ to the i-th 2D reference image as $\mathbf{p}^{ref}_{i}=(\mathbf{u}_{i},\mathbf{v}_{i})$. Since the talking head is dynamic, directly performing projection may bring some errors. Thus, an image warper is imposed to calibrate the 2D coordinate by predicting its offset $\Delta\mathbf{p}^{ref}_{i}$ on the feature plane. The calibrated coordinate is denoted by

$$\mathbf{p}^{{ref}^{\prime}}_{i}=\mathbf{p}^{ref}_{i}+\Delta\mathbf{p}^{ref}_{i}.$$

(2)

In a talking head video, if the speakers in the reference image and the target image have similar speech contents, they tend to have similar mouth shapes, as we mentioned above. Therefore, we introduce Audio Aware Aggregation module into the feature fusion process to make the reference image, whose audio is more similar to target, contribute more. Let $\mathbf{a}$, $\mathbf{a}^{ref}_{1...L}$ and $\mathbf{f}^{ref}_{1...L}$ be target audio feature, reference audio feature and reference image feature respectively, then we have

$\displaystyle\tilde{\mathbf{f}}=\operatorname{AAA}(\mathbf{a},\mathbf{a}^{ref}_{1...L},\mathbf{f}^{ref}_{1...L}).$

(3)

The Audio Aware Aggregation module utilizes a transformer structure with cross attention blocks to fuse the audio and the visual information. To be more concrete, $\mathbf{a}$, $\mathbf{a}^{ref}_{1...L}$, and $\mathbf{f}^{ref}_{1...L}$ are projected and reshaped to get their corresponding tokens $\mathbf{T}_{a}^{tar}\in\mathbb{R}^{N\times 128}$, $\mathbf{T}_{a}^{ref}\in\mathbb{R}^{NL\times 128}$ and $\mathbf{T}_{f}^{ref}\in\mathbb{R}^{NL\times 128}$, $N$ is the number of query point samples in a batch. The audio and image tokens are then modeled by the Multi-Head Cross Attention (MHCA), along with Layer Normalization (LN), Residual Connection (RC) and Feed Forward Network (FFN). The cross attention block, with $\mathbf{T}_{a}^{tar}$ as query, $\mathbf{T}_{a}^{ref}$ as key, $\mathbf{T}_{f}^{ref}$ as value, can be formalized as

		$\displaystyle\operatorname{MHCA}\left(\mathbf{T}_{a}^{tar},\mathbf{T}_{a}^{ref},\mathbf{T}_{f}^{ref}\right)=$		(4)
		$\displaystyle\operatorname{Softmax}\left[\frac{\mathbf{T}_{a}^{tar}\boldsymbol{W}^{Q}\left(\mathbf{T}_{a}^{ref}\boldsymbol{W}^{K}\right)^{T}}{\sqrt{d}}\right]\mathbf{T}_{f}^{ref}\boldsymbol{W}^{V},$

We also introduce the audio aware manner into the image warper module. The warper takes the target query point $\mathbf{x}$ and the target audio $\mathbf{a}$ as input, together with the corresponding audio $\mathbf{a}_{ref}$ and ray direction $\mathbf{d}_{ref}$ of the $i$-th reference, and outputs the coordinate offset, achieving more precise feature extraction:

$\displaystyle\Delta\mathbf{p}^{ref}_{i}=(\Delta\mathbf{u}_{i},\Delta\mathbf{v}_{i})=\operatorname{Warper}(\mathbf{x},\mathbf{a},\mathbf{a}^{ref}_{i},\mathbf{d},\mathbf{d}^{ref}_{i}).$

(5)

As stated before, a disentangled modeling of the audio related reigon and the audio decoupled region is of great significance. Our Audio-Aligned Face Generation strategy uses an audio associated NeRF and an audio independent NeRF to model these two regions separately, and only the audio associated NeRF conditions on audio feature. To merge the rendering results of the two NeRF models, we add an additional parsing branch to predict mask $\mathbf{m}_{aa}$ or $\mathbf{m}_{ai}$, where $\mathbf{m}_{aa}$ has 1 in audio related region (i.e., the lower half face ), and 0 in audio independent region, $\mathbf{m}_{ai}$ is otherwise. Then we have the following two formulations:

		$\displaystyle\mathcal{F}_{aa}({\mathbf{x}},\mathbf{d},\mathbf{a},\tilde{\mathbf{f}})=(\mathbf{c}_{aa},\sigma_{aa},\mathbf{m}_{aa})$		(6)
		$\displaystyle\mathcal{F}_{ai}({\mathbf{x}},\mathbf{d},\tilde{\mathbf{f}})=(\mathbf{c}_{ai},\sigma_{ai},\mathbf{m}_{ai}).$

For a query point $\mathbf{x}$, we can input it into two NeRFs and use the predicted mask to blend the colors or densities of both outputs, like (Ma et al. 2023). However, it will double the training time because a batch of rays has to go through the two NeRFs simultaneously.

Concretely, for a set of sampling rays $\Omega$, we use $\epsilon\times\Omega$ to represent a new set that has $\epsilon\times|\Omega|$ rays randomly sampled from $\Omega$ ($|\Omega|$ means the number of rays in $\Omega$). The set of the rays from the audio related region and the audio decoupled region are denoted as $\Omega_{ar}$ and $\Omega_{ad}$. An overlap region is defined as $\Omega_{overlap}=(\epsilon\times\Omega_{ar})\cup(\epsilon\times\Omega_{ad})$, where rays are fed into two NeRFs simultaneously. While the remaining parts $\Omega_{aa}=\Omega_{ar}\setminus\Omega_{overlap}$ and $\Omega_{ai}=\Omega_{ad}\setminus\Omega_{overlap}$ are rendered by their corresponding NeRFs separately. ($\cup$ and $\setminus$ denote sets union and subtraction operations). In practice, $\epsilon$ is set to 0.4 for the best lip generation result. More effects of this Regionwise Ray Sampling mechanism can be found in the supplementary material.

		$\displaystyle\mathbf{c}=\begin{cases}\mathbf{m}_{aa}\cdot\mathbf{c}_{aa}+\mathbf{m}_{ai}\cdot\mathbf{c}_{ai},&\mathbf{r}\in\Omega_{overlap}\\ \mathbf{c}_{aa},&\mathbf{r}\in\Omega_{aa}\\ \mathbf{c}_{ai},&\mathbf{r}\in\Omega_{ai}\end{cases}$		(7)
		$\displaystyle\sigma=\begin{cases}\sigma_{aa}+\sigma_{ai},&\mathbf{r}\in\Omega_{overlap}\\ \sigma_{aa},&\mathbf{r}\in\Omega_{aa}\\ \sigma_{ai}.&\mathbf{r}\in\Omega_{ai}\end{cases}$

During inference, the whole image are regarded as the overlap region, and rendered by two NeRFs at the same time, since there is no ground truth mask available. To get the predicted RGB pixel and the mask, we utilize classical volume rendering to accumulate the samples on the ray:

		$\displaystyle\hat{\mathbf{C}}(\mathbf{r})=\int_{t_{n}}^{t_{f}}T(t)\sigma(\mathbf{r}(t),\mathbf{a},\tilde{\mathbf{f}})\mathbf{c}(\mathbf{r}(t),\mathbf{d},\mathbf{a},\tilde{\mathbf{f}})dt,$		(8)
		$\displaystyle\hat{\mathbf{m}}(\mathbf{r})=\int_{t_{n}}^{t_{f}}T(t)\sigma(\mathbf{r}(t),\mathbf{a},\tilde{\mathbf{f}})\mathbf{m}(\mathbf{r}(t),\mathbf{d},\mathbf{a},\tilde{\mathbf{f}})dt,$

Following the original NeRF (Mildenhall et al. 2020), we use a reconstruction loss term to optimize the coarse and the fine network (we still take audio associated NeRF as example if not specified), which can be written as

		$\displaystyle\mathcal{L}_{\mathrm{p}}=\sum_{\mathbf{r}\in\mathcal{R}}\left[\left\|\hat{\mathbf{C}}_{c}\left(\mathbf{r}\right)-\mathbf{C}\left(\mathbf{r}\right)\right\|_{2}^{2}+\left\|\hat{\mathbf{C}}_{f}\left(\mathbf{r}\right)-\mathbf{C}\left(\mathbf{r}\right)\right\|_{2}^{2}\right],$		(9)
		$\displaystyle\mathcal{L}_{\mathrm{m}}=\sum_{\mathbf{r}\in\mathcal{R}}\left[\left\|\hat{\mathbf{m}}_{c}\left(\mathbf{r}\right)-\mathbf{m}\left(\mathbf{r}\right)\right\|_{2}^{2}+\left\|\hat{\mathbf{m}}_{f}\left(\mathbf{r}\right)-\mathbf{m}\left(\mathbf{r}\right)\right\|_{2}^{2}\right],$

Our final loss term can be given as

$$\mathcal{L}=\mathcal{L}_{\mathrm{p}}+\lambda_{m}\mathcal{L}_{\mathrm{m}}+\lambda_{o}\mathcal{L}_{\mathrm{o}},$$

(12)

where $\lambda_{m}$ and $\lambda_{o}$ are weight parameters.

Dataset Preparation. We use the videos provided by AD-NeRF (Guo et al. 2021), DFRF (Shen et al. 2022), and HDTF dataset (Zhang et al. 2021) to conduct our experiments. Videos are all resampled to 25 fps and resized to a resolution of $512\times 512$. For each video, the first half of it is used for training and the second half is used for inference.

To compare the quality of the generated talking head thoroughly, two different settings are taken into account: Self driven setting, where the video and the audio are from the same person. Cross driven setting, where the audio from one person is used to drive another identity. Each video is about two minutes in length.

Quantitative comparison results on AD-NeRF and DFRF videos are shown in Testset A part of Tab.1. Wav2Lip (Prajwal et al. 2020) uses a pretrained SyncNet (Chung and Zisserman 2017) as the optimization objective and it achieves a SyncNet score even better than the ground truth. However, the quality of its generated images is relatively low. NeRF-based methods have shown their superiority not only at the pixel level (PSNR and SSIM) but also at the feature level (LPIPS). AD-NeRF performs slightly worse than DFRF and our method in image quality metrics due to the head-torso misalignment. It can be seen that our AE-NeRF surpasses baseline methods by a large margin in SyncNet confidence and LMD, which indicates the effect of learning aggregated audio visual features for lip synthesis.

We compare our AE-NeRF with other NeRF based talking head methods under a more challenging setting, few shot talking head synthesis, which further validate the generation ability of our method.

An ablation study is conducted to show the function of each proposed module. We replace our Audio Aware Aggregation with a slot attention module (Locatello et al. 2020) which simply fuses the visual feature without considering audio signals, denoted as w/o AAA. We also test the performance of the model without Audio-Aligned Face Generation, where an audio associated NeRF is used to model the whole face area, denoted as w/o AAFG. The model trained without the warper is denoted as w/o Warper. The qualitative and quantitative results are shown in Fig. 6 and Tab. 4. The Audio Aware Aggregation module brings better quality in both pixel and feature level. Model obtained without this module has less awareness of the audio signal, resulting in inferior image quality. Our Audio-Aligned Face Generation strategy brings more correct lip shapes and more natural expressions. The warper can bring more accurate visual feature points, which can also affect the audio lip synchronization.

To further study the effect of the Audio Aware Aggregation module, we calculate the average attention scores between different references and the target feature in the first attention block of our Audio Aware Aggregation module, shown in Fig. 5. As a comparison, scores in the feature aggregation module used in DFRF are also taken into account. The inner product between target-reference pairs with similar audio visual contents is significantly higher than those with distinct contents in our cross attention block. This proves that the audio visual interaction process assists the rendering of the target pixel, resulting in better generation results.