[Paper 리뷰] EnCodec: High-Fidelity Neural Audio Compression

EnCodec: High Fidelity Neural Audio Compression

---

• Neural network를 사용하여 real-time, high-fidelity의 audio codec을 구성할 수 있음
• EnCodec
  • End-to-End 방식으로 학습된 quantized latent space를 가지는 streaming encoder-decoder architecture를 활용
  • Artifact를 줄이고 고품질의 sample을 합성하기 위해 multi-scale spectrogram adversary를 사용하여 training을 단순화하고 속도를 향상함
  • 이때 training을 stabilize 할 수 있는 loss balancer mechanism을 도입
• 논문 (Meta AI 2022) : Paper Link

---

1. Introduction

• Lossy signal compression에서는 sample의 bitrate를 최소화하면서 human perception과 관련된 distortion을 줄이는 것을 목표로 함
  • 이때 audio codec은 audio content의 redundancy를 제거하고 compact bitstream을 생성하기 위해 encoder-decoder로 이루어진 pipeline을 주로 활용
    - 기존에는 signal processing을 주로 활용했지만, neural network로 확장하여 구성할 수도 있음
  • BUT, lossy neural compression model은 다음의 2가지의 문제를 가지고 있음
    1. 모델은 training set에 overfitting 되지 않으면서 외부의 audio를 생성할 때 artifact가 발생하지 않는 wide-range의 signal을 represent 할 수 있어야 함
       
    2. 모델은 computing time, size 측면에서 효율적으로 compressing을 수행해야 함
       

-> 그래서 고품질의 lossy compression을 유지하면서 real-time으로 동작할 수 있는 neural codec인 EnCodec을 제안 

 

• EnCodec
  • Encoder-Decoder architecture를 기반으로 perceptural loss, quantization method를 적용한 end-to-end neural compression model
  • Neural encoder floating-point ouptut에 대한 residual vector quantization을 통해 computing time을 절감
  • 고품질의 audio를 위해 다양한 training set로 모델을 학습하고, perceptual loss 역할을 할 수 있는 discriminator network를 적용

< Overall of EnCodec >

• End-to-End 방식으로 학습된 quantized latent space를 가지는 streaming encoder-decoder architecture
• Artifact를 줄이고 고품질의 sample을 합성하기 위해 multi-scale spectrogram adversary를 사용하고 training을 stabilize 할 수 있는 loss balancer mechanism을 도입
• 결과적으로 기존 방식들보다 뛰어난 품질과 빠른 latency를 달성

2. Method

• Duration $d$의 audio signal은 $x\in [-1,1]^{C_{a}\times T}$의 sequence로 represent 될 수 있음
  • $C_{a}$ : audio channel 수, $T=d\cdot f_{sr}$ : 주어진 sampling rate $f_{sr}$에서의 audio sample 수
  • 이때 EnCodec은 3가지 component로 구성됨
    1. Encoder network $E$는 audio extract를 input 하여 latent representation $z$를 output 함
    2. Quantization layer $Q$는 vector quantization을 사용하여 compressed representation $z_{q}$를 생성함
    3. Decoder network $G$는 compressed representation $z_{q}$로부터 time-domain signal $\hat{x}$를 reconstruct 함
  • 전체적으로는 서로 다른 resolution에서 동작하는 discriminator 형태의 perceptual loss와 함께 time, frequency domain 모두에 적용되는 reconstruction loss를 최소화하도록 end-to-end training 됨

- Encoder and Decoder Architecture

• EnCodec은 latent representation에 대해 sequential modeling component를 가지는 simple streaming, convolutional-based encoder-decoder architecture
• Encoder-Decoder
  • Encoder $E$는 $C$ channel, 7 kernel size의 1D convolution과 이후의 $B$개의 convolution block으로 구성됨
    1. 여기서 각 convolution block은 single residual unit과 stride $S$의 2배인 kernel size $K$를 가지는 stride convolution으로 구성된 downsampling layer로 구성
       - Residual unit에는 kernel size가 3인 2개의 convolution과 skip-connection이 포함되고, downsampling이 수행될 때마다 channel 수는 2배로 늘어남
    2. Convolution block 다음에는 sequence modeling을 위해 two-layer LSTM과 kernel size가 7이고 output channel이 $D$인 final 1D convolution이 이어짐
       - 논문에서는 $C=32, B=4$, stride로써 $(2,4,5,8)$을 사용하고 ELU non-linear activation을 적용함
  • Decoder는 stride convolution 대신 transposed convolution을 사용하고, enocder의 stride를 역순으로 적용하여 최종적인 mono/stereo audio를 output 함
  • 이때 low-latency streamable setup과 high-fidelity non-streamable setup에 따라 EnCodec의 2가지 variant를 고려
    
• Non-Streamable
  • Non-streamable setup에서 EnCodec은 각 convolution에 대해 $K-S$의 total padding을 사용하고, first time step 이전과 last step 이후에 대해 동일하게 split 함
  • 이때 input을 1초의 chunk로 split 하고 10ms를 overlap 한 다음, 모델에 제공되기 전에 각 chunk를 normalize 하고 decoder의 output에 inverse operation을 적용함
    - 이를 통해 scale을 transmit 하는 negligible bandwidth overhead를 추가함
  • 해당 setup에서는 relative scale information을 유지하기 위해 time dimension을 포함한 statistics를 계산하는 layer normalization을 사용
• Streamable
  • Streamable setup의 경우, 모든 padding은 first time step 이전에 수행됨
  • Stride가 $s$인 transposed convolution은, $s$ first time step을 output 하고 next frame이 available 할 때까지 나머지 $s$ step을 memory에 유지하거나 stream의 마지막에 discarding 함
    - 해당 padding을 통해 EnCodec은 첫 320 sample (13ms)를 receive 하자마자 320 sample을 output 할 수 있음
  • 이때 layer normalization을 weight normalization으로 대체하여 time dimension에 대한 statistics를 계산함

Overall of EnCodec

- Residual Vector Quantization

• Encoder의 output을 quantize 하기 위해 Residual Vector Quantization (RVQ)를 도입
  • Vector quantization은 주어진 size의 codebook에서 가장 가까운 entry에 input vector를 projecting 함
    - 이때 RVQ는 quantization 이후 residual을 계산하고, second codebook을 사용하여 further quantizing 하는 방식으로 vector quantization process를 수행함
  • RVQ의 training procedure는
    1. 먼저 각 input에 대해 select 된 codebook entry는 0.99의 decay를 가지는 exponential moving average를 사용하여 update 되고, 사용되지 않은 entry는 current batch에서 sample 된 candidate로 대체됨
    2. 여기서 encoder의 gradient를 계산하기 위해 straight-through-estimator를 사용함
       - 이를 통해 quantization step이 backward phase에 대한 identity function인 것처럼 사용할 수 있음
    3. 최종적으로 input에 대해 계산된 gradient와 quantizer의 input/output 간의 MSE로 계산된 commitment loss를 전체 training loss에 추가됨
  • Training time에 대해 다양한 residual step을 선택함으로써 single model로도 multiple bandwidth를 지원할 수 있음
    - 이를 위해 EnCodec은 각각 1024개의 entry가 있는 최대 32 codebook을 사용하고, variable bandwidth training 시 4의 배수로 codebook 수를 randomly select 함
    
  • 결과적으로 $[B, D, T]$의 shape를 가지는 continuous latent representation이 encoder의 ouptut으로 주어지면, quantizer는 selected codebook 수 $N_{q}$를 사용하여 discrete index set $[B, N_{q}, T]$로 변환함
    - 이러한 discrete representation은, decoder에 제공되기 전에 해당 codebook entry와 summing 하여 vector로 변경될 수 있음

- Language Modeling and Entropy Coding

• Single CPU core에서 real time end-to-end compression/decompression 속도를 유지하기 위해 Transformer-based language model을 추가로 training 함 
  • 해당 transformer language model은 5개의 layer, 8개 head, 200개 channel, 800-dimensional feed-forward block으로 구성됨
  • Training 시에 bandwidth와 해당 codebook 수 $N_{q}$를 select 함
    1. Time step $t$에 대해 time $t-1$에서 얻은 discrete representation은 각 codebook에 하나씩 학습된 embedding table을 사용하여 continuous representation으로 변환되고 summed 됨 
       - $t=0$에 대해서는 special token을 대신 사용
    2. 이후 transformer의 output은 각 codebook cardinality만큼의 output channel을 가진 $N_{q}$ linear layer에 공급되어 time $t$에 대한 각 codebook의 추정 분포의 logit을 제공 
  • 즉, EnCodec은 single time step에서 codebook 간의 potential mutual information을 neglect 함
    - 이를 통해 final cross entropy에 대한 영향을 제한하면서 추론 속도를 높일 수 있음
  • 추가적으로 각 attention layer에는 3.5초의 causal receptive field가 존재하고, 더 긴 sequence를 emulate 하기 위해 sinusoidal position embedding으로 initial position을 random amount로 offset 함
    
• Entropy Encoding
  • EnCodec은 language model이 제공하는 추정 확률을 사용하기 위해 range-based arithmetic coder를 사용함 
  • 여기서 동일한 모델이라도 architecture나 floating point approximation으로 인해 다른 결과가 발생할 수 있음 
    - 즉, encoder-decoder가 정확히 동일한 code를 사용하지 않으면 decoding error가 발생함
  • 실제로 batch evaluation에서 decoder의 real-life streaming evaluation은 $10^{-8}$ 이상의 차이가 발생하는 것으로 나타남
    - 따라서 논문에서는 추정 확률을 $10^{-6}$의 precision으로 반올림하고, total range를 $2^{24}$, minimum range width를 $2$로 assign 함

- Training Objective

• EnCodec은 reconstruction loss, discriminator를 통한 perceptual loss, RVQ commitment loss를 결합하여 training objective를 구성함
• Reconstruction Loss
  • Reconstruction loss는 time-domain loss term과 frequency-domain loss term으로 이루어짐
  • Time-domain에서는 target과 compressed audio 간의 $L1$ distance를 최소화함:
    $\ell_{t}=||x-\hat{x}||_{1}$
  • Frequency-domain에서는 여러 time scale을 사용하여 mel-spectrogram에 대한 $L1, L2$ loss 간의 linear combination을 사용함:
    (Eq. 1) $\ell_{f}(x,\hat{x})=\frac{1}{|\alpha| \cdot |s|} \sum_{\alpha_{i}\in \alpha}\sum_{i\in e}|| \mathcal{S}_{i}(x)-\mathcal{S}_{i}(\hat{x})||_{1}+\alpha_{i}|| \mathcal{S}_{i}(x)-\mathcal{S}_{i}(\hat{x})||_{2}$
    - $\mathcal{S}_{i}$ : window size가 $2^{i}$, hop length가 $2^{i}/4$인 normalized STFT를 사용하는 64-bin mel-spectrogram
    - $e=5,....,11$ : scale set, $\alpha$ : $L1, L2$간의 balancing scalar coefficient
    - 논문에서는 $\alpha_{i}=1$ 사용
• Discriminative Loss
  • 생성된 sample의 품질을 더욱 향상하기 위해, 아래 그림과 같은 multi-scale STFT discriminator (MS-STFTD)를 기반으로 하는 perceptual loss를 도입
  • MS-STFTD는 실수부와 허수부가 concatenate 된 multi-scaled complex-valued STFT에서 동작하는 network로 구성됨
    1. 각 sub-network는 kernel size 3, 32 channel을 가지는 2D convolution layer로 구성되고, 1,2,4의 time dimension에서 dilation rate가 증가한 다음, frequency axis에 대한 stride가 2인 2D convolution이 추가됨
    2. 이후 kernel size가 $3\times 3$이고 stride가 $(1,1)$인 final 2D convolution layer로 최종적인 예측을 수행함
       - 이를 STFT window length가 $[2048,1024,512,256,128]$인 5가지 scale에 대해 적용
    3. Leaky ReLU를 non-linear activation으로 사용하고 discriminator network에 weight normalization을 적용함
  • 여기서 generator에 대한 adversarial loss는:
    $\ell_{g}(\hat{x})=\frac{1}{K}\sum_{k}\max(0,1-D_{k}(\hat{x}))$
    - $K$ : discriminator 수
  • 추가적으로 generator에 대한 다음의 relative feature matching loss를 적용함:
    (Eq. 2) $\ell_{feat}(x,\hat{x})=\frac{1}{KL}\sum_{k=1}^{K}\sum_{l=1}^{L}\frac{||D_{k}^{l}(x)-D_{k}^{l}(\hat{x})||_{1}}{\mathrm{mean}(||D_{k}^{l}(x)||_{1})}$
    - $\mathrm{mean}$ : 모든 dimension에 대해 계산되는 평균, $D_{k}^{l}$ : discriminator, $L$ : discriminator의 layer 수
    
  • 한편으로 discriminator는 다음의 hinge adversarial loss function을 최소화하도록 training 됨:
    $L_{d}(x,\hat{x})=\frac{1}{K}\sum_{k=1}^{K}\max(0,1-D_{k}(x))+\max(0,1+D_{k}(\hat{x}))$
    - 이때 discriminator는 decoder를 쉽게 overpower 할 수 있으므로, 24kHz에서는 $2/3$, 48kHz에서는 0.5의 확률로 wegiht를 update 함

MS-STFTD architecture

• Multi-bandwidth Training
  • 24kHz에서는 RVQ step에서 유지할 적절한 수의 codebook을 select 하여 bandwidth 1.5, 3, 6, 12 kpbs를 지원하도록 EnCodec을 training 함
    - 48kHz에서는 3, 6, 12, 24kpbs를 지원하도록 training 함
  • 이때 bandwidth 별로 dedicated discriminator를 사용하면 audio 품질 향상에 도움이 되는 것으로 나타남
  • 따라서 전체 batch에 대해 주어진 bandwidth를 select 하고 해당하는 discriminator만 evaluate 하고 update 함
• VQ Commitment Loss
  • Encoder output과 quantized value 사이에 commitment loss $l_{w}$를 추가하고, quantized value에 대한 gradient는 계산하지 않음
  • 각 residual step $c\in \{1,...,C\}$에서 $z_{c}$를 current residual, $q_{c}(z_{c})$를 해당 codebook에서 가장 가까운 entry라고 했을 때, commitment loss $\ell_{w}$는:
    (Eq. 3) $\ell_{w}=\sum_{c=1}^{C}||z_{c}-q_{c}(z_{c})||_{2}^{2}$
    - $C$ : current batch의 bandwidth target
  • 결과적으로 generator는 batch 전반에 걸쳐 summation 된 다음의 total loss를 최적화하도록 training 됨:
    (Eq. 4) $L_{G}=\lambda_{t}\cdot\ell_{t}(x,\hat{x})+\lambda_{f}\cdot\ell_{f}(x,\hat{x})+\lambda_{g}\cdot\ell_{g}(\hat{x})+\lambda_{feat}\cdot\ell_{feat}(x,\hat{x})+\lambda_{w}\cdot\ell_{w}(w)$
    - $\lambda_{t}, \lambda_{f}, \lambda_{g}, \lambda_{feat}, \lambda_{w}$ : balancing scalar coefficient
• Balancer
  • Discriminator로부터 발생하는 다양한 scale의 gradient에 의한 training을 stabilize 하기 위해 loss balancer를 도입함
    - 이때 balancer는 scale에 관계없이 다양한 loss weight에 대해 쉽게 얻어질 수 있음
  • 먼저 모델 $\hat{x}$의 output에만 의존하는 여러 loss $(\ell_{i})_{i}$가 있다고 하자
    1. 그러면 $g_{i}=\frac{\partial \ell_{i}}{\partial x}$와 last training batch에 대한 $g_{i}$의 exponential moving average $\langle ||g_{i}||_{2}\rangle_{\beta}$를 정의할 수 있음
    2. 이때 weight set $(\lambda_{i})$와 reference norm $R$이 주어지면 다음을 얻을 수 있음:
       (Eq. 5) $\tilde{g}_{i}=R\frac{\lambda_{i}}{\sum_{j}\lambda_{j}}\cdot \frac{g_{i}}{\langle ||g_{i}||_{2} \rangle_{\beta}}$
       - 논문에서는 $R=1$, $\beta=0.999$로 사용
  • 위를 사용하여 기존의 $\sum_{i}\lambda_{i}g_{i}$ 대신 $\sum_{i}\tilde{g}_{i}$로 backpropagate 할 수 있음
    - 이를 통해 최적화 문제를 변경하여, 각 loss의 natural scale에 관계없이 $\lambda_{i}$를 interpretable 하게 만듦
    - e.g.) $\sum_{i}\lambda_{i} = 1$인 경우, 각 weight는 해당 loss에서 나오는 모델 gradient의 fraction으로 interpret 될 수 있음
  • 모델의 output에 대해 정의되지 않은 commitment loss를 제외하고 (Eq. 4)의 모든 generator loss들은 balancer에 fit 됨

3. Experiments

- Settings

• Dataset : DNS challenge, Common Voice, AudioSet, FSD50K, Jamendo
• Comparisons : OPUS, EVS, Lyra v2

- Results

• $1.5, 3,6,12kps bandwidth에 대해 EnCodec의 성능을 비교해 보면
  • 동일한 bandwidth를 고려할 때 EnCodec은 가장 높은 MUSHRA score를 보임
  • 특히 3kpbs의 EnCodec은 6kbps의 OPUS나 12kbps의 Lyra v2 보다 더 나은 성능을 달성함

MUSHRA score 비교

• 한편으로 code에 대한 추가적인 language model을 고려하는 경우, bandwidth를 $25\text{~}40\%$까지 줄일 수 있음
  - e.g.) 3kpbs의 모델은 bandwidth를 1.9kbps로 줄이는 것이 가능

Bandwidth 별 모델 성능 비교

• The Effect of Discriminators Setup
  • 생성된 audio의 품질을 향상하기 위해 다양한 discriminator를 고려할 수 있음
  • 이때 EnCodec에서 Multi-Scale Discriminator (MSD), Multil-Period Discriminator (MPD), Multi-Scale STFT Discriminator (MS-STFTD), MS-STFTD+MPD를 각각 비교해 보면
  • MS-STFTD를 사용하는 것만으로도 EnCodec은 고품질의 audio를 충분히 생성할 수 있음
  •  

Discriminator 종류에 따른 EnCodec의 성능

• The Effect of the Streamable Modeling
  • EnCodec을 non-streamable에서 streamable 설정으로 전환하면 약간의 성능 저하가 발생함
  • BUT, EnCodec은 여전히 강력한 streamable 추론 성능을 보임

Streamable Setting에 따른 성능 비교

• The Effect of the Balancer
  • Balancer의 유무와 다양한 $\lambda_{t}, \lambda_{f}, \lambda_{g}, \lambda_{feat}$ 값에 대해 EnCodec의 성능을 비교해 보면
  • Balancer는 EnCodec의 training process를 상당히 stabilize 하는 것으로 나타남

Balancer의 영향 비교

• Stereo Evaluation
  • Monophonic setup 외에도 stereo compression 성능을 확인해 보면
  • EnCodec은 6kpbs에서 OPUS 보다 뛰어나고, 64kpbs의 MP3와 비교할만한 성능을 달성함
  • 추가적인 language model과 entropy coding을 사용하면 $20\text{~}30\%$의 variable gain을 얻을 수 있음

Stereophonic Setup에서의 성능 비교 결과

• Latency and Computation Time
  • 24kHz streaming EnCodec의 initial latency는 13.3ms이고, 48kHz non-streaming은 1초의 initial latency를 보임
    - Entropy coding을 사용하는 경우 각 frame을 flush 할 수 없기 때문에 initial latency가 늘어남
  • 24kHz EnCodec은 Lyra v2 보다는 느리지만 여전히 real-time보다 10배 빠르게 처리할 수 있음

Latency, Computation Time 비교