[논문] Contrastive Audio-Visual Masked Autoencoder

https://arxiv.org/abs/2210.07839

Contrastive Audio-Visual Masked Autoencoder

해당 논문을 보고 작성했습니다.

 

Abstract

이 논문에서 저자들은 single modality에서 동작하는 Masked Auto-Encoder (MAE) model을 audio-visual multi-modality로 확장한 첫 연구를 제안합니다. 저자들은 contrastive learning과 masked data modeling을 결합하여 공통 audio-visual representation을 학습하는 Contrastive Audio-Visual Masked Auto-Encoder (CAV-MAE)을 제안했습니다.

저자들의 실험을 통해 contrastive audio-visual correspondence learning objective는 model이 audio-visual retrieval task만 잘 수행할 뿐만 아니라 joint representation도 잘 학습함을 보여줍니다. 

 

Introduction

acoustic and visual modality는 다른 특성을 가지고 있지만, 사람들은 원활하게 연결하고 통합하여 인식할 수 있습니다. 이러한 능력을 모방하는 학습 알고리즘을 개발하는 것은 매우 흥미로운 연구 주제입니다. 수동으로 audio와 video에 annotation을 만드는 것은 비싸고 어렵기 때문에 web-scale unlabeled video data를 사용하는 self-supervised 연구들이 등장했습니다.

audio-visual self-supervised learning 연구의 주요 흐름 중 하나로 video에서 발견되는 자연스러운 audio-visual correspondence를 활용하는 것입니다. 그중 하나로 Contrastive Audio-Visual Learning이 있으며, 간단하지만 효과적인 방식입니다. 이는 mismatched sample보다 paired audio-visual sample에 대한 representation이 더 가깝도록 학습됩니다. 이러한 coordinated representation은 cross-modal retrieval과 같은 task에 특히 유용합니다.

또 다른 널리 사용되는 self-supervised learning framework로 Masked Data Modeling (MDM)이 있으며, 이는 손상된 input이나 feature로부터 original input 또는 feature를 복원하는 pretext task를 통해 의미 있는 representation을 학습합니다.

Audio Spectrogram Trasnformer와 Vision Transformer backbone 기반인 single-modal Masked Auto-Encoder (MAE)가 image and audio task 모두에서 SOTA performance를 달성했습니다. 이러한 성공에 영감을 받아, 저자들은 single-modal MAE를 unimodal signal을 융합하는 joint representation을 학습하는 것을 목표로 하는 Audio-Visual Masked Auto-Encoder (AV-MAE)로의 확장을 제안합니다. 

두 major self-supervised framework가 각각 널리 사용되어 왔지만, 저자들이 알기론 audio-visual learning에서는 한번도 결합된 적이 없었습니다. 저자들은 상호 보완적이라는 것을 발견했습니다. Contrastive auiod-visual learning은 명시적으로 매우 유용한 audio-visual pair information을 사용하지만, downstream task에서 유용한 modality-unique information은 손실할 수 있습니다.

이를 통해 저자들은 단일 model로 joint and coordinated audio-visual representation을 학습할 수 있는 contrastive learning and masked data modeling을 사용한 Contrastive Audio-Visual Masked Autoencoder (CAV-MAE)를 제안합니다.

저자들의 contribution을 요약하면 다음과 같습니다.

• cross-modal masked data modeling을 통해 audio-visual input을 fuse하여 self-supervised learning을 수행하는 multi-modal AV-MAE를 제안했습니다.
• 저자들은 contrastive audio-visual learning과 masked data modeling을 최적으로 결합하는 방법에 대한 연구를 수행했으며, CAV-MAE를 제안했습니다.
• 저자들은 contrastive and masked data modeling objective가 상호 보완적이라는 것을 입증했습니다.

Contrastvie Audio-Visual Masked Autoencoder

 

Preliminaries

- Audio and Image Pre-processing and Tokenization

 

위 그림과 같이 저자들은 pre-processing과 tokenization을 audio와 image input에 적용했습니다. audio에는 Audio Spectrogram Transformer (AST)를 적용하고, image에는 ViT를 적용하였습니다. 구체적으로 AudioSet과 VGGSound에 있는 10초 길이 video (with parallel audio)를 사용하여 model을 pretrain 하고 fine-tune 했습니다. audio의 경우, 각 10초 길이 audio waveform을 128차원 log Mel filterbank (fbank) feature sequence로 변환합니다. 결과로 1024(time) x 128(freq) spectrogram을 얻습니다. 그다음 spectrogram을 512, 16x16 square patch로 나누어 model의 input으로 사용합니다. Transformer model로 video를 처리하는 것은 비용이 많이 들고 연산 자원이 많이 필요로 합니다. 그래서 연산량을 저자들의 resource에 맞추기 위해, 저자들은 frame aggregation strategy를 사용했습니다. 구체적으로 각 10초 길이 video에서 10 RGB frame을 sample 합니다 (i.e., 1 FPS). 학습할 때 random 하게 RGB frame sample 중 하나를 선택하고 input으로 사용합니다. inference 할 때는 각 RGB frame에 대한 model prediction의 평균을 구해 video prediction으로 사용합니다. 여러 RGB frame을 concatenation 하여 Transformer의 input으로 사용하면 연산량이 제곱으로 증가되지만, frame aggregation은 inter-frame correlation을 고려하지 않기 때문에 time에 선형적으로 연산량이 증가됩니다. 각 RGB frame의 경우, 저자들은 resize and center crop 하여 224x224 sample을 얻으며, 이를 196, 16x16 square patch로 나눕니다.

 

- The Transformer Architecture

이 논문 전반에 걸쳐, 저자들은 standard Transformer를 사용했습니다. 각 Transformer layer는 multi-headed self-attention (MSA), layer normalization (LN), residual connection이 있는 multilayer perceptron (MLP) block으로 구성됩니다. 구체적으로 Transformer layer y = Transformer(x; MSA, LN11, LN2, MLP)는 다음과 같습니다.

 

MSA는 x의 각 element의 dot-product attention을 계산하며, x 크기의 제곱만큼 연산량이 필요합니다. 

 

- Contrastive Audio-Visual Learning (CAV)

video의 자연스러운 audio and visual information pairing은 self-supervision을 통해 audio-visual representation을 학습하기에 유용한 signal 입니다. conventional CAV model은 다음과 같습니다.

 

mini-batch로 N개 audio-visual pair sample을 사용하며, audio와 iamge를 pre-process하고 tokenize 합니다. 그리고 각 sample i에 대해 audio and visual token {a_i, v_i} sequence를 얻습니다. 그다음 a_i는 audio transformer E_a에 넣고 v_i는 visual transformer E_v에 넣습니다. 그다음 mean pooled audio representation c_i^a, visual representation c_i^v를 얻습니다. c_i^a = MeanPool(E_a(Proj_a(a_i)))이고 c_i^v = MeanPool(E_v(Proj_v(v_i)))이며, 여기서 Proj_a는 audio token을 768차원으로 mapping 하는 linear projection을 나타내고, Proj_v는 visual token을 768차원으로 mapping 하는 linear projection을 나타냅니다. c_i^a와 c_i^v에 contrastive loss를 적용했습니다.

 

- Single Modality Masked Autoencoder (MAE)

또 다른 self-supervised framework의 주요 line은 masked data modeling (MDM) 입니다. MDM의 다양한 변형들 중, masked auto-encoder (MAE)가 간단하지만 효과적입니다. tokenize 된 input sample x가 있을 때, MAE는 input x_{mask}의 일부를 mask 하고 unmasked token을 Transformer 기반 encoder-decoder model에 feed 합니다. model은 mean square error (MSE) loss를 minimize 하도록 masked token을 reconstruct 합니다. 이 과정에서 model은 input data의 의미 있는 representation을 학습합니다. MAE의 이점은 다양합니다. MAE는 original input을 prediction target으로 바로 사용하며, 이를 통해 학습 과정을 매우 단순화시킬 수 있습니다. MAE는 unmasked token을 encoder의 input으로 사용하고, maksing ratio가 높은 경우 MAE의 연산량은 상당히 줄어듭니다. 그리고 MAE는 audio modality와 visual modality 둘 다에서 single-modal task에 대해 좋은 성능을 보였습니다.

 

Vanilla Audio-Visual Masked Autoencoder (AV-MAE)

여태까지 MAE를 audio와 visual modality 각각에 적용되었지만, audio-visual multi-modality learning에 사용된 적은 없습니다. 이 연구의 첫 contribution으로, 저자들은 single modality를 audio-visual multi-modality로 MAE를 확장했으며 vanilla audio-visual autoencoder (AV-MAE)를 구성했습니다. 

 

위 그림과 같이 audio and image pair input들을 tokenize 한 후 768차원으로 project합니다. projection layer는 두 modality에 따라 다르며, 각 modality에 맞춰 2-D sinusoidal positional embedding을 더했습니다. 즉 a' = Proj_a(a) + E_a + E_a^P가 되며, 여기서 a는 audio input token을 나타내며, E_a는 audio라는 modality에 대한 type embedding이고, E_a^P는 2-D sinusoidal postional embedding을 의미합니다. image에도 똑같이 적용하여 v' = Proj_v(v) + E_v + E_v^P를 구합니다. a'과 v'을 concatenate 하고 joint embedding x = [a', v']를 만듭니다. 그다음 75% 정도 masking을 수행하고 unmasked token만 audio-visual joint encoder E_j()에 feed 하여 output x'_{unmask}를 얻습니다. 그다음 original position에 맞춰 x'_{unmask}에 학습 가능한 masked token을 padding합니다. 이제 audio-visual joint decoder D_j()에 feed 하기 전에, padding 된 token들을 가지고 각 modality의 2-D sinusoidal positional embedding을 더해주고 modality type embedding E'_a, E'_v를 더해줍니다. joint decoder D_j()에 값을 넣어 input을 reconstruct 하여 a^, v^ = D_j(x' + [E'_a, E'_v] + [E_a^P, E_v^P])를 얻고, mean square error (MSE)를 적용하여 이 error를 minimize 하도록 학습됩니다.

single-modal MAE와 비교했을 때, AV-MAE는 cross-modal masked data modelling objective를 사용하는데, 이는 model이 한 modality를 reconstruct하기 위해서 다른 modality의 정보도 사용하도록 만들어줍니다. 이를 통해 model이 audio-visual correlation을 학습하는데 도움을 줄 수 있습니다. 하지만 paired audio-visual correspondence에 대한 명시적인 objective가 없다면, vanilla AV-MAE는 audio-visual pairing 정보를 효율적으로 사용할 수 없게 됩니다. 또한 두 modality에 대한 joint encoder를 사용하면 cross-modal attention이 가능하지만, 매우 다른 두 modality가 동일한 weight로 처리됨을 의미하며 sub-optimal solution으로 수렴될 가능성이 존재하게 됩니다.

 

Contrastive Audio-Visual Masked Autoencoder (CAV-MAE)

contrastive audio-visual learning과 AV-MAE 각각 장점과 단점 모두 존재합니다. 저자들은 CAV와 AV-MAE의 상호보완적 장점 모두 얻고 싶어합니다. 이를 위해, 저자들은 Contrastive Audio-Visual Masked Autoencoder (CAV-MAE)를 제안합니다. 그림으로 보면 다음과 같습니다.

 

mini-batch로 N개 audio-visual pair sample을 사용하며, 저자들은 audio와 image에 pre-processing & tokenize를 적용하여 audio and visual token sequence {a_i, v_i}를 얻었습니다. 그다음 각 modality에 맞는 linear projection layer를 사용하여 768차원으로 project 합니다. 그리고 modality type embedding E_a, E_v를 더하고 각 modality에 맞는 2D sinusoidal positional embedding E_a^p, E_v^p를 더했습니다. 그다음 각 modality에 75% token들을 masking 합니다. 

 

식으로 나타내면 위와 같습니다. 그 다음 a_i^{unmask}, v_i^{unmask}를 각 audio Transformer encoders E_a(), visual Transformer encoders E_v()에 feed 하여 a'_i, v'_i를 얻습니다. 그다음 a'_i, v'_i에 multi-stream forward pass로 joint audio-visual encoder E_j (MSA, LN1, LN2, MLP)를 사용합니다. 구체적으로 input audio tokens a'_i, video tokens v'_i, concatenated audio-visual tokens [a'_i, v'_i]를 각각 독립적인 forward pass를 적용합니다. 각 stream에서 서로 다른 normalization layer LN1_{a, v, av}, LN2_{a, v, av}를 사용하지만, E_j의 나머지 다른 모든 weight들 (MSA, MLP의 weight)은 공유됩니다. 

 

audio and visual single modality stream output인 c_i^a, c_i^v를 contrastive learning에 사용하고 audio-visual multi-modal stream output x_i는 reconstruction task에 사용합니다. contrastive audio-visual learning의 경우, contrastive loss L_c는 다음과 같습니다. 

 

s_{i, j} = ||c_i^v||^T||c_i^a||를 나타냅니다.

reconstruction task의 경우, 저자들은 x_i에 학습 가능한 masked token을 사용하여 original position으로 padding을 수행합니다.  그 다음 modality type embedding E'_a, E'_v를 더한 후 각 modality에 대한 2D sinusoidal positional embedding을 더해 joint audio-visual decoder D_j()에 feed 합니다. 그렇게 input audio and image를 reconstruct 합니다. D_j()는 마지막 modal-specific projection layer를 제외한 모든 weight를 공유하여 audio and visual token을 처리합니다. 그다음 mean square error reconstruction loss L_r를 적용합니다.

 

N은 mini-batch size를 나타내고 hat이 없는 term들은 original masked patch를 나타내고, hat이 있는 term들은 predicted masked patch를 나타냅니다. |a_^{mask}|는 masked audio patch 수를 나타내고, |v_^{mask}|는 masked visual patch 수를 나타냅니다. 

최종적으로 contrastive loss L_c와 reconstruction loss L_r을 결합하여 CAV-MAE loss로 사용합니다. 

 

pre-train한 이후, 저자들은 decoder는 버리고 encoder만 downstream task에 사용합니다. 이 encoder를 통해 single-modality stream output과 multi-modal modality stream output을 더해서 사용할 수 있고, 또는 multi-modal stream output만 사용하여 finetuning을 수행해도 됩니다.

 

다음은 저자들이 CAV-MAE의 key design을 왜 위와 같이 구현했는지에 대한 이유입니다.

 

    - Multi-stream forward passes of the joint encoder

저자들은 contrastive audio-visual learning에 사용되는 representation들을 제한하는 것이 중요하다고 생각했습니다. c^a는 audio input에서만 등장해야 하고, c^v는 visual input에서만 등장해야만 하며, 만약 그렇제 않다면 contrastive objective는 collapse 될 수 있기 때문에, representation들을 제한하는 것이 중요하다고 생각했습니다. 동시에 저자들은 encoder가 reconstruction task와 downstream task를 위해 audio and visual information을 잘 결합하길 원했습니다. 그래서 위와 같이 multi-stream forward pass strategy를 사용해 설계했다고 합니다.

 

    - Modality-specific encoders and LN layers

통합된 network로 audio and visual modality를 동시에 처리하는 연구들이 등장했지만, audio and visual modality의 매우 다른 특성 때문에 modality-specific network가 아직까지 성능이 더 뛰어난 모습을 보인다는 결론을 얻었습니다. 그래서 저자들은 audio and visual input을 modality-specific encoder를 사용하여 encode 한 다음 joint encoder에 feed 하는 것을 선택했습니다. 동일한 이유로 joint encoder의 각 stream 마다 서로 다른 normalization statistic을 사용하도록 만들었습니다.

효율성 측면에서, 두 modality-specific encoder를 사용하는 것이 model size를 증가시키지만 input sequence length의 제곱으로 복잡도가 증가하는 Transformer이기 때문에 연산량을 줄일 수 있었다 합니다. 

 

    - Masked contrastive audio-visual learning

single-modality contrastive learning과 다르게, conventional contrastive audio-visual learning 은 augmentation과 masking을 하지 않습니다. 이 논문에서 저자들은 masked contrastive audio-visual learning을 제안하며, 즉 random하게 input의 일부를 masking 하여 contrastive learning에 사용합니다. 이를 통해 AV-MAE에 CAV를 결합할 뿐만 아니라, overfitting을 방지할 수 있습니다. masking ratio가 75%이고 bach size가 27일 때, pretraining 할 때 audio-visual matching accuracy가 약 72% 정도 나오고 이는 task가 불가능한 것도 아니지만 쉽지 않음을 보여줍니다.

 

- Implementation Details

기본적으로 모든 encoder Transformer layer는 768차원이고 12개 attention head를 가지고 있습니다. joint encoder로 Vanilla AV-MAE이고, 12-layer Transformer를 사용합니다.

CAV-MAE의 audio and visual encoder는 11개 layer로 구성된 Transformer (각 768차원)이고 joint encoder는 single-layer Transformer입니다. 즉, encoder의 총 layer 수는 12개이지만, modlaity-specific encoder를 사용하기 때문에 CAV 및 CAV-MAE는 더 큰 model입니다. AV-MAE와 CAV-MAE의 decoder는 512 embedding dimension & 16개 attention head를 가지는 8개 layer Transformer입니다. contrastive loss temperature τ = 0.05로 고정했습니다 .CAV-MAE의 경우, 저자들은 λ_c = 0.01로 설정했습니다. L_c가 L_r에 비해 더 큰 gradient를 보이기 때문에 λ_c를 상대적으로 작게 설정했으며, 이게 contrastive objectiv가 덜 중요하다는 것을 의미하는 건 아닙니다. encoder와 decoder는 각각 164M, 27M개 parameter를 갖습니다. 

audio-visual learning에 대한 일반적인 관행에 따라, 저자들은 ImageNet으로 pretrain된 weight를 model의 초기 weight로 사용합니다. 구체적으로 original vision MAE의 weight를 초기 weight로 사용함을 의미합니다. supervised pretrianed weight를 사용하는 이전 연구들과 다르게, 저자들은 self-supervised pretrained weight만 사용합니다. 이를 통해 가장 좋은 성능을 얻을 순 없지만, 모든 학습 과정일 self-supervised pipeline으로 만들 수 있었습니다. 

 

Self-Supervised Model Pretraining

저자들은 다음 model들을 가지고 pretrain을 수행하고 성능을 비교했습니다.

Audio-MAE/Visual-MAE

single-modal masked auto-encoder model입니다. model 구조는 Vanilla AV-MAE와 동일하지만 single modality data를 가지고 pretrain되었습니다.

CAV

reconstruction objective가 없는 contrastive audio-visual learning model입니다. 공평한 비교를 위해, 저자들은 CAV-MAE와 동일한 encoder 구조 (modal-specific encoders + joint encoder)를 사용했으며, reconstruction objective L_r은 제거했습니다.

Vanilla AV-MAE

joint encoder를 가지고 있고 contrastive objective가 없는 vanilla audio-visual masked auto-encoder입니다.

AV-MAE

two modal-specific encoder를 사용하는 audio-visual masked auto-encoder와 joint encoder입니다. CAV-MAE와 동일한 구조를 띄지만, λ_c = 0으로 설정하여 contrastive loss를 사용하지 않도록 만들었습니다. 저자들은 modal-specific encoder의 영향을 확인하기 위해 Vanilla AV-MAE와 비교할 때 이 model을 사용하며, contrastive objective의 영향을 확인하기 윟 ㅐCAV-MAE와 비교할 때 이 model을 사용합니다.

CAV-MAE

저자들이 제안하는 contrastive masked auto-encoder 구조입니다.

CAV-MAE^{scale+}

CAV-MAE와 동일한 mdoel이지만, 학습할 때 batch size를 108로 키웠습니다.

 

공평한 비교를 위해 CAV-MAE^{scale+}를 제외한 모든 model들은 동일한 pipeline에서 pretrain 되었습니다. pretrain 과정에서 저자들은 의도적으로 class balanced sampling을 사용하지 않았습니다. 모든 pre-training process는 fully self-supervised입니다. 

 

Audio-Visual Retrieval

CAV-MAE가 audio-visual retrieval을 수행할 수 있도록 audio-visual correspondence를 잘 capture 하는 cooredinated representation을 학습하는지 확인하겠습니다. 구체적으로 저자들은 AudioSet에서 1725개 audio-visual sample을 추출하고 VGGSound에서 1545개 audio-visual sample을 추출했습니다. 저자들은 audio와 image를 각 model에 independent forward pass에 따라 feed 하였습니다. 그다음 rank 1, 5, 10으로 retrieval recall을 계산했습니다. 이때 audio visual reprsentation에 cosine similarity를 사용하였습니다. 모든 model들은 self-supervised pretrained였지만 fine-tune 되지는 않았습니다. 

 

위 표를 통해 visual data로 audio를 retrieval 한 결과를 보여줍니다. 

 

위는 sample 예시를 보여줍니다. 저자들은 vanilla-MAE와 AV-MAE가 거의 random guess와 유사한 성능을 보인다는 것을 확인했고, contrastie objective가 audio-visual retrieval task에서 필수적임을 보였습니다. cross-modal masked data modeling objective는 retrieval 성능을 향상시켜줍니다. CAV-MAE가 CAV보다 더 나은 성능을 보여줍니다. batch size를 키우고 training epoch을 늘리는 것이 더 나은 retrieval 성능을 보여줍니다. pretrain dataset과 test dataset이 달라도 retrieval 성능이 경쟁력 있으며, audio-visual correspondence가 audio and visual representation에서 잘 동작한다는 것을 보여줍니다.

 

Conclusion

이 논문에서 저자들은 새로운 audio-visual learning model인 CAV-MAE를 제안합니다. 이 paper의 main idea는 간단합니다: masked data modeling and contrastive learning은 audio-visual self-supervised learning에서 서로 상호보완적입니다.

효율적으로 두 framework를 결합하고 representation collapse를 피하는 것은 multi-stream forward pass strategy, joint-specific encoder architecture, masked contrastive learning과 같이 세밀하게 design 된 몇 가지 module을 필요로 합니다. representation learning 관점에서 CAV-MAE는 joint and coordinated representation을 학습하고 audio-visual joint event classification & coordinated audio-visual retrieval task 모두에서 잘 동작합니다. 결과적으로 audio-visual event classification task에서 CAV-MAE가 SOTA model과 동일한 성능을 보이거나 더 뛰어난 성능을 보이기도 합니다. 마지막으로 CAV-MAE multi-modal pretraining이 strong single-modal representation을 학습할 수 있으며 audio-based event classification task에서 새로운 SOTA 성능을 보였습니다.