HyperAIHyperAI

Command Palette

Search for a command to run...

Diamond: A Sequence-to-Sequence Model for Speech Restoration via an Autoregressive RQ-Transformer over Neural Audio Codec Tokens

Almaz Zholdoshbek uulu Denis Pavlov Nursultan Bakashov Ulanbek Abdurazakov

Diamond-1.0: A Self-Autoregressive Speech Repair Model

Go to Notebook

Abstract

We present Diamond, a sequence-to-sequence speech-restoration model that turns degraded audio into studio (or near-studio) quality at 44.1 kHz. The degraded mel-spectrogram is processed by a bidirectional Transformer encoder; an autoregressive decoder with cross-attention predicts the tokens of a frozen neural audio codec (DAC, 9 RVQ codebooks). The decoder is an RQ-Transformer: a time-transformer models the frame sequence, while a compact depth-transformer models the codebook axis inside each frame, restoring the causal RVQ chain and distributing the output projection (read-out). Autoregressive models over acoustic tokens have been proposed before [41], but there the codebooks of a frame are read by parallel heads from a single vector; to the best of our knowledge, Diamond is the first speech-restoration model whose read-out is a two-transformer RQ decoder (a framelevel time-transformer plus a frame-local depth-transformer along the RVQ axis). Diamond is trained from scratch (≈166.6M parameters, no pretrained backbone) and reaches the level of the strongest current open systems for speech restoration, while remaining the best from-scratch restorer in its class.

One-sentence Summary

Researchers at nineninesix.ai propose Diamond, a sequence-to-sequence speech restoration model that turns degraded audio into studio-quality 44.1 kHz output using a bidirectional Transformer encoder and an autoregressive RQ-Transformer decoder over DAC tokens, which is the first restorer with a two-transformer read-out (a frame-level time-transformer and a frame-local depth-transformer along the RVQ axis), achieving state-of-the-art from-scratch performance.

Key Contributions

  • Diamond uses a two-transformer RQ decoder over neural audio codec tokens for speech restoration: a frame-level time-transformer and a frame-local depth-transformer along the RVQ axis restore the causal RVQ chain and distribute the read-out, replacing the parallel heads of prior autoregressive codec-token models.
  • The work diagnoses and resolves key instabilities of deep training over codec tokens from scratch, including layer-0 runaway, read-out inflation, and the exposure-bias cascade, with mechanistic explanations grounded in the literature.
  • A data and augmentation recipe (681.5 h with 28% native wideband content) is combined with a two-axis evaluation (DNSMOS-P.835 and CER) on real degraded recordings, showing second-place perceptual quality among open baselines (+0.255 OVRL, 657 of 750 clips improved) while openly measuring the content-preservation trade-off inherent to autoregressive generation.

Introduction

The authors tackle speech restoration, the task of recovering studio-quality 44.1 kHz audio from degraded recordings that have undergone aggressive lossy compression, band-limiting, clipping, and noise. Such recordings dominate real-world corpora but are unusable as training data because models reproduce the degradation. Restoration differs fundamentally from denoising: missing spectral content and transients must be generated, not just filtered, making it a generative problem. The authors introduce Diamond, a sequence-to-sequence model that leverages a frozen DAC and a two-transformer RQ decoder—a time-transformer plus a depth-transformer—to predict discrete codec tokens from mel spectrograms, trained entirely from scratch. They also diagnose and resolve key instabilities in deep codec-token training, and evaluate on both perceptual quality and content preservation, placing Diamond among the strongest open restorers without a pretrained backbone.

Dataset

The authors train Diamond on a studio speech dataset of 681.5 hours, split into 353,075 clips from approximately 2,500 speakers. A key property is that 28% of the data (≈190 hours) is native wideband at 44.1–48 kHz, providing real high-frequency content up to 22 kHz, not just upsampled 24 kHz targets.

Dataset composition and filtering

  • All clips are studio or near-studio recordings.
  • Clips are selected solely by duration: those between 2.5 and 30 seconds are kept. Shorter clips do not cover the training crop window, and longer ones inflate memory.
  • No perceptual quality filtering (e.g., DNSMOS, VQScore) is applied; the authors leave this for future work, noting that on less clean corpora it would improve results.
  • The manifest is shuffled on disk with a fixed seed.

How the data is used

  • The studio clips serve as clean targets. Targets are tokenized in advance.
  • The degraded input is synthesized on the fly inside the DataLoader, so the same target is seen under dozens of different degradations across epochs.
  • The augmentation pipeline is calibrated to match real-world degradations. It uses a reference set of 96,899 real degraded recordings to build a palette of 5,000 DSP parameter combinations via 8 channel features (SNR, noise floor, dynamic range, spectral rolloff, energy ratios in 3.4–8 kHz and above-8 kHz, and high-frequency instability).
  • The DSP chain: downsampling (24/32 kHz) → bass boost → microphone coloration → low‑pass (cutoff up to 12 kHz) → dynamic‑range compression → codec (MP3 57%, AAC 23%, OGG 15%, no codec 5%; bitrates 16–128 kbps; 1–3 passes, with 29% double and 9% triple re‑encodings) → optional high‑frequency instability (25%) → loudness levelling via tanh saturation to RMS –16 dB → real background noise transfer → peak normalization to 0.99. The noise consists of real background segments cut from live recordings with a VAD.
  • Sampling strategy: 60% k‑NN retrieval from the palette, 25% random palette combination, 15% freshly generated combination with wider spread; continuous parameters are jittered by ±40%. 5% of samples pass undegraded to teach the model to preserve clean audio.
  • Robustness track applied with probability 0.15 per sample: reverberation before the codec, broadband noise after the codec, packet loss (0.15), synthetic reverberation (0.10), and hard clipping (0.10). The order is deliberate — reverb survives encoding, so it is applied before the codec; additive noise is largely erased by the codec, so it is added after — ensuring the model learns dereverberation, denoising, and codec‑artifact removal simultaneously.

The training input is always generated fresh; no fixed degraded dataset is stored.

Method

The authors design a sequence-to-sequence Transformer architecture to remove a mixture of unknown degradations in a single pass. They explicitly avoid diffusion models because such models are typically conditioned on a specific degradation operator and become brittle when faced with the compound, non-differentiable, and unknown distortions found in real recordings. Instead, a bidirectional encoder with cross-attention allows a single conditional module to observe the entire degraded signal and jointly remove non-local artifacts.

To prevent the over-smoothing inherent in regressing to continuous targets, the model predicts discrete tokens rather than a spectrum or waveform. Predicting a continuous target under an L1L_1L1 or L2L_2L2 loss converges to the conditional mean, which muffs fine structures like sibilants and plosives. By contrast, predicting a token from a finite vocabulary is a classification task that selects one variant instead of averaging them. The decoder emits code indices from a frozen DAC (Descript Audio Codec) decoder, which then renders the codes into natural audio.

The model architecture processes the input through several specialized modules. The input is resampled to 24 kHz and converted into an 80-channel log-mel spectrogram. This front-end feeds into a twenty-layer bidirectional Transformer encoder. The encoder uses pre-norm RMSNorm, RoPE with learnable per-layer scales, and LayerScale on every residual branch to maintain stable training dynamics without temporal downsampling, preserving fine temporal details.

The decoder is a twenty-layer causal Transformer, referred to as the time-transformer. It emits one hidden vector per frame by attending to past frames via causal self-attention and to the full encoder output via cross-attention. To handle the nine codebooks of the RVQ codec, the authors introduce a compact depth-transformer. Instead of reading all nine codebooks in parallel from a single frame hidden vector, which dilutes capacity and breaks RVQ causality, the depth-transformer generates the nine codebooks sequentially within each frame. This restores the causal chain where each codebook refines the residual left by the previous ones.

Training is optimized using a per-frame weighted cross-entropy loss with label smoothing. The weights are assigned such that the coarse codebook carrying semantic content receives the largest weight, while a lifted tail keeps the fine codebooks from collapsing to the mode. A z-loss term is added to prevent logit inflation.

A critical component of the training pipeline is the data augmentation strategy, designed to emulate observed natural artifacts rather than applying arbitrary synthetic effects. The authors build a custom augmentor as a fitting problem. They extract channel features from a large pool of real degraded recordings and use inverse calibration to find DSP parameter combinations that statistically match the observed distribution.

The authors also analyze the structure of the artifact space to justify their augmentation strategy. They demonstrate that degradations do not fall into discrete classes but form a continuum. Standard clustering metrics fail to find a distinct number of groups, indicating that the model must cover the full density of the distribution rather than a handful of class representatives.

To ensure the synthesized degradations are realistic, the authors employ per-sample retrieval. For each training sample, a random real profile is drawn, and the nearest palette combination is retrieved. This strategy preserves the covariance structure of real transmission paths. Validation shows that the per-sample retrieval achieves a reproduction error lower than the intrinsic spread of the target corpus, effectively closing the gap between studio and real degraded recordings.

During inference, long autoregressive decoding is chunked to manage exposure bias. The model generates tokens sequentially, mitigated by a KV cache and batching of independent chunks. A repetition penalty is applied before every argmax to break autoregressive repetition loops, which is crucial for content preservation.

Experiment

The evaluation on Diamond-Bench (750 degraded recordings) with perceptual and content metrics shows that the autoregressive model, trained from scratch without pretrained backbones, delivers strong perceptual quality but exhibits a tail of content errors inherent to autoregressive generation. The design process revealed that transmission degradations occupy a continuous space without discrete classes, so per‑sample retrieval of real profiles preserves covariance better than random parameter sampling, and weighted sampling avoids over‑emphasizing unrealistic tails. The final recipe matches only channel features to prevent voice timbre transfer and mitigates exposure bias at inference rather than through time‑axis scheduled sampling, which caused gradient spikes.

The training dataset totals 681.5 hours from several sources, but only 28% consists of native wideband recordings; the remaining 72% are upsampled from 24 kHz, meaning the model must indirectly synthesize content above 12 kHz. This limited scale and bandwidth composition define the experimental budget, and the authors identify increasing the share of native wideband as the most direct lever for improvement. Only 28% of the 681.5-hour training set is native wideband; the rest is upsampled from 24 kHz, forcing the model to synthesize high-frequency content indirectly. The dataset is small by field standards (nearly four times smaller than the RE-USE training set), and the authors did not push data or compute to their limits, leaving headroom for future scaling.

A systematic set of experiments refuted seven intuitive hypotheses for speech restoration, including hand-crafted presets, degradation clustering, and uniform parameter sampling. The resulting Diamond model achieves strong perceptual quality (second of six systems) and median content preservation comparable to non-autoregressive restorers, but its autoregressive token generation introduces a thin tail of high-error clips, a structural limitation rather than a tuning defect. All seven manual design hypotheses were refuted or failed, steering the design away from hand-crafted presets and degradation classes toward a learned augmentation approach. Diamond improves 88% of clips with a mean OVRL gain of +0.255, but its mean CER (0.131) is inflated by rare derailments while median CER (0.028) matches top non-autoregressive systems.

Diamond delivers the second-highest perceptual quality (OVRL 3.829) among the evaluated restoration models, improving 88% of the degraded clips with a mean OVRL gain of 0.255. Its median CER of 0.028 matches other strong systems, but a thin tail of generation failures raises the mean CER to 0.131, a structural cost of its autoregressive design. All results were achieved with a small training budget and no pretrained backbone, making Diamond the best open restorer in its class. Diamond improves 657 of 750 degraded clips, lifting the overall perceptual quality above RE-USE, Resemble Enhance, UniSE, and VoiceFixer, while trailing only Sidon. The median CER is low (0.028) and comparable to UniSE and Resemble Enhance, but the mean CER (0.131) is dragged upward by a small number of clips where autoregressive generation derails, a pattern also seen in the other autoregressive system, UniSE.

Using a 681.5-hour training set with only 28% native wideband, systematic experiments disproved seven hand-crafted design hypotheses, steering the model toward a learned augmentation strategy. The resulting Diamond restorer achieves the second-highest perceptual quality among evaluated systems, improving 88% of degraded clips, but its autoregressive token generation introduces a thin tail of high-error clips that inflates the mean CER, a structural limitation rather than a tuning defect. Overall, Diamond is the best open restorer in its class, with further gains expected from a larger share of native wideband data.


Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing

HyperAI Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp