SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rithesh Kumar*  
Descript, Inc.
Alejandro Luebs  
Descript, Inc.
Prem Seetharaman*  
Descript, Inc.
Ishaan Kumar  
Descript, Inc.
Kundan Kumar  
Descript, Inc.
Language models have been successfully used to model natural signals, such as images, speech, and music. A key component of these models is a high quality neural compression model that can compress high-dimensional natural signals into lower dimensional discrete tokens. To that end, we introduce a high-fidelity universal neural audio compression algorithm that achieves 90x compression of  $44.1\mathrm{kHz}$  audio into tokens at just 8kbps bandwidth. We achieve this by combining advances in high-fidelity audio generation with better vector quantization techniques from the image domain, along with improved adversarial and reconstruction losses. We compress all domains (speech, environment, music, etc.) with a single universal model, making it widely applicable to generative modeling of all audio. We compare with competing audio compression algorithms, and find our method outperforms them significantly. We provide thorough ablations for every design choice, as well as open-source code and trained model weights. We hope our work can lay the foundation for the next generation of high-fidelity audio modeling.
Generative modeling of high-resolution audio is difficult due to high dimensionality ( $\sim 44,100$  samples per second of audio) [24, 19], and presence of structure at different time-scales with both short and long-term dependencies. To mitigate this problem, audio generation is typically divided into two stages: 1) predicting audio conditioned on some intermediate representation such as mel-spectrograms [24, 28, 19, 30] and 2) predicting the intermediate representation given some conditioning information, such as text [35, 34]. This can be interpreted as a hierarchical generative model, with observed intermediate variables. Naturally, an alternate formulation is to learn the intermediate variables using the variational auto-encoder (VAE) framework, with a learned conditional prior to predict the latent variables given some conditioning. This formulation, with continuous latent variables and training an expressive prior using normalizing flows has been quite successful for speech synthesis [17, 36].
A closely related idea is to train the same variational-autoencoder with discrete latent variables using VQ-VAE [38]. Arguably, discrete latent variables are a better choice since expressive priors can be trained using powerful autoregressive models that have been developed for modeling distributions over discrete variables [27]. Specifically, transformer language models [39] have already exhibited the capacity to scale with data and model capacity to learn arbitrarily complex distributions such as text[6], images[12, 44], audio [5, 41], music [1], etc. While modeling the prior is straightforward, modeling the discrete latent codes using a quantized auto-encoder remains a challenge.
*Equal contribution to this work. Address correspondence to papers@descript.com, or raise an issue at https://github.com/descriptinc/descript-audio-codec.
Learning these discrete codes can be interpreted as a lossy compression task, where the audio signal is compressed into a discrete latent space by vector-quantizing the representations of an autoencoder using a fixed length codebook. This audio compression model needs to satisfy the following properties: 1) Reconstruct audio with high fidelity and free of artifacts 2) Achieve high level of compression along with temporal downscaling to learn a compact representation that discards low-level imperceptible details while preserving high-level structure [38, 33] 3) Handle all types of audio such as speech, music, environmental sounds, different audio encodings (such as mp3) as well as different sampling rates using a single universal model.
While the recent neural audio compression algorithms such as SoundStream [46] and EnCodec [8] partially satisfy these properties, they often suffer from the same issues that plague GAN-based generation models. Specifically, such models exhibit audio artifacts such as tonal artifacts [29], pitch and periodicity artifacts [25] and imperfectly model high-frequencies leading to audio that are clearly distinguishable from originals. These models are often tailored to a specific type of audio signal such as speech or music and struggle to model generic sounds. We make the following contributions:
- We introduce Improved RVQGAN a high fidelity universal audio compression model, that can compress 44.1 KHz audio into discrete codes at 8 kbps bitrate ( $\sim$ 90x compression) with minimal loss in quality and fewer artifacts. Our model outperforms state-of-the-art methods by a large margin even at lower bitrates (higher compression), when evaluated with both quantitative metrics and qualitative listening tests.
- We identify a critical issue in existing models which don't utilize the full bandwidth due to codebook collapse (where a fraction of the codes are unused) and fix it using improved codebook learning techniques.
- We identify a side-effect of quantizer dropout - a technique designed to allow a single model to support variable bitrates, actually hurts the full-bandwidth audio quality and propose a solution to mitigate it.
- We make impactful design changes to existing neural audio CODEs by adding periodic inductive biases, multi-scale STFT discriminator, multi-scale mel loss and provide thorough ablations and intuitions to motivate them.
- Our proposed method is a universal audio compression model, capable of handling speech, music, environmental sounds, different sampling rates and audio encoding formats.
We provide code 1, models, and audio samples 2 that we encourage the reader to listen to.
High fidelity neural audio synthesis: Recently, generative adversarial networks (GANs) have emerged as a solution to generate high-quality audio with fast inference speeds, due to the feedforward (parallel) generator. MelGAN [19] successfully trains a GAN-based spectrogram inversion (neural vocoding) model. It introduces a multi-scale waveform discriminator (MSD) to penalize structure at different audio resolutions and a feature matching loss that minimizes L1 distance between discriminator feature maps of real and synthetic audio. HifiGAN [18] refines this recipe by introducing a multi-period waveform discriminator (MPD) for high fidelity synthesis, and adding an auxiliary mel-reconstruction loss for fast training. UnivNet [16] introduces a multi-resolution spectrogram discriminator (MRSD) to generate audio with sharp spectrograms. BigVGAN [21] extends the HifiGAN recipe by introducing a periodic inductive bias using the Snake activation function [47]. It also replaces the MSD in HifiGAN with the MRSD to improve audio quality and reduce pitch, periodicity artifacts [25]. While these the GAN-based learning techniques are used for VOCing, these recipes are readily applicable to neural audio compression. Our Improved RVQGAN model closely follows the BigVGAN training recipe, with a few key changes. Our model uses a new multi-band multi-scale STFT discriminator that alleviates aliasing artifacts, and a multi-scale mel-reconstruction loss that better model quick transients.
Neural audio compression models: VQ-VAEs [38] have been the dominant paradigm to train neural audio CODECs. The first VQ-VAE based speech codec was proposed in [13] operating at 1.6 kbps. This model used the original architecture from [38] with a convolutional encoder and an autoregressive
<https://github.com/descriptinc/descript-audio-codec
$^{2}$ https://descript.notion.site/Descript-Audio-Codec-11389fce0ce2419891d6591a68f814d5
wavenet [27] decoder. SoundStream [46] is one of the first universal compression models capable of handling diverse audio types, while supporting varying bitrates using a single model. They use a fully causal convolutional encoder and decoder network, and perform residual vector quantization (RVQ). The model is trained using the VQ-GAN [12] formulation, by adding adversarial and feature-matching losses along with the multi-scale spectral reconstruction loss. EnCodec [8] closely follows the SoundStream recipe, with a few modifications that lead to improved quality. EnCodec uses a multiscale STFT discriminator with a multi-scale spectral reconstruction loss. They use a loss balancer which adjusts loss weights based on the varying scale of gradients coming from the discriminator.
Our proposed method also uses a convolutional encoder-decoder architecture, residual vector quantization and adversarial, perceptual losses. However, our recipe has the following key differences: 1) We introduce a periodic inductive bias using Snake activations [47, 21] 2) We improve codebook learning by projecting the encodings into a low-dimensional space [44] 3) We obtain a stable training recipe using best practices for adversarial and perceptual loss design, with fixed loss weights and without requiring a sophisticated loss balancer. We find that our changes lead to a near-optimal effective bandwidth usage. This allows our model to outperform EnCodec even with  $3\mathrm{x}$  lower bitrate.
Language modeling of natural signals : Neural language models have demonstrated great success in diverse tasks such as open-ended text generation [6] with in-context learning capabilities. A key-component of these models is self-attention [39], which is capable of modeling complex and long-range dependencies but suffers from a quadratic computational cost with the length of the sequence. This cost is unacceptable for natural signals such as images and audio with very high dimensionality, requiring a compact mapping into a discrete representation space. This mapping is typically learnt using VQ-GANs [12, 44], followed by training an autoregressive Transformer on the discrete tokens. This approach has shown success across image [45, 44, 32], audio [5, 41], video and music [9, 1] domains. CODECs like SoundStream and EnCodec have already been used in generative audio models, like AudioLM [5], MusicLM [1], and VALL-E [41]. Our proposed model can serve as a drop-in replacement for the audio tokenization model used in these methods, allowing for highly superior audio fidelity, and more efficient learning due to our maximum entropy code representation.
3 The Improved RVQGAN Model
Our model is built on the framework of VQ-GANs, following the same pattern as SoundStream [46] and EnCodec [8]. Our model uses the fully convolutional encoder-decoder network from SoundStream, that performs temporal downscaling with a chosen striding factor. Following recent literature, we quantize the encodings using Residual Vector Quantization (RVQ), a method that recursively quantizes residuals following an initial quantization step with a distinct codebook. Quantizer dropout is applied during training to enable a single model that can operate at several target bitrates. Our model is similarly trained using a frequency domain reconstruction loss along with adversarial and perceptual losses.
Codec	Sampling rate (kHz)	Target bitrate (kbps)	Striding factor	Frame rate (Hz)	# of 10-bit codebooks	Compression factor
Proposed	44.1	8	512	86	9	91.16
EnCodec	24	24	320	75	32	16
48	24	320	150	16	32
SoundStream	24	6	320	75	8	64
Table 1: Comparison of compression approaches.
An audio signal with sampling rate  $f_{s}$  (Hz), encoder striding factor  $M$ , and  $N_{q}$  layers of RVQ produce a discrete code matrix of shape  $S \times N_{q}$ , where  $S$  is the frame rate defined as  $f_{s} / M$ . Table 1 compares our proposed model against baselines to contrast the compression factors and the frame rate of latent codes. Note that the target bitrate mentioned is an upper bound, since all models support variable bitrates. Our model achieves a higher compression factor compared to all baseline methods while outperforming them in audio quality, as we show later. Finally, a lower frame rate is desirable when training a language model on the discrete codes, as it results in shorter sequences.
3.1 Periodic activation function
Audio waveforms are known to exhibit high periodicity (especially in voiced components, music, etc.) While current non-autoregressive audio generation architectures are capable of generating high fidelity audio, they often exhibit jarring pitch and periodicity artifacts [25]. Moreover, common
neural network activations (such as Leaky ReLUs) are known to struggle with extrapolating periodic signals, and exhibit poor out-of-distribution generalization for audio synthesis [21].
To add a periodic inductive bias to the generator, we adopt the Snake activation function proposed by Liu et al. [47] and introduced to the audio domain in the BigVGAN neural vocoding model [21]. It is defined as  $\mathrm{snake}(x) = x + \frac{1}{\alpha}\sin^2 (\alpha x)$ , where  $\alpha$  controls the frequency of periodic component of the signal. In our experiments, we find replacing Leaky ReLU activations with Snake function to be an influential change that significantly improves audio fidelity (Table 2).
3.2 Improved residual vector quantization
Figure 1: Entropy for each codebook, computed using code usage statistics across a large test set.
While vector quantization (VQ) is a popular method to train discrete auto-encoder, there are many practical struggles when training them. Vanilla VQ-VAEs struggle from low codebook usage due to poor initialization, leading to a significant portion of the codebook being unused. This reduction in effective codebook size leads to an implicit reduction in target bitrate, which translates to poor reconstruction quality.
To mitigate this, recent audio codec methods use k-means clustering to initialize the codebook vectors, and manually employ randomized restarts [9] when certain codebooks are unused for several batches. However, we find that the EnCodec model trained at 24kbps target bitrate, as well as our proposed model with the same codebook learning method (Proposed
w/ EMA) still suffers from codebook under-utilization (Figure 1).
To address this issue, we use two key techniques introduced in the Improved VQGAN image model[44] to improve codebook usage: factorized codes and L2-normalized codes. Factorization decouples code lookup and code embedding, by performing code lookup in a low-dimensional space (8d or 32d) whereas the code embedding resides in a high dimensional space (1024d). Intuitively, this can be interpreted as a code lookup using only the principal components of the input vector that maximally explain the variance in the data. The L2-normalization of the encoded and codebook vectors converts euclidean distance to cosine similarity, which is helpful for stability and quality [44].
These two tricks along with the overall model recipe significantly improve codebook usage, and therefore bitrate efficiency (Figure 1) and reconstruction quality (Table 2), while being simpler to implement. Our model can be trained using the original VQ-VAE codebook and commitment losses [38], without k-means initialization or random restarts. The equations for the modified codebook learning procedure are written in Appendix A
3.3 Quantizer dropout rate
Quantizer dropout was introduced in SoundStream [46] to train a single compression model with variable bitrate. The number of quantizers  $N_{q}$  determine the bitrate, so for each input example we randomly sample  $n\sim \{1,2,\dots ,N_q\}$  and only use the first  $n_q$  quantizers while training. However, we noticed that applying quantizer dropout degrades the audio reconstruction quality at full bandwidth (Figure 2).
To address this problem, we instead apply quantizer dropout to each input example with some probability  $p$ . Interestingly, we find that dropout probability  $p = 0.5$  closely matches the reconstruction quality of baseline at lower bitrates, while closing the gap to full-bandwidth quality of a model trained without quantizer dropout ( $p = 0.0$ ).
Figure 2: Effect of quantizer dropout on audio quality vs bitrate.
Moreover, we provide additional insight into the practical behavior of quantizer dropout and it's interaction with RVQ. Firstly, we find that these techniques put together lead the quantized codes to learn most-significant to least significant bits of information with each additional quantizer. When the codes are reconstructed with  $1 \ldots N_{q}$  codebooks, we can see each codebook adds increasing amounts of fine-scale detail. We believe this interaction is beneficial when training hierarchical generative models on top of these codes [5, 41, 1], for example to partition the codes into "coarse" tokens (denoting the most significant codes) and "fine" tokens.
3.4 Discriminator design
Like prior work, we use multi-scale (MSD) and multi-period waveform discriminators (MPD) which lead to improved audio fidelity. However, spectrograms of generated audio can still appear blurry, exhibiting over-smoothing artifacts in high frequencies[16]. The multi-resolution spectrogram discriminator (MRSD) was proposed in UnivNet to fix these artifacts and BigVGAN [21] found that it also helps to reduce pitch and periodicity artifacts. However, using magnitude spectrograms discards phase information which could've been otherwise utilized by the discriminator to penalize phase modeling errors. Moreover, we find that high-frequency modeling is still challenging for these models especially at high sampling rates.
To address these issues, we use a complex STFT discriminator [46] at multiple time-scales [8] and find that it works better in practice and leads to improved phase modeling. Additionally we find that splitting the STFT into sub-bands slightly improves high frequency prediction and mitigates aliasing artifacts, since the discriminator can learn discriminative features about a specific sub-band and provide a stronger gradient signal to the generator. Multi-band processing was earlier proposed in [43] to predict audio in sub-bands which are subsequently summed to produce the full-band audio.
3.5 Loss functions
Frequency domain reconstruction loss: while the mel-reconstruction loss [18] is known to improve stability, fidelity and convergence speed, the multi-scale spectral losses[42, 11, 15] encourage modeling of frequencies in multiple time-scales. In our model, we combine both methods by using a L1 loss on mel-spectrograms computed with window lengths of [32, 64, 128, 256, 512, 1024, 2048] and hop length set to window_length / 4. We especially find that using the lowest hop size of 8 improves modeling of very quick transients that are especially common in the music domain.
EnCodec [8] uses a similar loss formulation, but with both L1 and L2 loss terms, and a fixed mel bin size of 64. We find that fixing mel bin size leads to holes in the spectrogram especially at low filter lengths. Therefore, we use mel bin sizes [5, 10, 20, 40, 80, 160, 320] corresponding to the above filter lengths which were verified to be correct by manual inspection.
Adversarial loss: our model uses the multi-period discriminator [18] for waveform discrimination, as well as the proposed multi-band multi-scale STFT discriminator for the frequency domain. We use the HingeGAN [22] adversarial loss formulation, and apply the L1 feature matching loss [19].
Codebook learning: we use the simple codebook and commitment losses with stop-gradients from the original VQ-VAE formulation [38], and backpropagate gradients through the codebook lookup using the straight-through estimator [3].
Loss weighting: we use the loss weightings of 15.0 for the multi-scale mel loss, 2.0 for the feature matching loss, 1.0 for the adversarial loss and 1.0, 0.25 for the codebook and commitment losses respectively. These loss weightings are in line with recent works [18, 21] (which use 45.0 weighting for the mel loss), but simply rescaled to account for the multiple scales and  $\log_{10}$  base we used for computing the mel loss. We don't use a loss balancer as proposed in EnCodec [8].
4.1 Data sources
We train our model on a large dataset compiled of speech, music, and environmental sounds. For speech, we use the DAPS dataset [26], the clean speech segments from DNS Challenge 4 [10], the Common Voice dataset [2], and the VCTK dataset [40]. For music, we use the MUSDB dataset
[31], and the Jamendo dataset [4]. Finally, for environmental sound, we use both the balanced and unbalanced train segments from AudioSet [14]. All audio is resampled to  $44\mathrm{kHz}$ .
During training, we extract short excerpts from each audio file, and normalize them to -24 dB LUFS. The only data augmentation we apply is to randomly shift the phase of the excerpt, uniformly. For evaluation, we use the evaluation segments from AudioSet [14], two speakers that are held out from DAPS [26] (F10, M10) for speech, and the test split of MUSDB [31]. We extract 3000 10-second segments (1000 from each domain), as our test set.
4.2 Balanced data sampling
We take special care in how we sample from our dataset. Though our dataset is resampled to  $44\mathrm{kHz}$ , the data within it may be band-limited in some way. That is, some audio may have had an original sampling rate much lower than  $44\mathrm{kHz}$ . This is especially prevalent in speech data, where the true sampling rates of the underlying data can vary greatly (e.g. the Common Voice data is commonly 8-16kHz). When we trained models on varying sampling rates, we found that the resultant model often would not reconstruct data above a certain frequency. When investigating, we found that this threshold frequency corresponded to the average true sampling rate of our dataset. To fix this, we introduce a balanced data sampling technique.
We first split our dataset into data sources that we know to be full-band - they are confirmed to contain energy in frequencies up to the desired Nyquist frequency (22.05kHz) of the codec - and data sources where we have no assurances of the max frequency. When sampling batches, we make sure that a full-band item is sampled. Finally, we ensure that in each batch, there are an equal number of items from each domain: speech, music, and environmental sound. In our ablation study, we examine how this balanced sampling technique affects model performance.
4.3 Model and training recipe
Our model consists of a convolutional encoder, a residual vector quantizer, and a convolutional decoder. The basic building block of our network is a convolutional layer which either upsamples or downsamples with some stride, followed by a residual layer consisting of convolutional layers interleaved with non-linear Snake activations. Our encoder has 4 of these layers, each of which downsamples the input audio waveform at rates [2, 4, 8, 8]. Our decoder has 4 corresponding layers, which upsample at rates [8, 8, 4, 2]. We set the decoder dimension to 1536. In total, our model has 76M parameters, with 22M in the encoder, and 54M in the decoder. We also examine decoder dimensions of 512 (31M parameters) and 1024 (49M parameters).
We use the multi-period discriminator [18], and a complex multi-scale STFT discriminator. For the first, we use periods of  $[2, 3, 5, 7, 11]$ , and for the second, we use window lengths [2048, 1024, 512], with a hop-length that is  $1/4$  the window length. For band-splitting of the STFT, we use the band-limits  $[0.0, 0.1, 0.25, 0.5, 0.75, 1.0]$ . For the reconstruction loss, we use distance between log-mel spectrograms with window lengths [32, 64, 128, 256, 512, 1024, 2048], with corresponding number of mels for each of  $[5, 10, 20, 40, 80, 160, 320]$ . The hop length is  $1/4$  of the window length. We use feature matching and codebook losses, as described in Section 3.5.
For our ablation study, we train each model with a batch size of 12 for  $250\mathrm{k}$  iterations. In practice, this takes about 30 hours to train on a single GPU. For our final model, we train with a batch size of 72 for  $400\mathrm{k}$  iterations. We train with excerpts of duration 0.38s. We use the AdamW optimizer [23] with a learning rate of  $1e - 4$ ,  $\beta_{1} = 0.8$ , and  $\beta_{2} = 0.9$ , for both the generator and the discriminator. We decay the learning rate at every step, with  $\gamma = 0.999996$ .
4.4 Objective and subjective metrics
To evaluate our models, we use the following objective metrics: