Geometric Iterative Retrieval for Neural Codec Resynthesis

Code accompanying the paper. Given the first RVQ codebook layer of a Descript Audio Codec (DAC) encoding, the models here predict the remaining residual layers and resynthesize the waveform. This repository contains everything needed to reproduce the data pipeline, training runs, and evaluation for our method and all baselines.

DAC configuration used throughout: 44.1 kHz, 9 codebook layers, codebook size 1024, codebook dimension 8, ~86 code stacks/second.

---

Methods and baselines

Each method is a PyTorch Lightning module. The short name in the first column is what conf/config.yaml:model_class (training) and --model_class (evaluation) expect, so a checkpoint trained under name X is evaluated with --model_class X.

Name	Module	Paradigm
attention	src/models/bert_latent_attention.py	Ours. Autoregressive over RVQ layers in continuous codebook space; cross-attention aggregation over previous layers; CLIP-style contrastive retrieval loss.
additive	src/models/bert_latent.py	Ablation of Ours: previous-layer embeddings are summed instead of cross-attended. Same per-layer contrastive objective.
full_residual	src/models/bert_latent_attention_full_residual.py	Variant of Ours trained against cumulative residual targets, with an extra magnitude head.
parallel	src/models/bert_parallel.py	One-step regression (OSR). Single forward pass predicts all residual codebook vectors at once, cosine-similarity objective.
osr_mse	src/models/osr_bert.py	MSE regression baseline. Predicts one 8-dim summed residual latent from layer 0; decomposed greedily into per-layer codebook indices at inference.
layer_grouped_ce	src/models/layer_grouped_tokenized_hf.py	Discrete token prediction (coarse-to-fine). Llama decoder with cross-entropy over flattened per-layer offset tokens, predicting layer k+1 frame by frame.
--baseline zeros	(no checkpoint)	Naive baseline. Decodes only the first n_given layers and drops the rest.

The first five share the DeBERTa-v3 backbone built in src/nn/deberta_builder.py; all of them read DAC's frozen codebooks and output projections via src/utils/dac_extraction.py.

Inference for every method lives in one place, src/benchmarks/inference.py (MODEL_CLASSES and predict_residual_codebooks), so the benchmark harness treats them uniformly.

---

Repository layout

conf/config.yaml               Hydra config: data, model, optimizer, logging
env.yaml                       Conda environment

src/
  data/tokenized_dataset.py    Fixed-length crops over pre-tokenized DAC codes;
                               weighted multi-dataset mixing; deterministic
                               crops for evaluation
  nn/                          Backbone builders (DeBERTa-v3, transformer blocks)
  models/                      Lightning modules — see table above
  utils/dac_extraction.py      Pull frozen codebooks / out_projs out of DAC
  train/
    train_bert.py              Trains attention | additive | full_residual |
                               parallel | osr_mse (select via model_class)
    train_layer_wise.py        Trains the CE token-prediction baseline
  benchmarks/
    data.py                    Val-split loaders, mirroring the training split
    metrics.py                 SI-SDR, log-spectral distance, FAD (VGGish)
    inference.py               Per-method prediction + DAC decoding
    codec_restoration.py       Main benchmark (SI-SDR / LSD / FAD vs n_given)
    layer_progression.py       Quality as a function of #predicted layers used
    aggregate_results.py       Campaign results -> paper-ready markdown tables
    generate_demo_samples.py   Per-method WAVs for listening tests
    generate_demo_clips.py     Long-form (chunked) per-method demo clips
    rank_demo_samples.py       Rank demo clips by improvement over naive
  scripts/
    processing/                Data download-side prep, tokenization, filtering
    job/                       SLURM wrappers for training / evaluation
    analysis/                  Codebook geometry analysis
  tests/                       Unit tests for the metrics

---

Setup

conda env create -f env.yaml
conda activate jupyternb

Place the DAC weights at data/model.pt (a torch.saved DAC model object, loaded with torch.load(..., weights_only=False)).

All entry points are run as modules from the repository root, e.g. python -m src.benchmarks.codec_restoration ....

---

Data pipeline

Three corpora are used: Jamendo-MTG (music), Common Voice (speech), and FMA (music). Each goes through resample → DAC-tokenize → filter, and is stored as a HuggingFace dataset on disk holding codes arrays of shape [9, T].

1. Resample to 44.1 kHz

   # Jamendo (and any HF audio dataset already on disk)
   python -m src.scripts.processing.preprocess \
       --input_dir  <hf_dataset_dir> \
       --output_dir data/jamendo-mtg-resampled \
       --sample_rate 44100
   
   # Common Voice, from the raw TSV + clips/ layout
   python -m src.scripts.processing.preprocess_common_voice \
       --input_dir  data/cv-corpus-24.0-2025-12-05/en \
       --output_dir data/common-voice-resampled \
       --split      train

2. Tokenize with DAC — chunked encoding with overlap to avoid boundary artifacts.

   python -m src.scripts.processing.tokenize_dataset \
       --input_dir  data/jamendo-mtg-resampled \
       --output_dir data/jamendo-tokenized \
       --model_path data/model.pt
   
   # FMA is tokenized in shards (streamed from HF), then merged
   python -m src.scripts.processing.tokenize_fma \
       --model_path data/model.pt \
       --output_dir data/fma_tokenized \
       --shard_size 5000
   
   python -m src.scripts.processing.merge_shards \
       --input_dir  data/fma_tokenized \
       --output_dir data/fma_tokenized_merged

3. Filter short samples — drop anything shorter than the training crop length, which would otherwise need padding.

   python -m src.scripts.processing.filter_dataset \
       --input_path  data/common-voice-tokenized \
       --output_path data/common-voice-tokenized-filtered \
       --min_length  128

Dataset locations are resolved from environment variables with the defaults in src/benchmarks/data.py:

Env var	Default
LOCAL_JAMENDO_DIR	data/jamendo-tokenized
LOCAL_CV_DIR	data/common-voice-tokenized-filtered
LOCAL_FMA_DIR	data/fma_tokenized_merged

Train/val splits are produced with train_test_split(test_size=cfg.val_ratio, seed=cfg.seed) per corpus. src/benchmarks/data.py reuses exactly that call, so evaluation runs on the same held-out samples the model was validated against — with deterministic_crop=True so every method sees identical time slices. Corpora are mixed during training with the dataset_weights ratio in the config ([10, 1, 2] for Jamendo/CV/FMA).

SLURM wrappers for these steps: src/scripts/processing/preprocess_job.sh, tokenize_job.sh, tokenize_fma_job.sh.

---

Training

Every DeBERTa-based method uses the same entry point; pick the method with the model_class config key:

python -m src.train.train_bert model_class=attention      # Ours
python -m src.train.train_bert model_class=additive       # aggregator ablation
python -m src.train.train_bert model_class=full_residual
python -m src.train.train_bert model_class=parallel       # one-step regression
python -m src.train.train_bert model_class=osr_mse        # MSE baseline

The CE token-prediction baseline has its own trainer (Llama backbone, FlashAttention-2, needs an Ampere+ GPU):

python -m src.train.train_layer_wise

Any config key can be overridden on the command line (Hydra), e.g. python -m src.train.train_bert model_class=parallel batch_size=32 learning_rate=1e-4. Set CKPT_PATH to resume; W&B run IDs are stored in the checkpoint and resumed automatically unless resume_wandb_run=false.

On SLURM: sbatch src/scripts/job/train_job.sh [overrides...]. That script stages the tokenized datasets from network storage onto node-local scratch, exports the LOCAL_*_DIR variables, runs the trainer (forwarding any Hydra overrides, e.g. sbatch src/scripts/job/train_job.sh model_class=parallel), and syncs checkpoints back at the end.

---

Evaluation

Main benchmark: SI-SDR, log-spectral distance (per pair) and Fréchet Audio Distance (set level, VGGish embeddings), against the full 9-layer DAC reconstruction, swept over the number of given layers:

python -m src.benchmarks.codec_restoration \
    --ckpt_path   data/checkpoints/<run>/last.ckpt \
    --model_class attention \
    --dac_path    data/model.pt \
    --n_given     1 2 3 \
    --per_dataset_samples 500 \
    --output_dir  logs/benchmarks

The naive baseline needs no checkpoint:

python -m src.benchmarks.codec_restoration --baseline zeros --n_given 1 2 3

Each run writes results.json (summary + per-dataset breakdown), per_file.csv (per-sample metrics), and run_info.txt (host, git SHA, argv).

Supporting analyses:

• python -m src.benchmarks.layer_progression --ckpt_path ... --model_class attention — decodes only the first K predicted layers for K = 1..9, showing whether higher predicted layers still help.
• python -m src.benchmarks.aggregate_results --manifest <campaign_manifest.json> --out_dir ... — collates a campaign into the paper's result tables.
• python -m src.scripts.analysis.codebook_cosine_sim --dac_model data/model.pt --audio_dir data/samples — per-layer cosine similarity between DAC latents and their nearest codebook entry.

On SLURM, src/scripts/job/eval_job.sh stages the DAC model and a checkpoint to local scratch and forwards any extra arguments to the benchmark. The exact submissions behind the paper's numbers are recorded in src/scripts/job/submit_sweep.sh (checkpoint selection) and src/scripts/job/submit_full.sh (final full evaluation) — including which checkpoint was used for each method.

Listening-test material

# Short per-method samples (Ours + OSR), aligned crops in one process
python -m src.benchmarks.generate_demo_samples --ours_ckpt ... --osr_ckpt ... --out_dir demo/samples

# Long-form clips for all five methods, chunked to each model's context length
python -m src.benchmarks.generate_demo_clips --ours_ckpt ... --osr_ckpt ... \
    --osr_mse_ckpt ... --ce_ckpt ... --clip_frames 860 --out_dir demo/clips_10s

# Rank clips by how much Ours improves on the naive baseline
python -m src.benchmarks.rank_demo_samples --samples_dir demo/samples

SLURM wrappers: src/scripts/job/demo_samples_job.sh, demo_clips_job.sh.