neural-audio-codec

A mini neural audio codec (SoundStream / EnCodec style) trained from scratch on real LibriSpeech speech — on a MacBook (Apple Silicon, MPS) in under 10 minutes.

Live in-browser demo (ONNX inference, no server): https://sugeerth.github.io/portfolio/ — see the audio-codec demo.

The v2 model is a 6.79M-parameter convolutional encoder/decoder with a 4-codebook residual vector quantizer (RVQ). Because it is trained with quantizer dropout, a single model serves three bitrates at inference time — just use fewer codebooks.

Results (v2, real LibriSpeech speech)

From assets_v2/metrics.json, evaluated on 6 held-out clips:

Codebooks	Bitrate	Compression vs PCM16	Mean SI-SDR
1	0.675 kbps	568.9x	-2.55 dB
2	1.35 kbps	284.4x	-1.96 dB
4	2.7 kbps	142.2x	-1.71 dB

Honest caveats: this is a deliberately tiny model trained for 1500 steps (~9.4 min on MPS) on a small slice of LibriSpeech dev-clean (73 utterances, ~8 min of audio). At these extreme bitrates the mean SI-SDR is negative — speech is intelligible but clearly lossy, and quality improves monotonically with codebook count (per-sample numbers in assets_v2/samples.json range from -7.2 to +3.3 dB). The point is the full pipeline: train → quantize → multi-bitrate decode → ONNX → browser.

Listen (original vs reconstruction A/B)

Held-out LibriSpeech clips, each rendered at 1, 2 and 4 codebooks:

Sample	Original	0.675 kbps	1.35 kbps	2.7 kbps
0	wav	wav	wav	wav
1	wav	wav	wav	wav
2	wav	wav	wav	wav
3	wav	wav	wav	wav
4	wav	wav	wav	wav
5	wav	wav	wav	wav

Transcripts and per-sample SI-SDR are in assets_v2/samples.json.

Architecture

codec.py — ~200 lines, three pieces:

• Encoder: 1-D conv stack with residual blocks (dilations 1 and 3), 4 strided downsampling stages (2x, 4x, 5x, 8x = 320x total). 24 kHz waveform → 128-dim latent at 75 Hz frame rate.
• Residual VQ: up to 4 codebooks of 512 entries (9 bits each), so bitrate = 75 Hz x 9 bits x n_codebooks → 675 / 1350 / 2700 bps.
  • EMA codebook updates (van den Oord 2017, as in EnCodec/DAC) instead of a codebook loss — keeps codebooks stable as the encoder distribution drifts.
  • Dead-entry restart: codes whose EMA cluster size decays below a threshold are re-seeded from random encoder outputs, preventing codebook collapse. Measured usage after training: 351–404 of 512 entries used per codebook, perplexity 203–291 (assets_v2/codebook_usage.json).
  • Quantizer dropout (SoundStream-style): each training step uses a random number of active codebooks, so one model decodes sensibly at any of the three bitrates.
• Decoder: mirror of the encoder with transposed convs, tanh output.

Training loss (losses.py): L1 waveform + multi-scale STFT (spectral convergence + log-magnitude at 4 FFT scales) + 0.25 x commitment loss. SI-SDR is tracked as the eval metric.

Repo layout

codec.py              encoder / RVQ (EMA + dead-entry restart) / decoder
losses.py             multi-scale STFT loss, SI-SDR
data.py               v1 synthetic "speech-like" signal generator
data_real.py          LibriSpeech clip loader (random crops from .npz cache)
cache_librispeech.py  download LibriSpeech dummy slice -> librispeech_cache.npz
train.py              v1 trainer (synthetic data, 2 codebooks) -> assets/
train_v2.py           v2 trainer (real speech, 4 codebooks, quantizer dropout) -> assets_v2/
export_onnx.py        export encoder.onnx + decoder.onnx + codebooks.bin for the browser demo
assets/               v1 artifacts (loss curves, metrics, one A/B wav pair, specs)
assets_v2/            v2 artifacts (metrics, curves, 6 A/B wav sets, specs, codebook usage)

Note: v1 (train.py, assets/metrics.json) is a 15.1M-param 2-codebook model trained on synthetic harmonic signals — its +6.98 dB SI-SDR is on that much easier synthetic data, not comparable to the v2 LibriSpeech numbers above.

How to train

Dependencies: torch numpy soundfile matplotlib scipy datasets (plus onnx for export).

# 1. Cache the LibriSpeech slice (hf-internal-testing/librispeech_asr_dummy, resampled to 24 kHz)
python3 cache_librispeech.py

# 2. Train v2 (1500 steps, ~10 min on Apple Silicon)
python3 train_v2.py

Device selection is automatic: MPS (Apple Silicon) → CUDA → CPU. The whole thing was developed and trained on MPS — no NVIDIA GPU required. Artifacts (metrics, curves, A/B wavs, spectrogram mosaic, codec.pt weights) land in assets_v2/.

Model weights (assets_v2/codec.pt, ~29 MB) and the data cache (librispeech_cache.npz, ~46 MB) are gitignored; both regenerate from the two commands above.

How to export to ONNX (browser inference)

python3 export_onnx.py

This exports encoder.onnx + decoder.onnx (opset 17, dynamic time axis) plus a raw codebooks.bin (4 x 512 x 128 float32) and a codec.meta.json. The RVQ itself runs in JavaScript — it's just per-frame argmin lookups — which avoids exporting one ONNX graph per bitrate. The script also sanity-checks the PyTorch RVQ against a NumPy reimplementation of the JS path. (Output paths are hardcoded to my local portfolio checkout; edit OUT_DIR / pt_path at the top for your setup.) That is exactly what powers the live demo at https://sugeerth.github.io/portfolio/.