[Paper 리뷰] EnCodecMAE: Leveraging Neural Codecs for Universal Audio Representation Learning

EnCodecMAE: Leveraging Neural Codecs for Universal Audio Representation Learning

---

• Universal audio representation learning은 다양한 downstream task에서 사용할 수 있는 foundational model을 요구함
• EnCodecMAE
  • BERT와 같은 self-supervised universal audio model에 적용하기 위해 audio signal을 discrete class set으로 mapping
  • 결과적으로 Masked AutoEncoder (MAE)를 기반으로 discrete target을 생성하는 neural codec을 구성
• 논문 (ICASSP 2024) : Paper Link

---

1. Introduction

• Universal audio model은 music, speech, environmental sound 등의 다양한 task에서 사용되는 것을 목표로 함
  • 대부분의 audio model은 Audioset과 같은 large-labeled dataset을 활용하여 superivsed 방식을 적용하고, unlabeled dataset의 경우 self-supervised learning (SSL)을 통해 pretrain 됨
  • 이때 SSL은 masked language modeling과 같이 utterance의 masked part를 예측하는 pretext task를 학습하거나 masked time step을 동일한 utterance의 distractor와 contrasting 하는 방식으로 동작함
    1. 결과적으로 pretext task는 frame-level에서 동작하면서 utterance의 relationship을 반영하고, contrastive learning은 instance-level embedding을 얻음
    2. 이러한 pretrained upstream model은 하나 이상의 layer에서 downstream network로 activation을 공급함으로써, 상대적으로 작은 dataset에 대한 downstream task를 수행함
       - 이때 pretrained model은 frozen 되거나 fine-tuning 됨
  • 특히 BERT 기반의 audio model은 audio domain에서 input signal을 masking 하고 이에 대응하는 discrete target을 예측하는 방식으로 동작함
    1. 대표적으로 HuBERT에서 discrete target은 MFCC의 $k$-means clustering과 second stage의 internal representation으로 얻어짐
    2. 여기서 masked language modeling의 성능을 향상하기 위해 EnCodec과 같은 neural codec으로 학습된 discretre unit을 target으로 사용할 수 있음

-> 그래서 masked audio model에서 discrete target을 mapping 하는 neural codec인 EnCodecMAE를 제안

 

• EnCodecMAE
  • EnCodec을 기반으로 audio signal을 represent 하고 masked language modeling pretext task에 대한 discrete target을 생성
  • HuBERT의 self-training stage는 speech를 비롯한 다양한 audio signal에 대해 적용됨
  • 이때 Masekd AutoEncoder (MAE) architecture는 spectrogram patch 대신 1D signal을 input으로 사용

< Overall of EnCodecMAE >

• MAE를 활용해 universal audio model을 학습하기 위한 discrete target을 생성하는 neural codec
• 결과적으로 speech, music, environmental sound 등의 다양한 downstream task에서 우수한 성능을 달성

2. Method

• EnCodecMAE는 EnCodec representation을 학습 중에 input과 target으로 사용하는 Masked AutoEncoder
  • 먼저 아래 그림과 같이 audio signal이 EnCodec encoder에 전달되어 embedding sequence를 생성함
    - 해당 sequence에는 sinusoidal positional embedding이 추가됨
  • 이후 embedding을 $M_{prop}$의 비율로 randomly mask 하고, replacement 없이 mask의 starting index를 sampling 한 다음, 모든 sampled index에 대해 subsequent $M_{gap}$ time step을 masking 함
    - 이때 wav2vec 2.0, HuBERT 등과 달리, masked embedding은 mask token으로 replace 되지 않고 sequence에서 discard 됨
    - 결과적으로 resulting sequence는 visible embedding만 포함하므로 original sequence보다 짧음
  • 그리고 masked sequence는 encoder를 통과한 다음, masked region에 mask token을 삽입하여 original size로 확장됨
  • 최종적으로 reconstruct 할 masked region의 position을 decoder에 inform 하기 위해, Sinusoidal positional embedding이 추가된 다음 resulting sequence를 decoder에 input으로 사용함
    - 이때 EnCodecMAE는 decoder에 의한 posterior output와 Residual Vector Quantization (RVQ) layer에 의해 생성된 discrete target 간의 weighted cross entropy를 최소화하도록 training 됨

Overall of EnCodecMAE

- EnCodec

• EnCodec은 quantized latent space를 가지는 encoder-decoder architecture로 구성된 neural codec
  • EnCodec은 높은 compression rate를 달성하면서 decoded audio의 perceptible distortion을 최소화할 수 있음
    - 실제로 동일한 bitrate의 SoundStream 보다 높은 MUSHRA score를 보임
  • 즉, EnCodec의 quantized latent pace는 complex signal을 reconstruct 하기 위한 대부분의 perceptually-related information을 포함하고 있음
  • EnCodec architecture는,
    1. Audio signal을 24kHz에서 75kHz로 downsampling 하는 4개의 convolution block과 2개의 LSTM layer, final convolution layer를 가지는 CNN encoder로 구성됨
    2. 여기서 encoder output은 128-dimensional vectoe sequence로써, 각각 1024개의 possible value를 가지는 32개 codebook이 있는 Residual Vector Quantization (RVQ) block을 통해 얻어짐
       - RVQ는 input을 first codebook의 closest entry에 mapping 한 다음, residual을 계산하여 second codebook의 closest entry에 mapping 하는 방식으로 동작함
       - 해당 process는 전체 codebook에 대해 repeat 되어 quantized output을 refining 함
    3. 최종적으로 decoder는 transposed convolution을 통해 encoded signal을 upsampling 하고 original audio를 reconstruct 함

- Masked Autoencoders

• Masked AutoEncoder (MAE)는 self-supervised model로써 spectrogram patch를 활용하여 audio representation learning에 적용되고 있음
  • MAE는 input signal의 일부를 masking 하고 unmasked signal에서 masked part를 예측하는 방식으로 BERT와 유사하게 training 됨
  • 이때 masked region을 mask token으로 대체하는 BERT와 달리, MAE는 encoder에 sequence를 전달하기 전에 masked region을 discard 하고 decoding 시에 mask token을 추가하여 encoder output을 original length로 확장함
    
  • 논문에서는 MAE로써 lightweight decoder와 larger encoder를 가지는 assymetric encoder-decoder를 활용함
    1. 이때 MAE는 input signal의 estimation을 직접 생성하지 않고, input signal를 reconstruction 할 수 있는 codebook set에 대한 index의 posterior를 생성함
    2. 먼저 MAE에서 masked region은 discard 되므로 encoder의 input seqeunce legnth는 original length 보다 짧음
       - 이를 통해 효율적인 학습을 지원하면서 encoder에서 mask token을 unseen 시켜 training, inference stage 간의 mismatch 문제를 해결
       - 한편으로 MAE는 transformer와 비슷하게 positional embedding이 encoder input에 추가되고, remaining embedding이 sequence의 absolute location을 전달할 수 있도록 함
    3. Decoder는 transformer layer stack과 softmax activation이 있는 classification layer로 구성됨
       - 이때 classification layer는 input audio의 neural codec representation에 해당하는 EnCodec RVQ에서 discrete target을 예측
       - 즉, decoder는 기존과 달리, masking 전 encoder input의 esitmation가 아닌 discrete version을 output 함

- Training Procedure

• Training procedure는 two-stage로 구성됨
  1. 먼저 model의 target $y_{q}(t)$은 EnCodec RVQ layer에 의해 생성되어 각 time step $t$와 codebook $q$에 대해 1부터 1024까지의 index로 구성됨
     • 그러면 loss function은 masked time step과 unmasked time step, 각 individual codebook에 대해 tunable weight를 가지는 weighted cross-entropy:
       (Eq. 1) $L=\sum_{q=0}^{Q}\gamma_{q}\left[\alpha\sum_{t\in M}C(y_{q}(t),z_{q}(t))+\beta\sum_{t\notin M}C(y_{q}(t),z_{q}(t))\right]$
       - $z_{q}(t)$ : time step $t$, codebook $q$에 대해 decoder에 의해 생성된 1024개 posterior의 vector
       - $M$ : masked index set, $\alpha=\delta/|M|, \beta = (1-\delta)/(T-|M|)$ : 각각 masked, unmasked region에 대한 weight
       - $\delta$ : tunable parameter, $T$ : sequence length
     • 여기서 각 codebook stream $q$에 해당하는 loss는 순차적으로 $\gamma_{q}$에 의해 weight 됨
     • $C(y_{q}(t),z_{q}(t))=-\log (z_{q,y_{q}(t)}(t))$는 target index에 대해 evaluate 된 $z_{q}(t)$의 negative logarithm으로써, 각 codebook과 time step에 대한 loss로 주어짐
  2. Second-stage에서는 HuBERT와 마찬가지로 500k training step 이후, self-training stage를 적용함
     • 이때 randomly-sampled 10k audio signal에 대해 last encoder layer에서 embedding을 추출하고, $k=1024$인 $k$-means model을 적용함
     • 그리고 pretraining dataset을 quantize 한 다음, resulting stream을 $Q$ quantizer에 추가함
       - 해당 stream에는 $\sum_{q}\gamma_{q}$와 동일한 weight가 부여되고, 150k의 additional step에 대해 학습됨
     • 이때 large model을 위해, self-training 이후 base model의 activation을 사용하여 $k$-means model을 train 함

3. Experiments

- Settings

• Dataset
  - Pretraining : Audioset, Free Music Archive, LibriLight
  - Evaluation : NSynth Pitch, GTZAN, Google Speech Commands, CREMA-D, Freesound Dataset, ESC-50
• Comparisons : Fuse H, Fuse W, BYOL-S, BYOL-A, Cat HWC, PaSST

- Results

• 단순히 EnCodec (EC)만을 사용하는 경우 모든 task에서 가장 낮은 성능을 보임
  • 한편으로 MAE를 추가하는 경우, randomly initialize 된 ECMAE-R을 포함하여 향상된 성능을 얻을 수 있음
    - 특히 MAE를 pretraining 하거나 self-training (ST)를 추가하면 추가적인 성능 개선이 가능함
    
  • Fuse H, BYOL-A 등의 기존 모델들과 비교해 보면, ECMAE는 모든 task에 대해서 범용적인 성능을 보임

모델 성능 비교