WaLLoC — Stereo Audio

@inproceedings{jacobellis2025learned,
  title={Learned Compression for Compressed Learning},
  author={Jacobellis, Dan and Yadwadkar, Neeraja J.},
  booktitle={Data Compression Conference},
  year={2025},
  organization={IEEE}
}

arXiv

Github

Pre-trained model weights

The audio sibling of the image WaLLoC tutorial, with the same asymmetric design: an invertible wavelet packet transform does the heavy lifting, a tiny learned encoder maps the subbands to a low-dimensional integer latent, and an off-the-shelf lossless image codec (WebP) serves as the entropy coder. This walkthrough uses a 5 s excerpt from the danjacobellis/music example set and the single stereo_5x operating point — there is no quality knob.

import io
import torch
import torchaudio
import matplotlib.pyplot as plt
import PIL.Image
import PIL.ImageOps
from IPython.display import Audio
from torchvision.transforms import ToPILImage
from datasets import load_dataset
from compressors.audio_eval import normalize_audio, pad_audio, psnr_db
from compressors.walloc_audio._codec import (
    load_codec,
    encode_to_latent,
    latent_to_webp_bytes,
    decode_from_webp_bytes,
)

torch.set_num_threads(8)
device = 'cpu'

Skipping import of cpp extensions due to incompatible torch version 2.11.0+cu130 for torchao version 0.15.0             Please see https://github.com/pytorch/ao/issues/2919 for more info

/home/dgj335/g/lib/python3.12/site-packages/pytorch_wavelets/dtcwt/coeffs.py:7: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
  from pkg_resources import resource_stream

Load the codec

codec, config = load_codec(device)

print(f'channels    = {config.channels}')
print(f'J           = {config.J}  (wavelet packet depth -> {config.channels * 2**config.J} subband channels)')
print(f'latent_dim  = {config.latent_dim}')
print(f'latent_bits = {config.latent_bits}')
print(f'post_filter = {config.post_filter}')

channels    = 2
J           = 8  (wavelet packet depth -> 512 subband channels)
latent_dim  = 108
latent_bits = 8
post_filter = True

/home/dgj335/g/lib/python3.12/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)

Load an example clip

The danjacobellis/music dataset holds 24 example music files (the source files are themselves MP3 at ~258 kbps); each row’s audio is a torchcodec AudioDecoder. We take a 5 s stereo 44.1 kHz excerpt and normalize it with the harness convention: zero-mean float in \([-0.5, 0.5]\). The time axis is padded to a multiple of \(2^{16}\) before analysis, matching the evaluation harness.

CLIP_INDEX = 16        # which of the 24 example files
START_SECONDS = 150.0  # offset into the file, rounded to the nearest sample
DURATION_SECONDS = 5.0 # excerpt length in seconds

ds = load_dataset("danjacobellis/music", split="train")
samples = ds[CLIP_INDEX]["audio"].get_all_samples()
fs = samples.sample_rate
i0 = round(START_SECONDS * fs)
x = samples.data[:, i0 : i0 + round(DURATION_SECONDS * fs)]
x = normalize_audio(x)
n_samples = x.shape[-1]
x_padded = pad_audio(x.unsqueeze(0), 2**16).to(device)
print(f'x shape        = {tuple(x.shape)}  ({n_samples / fs:.1f} s at {fs} Hz)')
print(f'x range        = [{x.min():.3f}, {x.max():.3f}]')
print(f'x_padded shape = {tuple(x_padded.shape)}  '
      f'(padded to a multiple of 2**16: +{x_padded.shape[-1] - n_samples} zero samples)')

x shape        = (2, 220500)  (5.0 s at 44100 Hz)
x range        = [-0.500, 0.477]
x_padded shape = (1, 2, 262144)  (padded to a multiple of 2**16: +41644 zero samples)

A log-mel spectrogram of each channel (the same make_spectrogram view used across all of the audio tutorials), rendered directly as a PIL image with low frequencies at the bottom. The spectrogram is taken on the unpadded excerpt — the padded silence tail would just add a black band on the right.

mel = torchaudio.transforms.MelSpectrogram(sample_rate=fs, n_fft=4096, hop_length=512)

def make_spectrogram(x_1ch):
    S = mel(x_1ch).log()
    S = ((S - S.mean()) / (3 * S.std()) + 0.5).clamp(0, 1)
    return ToPILImage()(S.flip(0))  # low frequencies at the bottom

display(make_spectrogram(x[0]))
display(make_spectrogram(x[1]))

Audio(x.numpy(), rate=fs, normalize=False)

Wavelet packet analysis

Eight levels of wavelet packet decomposition turn the stereo waveform \((1, 2, L)\) into \((1, 2 \cdot 2^J, L/2^J)\) — 512 subband channels at 1/256 the time resolution. The transform is invertible, so this stage loses nothing; its job is energy compaction, concentrating the signal into the low-index subbands that the learned encoder can then exploit.

with torch.inference_mode():
    X = codec.wavelet_analysis(x_padded, codec.J)
print(f'X shape = {tuple(X.shape)}  '
      f'(expected: 1, 2 * 2**J = {2 * 2**config.J}, L / 2**J = {x_padded.shape[-1] // 2**config.J})')

X shape = (1, 512, 1024)  (expected: 1, 2 * 2**J = 512, L / 2**J = 1024)

plt.figure(figsize=(8, 4), dpi=120)
plt.imshow(X[0].abs().log1p().numpy(), aspect='auto', origin='lower', cmap='magma',
           extent=[0, x_padded.shape[-1] / fs, 0, X.shape[1]])
plt.xlabel('time [s]')
plt.ylabel('subband channel')
plt.colorbar(label='log1p |X|')
plt.show()

Encode and round

The lightweight encoder is a single 1×1 Conv1d followed by the entropy bottleneck, which compands and (in eval mode) rounds to the nearest integer. Note the latent is two samples longer than the subband signal — an upstream Conv1d(kernel_size=1, padding=1) quirk. The byte path carries this length around as side info.

with torch.inference_mode():
    z = codec.encoder[0:2](X)        # continuous companded latent
    z_hat = codec.encoder[2](z)      # rounded integer latent

print(f'z_hat shape = {tuple(z_hat.shape)}  '
      f'(1, latent_dim, L/2**J + 2 = {x_padded.shape[-1] // 2**config.J + 2})')
print(f'z_hat range = [{z_hat.min():.0f}, {z_hat.max():.0f}]  '
      f'(integer-valued, dtype={z_hat.dtype})')
print(f'dimensionality reduction = {x_padded.numel() / z_hat.numel():.2f}x')

z_hat shape = (1, 108, 1026)  (1, latent_dim, L/2**J + 2 = 1026)
z_hat range = [-32, 28]  (integer-valued, dtype=torch.float32)
dimensionality reduction = 4.73x

plt.figure(figsize=(5, 3), dpi=120)
plt.hist(z.flatten().numpy(), bins=151, range=(-25, 25), density=True)
plt.title('Histogram of pre-rounding latents (z)')
plt.xlabel('latent value')
plt.ylabel('density')
plt.grid(alpha=0.3)
plt.show()

Pack to bytes (WebP-lossless)

latent_to_webp_bytes pads the latent to a multiple of 128 along time, folds it to a 2-D grid (b c (w h) -> b c w h, h=128), packs 3 latent channels per RGB pixel plane, and saves the result as lossless WebP: an off-the-shelf lossless image codec is the audio codec’s entropy coder. The 4.74× dimensionality reduction becomes ~9.5× by storing 8-bit latents instead of 16-bit samples, and WebP-lossless adds another ~1.7× on this clip. At this 5 s clip length the harness padding is material: the codec actually codes \(2^{18} = 262{,}144\) samples, ~19% more than the 220,500-sample excerpt, so the rate below carries some overhead from coding the padded silence.

blob, latent_length = latent_to_webp_bytes(codec, z_hat.cpu())
buff = io.BytesIO(blob)
kbps = 8 * len(buff.getbuffer()) / (n_samples / fs) / 1000
CR = 16 * 2 * n_samples / (8 * len(buff.getbuffer()))
print(f'latent_length     = {latent_length}  (side info, carried alongside the blob)')
print(f'len(buff)         = {len(buff.getbuffer())} bytes')
print(f'rate              = {kbps:.1f} kbps')
print(f'compression ratio = {CR:.2f}x  (vs 16-bit stereo PCM)')

latent_length     = 1026  (side info, carried alongside the blob)
len(buff)         = 53254 bytes
rate              = 85.2 kbps
compression ratio = 16.56x  (vs 16-bit stereo PCM)

Visualize the packed latent

def scale_for_display(img, target=512):
    """Autocontrast, then integer-upscale (nearest) so the latent image is visible."""
    img = PIL.ImageOps.autocontrast(img)
    scale = max(1, target // max(img.size))
    return img.resize((img.width * scale, img.height * scale), PIL.Image.NEAREST)

latent_pil = PIL.Image.open(io.BytesIO(blob))
print(f'latent image: size = {latent_pil.size}, mode = {latent_pil.mode}')
scale_for_display(latent_pil)

latent image: size = (768, 54), mode = RGB

Decode from the buffer

The decoder gets only the WebP bytes in buff, plus the latent length and clip length as side info — never the in-memory z_hat. WebP-lossless makes the round trip exact, but the structure is the point: everything the receiver needs crosses the byte boundary. decode_from_webp_bytes inverts every encode stage in reverse order — WebP → PIL → latent unpack, unfold and trim to latent_length, learned decoder, wavelet synthesis, post filter, clamp — and trims to n_samples.

with torch.inference_mode():
    x_hat = decode_from_webp_bytes(codec, buff.getvalue(), latent_length, n_samples, device)
print(f'x_hat shape = {tuple(x_hat.shape)}')
print(f'x_hat range = [{x_hat.min():.3f}, {x_hat.max():.3f}]')
print(f'PSNR = {psnr_db(x, x_hat):.2f} dB')

x_hat shape = (2, 220500)
x_hat range = [-0.500, 0.468]
PSNR = 50.25 dB

Mel spectrograms and reconstruction

The same make_spectrogram view, channel 0 of the original next to the reconstruction. This excerpt keeps nearly all of its energy in the lower part of the image — 95% of it sits below ~1.7 kHz — and over that range the two images are visually indistinguishable. The only difference is at the very top: above ~8 kHz the original holds almost no energy to begin with, and the reconstruction attenuates that already-faint band further. That is consistent with the ~50 dB PSNR above.

print('original (channel 0)')
display(make_spectrogram(x[0]))
print('WaLLoC reconstruction (channel 0)')
display(make_spectrogram(x_hat[0].float().cpu()))

original (channel 0)

WaLLoC reconstruction (channel 0)

Audio(x_hat.numpy(), rate=fs, normalize=False)

Spatial quality (SSDR / SRDR)

WaLLoC codes the two stereo channels jointly, so spatial quality is worth checking: SSDR / SRDR (Watcharasupat & Lerch, ICASSP 2024) separate spatial distortion from residual distortion.

from compressors.spatial_audio_quality import ssdr_srdr

spatial = ssdr_srdr(x, x_hat, fs)
print(f"SSDR = {spatial['SSDR']:.2f} dB")
print(f"SRDR = {spatial['SRDR']:.2f} dB")

SSDR = 43.67 dB
SRDR = 23.24 dB

Sanity: codec.forward matches the manual pipeline

with torch.inference_mode():
    z_helper = encode_to_latent(codec, x_padded)
    x_hat_fwd, _, _ = codec(x_padded)
print(f'encode_to_latent matches manual encode: {torch.equal(z_helper, z_hat)}')
max_diff = (x_hat_fwd[0, :, :n_samples] - x_hat).abs().max().item()
print(f'max | codec.forward - buffer decode | = {max_diff:.6e}')

encode_to_latent matches manual encode: True
max | codec.forward - buffer decode | = 0.000000e+00