[Paper 리뷰] LightCodec: A High Fidelity Neural Audio Codec with Low Computation Complexity

LightCodec: A High Fidelity Neural Audio Codec with Low Computation Complexity

---

• Neural codec은 높은 computational complexity의 한계를 가지고 있음
  - 즉, complexity를 줄이는 경우 성능이 현저하게 저하되므로 low computation resource에서 사용하기 어려움
• LightCodec
  • 높은 품질을 유지하면서 낮은 complexity를 가지는 neural audio codec
  • Frequency band division에 기반한 structure를 도입하고 Within Band-Across Band Interaction (WBABI) module을 통해 subband에 대한 feature를 학습하도록 함
  • Quantization-compensation module을 통해 quantization error를 90% 감소시킴
• 논문 (ICASSP 2024) : Paper Link

---

1. Introduction

• Audio codec은 continuous audio를 discrete representation으로 compress 하고, distortion을 최소화하면서 discrete representation에서 original audio를 recover 하는 것을 목표로 함
  • 기존의 parametric codec은 psycho-acoustics나 signal processing piepline을 활용했음
  • 한편으로 최근의 neural codec은 hybrid codec과 end-to-end codec으로 나누어짐
    1. Hybrid codec에서 encoder는 input speech에서 acoustic feature를 추출하고, 해당 feature를 integer type으로 변환해 bitstream transmission을 수행함
       - 이때 vocoder는 transmitted feature를 활용해 waveform을 변환하므로 적합한 vocoder를 구축하는 것이 중요한 요소임
    2. End-to-End codec은 input audio를 high-dimensional space에 mapping 하여 encoder-decoder architecture를 기반으로 reconstruction 함
       - 대표적으로 SoundStream은 residual vector quantization (RVQ)를 활용해 encoder output을 discretize 하는 방식을 사용함
       
  • End-to-End codec은 일반적으로 hybrid codec이나 기존 방법들보다 뛰어난 성능을 보임
    1. BUT, quantization module로 인해 error가 발생하므로 original과 decoded waveform 간의 차이가 발생할 수 있음
    2. 추가적으로 high-dimensional feature 추출과 quantization error repair 과정에서 상당한 computational complexity가 필요함
       - 결과적으로 complexity가 줄어들면 codec의 성능도 급격히 저하됨

-> 그래서 낮은 complexity을 가지면서 고품질의 음성을 유지할 수 있는 lightweight neural codec인 LightCodec을 제안

 

• LightCodec
  • Input signal을 여러 subband로 divide 하고, subband의 withing/across에 걸쳐 information을 학습
    - 이를 통해 low computational complexity로 deep feature extraction을 달성할 수 있음
  • Quantizer의 quantization error를 shallow network를 통해 reduce 하는 quantization-compensation module을 도입 

< Overall of LightCodec >

• 높은 품질을 유지하면서 낮은 complexity를 가지는 neural audio codec
• Within Band-Across Band Interaction module을 통해 subband에 대한 feature를 학습하고 quantization-compensation module을 통해 quantization error를 90% 감소시킴
• 결과적으로 24kHz audio를 처리하는데 0.8 GMACs의 complexity만을 사용하여 기존보다 뛰어난 품질을 달성

2. Method

• LightCodec은 encoder, quantizer, decoder의 3가지 module로 구성됨
  • Encoder는 band information interaction을 통해 input speech를 high-dimensional feature로 mapping 하고 decoder는 그 반대의 process를 수행함
  • Quantizer는 encoder에서 continuous feature를 discretize 하고 해당 codebook feature의 index를 transmit 함
  • LightCodec에서 time과 time-frequency domain loss는 training 과정에서 original waveform을 reconstruct 하기 위한 constraint로 사용되고, adversarial training 시에는 discriminator가 사용됨
    - 아래 그림에서, $n$ : output channel, $k$ : kernel size, $s$ : CNN layer의 stride

Overall of LightCodec

- Encoder and Decoder

• Encoder와 decoder는 neural codec에 있어 가장 computational intensive 한 부분
  • 따라서 LightCodec은 subband-based codec을 활용하여, subband signal processing에 중점을 두고 low complexity에서 high performance를 달성함
    1. 먼저 MB-MelGAN은 multi-band merging을 위해 pseudo quadrature mirror filter bank (PQMF)를 사용하여 low complexity의 vocoder를 구성했음
    2. 따라서 LightCodec도 이와 비슷하게 PQMF를 사용하여 encoder 이전에 original waveform을 여러 subband signal로 divide 하고, decoder 이후에 생성된 subband signal을 full-band waveform으로 합성함
  • 이때 end-to-end codec은 signal compressing을 위해 상당한 complexity를 가지는 stacking CNN layer를 활용함
    - 해당 구조로 발생하는 complexity를 줄이기 위해, LightCodec은 Within Band-Across Band Interaction (WBABI) module을 통해 latent feature를 학습함
  • 여기서 WBABI module은 구조적으로
    1. 각 subband signal은 먼저 해당 subband 내 (within)에서 학습되고, 이는 shallow CNN과 ELU layer로써 구현됨
       - 특히 각 subband 간에는 correlation이 존재하고, 해당 relationship은 latent feature에 반영되어야 함
       - 따라서 단순히 within information을 학습하는 것 외에도 subband 전반 (across)에 대한 information interaction을 모델링해야 함 
    2. 결과적으로 CNN의 output은 subband across information을 학습하기 위해 $1\times 1$ convolution layer에 제공된 다음, 두 feature를 add 하고 input을 residual로 사용
    3. 이러한 WBABI를 통해 encoder는 두 shallow layer를 통해 within/across subband에 대한 characteristic을 얻고, frequency band의 global feature와 local feature를 학습
  • Encoder의 구조는 dilated CNN과 long-short term memory (LSTM)이 long-time sequence modeling에 사용된다는 점에서 EnCodec과 유사함
    1. 이때 LightCodec은 EnCodec의 residual unit을 앞선 WBABI module로 대체하고, downsampling rate를 $[5,4,2,2]$로 설정함
    2. 이를 통해 encoder에서 각 subband feature는 interactive/independent learning을 continuously realize 하여 high-dimensional vector $Z$로 mapping 됨
  • Decoder 구조는 encoder의 반대와 같고, transposed convolution을 통해 upsampling 함
    - Decoder output은 4-channel subband signal을 생성하고, 이를 PQMF를 통해 waveform으로 합성

- Quantization-Compensation Module

• End-to-End codec은 encoder/decoder가 서로 symmetric 하므로 lossless transmission에서 가장 높은 성능을 보임
  • BUT, quantizer는 quantization error를 필연적으로 발생시킴
    - 실제로 RVQ와 GRVQ는 quantization error를 생성하고 해당 error는 decoder의 성능에 영향을 미침
    - 따라서 decoder는 waveform reconstruction을 수행하면서 해당 error를 최대한 restore 해야 함
  • 결과적으로 quantizer가 error를 최소한으로 발생시킨다면, encoder/decoder의 성능을 향상할 수 있음
    - 이를 위해 LightCodec은 encoder/decoder의 시작 부분에 있는 2개의 shallow CNN에 의한 quantization error를 예측하는 quantization-compensation module을 도입
    - 해당 module은 quantizer로 발생하는 error를 최소화하는 것을 목표로 함
  • Quantization-compensation module에서,
    1. 먼저 encoder output은 GRVQ에 의해 quantize 되고 communication channel을 통해 transmit 됨
    2. Receiver에서는 quantization result가 compensation module (CM)을 통과함
       - 해당 module은 kernel size가 각각 7, 1인 2개의 CNN layer와 사이의 ReLU activation으로 구성
    3. 최종적으로 output result는 예측된 quantization error로써, 해당 error를 GRVQ의 dequantization result에 추가하여 decoder input으로 사용
       - 여기서 compensation module이 quantization error를 학습할 수 있도록, loss function을 사용하여 CM output을 constraint 함

- Discriminator

• Discriminator로써 MelGAN의 multi-scale discriminator (MSD)와 HiFi-Codec의 multi-scale STFT (MS-STFT) discriminator를 사용함
  • MSD는 input waveform을 downsampling 하여 다양한 time scale에 대해 fake/real을 discriminate 함
  • MS-STFT는 다양한 time-frequency domain resolution에 대한 complex spectra를 discriminate 함
    - 이때 $[2048, 1024, 512]$의 FFT size, window length를 사용하고, hop size는 window length의 $1/4$로 설정

- Loss Function

• Training 과정에서 LightCodec은 generator와 discriminator를 모두 최적화함
  • Generator는 reconstruction loss, quantization-compensation module의 commitment loss, adversarial loss를 jointly optimize 하고, Discriminator는 adversarial hinge-loss 기반의 discriminator loss를 사용함
  • 먼저 LightCodec은, full-band/subband에 대한 multi-resolution STFT ($mr\text{-}STFT$) loss와 time-domain multi-band $L1$ ($mb\text{-}L1$) loss의 2가지 reconstruction loss를 사용
    1. $mb\text{-}L1$ loss는 서로 다른 frequency band에서 모델을 최적화하는 것을 목표로 하고, 이를 위해 full-band/subband 각각에서 $L1$ loss를 계산함:
       (Eq. 1) $L_{r}=mb\text{-}L_{1}(x_{r},x_{g})+mr\text{-}STFT(x_{r},x_{g})$
       (Eq. 2) $mb\text{-}L1(x_{r},x_{g})=L1(x_{r\text{-}sb},x_{g\text{-}sb})+L1(x_{r},x_{g})$
       - $x_{r}$ : original waveform, $x_{g}$ : generated waveform
       - $x_{r\text{-}sb}, x_{g\text{-}sb}$ : 각각 $x_{r}, x_{g}$의 subband signal
       - Full-band $mr\text{-}STFT$ loss에 사용된 FFT point와 window length는 $[512, 1024,2048]$이고, subband는 $[384,683,171]$
    2. Commitment loss는:
       (Eq. 3) $L_{c}=MSE(Z, Q(Z))+MSE(Z,Comp(Q(Z))+Q(Z))$
       - $Z$ : encoder의 continuous output, $Q$ : GRVQ layer, $Comp$ : compensation module
       - 결과적으로 $L_{c}$는 quantizer의 embedding space와 encoder result를 가깝게 만들고, compensation module이 quantization error에 잘 fit 되도록 함
    3. Adversarial loss $L_{adv}$는 discriminator를 fake 하는 데 사용됨:
       (Eq. 4) $L_{adv}=\frac{1}{k}\sum_{i=1}^{k}(1-D_{k}(x_{g}))^{2}$
       - $k$ : discriminator 수
    4. 종합하면 generator의 loss function은:
       (Eq. 5) $L_{g}=\lambda_{1}L_{r}+\lambda_{2}L_{c}+\lambda_{3}L_{adv}$
       - $\lambda_{1}, \lambda_{2}, \lambda_{3}$ : hyperparameter
  • Discriminator의 training에는 다음의 objective를 사용함:
    (Eq. 6) $L_{d}=\frac{1}{k}\sum_{i=1}^{k}((1-D_{k}(x_{r}))^{2}+(D_{k}(x_{g}))^{2})$

3. Experiments

- Settings

• Dataset : AISHELL3, VCTK, LibriTTS
• Comparisons : SoundStream, EnCodec

- Results

• 정량적인 지표 측면에서 성능을 비교해 보면, LigthCodec은 0.8 GMACs만을 사용하여 다른 모델들보다 비슷하거나 더 높은 성능을 보임
  
  • 특히 low bitrate에서 LightCodec은 기존의 EVS, Opus 보다 뛰어남
  • 즉, LightCodec은 낮은 complexity로 우수한 성능을 달성할 수 있음

모델 성능 비교

• MUSHRA 측면에서 비교해 보면, 마찬가지로 LightCodec은 낮은 complexity로 가장 뛰어난 품질을 보임

주관적 품질 비교 결과

• LightCodec의 compensation module의 효과를 알아보면
  • Simple complexity reduction을 사용하는 경우 모델 성능이 크게 저하됨
  • 반면 compensation module을 사용하는 경우 commitment loss (CL)를 90% 줄여 낮은 complexity에서 모델의 reconstruction ability를 크게 향상함

Compensation Module의 효과