TriBERT: Human-centric Audio-visual Representation Learning

The recent success of transformer models in language , such as BERT , has motivated the use of such architectures for multi-modal feature learning and tasks . However , most multi-modal variants ( e.g. , ViLBERT ) have limited themselves to visuallinguistic data . Relatively few have explored its use in audio-visual modalities , and none , to our knowledge , illustrate them in the context of granular audio-visual detection or segmentation tasks such as sound source separation and localization . In this work , we introduce TriBERT – a transformer-based architecture , inspired by ViLBERT , which enables contextual feature learning across three modalities : vision , pose , and audio , with the use of flexible co-attention . The use of pose keypoints is inspired by recent works that illustrate that such representations can significantly boost performance in many audio-visual scenarios where often one or more persons are responsible for the sound explicitly ( e.g. , talking ) or implicitly ( e.g. , sound produced as a function of human manipulating an object ) . From a technical perspective , as part of the TriBERT architecture , we introduce a learned visual tokenization scheme based on spatial attention and leverage weak-supervision to allow granular cross-modal interactions for visual and pose modalities . Further , we supplement learning with sound-source separation loss formulated across all three streams . We pre-train our model on the large MUSIC21 dataset and demonstrate improved performance in audio-visual sound source separation on that dataset as well as other datasets through fine-tuning . In addition , we show that the learned TriBERT representations are generic and significantly improve performance on other audio-visual tasks such as cross-modal audio-visual-pose retrieval by as much as 66.7 % in top-1 accuracy . 1 Introduction . Multi-modal audio-visual learning [ 57 ] , which explores and leverages the relationship between visual and auditory modalities , has started to emerge as an important sub-field of machine learning and computer vision . Examples of typical tasks include : audio-visual separation and localization , where the goal is to segment sounds produced by individual objects in an audio and/or to localize those objects in a visual scene [ 15 , 16 , 42 , 55 ] ; and audio-visual correspondence , where the goal is often audio-visual retrieval [ 23 , 47 , 53 ] . Notably , some of the most recent audio-visual methods [ 15 ] leverage human pose keypoints , or landmarks , as an intermediate or contextual representation . This tends to improve the overall performance of sound separation , as pose and motion are important cues for characterising both the type of instrument being played and , potentially , over time , the rhythm of the individual piece [ 15 ] . It can also serve as an intermediate representation when generating video from acoustic signals [ 8 , 44 ] for example . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Most of the existing architectures tend to extract features from the necessary modalities using pre-trained backbones ( e.g. , CNNs applied to video frames [ 55 ] , object regions [ 16 ] , and audio spectrograms ; and/or graph CNN for human pose [ 15 ] ) and then construct problem-specific architectures that often utilize simple late fusion for cross-modal integration in decoding ( e.g. , to produce spectrogram masks [ 15 , 16 , 55 ] ) . This is contrary to current trends in other multi-modal problem domains , where over the past few years , approaches have largely consolidated around generic multi-modal feature learning architectures that are task agnostic to produce contextualized feature representations and then fine-tune those representations to a variety of tasks ( e.g. , visual question answering ( VQA ) or reasoning ( VCR ) ) and datasets . Examples of such architectures include ViLBERT [ 33 ] , VL-BERT [ 46 ] , and Unicoder-VL [ 31 ] , all designed specifically for visual-linguistic tasks . Audio-visual representation learning has , in comparison , received much less attention . Most prior works [ 51 ] assume a single sound source per video and rely on audio-visual alignment objectives . Exceptions include [ 39 ] , which relies on proposal mechanisms and multiple-instance learning [ 49 ] or co-clustering [ 25 ] . These approaches tend to integrate multi-modal features extracted using pre-trained feature extractors ( e.g. , CNNs ) at a somewhat shallow level . The very recent variants [ 6 , 28 , 35 ] leverage transformers for audio-visual representation learning through simple classification [ 6 ] and self-supervised [ 28 ] or contrastive [ 35 ] learning objectives while only illustrating performance on video-level audio-visual action classification . To the best of our knowledge , no audio-visual representation learning approach to date has explored pose as one of the constituent modalities ; nor has shown that feature integration and contextualization at a hierarchy of levels , as is the case for BERT-like architectures , can lead to improvements on granular audio-visual tasks such as audio-visual sound source separation . To address the aforementioned limitations , we formulate a human-centric audio-visual representation learning architecture , inspired by ViLBERT [ 33 ] and other transformer-based designs , with an explicit goal of improving the state-of-the-art in audio-visual sound source separation . Our transformer model takes three streams of information : video , audio , and ( pose ) keypoints and co-attends among those three modalities to arrive at enriched representations that can then be used for the final audiovisual sound separation task . We illustrate that these representations are general and also improve performance on other auxiliary tasks ( e.g. , forms of cross-modal audio-visual-pose retrieval ) . From a technical perspective , unlike ViLBERT and others , our model does not rely on global frame-wise features nor an external proposal mechanism . Instead , we leverage a learned attention to form visual tokens , akin to [ 42 ] , and leverage weakly supervised objectives that avoid single sound-source assumptions for learning . In addition , we introduce spectrogram mask prediction as one of our pre-training tasks to enable the network to better learn task-specific contextualized features . Contributions : Foremost , we introduce a tri-modal VilBERT-inspired model , which we call TriBERT , that co-attends among visual , pose keypoint , and audio modalities to produce highly contextualized representations . We show that these representations , obtained by optimizing the model with respect to uni-modal ( weakly-supervised ) classification and sound separation pretraining objectives , produce features that improve audio-visual sound source separation at large and also work well on other downstream tasks . Further , to avoid reliance on the image proposal mechanisms , we formulate tokenization in the image stream in terms of learned attentional pooling , which is learned jointly . This alleviates the need for externally trained detection mechanisms , such as Faster R-CNN and variants . We illustrate competitive performance on a number of granular audio-visual tasks both by using the TriBERT model directly , using it as a feature extractor , or by fine-tuning it . 2 Related works . Audio-visual Tasks . There exists a close relationship between visual scenes and the sounds that they produce . This relationship has been leveraged to complete various audio-visual tasks . Based on [ 57 ] ’ s survey of audio-visual deep learning , these tasks can be categorized into four subfields , three of which are addressed in this paper and described in the following three subsections . Audio-visual Sound Source Separation and Localization . Sound source separation and the related task of sound source localization have been studied quite extensively . Previous works studying separation , also known as the cocktail party problem [ 19 ] , leverage multi-modal audio-visual information [ 11 , 14 ] to help improve performance with respect to their audio-only counterparts [ 26 , 34 ] . Examples include learning correlations between optical flow and masked frequencies [ 9 , 13 ] , using graphical models [ 21 ] , detecting salient motion signals that correspond to audio events [ 30 , 40 ] , and extracting pose keypoints to model human movements [ 15 ] . A close connection between separation and localization has also been illustrated [ 40 , 43 , 55 , 56 ] . For example , [ 16 , 42 ] both formulate the task as one of auditory and visual co-segmentation , either with pre-trained object regions obtained by the detector [ 16 ] or directly from the image [ 42 ] . All of these approaches contain highly specialized architectures with custom fusion schemes . We aim to leverage the flexibility of transformer models to create generalized multi-modal representations that improve on audio-visual tasks . Audio-visual Representation Learning . The goal is typically to learn aligned representations . The quality of these representations has been shown to greatly impact the overall performance of tasks downstream [ 4 ] . A common strategy for representation learning is to introduce a proxy task . In the audio-visual space , past works [ 1 , 2 , 38 ] have trained networks by having them watch and listen to a large amount of unlabeled videos containing both positive samples of matching audio and visual pairs and negative samples of mismatched pairs ; the proxy task is binary classification of whether or not the audio and visual match each other . Other proxy tasks include determining whether or not an audio-visual pair is time synchronized [ 27 ] ; and [ 29 ] uses a classification task to identify the correct visual clip or audio stream from a set with negative samples . However , these works rely on the assumption that only one main sound source occurs at a time and everything else is background noise . Our model uses a weakly supervised proxy objective to learn representations for multiple sources of sound ( two in experiments ) occurring simultaneously and also learns to incorporate pose features . Audio-visual Correspondence Learning . One of the fundamental tasks in correspondence learning related to our work is cross-modality retrieval . Most prior works focus on audio-visual retrieval [ 24 , 36 , 48 ] and propose learning a joint embedding space where both modalities can be mapped to . In this space , semantically related embeddings are close to each other and thus retrieval can be performed by selecting the closest embedding to the query from the alternate modality . In our work , we demonstrate that enhanced feature representations obtained by our pretrained model capture aligned semantics and lead to much better cross-modal retrieval than baseline representations . Visiolinguistic Representation Learning . Our model is inspired by the recent successes of visiolinguistic representations . Most such approaches leverage a combination of uni-modal and cross-modal transformer modules to pre-train generic visiolinguistic representations on masked language and/or multi-modal alignment tasks . For example , [ 33 ] proposes separate streams for each modality that communicate with each other through co-attention , while [ 46 ] uses a single-stream model that takes both visual and linguistic embeddings as input . In our work , we also leverage co-attention modules to learn joint representations between audio , pose , and vision modalities . However , in addition to extending co-attention , we also focus on reformulating image tokenization and demonstrate the ability to learn with weakly-supervised classification objectives as opposed to masked token predictions . 3 Approach . We introduce TriBERT , a network that learns a joint representation of three modalities : vision , pose , and audio . We briefly review ViLBERT , the architecture that inspired TriBERT , in Section 3.1 . We then describe our TriBERT architecture in Section 3.2 , including pretraining tasks and objectives . 3.1 Reviewing Vision-and-Language BERT ( ViLBERT ) . Motivated by the recent success of the BERT architecture for transfer learning in language modeling , Lu et al . [ 33 ] proposed ViLBERT to represent text and visual content jointly . ViLBERT is a twostream model for image regions and text segments . Each stream is similar to the BERT architecture , containing a series of transformer blocks ( TRM ) [ 50 ] . Given an image I with corresponding regions-of-interest ( RoIs ) or bounding boxes v0 , v1 , ... vN and an input sentence S with word tokens w0 , w1 , ... wT , the final output representations are hv0 , hv1 , ... , hvN and hw0 , hw1 , ... , hwT for the visual and linguistic features , respectively . To exchange information between the two modalities , the authors introduced a co-attentional transformer layer which computes query ( Q ) , key ( K ) , and value ( V ) pairs like a standard transformer block . The keys and values from each modality are then fed to the multi-headed attention block of the other modality . The attention block in each stream generates attention-pooled features conditioned on the other modality and outputs a multi-modal joint representation which outperforms single-stream models across multiple vision-and-language tasks .

Summary: The paper presents TriBERT a BERT-like model for learning multi-model representations of three modalities simultaneously, namely: vision, pose and audio. The architecture of TriBERT takes inspiration from the visual-linguistic model VilBERT (notably known for the applications in VQA) in terms of its Transformer architecture and of its multi-modal attention (the latter is extended in the paper to 3 modalities instead of 2).  TriBERT training comprises four losses to optimize. In particular, after the projection of the 3 unimodal feature sequences into the joint feature space by the multimodal transformer, there are 3 sequences of features, namely: <h_v>, <h_p> and <h_a> for respectively enriched visual, pose and audio features. Then, the <h_v> and <h_p> features are (independently) trained to localize and predict the sound source (the musical instrument(s) which is (are) played in the recording) using 2 (independent) BCE losses. It should be noted that the localization part is learnt in a weakly-supervised manner as only global instrument classification annotations are available in the training data. The <h_a> audio features are trained to classify the sound source (the 3rd BCE loss).  Finally, the 4th supervision task is the multi-modal sound source separation one. The original audio spectrogram is mixed with other audio spectrogram(s) and the task consists in predicting a spectrogram mask which would separate the original audio from the others (the 4th BCE loss). Following [37], for this last task, a U-Net neural architecture is used. U-Net takes the mixed audio spectrogram at its input and the fusion of <h_v>, <h_p> and <h_a> at its bottleneck representation to predict the desired spectrogram mask.  Once trained, TriBERT is evaluated on the MUSIC21 dataset for the sound separation task, and it also fine-tuned and evaluated for the multi-modal retrieval downstream task.