E2A2 — clean 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 clean speech is the 16 kHz mono speech codec trained clean→clean (an ordinary speech codec; LibriSpeech-only training pool). Ladder ps 256/128/64/32/16 × 4/6/6/12/8 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, clean_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 = 'clean_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/clean_speech   (hub format e2a2_hub1, source checkpoint rG1c_g4gs12_g5.pth)
modality      = audio   spatial rank = 1   input channels = 1
ps ladder     = [256, 128, 64, 32, 16] x [4, 6, 6, 12, 8] channels
boundaries    = [4, 10, 16, 28, 36]  (cumulative channels = operating points)
  op  4: decoder ps=128 dim=1024 depth=24
  op 10: decoder ps=128 dim=1024 depth=24
  op 16: decoder ps=128 dim=1024 depth=24
  op 28: decoder ps=128 dim=1024 depth=24
  op 36: decoder ps=128 dim=1024 depth=24
  scale 1: ps= 256  channels 1-4
  scale 2: ps= 128  channels 5-10
  scale 3: ps=  64  channels 11-16
  scale 4: ps=  32  channels 17-28
  scale 5: ps=  16  channels 29-36

entropy coder: GGDLPC, stream_mode=all_detail, macroregion=256 samples, no global stream
  alphabet [-31, 31]; persistent model = 216 law scalars + 108 predictor ints for 36 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 clean column; normalization follows the trainer’s audio contract: zero-mean, peak-normalize the INPUT to [−1, 1]; 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['clean']['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['clean']['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.50]

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 (clean)')
display(make_spectrogram(x[0]))
Audio(x.numpy(), rate=fs)

input (clean)

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 -> 165000 latent values (1.031x)
  scale 1 (ps= 256): latent (1, 4, 625)  range [-6.6, 7.0]
  scale 2 (ps= 128): latent (1, 6, 1250)  range [-5.8, 8.6]
  scale 3 (ps=  64): latent (1, 6, 2500)  range [-9.2, 8.2]
  scale 4 (ps=  32): latent (1, 12, 5000)  range [-10.4, 10.1]
  scale 5 (ps=  16): latent (1, 8, 10000)  range [-9.8, 9.2]
clipped: 0/165000

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 21,552 bytes total (per region min/median/max = 7/28/146)

  op  stream bits  bits/sample    kbps       CR  bits/latent
   4        8,416     0.052600    0.84   304.2x        3.366
  10       20,992     0.131200    2.10   122.0x        2.099
  16       33,824     0.211400    3.38    75.7x        1.353
  28      125,512     0.784450   12.55    20.4x        1.477
  36      172,416     1.077600   17.24    14.8x        1.045

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, 10, 16, 28, 36]: every op decoded from its truncated bytes alone, bit-exact

Reconstructions at each operating point

Each op decodes its byte-prefix with its own merged decoder; PSNR vs the input (= the clean reference for this codec).

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.052600 bits/sample    0.84 kbps   PSNR vs clean = 34.43 dB   PSNR vs input = 34.43 dB

10 ch:  0.131200 bits/sample    2.10 kbps   PSNR vs clean = 39.39 dB   PSNR vs input = 39.39 dB

16 ch:  0.211400 bits/sample    3.38 kbps   PSNR vs clean = 42.12 dB   PSNR vs input = 42.12 dB

28 ch:  0.784450 bits/sample   12.55 kbps   PSNR vs clean = 47.31 dB   PSNR vs input = 47.31 dB

36 ch:  1.077600 bits/sample   17.24 kbps   PSNR vs clean = 48.61 dB   PSNR vs input = 48.61 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_clean_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).

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)')

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 clean speech — validation set')
plt.grid(True, which='both', alpha=0.3); plt.legend(fontsize=7); plt.tight_layout(); plt.show()

print('E2A2 clean_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 clean_speech, validation average (verified streams):
   4 ch:  0.052983 bits/sample  (  0.85 kbps)   PSNR=28.79 dB
  10 ch:  0.159882 bits/sample  (  2.56 kbps)   PSNR=33.10 dB
  16 ch:  0.292507 bits/sample  (  4.68 kbps)   PSNR=35.43 dB
  28 ch:  1.090351 bits/sample  ( 17.45 kbps)   PSNR=42.07 dB
  36 ch:  1.676055 bits/sample  ( 26.82 kbps)   PSNR=46.22 dB