XKD: Cross-modal Knowledge Distillation with Domain Alignment for Video Representation Learning

Figure 1 presents an overview of our framework. Our method consists of two sets of autoencoders for audio ($\theta_{ae}^{a}$) and video ($\theta_{ae}^{v}$). First, we train these autoencoders to reconstruct from the masked inputs, which helps the encoders ($\theta_{e}^{a}$ and $\theta_{e}^{v}$) to learn modality-specific information. Next, to transfer knowledge between modalities to learn complementary information, we align the two domains by identifying the most transferable features and minimizing the domain gaps. Finally, audio ($\theta_{t}^{a}$) and video ($\theta_{t}^{v}$) teachers are employed to provide cross-modal supervision to the opposite modalities. The details of our proposed method are mentioned below.

Let be given video clip $x$, where the visual frames and audio spectrograms are denoted as $x_{v}$ and $x_{a}$, respectively. We create ‘global’ and ‘local’ views from each modality which are then used by the teachers and students respectively. First, we apply augmentations on $x_{v}$ and $x_{a}$ to generate the global views as $x_{v}^{g}$ and $x_{a}^{g}$. Next, we generate $n$ local views $x_{v}^{l}$ and $x_{a}^{l}$ from $x_{v}$ and $x_{a}$ respectively, where $x_{v}^{l}=\{x_{v1}^{l},\dots,x_{vn}^{l}\}$ and $x_{a}^{l}=\{x_{a1}^{l},\dots,x_{an}^{l}\}$. Specifically, $x_{ai}^{l}$ is an augmented time-frequency crop of $x_{a}$ and $x_{vi}^{l}$ is an augmented spatio-temporal crop of $x_{v}$. To further elaborate, $x_{v}^{l}$ and $x_{a}^{l}$ are differently augmented than $x_{v}^{g}$ and $x_{a}^{g}$. Both local and global views of audio and visual inputs are projected into an embedding space. For example, $x_{a}\in\mathbb{R}^{F\times T_{a}}$ are reshaped into $N_{a}$ smaller patches of size $f\times t_{a}$, where $N_{a}=F/f\times T_{a}/t_{a}$. Similarly, we reshape the videos $x_{v}\in\mathbb{R}^{T_{v}\times H\times W\times C}$ into $N_{v}$ smaller cuboids of size $t_{v}\times\!h\times\!w\times\!C$, where $N_{v}=T_{v}/t_{v}\!\times\!H/h\!\times\!W/w\!\times\!C$. Finally, the spectrogram patches and visual cuboids are flattened into vectors and linearly projected onto the embedding space, which are then fed to the encoders.

Inspired by the recent success and scalability of pretraining with masked reconstruction in different domains (Devlin et al. 2018; Bao, Dong, and Wei 2021; Wang et al. 2021; Gong, Chung, and Glass 2021a; Baevski et al. 2022; He et al. 2021; Niizumi et al. 2022b; Tong et al. 2022), we adopt masked data modelling in our framework to learn modality-specific representations. The masked reconstruction employs an autoencoder $\theta_{ae}$, which consists of an encoder $\theta_{e}$ and a decoder $\theta_{d}$. Let $x$ be the input, which can be further expressed as a sequence of tokens $\{x_{i}\}_{i=1}^{N}$, as described in Section 3.2. Here, we randomly mask some of the input tokens with $m\in\{0,1\}^{N}$, hence the masked tokens $x^{[m]}$ are represented as $\{x_{i}|m_{i}=1\}_{i=1}^{N}$ while the corrupted inputs $\hat{x}$ are represented as $\{x_{i}|m_{i}=0\}_{i=1}^{N}$. Further, we drop the masked tokens $x^{[m]}$ before feeding the input to $\theta_{ae}$ for computational efficiency (Akbari et al. 2021; He et al. 2021). We train $\theta_{ae}$ to reconstruct $x$, based on input $\hat{x}$. Here, $\theta_{ae}$ is trained to minimize the reconstruction loss $\mathcal{L}_{\mathrm{recon}}$ as:

In particular, using Equation 1, we define the video and audio reconstruction losses as $\mathcal{L}_{\mathrm{recon}}^{v}$ and $\mathcal{L}_{\mathrm{recon}}^{a}$ for given inputs $x_{v}^{g}$ and $x_{a}^{g}$, to train $\theta_{ae}^{v}$ and $\theta_{ae}^{a}$, respectively. Here $\theta_{ae}^{v}$ and $\theta_{ae}^{a}$ denote video and audio autoencoders. To jointly train the audio-visual masked autoencoder, we define the final reconstruction loss $\mathcal{L}_{\mathrm{ae}}$ as:

To facilitate cross-modal knowledge sharing, we adopt a teacher-student setup (Tarvainen and Valpola 2017). The teacher ($\theta_{t}$) and student ($\theta_{s}$) are comprised of a backbone and a projector head ($\theta_{h}$), where the teacher and student network architectures are the same but differently parameterized. Moreover, we parameterize $\theta_{s}$ as $\{\theta_{e}$, $\theta_{h}\}$, where $\theta_{e}$ is the same encoder used in reconstruction (explained in the previous subsection) and $\theta_{h}$ is a newly added projector head. We define $\theta_{s}^{v}$ and $\theta_{s}^{a}$ as video and audio students, whereas, $\theta_{t}^{v}$ and $\theta_{t}^{a}$ are denoted as video and audio teachers.

As mentioned earlier, cross-modal knowledge distillation between audio and video is particularly difficult due to the inherent diversity, complexity, and domain-specific nature of these two modalities. To tackle this, we perform domain alignment by identifying the most transferable features and minimizing the domain discrepancy. The proposed domain alignment ensures meaningful target distributions are set by the teachers in order to perform successful cross-modal knowledge distillation.

To train XKD, we define the final loss function as the combination of reconstruction, domain alignment, and cross-modal knowledge distillation losses expressed as:

Here, $\lambda_{\mathrm{ae}}$, $\lambda_{\mathrm{da}}$, and $\lambda_{\mathrm{kd}}$ are the loss coefficients corresponding to the three loss terms, respectively. Empirically, we set $\lambda_{\mathrm{ae}}$, $\lambda_{\mathrm{da}}$, and $\lambda_{\mathrm{kd}}$ as $5$, $1$, and $1$ respectively. It should be noted that $\mathcal{L}_{\mathrm{xkd}}$ is only used to train $\theta_{s}$ and $\theta_{ae}$, not $\theta_{t}$. EMA (Grill et al. 2020; Tarvainen and Valpola 2017; Chen, Xie, and He 2021) is used to slowly update $\theta_{t}^{a}$ and $\theta_{t}^{v}$ as:

where $\lambda_{a}$ and $\lambda_{v}$ are the EMA coefficients corresponding to $\theta_{t}^{a}$ and $\theta_{t}^{v}$. We present the proposed algorithm in Alg. 1.

We take our proposed approach a step forward and attempt to train it in a modality-agnostic fashion with the goal of developing a ‘general’ network capable of handling both modalities. This is a very challenging task in the context of our work given the very diverse nature of audio and video streams. We introduce two modality-agnostic variants XKD-MATS and XKD-MAS. As the name suggests, in XKD-MATS the audio and visual teachers share their backbones, and so do the audio and visual students. In the case of XKD-MAS, the audio and visual students share their backbones, while the audio and visual teachers use modality-specific backbones. Please note that in all the setups, we use the Equation 13 to update the teachers using their respective students. Moreover, given the need to reconstruct different modalities, all of the variants use modality-specific decoders and input projection layers. The rest of the setups remain the same as the original XKD. Both variants are trained with $\mathcal{L}_{\mathrm{xkd}}$ (see Equation 12).

As discussed, the proposed XKD is pretrained to solve masked data modelling and cross-modal knowledge distillation. Therefore, to obtain an accurate understanding of the impact of cross-modal knowledge distillation, we compare XKD ($\mathbf{\mathcal{L}_{\mathrm{\bf xkd}}}$) with respect to the following baselines:

• •

  masked video reconstruction ($\mathbf{\mathcal{L}_{\mathrm{\bf recon}}^{v}}$)

• •

  masked audio reconstruction ($\mathbf{\mathcal{L}_{\mathrm{\bf recon}}^{a}}$)

• •

  audio-video masked autoencoder ($\mathbf{\mathcal{L}_{\mathrm{\bf ae}}}$).

We pretrain the above variants for the full training schedule of $800$ epochs and report linear evaluation results on the split-1 of downstream benchmarks.

We conduct a thorough ablation study investigating the impact of domain alignment in our proposed framework, as presented in Table 3. First, without domain alignment, the model fails to perform cross-modal knowledge distillation due to domain discrepancy, and the model behaves as an audio-visual masked autoencoder. Second, quite expectedly, naively aligning the two domains without knowledge distillation has a minor impact. Last, our results exhibit that the proposed $\mathcal{L}_{da}$ and $\mathcal{L}_{kd}$ are complementary to each other, and their combined effect significantly improves the performance, e.g., by $8.1-14.2\%$ on UCF101, HMDB51, and Kinetics400. In addition to the absence of domain alignment, we identify two more factors that could cause training instability, discussed in the Suppl. Mat. Sec. B.5 and B.6. Please see alternative design choices for domain alignment in the Suppl. Mat. Sec. B.7.

To study the impact of $\mathrm{refine}$ in cross-modal knowledge distillation, we modify ${\mathcal{L}_{\mathrm{da}}}$ in the final loss function (Equation 12) as $\mathcal{L}_{\mathrm{mmd}}(\theta_{t}^{a}(x_{a}^{g}),\theta_{t}^{v}(x_{v}^{g}))+\mathcal{L}_{\mathrm{mmd}}(f_{s}^{v},f_{s}^{a})$. The results presented in Table 4 demonstrate that $\mathrm{refine}$ improves downstream performance by $5.8\%$, $3.7\%$, and $3.0\%$ on HMDB51, UCF101, and ESC50 respectively. In Figure 3, we visualize the attention from the last layer with and without $\mathrm{refine}$, which further confirms the ability of $\mathrm{refine}$ in identifying the most transferable and key features for enhanced downstream performance. Please see alternative design choices of feature refinement Suppl. Mat. Sec. B.8.

In Table 6, we compare the performance of modality-agnostic (MA) variants with the modality-specific (MS) one. Amongst the MA variants, XKD-MATS works slightly better on ESC50, whereas, XKD-MAS shows better performance on both UCF101 and HMDB51. The results are promising, as these variants show minor performance drops (e.g., $0.7\%-4.1\%$) in comparison to the default XKD. Please note that such performance drops are expected as we keep the encoder size fixed for both MA and MS variants (Akbari et al. 2021). To further elaborate, while we dedicate $2$ separate backbones of $87$M parameters to learn audio and visual representations for the MS variant, we use just one backbone of $87$M parameters for both audio and visual modalities in the MA variants. Therefore, the network parameters or model weights become saturated relatively quickly, limiting their performance. A simple solution could be to use a larger backbone for MA variants, which could be explored in future.

In Table 6, we present the comparison between teachers and students. Our results exhibit that due to the slow weight update schedule using EMA (see Equation 13), the teachers become slightly stronger learners compared to the students. We study the impact of different EMA schedules (presented in the Suppl. Mat. Sec. B.2) and find that EMA coefficients as $0.997$ work best in our setup. As presented in Table 6, the video teacher outperforms the video student by $0.3\%$ and $1.5\%$ on UCF101 and HMDB51. Next, the audio teacher outperforms the audio student by $0.7\%$ on ESC50. Please note that by default, we use the teachers to perform downstream tasks. We present a detailed comparison between the modality-specific and modality-agnostic variants using both teacher and student encoders in the Suppl. Mat. Sec. B.9.

We study the scalability of XKD on $3$ pretraining datasets of different sizes, similar to (Sarkar and Etemad 2023; Ma et al. 2020), i.e., Kinetics-Sound ($22$K), Kinetics400 ($240$K), and AudioSet ($1.8$M). Additionally, we experiment with $2$ different sizes of ViT variants, i.e., ViT-B ($87$M) and ViT-L ($305$M). We try both ViT-B and ViT-L as the video backbone when pretrained on AudioSet. We report the linear evaluation top-1 accuracy averaged over all the splits on UCF101, HMDB51, and ESC50. Figure 4 shows that XKD continues to perform better as we increase the number of training samples and/or the size of the network. Such scalability is a much-desired property, which shows XKD can likely be scaled on even larger datasets like HowTo100M (Miech et al. 2019) or larger networks like ViT-22B (Dehghani et al. 2023). We also study the effect of longer pretraining of XKD, presented in the Suppl. Mat. Sec. B.10.

Following the standard practice in (Alwassel et al. 2020; Sarkar and Etemad 2023; Ma et al. 2020), we compare XKD with other leading audio-visual self-supervised frameworks in Table 7 in both linear evaluation (Lin.) and finetuning (FT.) on UCF101, HMDB51, and Kinetics400. We note a large variability in the experiment setups amongst the prior works, however, they are included for a more inclusive and thorough comparison. We report top-1 accuracy averaged over all the splits for both UCF101 and HMDB51.

The results presented in Table 7 show that XKD outperforms or achieves competitive performance amongst the leading audio-visual self-supervised methods. For example, XKD pretrained on Kinetics400 outperforms AVID and CrissCross in linear evaluation on Kinetics400, by a very large margin. Moreover, XKD outperforms powerful state-of-the-art models like STiCA, BraVe, and CrissCross among several others when finetuned on UCF101. Additionally, XKD outperforms prior works that are pretrained with massive datasets and tri-modalities like Elo pretrained with video, audio, and optical flow from YouTube8M (YT) (Abu-El-Haija et al. 2016), compared to XKD pretrained with audio and video from AudioSet 2M. Our modality-agnostic variants also achieve encouraging results, e.g., a minor performance drop is noticed between XKD-MAS and XKD, e.g., $0.7\%$ on UCF101 and $1.7\%$ on Kinetics400, when finetuned.

In Table 8, we compare the performance of our proposed method using linear evaluation (Lin.) and finetuning (FT.) on sound classification using $2$ popular audio benchmarks ESC50 and FSD50K. Following (Fonseca et al. 2022; Piczak 2015), we report top-1 accuracy averaged over all the splits on ESC50 and mean average precision on FSD50K. XKD outperforms prior works like XDC, AVID, and CrissCross on ESC50 in both linear evaluation and finetuning. Moreover, XKD and its MA variants outperform VATT and VATT-MA on ESC50 by $4.7\%$ and $7.5\%$, even though VATT is pretrained with $136$M videos, compared to XKD which is only trained on $240$K samples. Additionally, when evaluated on FSD50K, XKD outperforms BYOL-A, AudioTransformer, and PSLA among others. Lastly, XKD shows state-of-the-art performance on ESC50, achieving top-1 finetuned accuracy of $96.5\%$ when pretrained with AudioSet.

Following (Sarkar and Etemad 2023), we evaluate XKD in multimodal action classification using Kinetics-Sound. We extract fixed audio and visual embeddings from the pretrained encoders and concatenate them together (i.e., late fusion), followed by a linear SVM classifier is trained and top-1 accuracy is reported. The results presented in Table 9 show that XKD outperforms CrissCross by $14.5\%$ and baseline audio-visual masked autoencoder ($\mathbf{\mathcal{L}_{\bf ae}}$) by $5.5\%$.

We present reconstruction examples in Figure 5, which shows that XKD retains its reconstruction ability even when a very high masking ratio is applied, for both audio and video modalities. This makes XKD also suitable for in-painting tasks. More examples are in the Suppl. Mat. Sec. C.2.