Test-time beam search cuts RVQ codec quantization error without retraining EnCodec or HiFi-Codec

Improving Test-Time Performance of RVQ-based Neural Codecs

Hyeongju Kim, Junhyeok Lee, Jacob Morton, Juheon Lee, Jinhyeok Yang

eess.AS, cs.SD

2025-09-24

Replacing RVQ's greedy per-level code selection with test-time beam search needs no retraining: on EnCodec, quantization error drops 5.10→4.63 and PESQ rises 2.73→2.85.

What problem this solves

Neural audio codecs like EnCodec, SoundStream, and HiFi-Codec compress audio into a few kbps of discrete tokens. They double as the audio tokenizer for generation models such as AudioLM and VALL-E. Almost all of them quantize continuous features with RVQ (residual vector quantization): the first codebook approximates the input, the second approximates the leftover residual, and so on, so shallow codebooks capture coarse structure and deeper ones fill in detail.

RVQ quantizes greedily, one level at a time: each level picks the single nearest code to the previous level's residual. The paper nails the problem with a toy three-codebook example. An input of 2.13, quantized greedily, lands at 3.0 with error 0.87. Pick the codes [1, 1, 0.1] instead and it lands at 2.1 with error 0.03. Per-level greedy is not the global optimum, and greedy RVQ has been quietly leaving that gap on the table.

Method

The fix is a test-time beam-search encoder (Algorithm 1). Instead of keeping one code per level, it keeps B candidates:

Conventional greedy RVQ is the B=1, k=1 special case. The point is to fold future levels into each level's decision instead of minimizing only the immediate residual. Exhaustive search costs O(S^L) (S codebook size, L levels); beam search bounds it with a fixed B.

Crucially this changes only the encode step. No decoder change, no retraining, drop-in on official checkpoints.

Results

Evaluated on EnCodec (Meta, default 6 kbps) and HiFi-Codec, with speech from LibriTTS, music from MUSDB18, plus an internal non-verbal-vocalization set (laughs, sighs). Metrics: quantization error (L2), mel distance, SI-SNR, PESQ, STOI, NISQA.

EnCodec at 6 kbps on speech shows a clean monotonic gain:

Beam sizeMel dist ↓SI-SNR ↑PESQ ↑NISQA ↑
B=1 (greedy)1.8934.2392.7263.491
B=161.8424.6492.8503.588

Quantization error falls 5.096 → 4.625 (B=16), about 9%. On music, SI-SNR rises 5.396 → 5.835. HiFi-Codec gains far less: PESQ moves only 3.013 → 3.034, with mel distance and SI-SNR shifting in the second decimal.

The variable-bitrate sweep tells the real story. At 3 kbps, PESQ goes 2.053 → 2.139, the largest win. At 24 kbps (32 codebooks) it shrinks to 3.670 → 3.691. More codebooks leave greedy less to correct.

Cost decides whether this is usable. A naive CPU implementation adds 229% latency when beam grows from 4 to 16. But the authors parallelize the inner loop (lines 7–16) onto GPU: the same 4-to-16 increase costs only 1.6%, and 5 seconds of audio on an RTX 4090 goes from 6.78 ms to 7.37 ms. On GPU the method is essentially free.

Why it matters

For TTS and audio-generation teams this is a cheap, free upgrade: no retraining, no model swap, just a beam search on the encode side for a steady objective-metric gain, biggest at low bitrates. Neural codecs are now both compression tools and tokenizers for generative models, so encode quality propagates straight into downstream synthesis. A few points of quantization error is not a trivial number in that chain.

It also reframes RVQ at the design level: per-level greedy is an engineering convenience, not an information-theoretic optimum. Where latency is not binding, batch encoding of an asset library for instance, beam search is a ready lever to pull.

Limitations

No subjective listening tests (MUSHRA, human MOS); everything rests on objective metrics. NISQA is a neural network that estimates MOS, not real listener scores, and PESQ and STOI target intelligibility more than quality. Quantization error is what the method optimizes directly, and its link to perceived quality is assumed rather than proven. Error drops 9% and PESQ rises 0.12, and the paper gives no theory for that exchange rate.

Validation covers only EnCodec and HiFi-Codec, not SoundStream, Descript Audio Codec (DAC), or newer codecs. Gains collapse at high bitrate, where 24 kbps is nearly negligible, so the method serves bandwidth-constrained low-bitrate use. HiFi-Codec's smaller gain is left unexplained; a plausible reason is its four-codebook group-residual design leaves little search space, but the authors do not analyze this. And the 229% CPU overhead means GPU-less, real-time deployments (on-device low-latency TTS) cannot use a large beam and must stop around B=4.

Terms

Source

What people are saying

Related papers

All paper explainers