[Paper 리뷰] Self-Guidance: Enhancing Neural Codecs via Decoder Manifold Alignment

Self-Guidance: Enhancing Neural Codecs via Decoder Manifold Alignment

---

• Neural codec은 quantization error로 인한 reconstruction fidelity의 한계가 존재함
• Self-Guidance
  • Quantized token, continuous embedding을 process 할 때 internal decoder feature manifold를 align
  • 이를 위해 lightweight feature mapping loss를 도입
• 논문 (ICML 2026) : Paper Link

---

1. Introduction

• SoundStream, EnCodec과 같은 Vector-Quantized Variational AutoEncoder (VQ-VAE) 기반의 neural codec은 high compression ratio에서도 high-fidelity reconstruction이 가능함
  • 여기서 VQ-VAE는 latent vector quantizer를 활용해 continuous audio를 discrete token으로 transform 함
  • BUT, quantization process는 inherently lossy 하므로 decoder의 reconstruction에는 한계가 있음
    - 해당 quantization error를 suppress 하기 위해 hierarchical quantization이나 multiple residual codebook을 활용할 수 있지만, downstream model에서 autoregressive modeling complexity가 증가하게 됨

-> 그래서 downstream modeling burden 없이 neural audio codec의 quantization error를 줄일 수 있는 Self-Guidance를 제안

 

• Self-Guidance
  • Decoder의 internal manifold를 align 하여 quantization error를 counteract 
  • 특히 quantized token embedding, continuous pre-quantized latent vector에 대한 intermediate decoder feature의 feature mapping loss를 활용

< Overall of Self-Guidance >

• VQ-VAE-based codec backbone의 quantization error를 줄이는 guidance mechanism 
• 결과적으로 기존보다 우수한 성능을 달성

2. Preliminary

- Revisiting the VQ-VAE Framework

• Vector-Quantized Variational AutoEncoder (VQ-VAE)는 neural audio codec의 foundation으로 사용됨
  • 구조적으로 VQ-VAE는 encoder-vector quantizer-decoder로 구성됨
    1. Encoder는 latent dimension $d_{e}$에 대해, input audio signal $x$를 continuous latent embedding sequence $z_{e}\in\mathbb{R}^{d_{e}}$로 process 함
    2. Vector quantizer는 $z_{e}$의 각 embedding을 finite codebook $\mathcal{Q}\subset \mathbb{R}^{d_{e}}$의 nearest entry에 mapping 함
    3. Decoder는 quantized token embedding sequence $z_{q}$로부터 audio signal $\hat{x}$를 reconstruct 함
  • Training 시에는 Straight-Through Estimation (STE)를 사용해 gradient를 $z_{q}$에서 $z_{e}$로 directly copy 하여 propagate 함
  • 해당 quantization process는 continuous latent vector를 discrete codebook에 project 하므로 quantization error가 존재하고, 해당 error는 discretization 시 발생한 information loss를 represent 함:
    (Eq. 1) $e_{q}=||z_{e}-z_{q}||_{2}$
    

- Observation of Quantization Artifacts

• Quantization error $e_{q}$의 information loss로 인해 quantization artifact가 나타남
  • 실제로 아래 표와 같이 decoder가 quantized token $z_{q}$ 대신 continuous pre-quantized latent $z_{e}$를 process 하는 경우 더 나은 reconstruction quality를 얻을 수 있음
  • 즉, 더 나은 reconstruction을 위해서는 quantization error에 robust 한 decoder가 필요함

기존 Neural Codec의 Reconstruction 성능

3. Method

• 논문은 neural speech codec에서 quantization artifact를 mitigate 하는 것을 목표로 함

- Self-Guidance Mechanism

• Self-Guidance는 input token $z_{q}$에서 quantization error로 인해 발생하는 information loss를 compensate 함
  • 즉, decoder가 quantized token $z_{q}$와 pre-quantized latent $z_{e}$ 모두에서 similar output을 produce 해야 함
  • Vanilla VQ-VAE는 original input $x$를 target으로 사용하여 decoder를 implicitly guide 할 수 있지만 quantization artifact를 fully address 하기에는 insufficient 함
    - 따라서 논문은 explicit guide를 위한 internal guidance로 pre-quantized latent $z_{e}$를 사용함
  • 먼저 training 시 $z_{e}$를 decoder에 feed 하는 additional forward pass를 추가함
    1. 이후 $z_{e}$ branch에서 $h_{e}$, $z_{q}$ branch에서 $h_{q}$의 intermediate hidden feature를 추출함
       - Feature $h$는 X-Codec decoder의 final Transformer block output으로 얻어지고, iSTFT head에 feed 되어 waveform을 reconstruct 함
    2. 여기서 논문은 해당 feature를 align 하기 위해 feature mapping loss $\mathcal{L}_{guide}$를 도입하고, 해당 loss term은 VQ-VAE objective에 add되어 end-to-end self-supervised training을 구성함:
       (Eq. 2) $\mathcal{L}_{guide}=||\text{sg}(h_{e})-h_{q}||_{2}^{2}$
       - $\text{sg}(\cdot)$ : stop-gradient operation 
  • Self-Guidance mechanism은 minimal computational, architectural overhead를 가짐
    - Training 시에는 $z_{e}$에 대한 additional forward pass만 필요하고, 추론 시 decoder는 standard VQ-VAE와 같이 $z_{q}$만 사용함

Overview

- Neural Speech Codec Model

• 논문은 Self-Guidance를 X-Codec2에 적용하여 low-bitrate encoding과 speech generation task를 수행함
  • 구조적으로 X-Codec2는 convolutional encoder, single-layer Finite Scalar Quantizer (FSQ), acoustic decoder로 구성되고, W2V-BERT feature 기반의 auxiliary autoencoder를 가짐
    
  • 특히 acoustic decoder는 Transformer backbone과 iSTFT head를 가지므로 해당 Transformer backbone output을 $\mathcal{L}_{guide}$를 compute 하는 데 사용할 수 있음
  • 결과적으로 training objective는 다음과 같이 구성됨:
    (Eq. 3) $\mathcal{L}_{total}=\lambda_{guide}\mathcal{L}_{guide}+\mathcal{L}_{semantic}+\mathcal{L}_{acoustic}+\mathcal{L}_{adv}$
    - $\mathcal{L}_{guide}$ : (Eq. 2)의 self-guiding feature mapping loss, $\mathcal{L}_{semantic}$ : semantic feature MSE loss, $\mathcal{L}_{acoustic}$ : multi-scale mel-spectrogram $L1$ loss, $\mathcal{L}_{adv}$ : adversarial loss, $\lambda_{guide}$ : weight 

4. Experiments

- Settings

• Dataset : LibriSpeech
• Comparisons : DAC, WavTokenizer, XCodec, TS3Codec, BigCodec

- Results

• Self-Guidance를 적용하면 더 나은 성능을 달성할 수 있음

Model 성능 비교

• Reconstruction 측면에서도 우수한 성능을 보임

Reconstruction 성능

• Feature Alignment
  • Self-Guidance는 hidden feature alignment error $||h_{e}-h_{q}||_{2}^{2}$를 suppress 함

Feature Alignment

• Decoder Manifold Alignment
  • Self-Guidance는 $k$NN Jaccard similarity, Procrustes residual을 크게 개선함

Manifold Alignment 성능

• Self-Guidance가 activate 된 경우 hidden feature가 token ID 별로 cluster 되어 나타남

$t$-SNE

• Self-Guidance는 decoder 전체에 대한 alignment를 개선함

Blockwise Alignment

• VQ-VAE Configs
  • Self-Guidance는 codebook size에 관계없이 성능 향상이 가능함

Codebook Size 별 성능

• 여러 codec에 대해서도 consistent 한 결과를 보임

Codec 별 성능

• Downstream Autoregressive TTS
  • Text-to-Speech (TTS) task에서도 Self-Guidance를 활용하면 더 나은 결과를 얻을 수 있음

TTS 성능