8-cut/docs/plans/2026-04-19-audio-pipeline-improvements-design.md

Audio Pipeline Improvements Design

Date: 2026-04-19

Goal

Improve audio scan classification accuracy, especially for non-speech sounds (suction, gagging, impacts), through three changes:

1. Multi-layer feature extraction from existing HuBERT/Wav2Vec2 models
2. Two new embedding models: AST (AudioSet-supervised) and EAT (self-supervised + AudioSet finetuned)
3. Calibrated classifier for better threshold behavior

1. Multi-Layer Feature Extraction

Current behavior

model(waveforms) extracts embeddings from the last transformer layer only.

Change

Use model.extract_features(waveforms) (torchaudio API) to get all layer outputs. Select layers at quartile boundaries, mean-pool each over time, concatenate.

Model	Layers	Single-layer dim	Multi-layer dim (4 quartiles)
HUBERT_XLARGE	48	1280	5120
HUBERT_LARGE	24	1024	4096
HUBERT_BASE	12	768	3072
WAV2VEC2_BASE	12	768	3072

Implementation

• New entries in _EMBED_MODELS: "HUBERT_XLARGE_ML" -> 5120, etc.
• _extract_w2v_windows: when model name ends with _ML, call extract_features() instead of model(), select quartile layers, concat
• Cache key: model name includes _ML suffix -> separate cache files
• No change to classifier or training pipeline (HistGBT handles high-dim fine)

2. AST (Audio Spectrogram Transformer)

What

MIT/ast-finetuned-audioset-10-10-0.4593 via HuggingFace transformers. 86M params, 768-dim, supervised on AudioSet 527 sound classes.

Integration

• Load: ASTModel.from_pretrained() + ASTFeatureExtractor
• Preprocessing: ASTFeatureExtractor handles mel spectrogram from 16kHz raw audio
• Batching: prepare input_values per window, stack into batch, forward through model
• Multi-layer: output_hidden_states=True returns 13 layers; AST_ML variant concats quartile layers -> 3072-dim
• Model cached via _get_w2v_model() same lazy-load pattern

Entries

• "AST" -> 768
• "AST_ML" -> 3072

3. EAT (Efficient Audio Transformer)

What

worstchan/EAT-base_epoch30_finetune_AS2M via HuggingFace with trust_remote_code=True. 88M params, 768-dim, self-supervised + AudioSet finetuned.

Integration

• Load: AutoModel.from_pretrained(..., trust_remote_code=True)
• Preprocessing: manual 128-bin Kaldi fbank mel spectrogram via torchaudio, normalize with EAT constants (mel - (-4.268)) / (4.569 * 2), reshape to [B, 1, T, 128]
• Feature extraction: model.extract_features(mel) returns [B, seq, 768]; CLS token [:, 0, :] for utterance-level, or mean-pool [:, 1:, :] for frame-level. Use mean-pool for consistency with other models.
• Multi-layer: not natively supported, skip for now

Entry

• "EAT" -> 768

4. Calibrated Classifier

Wrap HistGradientBoostingClassifier in CalibratedClassifierCV(clf, cv=3, method='isotonic') after fitting. Gives well-calibrated probabilities -> threshold slider maps more linearly to precision/recall.

One change in train_classifier(), no UI changes needed.

5. Requirements

Add to requirements.txt:

transformers>=4.30
timm>=0.9

Both AST and EAT need transformers. EAT additionally needs timm (used internally by its custom model code). Both setup scripts (setup_env.sh, setup-windows.ps1) install from requirements.txt so no changes needed there.

Cache Compatibility

• All new model variants get distinct cache keys via model name in the hash
• Existing caches for HUBERT_XLARGE, BEATs, etc. remain valid and untouched
• New models create new .npz files in the same cache/w2v/ directory

UI Changes

• _EMBED_MODELS dict additions appear automatically in Train dialog model dropdown and scan model dropdown
• No other UI changes needed