[Paper 리뷰] High-Fidelity Audio Compression with Improved RVQGAN

High-Fidelity Audio Compression with Improved RVQGAN

---

• Language model의 핵심 component는 high-dimensional natural signal을 low-dimensional discrete token으로 compress 하는 neural codec임
• Improved RVQGAN
  • Adversarial, reconstruction loss와 vector quantization technique을 도입하여 high-fidelity의 audio compression을 보장
  • 추가적으로 speech, environment, music 등의 다양한 domain에 대한 universal compression을 지원
• 논문 (NeruIPS 2023) : Paper Link

---

1. Introduction

• High-resolution audio의 generative modeling은 high-dimensionality와 short/long-term dependency를 모두 가지는 time-scale structure로 인해 어려움이 있음
  • 해당 문제를 완화하기 위해 audio generation은 일반적으로 two-stage 방식을 활용함 
    
    1. Text와 같은 conditioning information이 주어진 intermediate representation을 예측한 다음,
    2. 해당 representation을 audio로 변환하는 것
  • 즉, 해당 과정은 intermediate variable이 obeserve 되는 hierarchical model로 볼 수 있고, Variational AutoEncoder (VAE), Flow 등을 사용하여 합성을 수행할 수 있음
  • 이때 VQ-VAE로 얻어지는 discrete latent variable로 VAE를 training 할 수도 있음
    
    1. 특히 discrete latent variable은 distribution modeling을 위해 powerful autoregressive model을 사용하여 expressive prior를 training 할 수 있음
    2. 실제로 MusicLM과 같은 transformer language model은 해당 방식을 사용하여 audio, music 등의 complex distribution을 효과적으로 모델링했음
       - BUT, prior 활용의 장점에도 불구하고, quantized autoencoder를 통해 discrete latent code를 모델링하는 것에는 여전히 어려움
  • 한편으로 discrete code를 학습하는 것은 lossy compression task로 취급할 수 있음 
    - 여기서 audio signal은 fixed length codebook을 사용하여 autoencoder의 representation을 vector-quantizing 하고 discrete latent space로 compress 됨
    
  • 이때 compression model은 다음의 property를 만족해야 함:
    • Audio codec은 high-fidelity, artifact-free 한 audio를 reconstruction 해야 함 
    • High-level structure는 유지하면서 low-level impreceptible detail은 discard 하여 temporal downsampling과 high-level compression이 가능해야 함 
    • Speech, music, environmental 등에 대한 다양한 type의 audio를 universal 하게 지원할 수 있어야 함 
  • BUT, SoundStream, EnCodec과 같은 neural codec은 기존의 GAN-based model과 비슷하게 tonal artifact, pitch/periodicity artifact 등의 문제를 겪으므로 high-freqeuncy를 불완전하게 모델링하는 문제가 있음

-> 그래서 quality loss와 artifact를 최소화하면서 universal compression을 지원하는 neural codec인 Improved RVQGAN을 제안

 

• Improved RVQGAN
  • 44.1kHz audio를 8kpbs의 bitrate의 discrete code로 compress 하는 universal neural codec
    
  • Codebook collapse 문제를 해결하기 위해, full bandwidth를 사용할 수 있는 improved codebook learning technique을 도입
  • Single 모델이 다양한 bitrate를 지원할 수 있도록, quantizer dropout을 적용
  • Periodic inductive bias, multi-scale STFT discriminator, multi-scale mel loss를 추가

< Overall of Improved RVQGAN >

• Vector quantization, adversarial training을 활용하여 44.1kHz audio에 대해 90배의 compression 성능을 달성하는 neural codec
• 결과적으로 speech, music, environment 등의 다양한 sound에 대해서 high-fidelity의 품질과 효과적인 compression 성능을 달성

2. Method

• Improved RVQGAN은 SoundStream, EnCodec과 유사한 VQ-GAN framework를 기반으로 함
  • 구조적으로,
    1. SoundStream의 fully convolutional encoder-decoder network를 활용하여 chosen striding factor로 temporal downsampling을 수행함
    2. 특히 initial quantization step 이후, distinct codebook을 사용하여 residual을 recursively quantize 하는 Residual Vector Quantization (RVQ)을 적용하여 encoding을 quantize 함
    3. Improved RVQGAN에서는 다양한 target bitrate에서 동작하도록 training 중에 quantizer dropout을 도입함 
    4. 최종적으로 Improved RVQGAN은 adversarial loss와 함께 frequency-domain reconstruction loss를 사용하여 training 됨 
  • 여기서 $f_{s}$의 sampling rate를 가지는 audio signal에 대해, $M$의 encoder striding factor, $N_{q}$ RVQ layer는 $S\times N_{q}$ shape의 discrete code matrix를 생성함
    - 여기서 $S=f_{s}/M$으로 정의된 frame rate
  • 요약적으로 아래 표와 같이 기존 codec들의 compression factor와 latent code의 frame rate를 비교했을 때, 제안하는 Improved RVQGAN은 모든 기존 codec들에 비해 더 높은 compression factor를 달성할 수 있음

기존 Neural Codec 비교

- Periodic Activation Function

• Audio waveform은 high periodicity를 가지므로, non-autoregressive architecture에서 jarring pitch와 periodicity artifact가 종종 발생함
  • 특히, Leaky ReLU와 같은 일반적인 neural activation은 periodic signal을 extrapolating 하는데 한계가 있고, audio synthesis에 대해 poor 한 out-of-distribution generalization을 보임
  • 따라서 Improved RVQGAN에서는 generator에 periodic inductive bias를 반영하기 위해, Snake activation을 도입함
    - 이때, BigVGAN과 마찬가지로 $\mathrm{snake}(x)=x+\frac{1}{\alpha}\sin^{2}(\alpha x)$로 정의됨
    - $\alpha$는 signal의 periodic component에 대한 frequency를 control 하는 역할
  • 결과적으로 Sanke function은 Leaky ReLU보다 audio fidelity를 더욱 향상할 수 있는 것으로 나타남

- Improved Residual Vector Quantization

• Vector Quantization은 discrete autoencoder를 training 하는데 자주 사용되는 방법임
  • BUT, vanilla VQ-VAE는 poor initialization과 low codebook usage로 인해 codebook의 상당 부분이 사용되지 않음
    - 결과적으로 이러한 effective codebook size의 reduction은 target bitrate에 대한 implicit reduction으로 이어지므로, 전체적인 reconstruction 저하가 발생함
  • 해당 문제를 해결하기 위해 $k$-means clustering을 활용하여 codebook vector를 initialize 하고, 특정 codebook이 several batch 동안 사용되지 않는 경우 randomized restart를 적용할 수 있음
    - BUT, 해당 방법을 도입하더라도 24kpbs의 target bitrate로 training 된 EnCodec은 여전히 낮은 codebook utilization을 보임
  • 따라서 Improved RVQGAN에서는 codebook utilization 문제를 해결하기 위해, factorized code와 $L2$-normalized code를 도입함
    1. 먼저 factorization은 code lookup과 code embedding을 decouple 함
       - 여기서 code lookup은 low-dimensional space ($8d, 32d$)에서 동작하고 code embedding은 high-dimensional space ($1024d$)에서 동작함
       - 즉, data의 variance를 maximally explain 하는 input vector의 principal component를 사용하는 code lookup으로 볼 수 있음
    2. 한편으로 encoded vector와 codebook vector에 대한 $L2$-normalization은 euclidean distance를 cosine similarity로 변환하여 stability와 quality를 향상함
  • 결과적으로 Improved RVQGAN은 위 두 가지 trick을 활용하여 bitrate efficiency와 reconstruction quality를 개선함
    - 특히, 모델은 기존의 $k$-means initialization이나 random restart 없이도 VQ-VAE codebook, commitment loss를 사용해 training 됨
• Modified Codebook Learning Algorithm
  • 논문에서는 다음의 modified quantization operation을 사용함:
    (Eq. 1) $z_{q}(x)=W_{out}e_{k}$
    - $k=\arg\min_{j}|| \ell_{2}(W_{in}z_{e}(x))-\ell_{2}(e_{j})||_{2}$
    - $W_{in}, W_{out}$ : projection matrix
  • 여기서 $W_{in}$은 encoder output을 intermediate representation으로 mapping 하고, $W_{out}$은 해당 representation을 quantized representation $z_{q}(x)$에 mapping 함:
    (Eq. 2) $W_{in}\in\mathbb{R}^{D\times M}, \,\,\, W_{out}\in\mathbb{R}^{M\times D}$
    - $D$ : encoder의 output dimension, $M$ : codebook dimension ($M\ll D$)
  • 그러면 vector quantizer loss function은:
    (Eq. 3) $z_{proj}(x)=W_{in}z_{e}(x)$
    (Eq. 4) $\mathcal{L}_{VQ}=||\text{sg}[\ell_{2}(z_{proj}(x))]-\ell_{2}(e_{k})||_{2}^{2}+\beta|| \ell_{2}(z_{proj}(x))-\text{sg}[\ell_{2}(e_{k})]||_{2}^{2}$
    - $\text{sg}$ : $e_{k}$를 통한 gradient back-propagation을 prevent 하는 stop-gradient operator
    - $\beta$ : balancing parameter

- Quantizer Dropout Rate

• Variable bitrate를 가지는 single compression model을 training 하기 위해, quantizer dropout을 도입함
  • Quantizer 수 $N_{q}$는 bitrate를 결정하므로, 각 input example에 대해 $n\sim\{1,2,...,N_{q}\}$를 randomly sample 하고 training 중에 first $n_{q}$ quantizer만 사용함
  • BUT, 해당 quantizer dropout을 단순히 적용하면 full bandwidth에서 audio reconstruction 성능이 저하됨
    1. 이를 해결하기 위해, 논문은 probability $p$를 사용하여 각 input example에 quantizer dropout을 적용하는 방식을 채택함
    2. 결과적으로 $p=0.5$의 dropout probability는 아래 그림과 같이, lower bitrate에서 baseline의 reconstruction quality와 거의 일치하는 것으로 나타남
  • 추가적으로 quantizer dropout과 RVQ 간의 interaction을 알아보면,
    1. 각 technique을 결합했을 때, quantized code는 most-significant에서 least-significant bit로 information을 학습할 수 있음
    2. Code가 $1...N_{q}$ codebook으로 reconstruct 되는 경우, 각 codebook은 점점 fine-scale detail을 추가하는 것으로 나타남
       - 해당 interaction은 AudioLM, MusicLM의 coarse token/fine token과 같이 hierarchical generative model을 training 할 때 유용함 

Quantization Dropout의 효과

- Discriminator Design

• Improved RVQGAN은 Multi-Scale Discriminator (MSD)와 Multi-Period Discriminator (MPD)를 도입하여 fidelity를 향상함
  • 이때 생성된 spectrogram은 blurry 하고 high-frequency에서 over-smoothing artifact가 나타날 수 있음
    - 한편으로 UnivNet의 Multi-Resolution Spectrogram Discriminator (MRSD)를 사용하여 해당 pitch/periodic artifact를 완화함
    - BUT, 해당 방식을 사용하면 magnitude spectrogram으로 인해 phase information이 discard 되므로 discriminator의 모델링에 영향을 줌
  • 따라서 논문은 multiple time-scale에서 동작하는 SoundStream의 complex STFT discriminator를 활용하여 phase modeling을 개선함 
    1. 이와 같이 STFT를 sub-band로 split 하면 high-frequency 예측이 향상되고 aliasing artifact가 완화됨
       - Discriminator가 specific sub-band에 대한 discriminative feature를 학습하고 generator에 stronger gradient signal을 제공할 수 있기 때문
       
    2. 추가적으로 Multi-band MelGAN에서 제안된 multi-band processing은 full-band audio를 생성하기 위해 summation 되는 sub-band audio를 예측하기 위해 도입됨

- Loss Function

• Frequency Domain Reconstruction Loss
  • Mel-reconstruction loss는 stability와 fidelity를 향상할 수 있지만, multi-scale spectral loss는 multiple time-scale에 대한 frequency modeling을 지원할 수 있음
  • 따라서 논문에서는 $[32, 64,128,256,512,1024,2048]$의 window length와 $\text{window_length}/4$의 hop size로 계산된 mel-spectrogram에서 $L1$ loss를 사용하여 두 방식을 결합함
    - 특히 8의 hop size를 사용하면 music domain에서 quick transient에 대한 모델링을 향상할 수 있음
  • 한편 EnCodec의 경우, $L1, L2$ loss와 64 size의 fixed mel bin을 사용함
    - 이와 같이 mel bin size를 fix 하는 경우, low filter length에서 spectrogram에 hole이 발생하는 것으로 나타남
    - 따라서, 논문은 $[5,10,20,40,80,160,320]$의 mel bin size를 사용
• Adversarial Loss
  • Improved RVQGAN은 waveform discrimination을 위해 MPD와 frequency domain에 대한 multi-band multi-scale STFT discriminator를 사용함
  • 이를 위해 HingeGAN의 adversarial loss와 $L1$ feature matching loss를 적용
• Codebook Learning
  • VQ-VAE formulation의 stop-gradient를 가지는 codebook, commitment loss를 활용함
  • 이때 straight-through estimator를 사용한 codebook lookup을 통해 gradient를 backpropagate 함
• Loss Weighting
  • Multi-scale mel loss는 $15.0$, feature matching loss는 $2.0$, adversarial loss는 $1.0$, codebook commitment loss에는 각각 $1.0, 0.25$의 weight를 적용함
  • 해당 loss weighting은 mel loss를 계산하는 데 사용된 multiple scale과 $\log_{10}$ base를 고려하여 rescale 된 값

3. Experiments

- Settings

• Dataset
  - Speech : DAPS, DNS Challenge, Common Voice, VCTK
  - Music : MUSDB, Jamendo
  - Environment : AudioSet
• Comparisons : EnCodec, Lyra, Opus

- Results

• Ablation Study
  • 먼저 architecture 측면에서 decoder dimension을 변경하면 성능 저하가 발생하는 것으로 나타남
    - 특히 ReLU activation을 사용할 때 보다 Snake function을 사용하는 것이 더 나은 성능을 보임
  • Discriminator 측면에서, multi-band discriminator는 high-frequency aliasing을 완화하는 것으로 나타남
    
  • Low-hop reconstruction 역시 waveform loss와 high-frequency modeling에 중요함
  • Codebook의 latent dimension은 bitrate efficiency와 reconstruction 품질에 영향을 미침
    - 너무 낮거나 높게 설정하는 경우, bitrate efficiency가 나빠지므로 성능 저하가 발생함
    - 경험적으로 8의 latent dimension을 사용하는 것이 효과적임
    

Ablation Study 결과

• 한편으로 exponential moving average로 codebook을 학습하면 SI-SDR에 대한 성능 저하가 발생하고, 아래 그림과 같이 모든 codebook에서 codebook utilization이 떨어짐
  • 추가적으로 quantization dropout 역시 성능에 영향을 미치는데, dropout이 0인 경우 codebook 수가 적어 낮은 reconstruction 성능을 보임
    - 경험적으로 0.5의 dropout rate를 사용
  • 그 외에도 balanced data sampling을 제거하는 경우에도 성능 저하가 발생했음

각 Codebook의 Entropy

• Comparison to Other Models
  • 다른 codec들과 비교해 보면, Improved RVQGAN은 정량적, 정성적 지표 모두에서 가장 우수한 성능을 보임 

모델 성능 비교

• 특히 24kHz sampling rate에서 EnCodec과의 비교해 보면, Improved RVQGAN은 지속적으로 우수한 성능을 달성함

24kHz에서 EnCodec과의 비교

• 각 sound category 별 성능에서도 마찬가지로 우수한 결과를 보임

Sound Category 별 성능