Efficient Evaluation of Quantization-Effects in Neural Codecs

TL;DR

The paper introduces an efficient framework for evaluating quantization effects in neural codecs by using surrogate data with a fixed bit budget and a lightweight surrogate codec to mimic non-linearities. It formalizes encoder/quantizer/decoder interactions, analyzes gradient-tracking challenges through $\mathcal{Q}$, and surveys approaches like STE, soft quantization, and noise emulation. A Modified Straight-Through Estimator (mSTE) is proposed to stabilize training by tying quantization noise to the computational graph via a noise-mcale factor, reducing encoder growth and divergence observed with standard STE. The framework enables rapid prototyping (training costs under 1 hour on a GPU with <400 MB memory) and is validated against an internal audio codec and the descript-audio-codec, demonstrating improved stability and performance. This work accelerates analysis of quantization in neural codecs and provides a practical path for future exploration of gradient estimators and architectural variations.

Abstract

Neural codecs, comprising an encoder, quantizer, and decoder, enable signal transmission at exceptionally low bitrates. Training these systems requires techniques like the straight-through estimator, soft-to-hard annealing, or statistical quantizer emulation to allow a non-zero gradient across the quantizer. Evaluating the effect of quantization in neural codecs, like the influence of gradient passing techniques on the whole system, is often costly and time-consuming due to training demands and the lack of affordable and reliable metrics. This paper proposes an efficient evaluation framework for neural codecs using simulated data with a defined number of bits and low-complexity neural encoders/decoders to emulate the non-linear behavior in larger networks. Our system is highly efficient in terms of training time and computational and hardware requirements, allowing us to uncover distinct behaviors in neural codecs. We propose a modification to stabilize training with the straight-through estimator based on our findings. We validate our findings against an internal neural audio codec and against the state-of-the-art descript-audio-codec.

Figures (8)

• Figure 1: A neural codec system with an encoder $(\mathcal{E})$ mapping the input to embeddings $E$, a quantizer ($\mathcal{Q}$) mapping $E$ to the quantized version $E_q$. The decoder $(\mathcal{D})$ maps $E_q$ to the output.
• Figure 2: Proposed data generation pipeline. A Gaussian noise process with a $P\times P$ identity covariance matrix is sampled to obtain $X$. Each element in $X$ is quantized via scalar quantization to obtain the network target $X_q$. A rotation matrix is applied to $X_q$ to obtain the network input $Y$.
• Figure 3: Training and when using a no quantizer and the standard training using the with .
• Figure 4: Training and when using or the for training with and without . Note that the curves of NA and $\text{NA}^\text{CL}$ are overlapping.
• Figure 5: Training and when using for training with and without . We compare training using detached and attached noise as in (\ref{['equ:NAdet']}) and (\ref{['equ:classicNA']}), respectively. In the , NA and $\text{NA}_\text{det}$ overlap.