E2A2 — denoising speech

E2A2 (E²A², Encoding Efficient Asymmetric Autoencoders, github.com/danjacobellis/E2A2) is the successor of FRAPPE v3: the same split trainer (frozen per-group encoders — one learned strided projection per latent channel — SC6 softsign companding onto the coder alphabet [−31, 31], one merged decoder per operating point) with the entropy stage replaced by GGDLPC (a learned FIR predictor per channel, two affine laws mapping a causal statistic to a folded-GGD Huffman table, plus a derived zero-run mode; the C engine builds itself on import). Every reported rate is real byte-padded stream bits, round-trip verified in this notebook by decoding the latents back from the truncated bytes.

E2A2 denoising speech is the 16 kHz mono speech codec trained noisy→clean: the input is speech + noise (+ reverb) and the reconstruction target is the clean speech, so the decoder outputs an estimate of the clean signal from the coded noisy input. Ladder ps 256/128/64/32 × 4/8/8/12 channels, all_detail streams at macroregion 256 samples (16 ms — streaming-capable).

Stream layout all_detail: no global stream — every channel is coded per macroregion, and every macroregion is decodable from its own bytes alone (one byte-pad each). The stream for operating point n is the literal bit-prefix covering channels ≤ n.

The codec is loaded from the Hugging Face hub (danjacobellis/E2A2, denoising_speech/ subdirectory) through the compressors.e2a2 public API — there is no local checkpoint or training code in this notebook.

import json, glob
import torch, numpy as np, matplotlib.pyplot as plt
import torchaudio
from IPython.display import Audio, display
from datasets import load_dataset
from torchvision.transforms import ToPILImage
from compressors import e2a2

Load the codec and its entropy coder

One MergedAutoencoder per operating point, all from the hub. Each operating point has its own merged decoder (multidec); the encoders froze at merge, so every operating point’s encoder is a channel prefix of the full one. entropy.json is the GGDLPC blob (per-channel predictor taps + law scalars); the Huffman tables are generated from it when the coder loads.

NAME = 'denoising_speech'
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
config, models, coder = e2a2.load_all_codecs(NAME, device=device)
ops = list(config.cumulative_channels)
model = models[ops[-1]]
rom = e2a2.rom_facts(coder.blob)

print(f'codec         = {config.codec}/{NAME}   (hub format {config.format}, source checkpoint {config.source_checkpoint})')
print(f'modality      = {config.modality}   spatial rank = {config.dim}   input channels = {config.input_channels}')
print(f'ps ladder     = {list(config.ps)} x {list(config.group_sizes)} channels')
print(f'boundaries    = {ops}  (cumulative channels = operating points)')
for n in ops:
    d = config.decoder_configs[ops.index(n)]
    print(f'  op {n:>2}: decoder ps={d["decoder_ps"]} dim={d["decoder_dim"]} depth={d["decoder_depth"]}')
for s, (ps_s, start, end) in enumerate(model.scale_groups):
    print(f'  scale {s+1}: ps={ps_s:>4}  channels {start+1}-{end}')
print(f'\nentropy coder: GGDLPC, stream_mode={coder.stream_mode}, macroregion={coder.macroregion} samples'
      + (f', global stream = channels 1-{coder.n_coarse}' if coder.n_coarse else ', no global stream'))
print(f'  alphabet [-31, 31]; persistent model = {rom["law_scalars"]} law scalars + '
      f'{rom["predictor_ints"]} predictor ints for {rom["n_channels"]} channels')
print(f'  {rom["note"]}')

codec         = E2A2/denoising_speech   (hub format e2a2_hub1, source checkpoint denoise_g4gs12.pth)
modality      = audio   spatial rank = 1   input channels = 1
ps ladder     = [256, 128, 64, 32] x [4, 8, 8, 12] channels
boundaries    = [4, 12, 20, 32]  (cumulative channels = operating points)
  op  4: decoder ps=128 dim=1024 depth=12
  op 12: decoder ps=128 dim=1024 depth=12
  op 20: decoder ps=128 dim=1024 depth=24
  op 32: decoder ps=128 dim=1024 depth=24
  scale 1: ps= 256  channels 1-4
  scale 2: ps= 128  channels 5-12
  scale 3: ps=  64  channels 13-20
  scale 4: ps=  32  channels 21-32

entropy coder: GGDLPC, stream_mode=all_detail, macroregion=256 samples, no global stream
  alphabet [-31, 31]; persistent model = 192 law scalars + 96 predictor ints for 32 channels
  edges, folded GGD pmfs, quotient split, canonical Huffman tables and the run mode are GENERATED from the per-channel scalars at load (GGDLPC.tables)

Analysis filterbank

One learned ps-tap projection per latent channel: impulse response (left) and magnitude response (right, Nyquist 8 kHz); ≤ 6 channels per panel.

def chunk_sizes(n, max_per=6):
    k = max(1, -(-n // max_per))
    return [n // k + (1 if i < n % k else 0) for i in range(k)]

fs_hz = config.sample_rate
N = 2048
freqs_khz = np.fft.rfftfreq(N) * fs_hz / 1e3
panels = []
for s, (ps_s, start, end) in enumerate(model.scale_groups):
    W = model.encoders[s][0].weight.data.cpu().numpy()[:, 0]
    off = 0
    for sz in chunk_sizes(end - start):
        panels.append((ps_s, start + off + 1, start + off + sz, W[off:off + sz]))
        off += sz

fig, axes = plt.subplots(len(panels), 2, figsize=(10, 1.6 * len(panels)), dpi=100, constrained_layout=True)
axes = np.atleast_2d(axes)
for (ax_imp, ax_mag), (ps_s, ch_lo, ch_hi, filt) in zip(axes, panels):
    t_ms = np.arange(filt.shape[1]) / fs_hz * 1e3
    for c, h in enumerate(filt, start=ch_lo):
        H = 20 * np.log10(np.abs(np.fft.rfft(h, N)) + 1e-9)
        ax_imp.plot(t_ms, h, lw=0.8, label=f'ch{c}')
        ax_mag.plot(freqs_khz, H, lw=0.8, label=f'ch{c}')
    ax_imp.set_title(f'ps={ps_s}   channels {ch_lo}-{ch_hi}  (impulse response)', fontsize=8)
    ax_mag.set_title(f'ps={ps_s}   channels {ch_lo}-{ch_hi}', fontsize=8)
    ax_mag.set_xlim(0, fs_hz / 2e3); ax_mag.set_ylim(-60, 10)
    ax_imp.set_ylabel('h(t)'); ax_mag.set_ylabel('|H(f)| (dB)')
    ax_mag.legend(fontsize=6, loc='upper right', framealpha=0.6, ncol=3)
axes[-1, 0].set_xlabel('time (ms)'); axes[-1, 1].set_xlabel('frequency (kHz)')
plt.show()

Load an example clip

One row of the frozen validation set danjacobellis/e2a2_speech_val (Audio columns clean / noisy / target, 16 kHz, one row per validation clip — the trainer’s frozen sets). The codec input is the noisy column; normalization follows the trainer’s audio contract: zero-mean, peak-normalize the INPUT to [−1, 1], and apply the same affine map to the reference. The reference is the target column — the speech component of noisy under the mixture’s own gain, reverbed in the rows where the mixture is (the codec removes the noise and keeps the room); this is what the trainer and the hub rd.json scored against, not the dry clean column; cropped (not looped) to a whole number of macroregions.

ROW_INDEX = 9
ds = load_dataset('danjacobellis/e2a2_speech_val', split='validation')
row = ds[ROW_INDEX]
fs = config.sample_rate
x = torch.from_numpy(row['noisy']['array'].astype(np.float32)).view(1, -1)
mean = x.mean(); scale = 1.0 / ((x - mean).abs().max() + 1e-8)
x = (x - mean) * scale
unit = max(coder.macroregion, max(config.ps))
L = unit * (x.shape[-1] // unit)
x = x[..., :L]
n_samples = x.shape[-1]
x_in = x.unsqueeze(0).to(device)
x_ref = torch.from_numpy(row['target']['array'].astype(np.float32)).view(1, -1)
x_ref = ((x_ref - mean) * scale)[..., :L]
print(f'fs   = {fs} Hz')
print(f'clip = {tuple(x.shape)}  ({n_samples / fs:.2f} s = {n_samples // coder.macroregion} macroregions of {coder.macroregion} samples),  range [{x.min():.2f}, {x.max():.2f}]')

fs   = 16000 Hz
clip = (1, 160000)  (10.00 s = 625 macroregions of 256 samples),  range [-1.00, 0.49]

mel = torchaudio.transforms.MelSpectrogram(sample_rate=fs, n_fft=1024, hop_length=256, n_mels=80)

def make_spectrogram(x_1ch):
    S = (mel(x_1ch) + 1e-8).log()
    S = ((S - S.mean()) / (3 * S.std()) + 0.5).clamp(0, 1)
    return ToPILImage()(S.flip(0)).resize((600, 160))

print('input (noisy)')
display(make_spectrogram(x[0]))
Audio(x.numpy(), rate=fs)

input (noisy)

print('clean reference (same normalization)')
display(make_spectrogram(x_ref[0]))
Audio(x_ref.numpy(), rate=fs)

clean reference (same normalization)

Analysis transform and rounding

Encode to the per-scale latents (companding included) and round to the coder’s [-31, 31] grid.

with torch.no_grad():
    latents = model.encode(x_in)
n_latent_values = sum(z.numel() for z in latents)
print(f'{x_in.numel()} input samples -> {n_latent_values} latent values ({n_latent_values / x_in.numel():.3f}x)')
for s, z in enumerate(latents):
    print(f'  scale {s+1} (ps={model.scale_groups[s][0]:>4}): latent {tuple(z.shape)}  range [{z.min():.1f}, {z.max():.1f}]')
n_clip = sum(int((z.round().abs() > 31).sum()) for z in latents)
latents_q = e2a2.encode_to_latents(model, x_in)
latents_np = [z[0].numpy() for z in latents_q]
print(f'clipped: {n_clip}/{n_latent_values}')

160000 input samples -> 92500 latent values (0.578x)
  scale 1 (ps= 256): latent (1, 4, 625)  range [-4.1, 3.7]
  scale 2 (ps= 128): latent (1, 8, 1250)  range [-3.7, 4.2]
  scale 3 (ps=  64): latent (1, 8, 2500)  range [-2.7, 2.3]
  scale 4 (ps=  32): latent (1, 12, 5000)  range [-2.8, 4.2]
clipped: 0/92500

GGDLPC all-detail entropy coding

encode_latents produces one byte-padded stream per 16 ms macroregion and nothing else — every region decodable from its own bytes alone; the stream for operating point n is the literal bit-prefix covering channels ≤ n. CR is against 16-bit PCM.

enc = e2a2.encode_latents(coder, latents_q)
assert enc.gblob == b''
region_bytes = (enc.prefix[:, -1] + 7) // 8
print(f'{len(enc.streams)} macroregion streams; at max detail {int(region_bytes.sum()):,} bytes total '
      f'(per region min/median/max = {int(region_bytes.min())}/{int(np.median(region_bytes))}/{int(region_bytes.max())})')
print(f'\n{"op":>4} {"stream bits":>12} {"bits/sample":>12} {"kbps":>7} {"CR":>8} {"bits/latent":>12}')
for n in ops:
    bits = e2a2.op_bits(coder, enc, n)
    nl = sum(min(max(n - start, 0), end - start) * z.shape[1]
             for z, (ps_s, start, end) in zip(latents_np, model.scale_groups))
    print(f'{n:>4} {bits:>12,} {bits / n_samples:>12.6f} {bits / (n_samples / fs) / 1000:>7.2f} '
          f'{16 * n_samples / bits:>7.1f}x {bits / nl:>12.3f}')

625 macroregion streams; at max detail 5,443 bytes total (per region min/median/max = 4/6/63)

  op  stream bits  bits/sample    kbps       CR  bits/latent
   4        7,376     0.046100    0.74   347.1x        2.950
  12       21,352     0.133450    2.14   119.9x        1.708
  20       30,520     0.190750    3.05    83.9x        0.939
  32       43,544     0.272150    4.35    58.8x        0.471

Round-trip verification

verify_latents decodes the latents back from the truncated bytes alone at every operating point’s literal bit-prefix (the global stream from its own bytes; every macroregion from its own bytes, no cross-region input) and asserts them bit-exact against the encoder’s. The reconstructions below are decoded from those byte-decoded latents (decode_latents), not from the encoder’s copy.

e2a2.verify_latents(coder, enc, latents_q, ops)
print(f'round trip verified at ops {ops}: every op decoded from its truncated bytes alone, bit-exact')

round trip verified at ops [4, 12, 20, 32]: every op decoded from its truncated bytes alone, bit-exact

Reconstructions at each operating point

The decoder outputs an estimate of the CLEAN speech from the coded noisy input: PSNR vs the clean reference (the objective) and, for reference, vs the noisy input (which this codec is not trying to match). Each op decodes its byte-prefix with its own merged decoder.

rd_clip = []
for n in ops:
    lts = e2a2.decode_latents(coder, enc, n)
    with torch.no_grad():
        xh = models[n].decode([z.to(device) for z in lts]).clamp(-1, 1)[0].float().cpu()
    bits = e2a2.op_bits(coder, enc, n)
    psnr_in = -10 * torch.nn.functional.mse_loss(x / 2 + 0.5, xh / 2 + 0.5).log10().item()
    psnr_ref = -10 * torch.nn.functional.mse_loss(x_ref / 2 + 0.5, xh / 2 + 0.5).log10().item()
    rd_clip.append((bits / n_samples, psnr_ref, n))
    print(f'{n:>2} ch:  {bits / n_samples:.6f} bits/sample  {bits / (n_samples / fs) / 1000:6.2f} kbps   '
          f'PSNR vs clean = {psnr_ref:.2f} dB   PSNR vs input = {psnr_in:.2f} dB')
    display(make_spectrogram(xh[0]))
    display(Audio(xh.numpy(), rate=fs))

 4 ch:  0.046100 bits/sample    0.74 kbps   PSNR vs clean = 33.10 dB   PSNR vs input = 33.04 dB

12 ch:  0.133450 bits/sample    2.14 kbps   PSNR vs clean = 37.73 dB   PSNR vs input = 37.54 dB

20 ch:  0.190750 bits/sample    3.05 kbps   PSNR vs clean = 39.72 dB   PSNR vs input = 39.46 dB

32 ch:  0.272150 bits/sample    4.35 kbps   PSNR vs clean = 40.24 dB   PSNR vs input = 39.95 dB

Per-macroregion rate profile

Each 16 ms macroregion carries its own byte-padded stream — the rate allocation over time is directly visible, and any region can be transmitted, dropped, or upgraded independently.

t_edges = np.arange(len(enc.streams) + 1) * coder.macroregion / fs
fig, axes = plt.subplots(2, 1, figsize=(8, 4), dpi=120, sharex=True, gridspec_kw={'height_ratios': [1, 1.6]})
axes[0].plot(np.arange(n_samples) / fs, x[0].numpy(), lw=0.3, color='0.4')
axes[0].set_ylabel('amplitude'); axes[0].set_title('input', fontsize=9)
for n in ops:
    rb = 8 * ((enc.prefix[:, n - 1] + 7) // 8)
    axes[1].stairs(rb / (coder.macroregion / fs) / 1000, t_edges, label=f'{n} ch')
axes[1].set_xlabel('time (s)'); axes[1].set_ylabel('kbps')
axes[1].set_title(f'per-macroregion rate ({coder.macroregion} samples = {coder.macroregion/fs:.3f} s each)', fontsize=9)
axes[1].legend(fontsize=8); axes[1].grid(True, alpha=0.3)
plt.tight_layout(); plt.show()

Rate–distortion

This clip’s operating points over the whole-validation-set curve from results/e2a2_denoising_speech/ (PSNR vs the clean reference; falls back to the hub rd.json). There is no in-repo baseline for this modality; the reference lines are the trainer ledger’s silence baseline (24.70 dB vs clean) and the identity ceiling (the noisy input itself vs clean: 34.55 dB).

def latest(pattern):
    paths = sorted(glob.glob(pattern))
    return paths[-1] if paths else None

rd_path = latest(f'results/e2a2_{NAME}/rate_distortion_*.json')
if rd_path is not None:
    rd = json.load(open(rd_path))
    key = 'channel_counts' if 'channel_counts' in rd else 'quality_values'
    rd_ops = list(rd[key])
    rd_label = f"E2A2 {NAME} — whole validation split ({rd.get('n_images', rd.get('n_samples'))} samples, compressors harness, verified streams)"
else:
    rd = None
    hub_rd = e2a2.hub_rd_curve(NAME)
    rd_ops = [p['n_ch'] for p in hub_rd['rd_curve']]
    rd_label = f"E2A2 {NAME} — whole validation split (hub rd.json, trainer's verified streams)"

if rd is not None:
    val_bps = [rd['results'][str(n)]['mean']['kbps'] * 1000 / rd['sample_rate'] for n in rd_ops]
    val_psnr = [rd['results'][str(n)]['mean']['PSNR_dB'] for n in rd_ops]
else:
    val_bps = [p['bpp'] for p in hub_rd['rd_curve']]; val_psnr = [p['psnr'] for p in hub_rd['rd_curve']]

plt.figure(figsize=(6.5, 4), dpi=120)
plt.semilogx(val_bps, val_psnr, 'o-', color='tab:blue', label=rd_label)
plt.semilogx([p[0] for p in rd_clip], [p[1] for p in rd_clip], 's--', color='tab:blue', alpha=0.5, label=f'E2A2 {NAME} — this clip (row {ROW_INDEX})')
plt.axhline(24.70, color='0.5', ls=':', lw=1, label='silence vs clean (24.70 dB, trainer ledger)')
plt.axhline(34.55, color='0.3', ls='--', lw=1, label='identity ceiling: noisy input vs clean (34.55 dB)')
for n, bv, psv in zip(rd_ops, val_bps, val_psnr):
    plt.annotate(f'{n}ch', (bv, psv), fontsize=7, color='tab:blue', xytext=(3, -9), textcoords='offset points')
plt.xlabel('bits per sample (16 kHz mono)'); plt.ylabel('PSNR vs clean (dB)')
plt.title('E2A2 denoising speech — validation set')
plt.grid(True, which='both', alpha=0.3); plt.legend(fontsize=7); plt.tight_layout(); plt.show()

print('E2A2 denoising_speech, validation average (verified streams):')
for n, bv, psv in zip(rd_ops, val_bps, val_psnr):
    print(f'  {n:>2} ch:  {bv:.6f} bits/sample  ({bv * fs / 1000:6.2f} kbps)   PSNR={psv:.2f} dB')

E2A2 denoising_speech, validation average (verified streams):
   4 ch:  0.046045 bits/sample  (  0.74 kbps)   PSNR=28.81 dB
  12 ch:  0.165305 bits/sample  (  2.64 kbps)   PSNR=32.19 dB
  20 ch:  0.293591 bits/sample  (  4.70 kbps)   PSNR=33.17 dB
  32 ch:  0.623112 bits/sample  (  9.97 kbps)   PSNR=33.76 dB