mrothroc/sc35-s4d-lin-mixlab

S4D-Lin on Speech Commands (35-way, raw waveform), reproduced in mixlab

A 306K-parameter state-space model that classifies spoken keywords directly from the raw 16 kHz waveform, with no spectrogram, no MFCC, no feature extraction of any kind. It is a reproduction of the S4D-Lin result from On the Parameterization and Initialization of Diagonal State Space Models (Gu et al., 2022), trained in mixlab.

A 1-second waveform is a 16,000-step sequence, and this model handles it with 306K parameters. For comparison, a 26M-parameter ConvNet scores lower on the same task. The training recipe is defined in mixlab json and it was trained on a single Apple M4. You are welcome to use it as-is, but the true purpose is to provide a starting point for tinkering.

The architecture is a JSON config, so swapping the sequence mixer is an edit rather than a rewrite. The same file runs on an Apple Silicon laptop and on a big CUDA machine with no changes, there are details in the cookbook entry. You can develop locally where iteration is cheap, then ship the file to bigger hardware when the run gets long.

Further details on some of the tinkering we did on this job are in the cookbook entry

How close is it

This is a representative reproduction, not a replication of the authors' experiment. The goal was to express the benchmark architecture as a mixlab config and see whether it lands where the paper says it should. It does.

Official testing_list.txt split, 11,005 utterances, single seed:

	test accuracy
this model	96.14% (10,580 / 11,005)
S4D-Lin, published (Table 11)	96.25% (±0.03)

0.11 points under the paper, at 306,083 parameters against their "306K".

For context from that same table: ResNet-18 (216K) scores 77.86, and a 26.2M-parameter ConvNet, 85x larger, scores 95.51.

Implementation notes

• Single seed. The paper reports +/-0.03 over multiple seeds; we ran one, so our number is outside their printed interval and 0.11 is not a precise gap.
• Two details are declared assumptions, not things the paper states: n_ssm: 2 (SSM parameter sharing), and the checkpoint-selection rule. This checkpoint is the best devel checkpoint, epoch 36 of 40 (step 190,872), scored once on test, which is the selection rule the reference uses.
• The paper's 8 kHz zero-shot transfer result is not included. It needs rate-aware rediscretisation, which mixlab does not implement at the time of this run.
• This is a keyword spotter for 35 fixed words, trained on Speech Commands v0.02. It is not a general speech model.
• It classifies whole clips, and is not a streaming detector. The blocks are bidirectional, so it reads the entire one second window before deciding.
• There is no unknown or silence class. Any input gets one of the 35 labels, including background noise and words it has never heard. If you need rejection, you have to add it.
• Inputs longer than one second are truncated to the first 16,000 samples.
• CPU inference latency and training memory are unmeasured; we have not profiled either.

Contents

file	what it is
mixlab_checkpoint.safetensors	the trained weights (mixlab native format)
config.json	the mixlab config that produced them; not a transformers config
labels.json	output index to word. The order is not alphabetical
modeling_s4d.py	standalone PyTorch port of the architecture. Hardcodes 6 layers / dim 128 / n_ssm 2, so it loads this checkpoint, not a re-architected one
mixlab_ckpt.py	reads mixlab's .st container, which the strict safetensors loader rejects
verify_labels.py	checks the label mapping against real audio

Usage

This is not a transformers model. AutoModel, pipeline("audio-classification", ...) and the Hub inference widget will not work on it: the architecture has no transformers equivalent, and config.json here is a mixlab config, not a transformers one. Use the code below instead. A separate PyTorch port of the block matched mixlab's forward outputs on 35 inputs to 6.251e-05 max absolute logit difference with 35/35 argmax agreement, measured on a one-step checkpoint. That is evidence the layer maths and weight layout line up; it is not a re-verification of the trained weights below.

Needs torch and numpy. Run it from inside this directory, so mixlab_ckpt and modeling_s4d import. You also need a clip to classify: any 16 kHz mono 16-bit wav of one of the 35 keywords. The path below names a file from Speech Commands v0.02 (2.3 GB tarball); substitute your own recording if you would rather not download it.

import json, wave
import numpy as np, torch
import mixlab_ckpt
from modeling_s4d import SC35Model

MEAN, STD, L = -2.791959, 2818.166625, 16000

def load_wav(path):                      # mono, 16 kHz, 16-bit PCM
    with wave.open(path, "rb") as w:
        a = np.frombuffer(w.readframes(w.getnframes()), dtype="<i2").astype(np.float32)
    x = np.zeros(L, np.float32)
    x[:min(len(a), L)] = a[:L]           # right-pad to 16,000
    return (x - MEAN) / (STD + 1e-5)     # PCM units, see below

model  = SC35Model(mixlab_ckpt.load("mixlab_checkpoint.safetensors")).eval()
labels = json.load(open("labels.json"))["id2label"]

x = torch.from_numpy(load_wav("yes/004ae714_nohash_0.wav")).reshape(1, L, 1)   # <- your clip
with torch.no_grad():
    logits = model(x)                    # -> [1, 35]
print(labels[str(int(logits.argmax(-1)))])

Do not smoke-test it with torch.randn: noise contains no spoken word, so the answer is meaningless and will not look obviously wrong. Use a real clip.

Use labels.json for the output mapping. Do not sort the class names yourself. The index order is two concatenated alphabetical runs, not one, so sorting the 35 Speech Commands folder names gives the wrong word for 34 of the 35 classes, silently.

verify_labels.py checks the mapping against real audio rather than asking you to take it on trust: point it at an extracted Speech Commands directory and it classifies one clip per word and compares against the folder name. Ten classes, ten correct, confidences 0.846 to 1.000 when we ran it.

The PyTorch port is for inference on CPU, and is not a normal nn.Module. It holds the BatchNorm affine weights and running statistics in a plain dict rather than registering them, which has three consequences: .to("cuda") or .to("mps") will not move them and inference fails on a device mismatch; state_dict() omits them, so saving and reloading the module loses model state; and parameters() reports 304,291 of the 306,083 learned values, so optimizing it would silently leave BatchNorm frozen. It is a faithful reader for these weights, not a training harness. Use mixlab for anything beyond checking outputs.

mixlab_ckpt.py is included because the strict safetensors loader rejects this file with SafetensorError: Error while deserializing header: MetadataIncompleteBuffer. The container carries 4 trailing bytes the header does not account for, so it is slightly out of spec rather than merely awkward. mixlab_ckpt.py parses the header directly and returns a plain dict of arrays, which is what SC35Model takes.

Normalization used in training: (x - mean) / (std + 1e-5) with mean -2.791959, std 2818.166625, computed over the training split in float64.

These constants are in signed 16-bit PCM amplitude units, i.e. roughly [-32768, 32767], not the [-1, 1] range most audio loaders return by default. If your loader normalises to [-1, 1], multiply by 32768 before applying them, or your inputs will be scaled wrong by four orders of magnitude.

Also required: audio must be mono at 16 kHz, and short clips must be right-padded with zeros, not left-padded. Most Speech Commands clips are under a second, so this applies to nearly all of them.

We measured which of these actually matter, on ten real clips:

mistake	effect
[-1,1] units instead of PCM	0/10 correct. Collapses to one class for every input
no normalization	0/10
8 kHz audio fed as 16 kHz	0/10
left-pad instead of right-pad	fine at native lengths; 4/10 predictions change once clips are half padding
pad before vs after normalizing	no effect (max logit delta 0.005, even at 75% padding)

The scale warning is important: torchaudio.load and librosa both return [-1,1], and the failure is uniform rather than noisy, so it looks like a broken checkpoint instead of a scaling bug. Padding order matched training but made no difference to any of the ten predictions.

Reproducing or changing it

config.json is the exact file that produced this checkpoint: 40 epochs, batch 16, seed 2222, lr 0.01 with state_lr 0.001 on the SSM parameters. About 33 hours on an M4 Max, on the laptop.

The same file runs unchanged on NVIDIA. We used that: the shorter runs and the mixer comparison were done on rented CUDA (RTX 4090 and L4) with no edits to the config, because the comparison needed two arms of 20 to 34 hours each and that is not laptop work. mixlab selects the native Metal S4D kernel on the Mac and the CUDA path on NVIDIA.

There is also a 10-epoch recipe that reaches 95.06% devel in roughly 11 hours on a single consumer GPU, if you want the cheap path in. Both configs, both learning curves, and the full write-up are in the cookbook entry.

Swapping the mixer is a config edit. Replace the six s4d blocks with mamba3-canonical and change the learning rate. We ran that arm too: it reaches 90.74% devel at a matched 6 epoch budget against s4d's 94.98% devel, and will not train at all at S4D's lr 0.01 (non-finite six times out of six, mostly during warmup).

That is roughly the expected direction, since raw audio favors a time invariant model. Treat it as a demonstration that the swap is cheap to try, not as an architecture verdict: the two arms did not get equal tuning effort, s4d's learning rate being the original authors' and ours for mamba3 being a value we found that trains. The details are in the cookbook entry.

For the story behind it, read I Built an ML Architecture Lab in Go.

Citation

The result being reproduced:

@inproceedings{gu2022parameterization,
  title     = {On the Parameterization and Initialization of Diagonal State Space Models},
  author    = {Gu, Albert and Gupta, Ankit and Goel, Karan and R{\'e}, Christopher},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2022}
}