[Paper 리뷰] DS-Codec: Dual-Stage Training with Mirror-to-NonMirror Architecture Switching for Speech Codec

DS-Codec: Dual-Stage Training with Mirror-to-NonMirror Architecture Switching for Speech Codec

---

• High-quality speech tokenizer가 필요함
• DS-Codec
  • Mirror-NonMirror architecture switching을 활용한 dual-stage training framework를 도입
  • Mirrored architecture를 통해 learned codebook의 robustness를 향상하고 Mirror-NonMirror structure를 통해 training을 balance
• 논문 (INTERSPEECH 2025) : Paper Link

---

1. Introduction

• 최근 VALL-E, AudioLM, AudioGen 등과 같이 Text-to-Speech (TTS)를 위해 Large Language Model (LLM)이 자주 활용되고 있음
  • 이를 위해서는 speech signal을 discrete representation으로 변환하는 neural codec이 필요함
  • 특히 대부분의 speech codec은 encoder, quantization module, mirrored decoder로 구성됨
    1. Encoder는 speech signal을 latent representation으로 encode 하고, quantization module은 주로 Residual Vector Quantization(RVQ)를 통해 latent representation을 discrete token으로 변환함
       
    2. Mirror decoder는 discrete token으로부터 speech signal을 reconstruct 함
  • BUT, 기존의 RVQ-based speech codec은 multiple token sequence가 필요하므로 additional non-autoregressive stage가 요구됨
    - 이를 위해서는 WavTokenizer, Single-Codec과 같은 single-codebook neural codec을 구성해야 함
  • 추가적으로 EnCodec, DAC와 같은 기존 mirrored structure 대신 decoder를 emphasize 하는 non-mirrored decoder upsampling structure를 활용하면 reconstruction ability를 향상할 수 있음

-> 그래서 non-mirrored structure를 기반으로 한 single-codebook codec인 DS-Codec을 제안

 

• DS-Codec
  • Multiple codebook 대신 single codebook을 활용하고 Product Quantization (PQ)를 적용
  • Mirror, non-mirror architecture를 switching 하는 two-stage training strategy를 도입

< Overall of DS-Codec >

• Sinlge codebook을 기반으로 Mirror-NonMirror switching training을 적용한 neural codec
• 결과적으로 기존보다 뛰어난 reconstruction 성능을 달성

2. Method

- Model Architecture

• DS-Codec은 non-mirrored architecture를 활용함
  • Encoder는 residual CNN block으로 구성되고, 각 block은 snake activation function과 2-layer unidirectional LSTM network를 가짐
    - CNN block은 specific factor로 input waveform을 downsampling 하고, 다양한 dilation rate를 가진 multiple convolutional layer를 사용하여 sequential data pattern을 capture 함
  • Decoder는 transpose convolution을 사용하여 upsampling을 수행함
    1. 이때 speech reconstruction capability를 향상하기 위해 논문은 LLaMa decoder layer와 같은 Transformer block을 도입함
       
    2. 해당 Transformer block은 residual attention과 residual SwiGLU, RMSNorm으로 구성됨
  • 추가적으로 discriminator로는 HiFi-GAN의 Multi-Period Discriminator (MPD)와 EnCodec의 Multi-Scale Short-Time Fourier Transform (MS-STFT) Discriminator를 채택함

Overview

- Quantization Module

• Vector Quantization
  • Vector Quantization (VQ) module은 DAC를 따라 $8192$ code를 포함하는 fixed-size single-codebook을 사용하여 latent representation을 discrete vector로 mapping 함
  • 이때 codebook utilization을 optimize 하기 위해 quantization 이전에 latent representation에 dimensionality reduction process를 적용함
    1. 먼저 low-dimensional space로 project 한 다음,
    2. Quantize 하고 original dimensionality로 project back 함
  • 추가적으로 latent variable과 codebook vector는 $L2$-normalize 되어 quantization process의 efficiency와 accuracy를 향상함
• Product Quantization
  • Product Quantization (PQ)는 multiple VQ module을 사용함
    1. 이때 latent representation은 channel dimension을 따라 non-overlapping segment로 divide 되고, 각 segment는 dedicated VQ module에 assign 됨
       
    2. 각 segment는 해당하는 VQ module을 통해 independently quantize 되고 resulting quantized vector는 final PQ output $\hat{y}$를 생성하기 위해 concatenate 됨
  • Training, inference 시에는 [Algorithm 1]과 같이 single-codebook으로 취급할 수 있음
  • 결과적으로 DS-Codec-PQ는 size가 $[16,16,16,16]$인 4개의 codebook을 통해 $65536$의 codebook size를 생성함
    - 여기서 VQ loss는 각 VQ module의 VQ loss summation으로 얻어짐

Product Quantization

- Training Strategy

• 논문은 다음과 같은 two-stage training strategy를 도입함
  1. Mirror Training
     • First stage는 standard mirrored speech codec framework를 활용한 joint training에 해당함
       - 여기서 model은 encoder, quantization module, mirrored decoder를 통해 training 됨
     • 해당 stage는 mirrored design을 활용하여 robust codebook을 얻는 것을 목표로 함
  2. Decoder Training
     • Second stage에서는 decoder capability를 향상하는 것을 목표로 함
     • 이를 위해 encoder, quantizer의 parameter를 freeze 하고, Transformer block을 integrate 하고, non-mirrored architecture로 switch 한 다음, discriminator parameter를 reinitialize 함
       - 이를 통해 efficient, stable training process를 보장하여 reconstruction ability를 향상할 수 있음
     • Decoder parameter는 reinitialize 하지 않는 대신 first training stage의 decoder parameter weight를 retain 함
       - 이는 unnecessary adjustment를 minimize 하여 training process를 accelerate 함

3. Experiments

- Settings

• Dataset : LibriSpeech
• Comparisons : DAC, EnCodec, TiCodec, WavTokenizer, BigCodec

- Results

• 전체적으로 DS-Codec의 성능이 가장 우수함

Model 성능 비교

• LJSpeech dataset에 대해서도 DS-Codec은 우수한 generalization ability를 보임

LJSpeech Dataset에서의 성능

• Mirror vs. Non-Mirror
  • Mirrored structure는 non-mirrored structure에 비해 높은 VQ loss를 보임

Training Loss 비교

• Ablation Study
  • Two-stage training을 활용하면 더 나은 성능을 달성할 수 있음

Ablation Study