VENOM: Very Efficient Neural Optimal-filter for MKIDs

A selective state space model (Mamba architecture) for photon energy estimation from raw MKID IQ timestream data. Replaces the traditional coordinate transform + optimal filter pipeline with a learned model that operates directly on I(t), Q(t). Designed to fit inside an FPGA-class resource budget on the MKIDGen3 RFSoC readout platform.

Runs on PyTorch: CUDA when available, otherwise CPU. The selective scan accumulates in float64 so the parallel (training) and recurrent (FPGA) forms agree, which rules out Apple GPUs: MLX has float64 only on its CPU device and PyTorch's MPS device has none. The MLX backend was removed for that reason.

Headline results

A single 3,314-parameter model, trained per detector from the same recipe and scored on a 15% stratified validation hold-out, averaged over three training seeds:

Detector	N wavelengths	Range	VENOM mean KDE R	Optimal filter mean R
InHf bilayer (Zobrist 2022)	10	254–1310 nm	26.4 ± 0.2	24.2
PtSi (Zobrist 2019)	5	808–1310 nm	10.3 ± 0.04	8.8

The InHf optimal filter values are the published per-wavelength results of Zobrist 2022, which use a separate template for each wavelength (mkidcalculator.filter_pulses(template_mask=True)); InHfData/metadata.json holds a later loose-cut rerun of that pipeline. Choosing a template needs the photon's wavelength, so that filter exists only for calibration data. VENOM is one model with no per-wavelength tuning. On InHf it leads the filter from 254 to 814 nm, is within 5% of it at 917, 979 and 1310 nm, and trails it at 1110 nm (−9%). The PtSi optimal filter is the published shared 920 nm template analysis at 808, 980 and 1310 nm, and scripts/ptsi_of/ repeats that analysis on the traces for 920 and 1120 nm. The PtSi edge lines (808, 1310 nm) depend on how far past them the synthetic energies extend, so the fair PtSi comparison is the interior lines: VENOM 9.3 / 8.3 / 9.7 against 8.8 / 9.0 / 8.8.

scripts/paper_campaign.py runs every training behind the paper, scripts/paper_tables.py builds its tables, and scripts/generate_paper_figures.py its figures.

Architecture

VENOM uses a two-tier design motivated by the eventual FPGA deployment target:

Tier 1 — SSM Backbone (SSMBackbone): Processes raw IQ samples through a single Mamba block (LayerNorm, causal depthwise Conv1D of width 4, selective SSM with input-dependent B, C, Δ, SiLU gate). Diagonal S4D-real state matrix stored as log_neg_A for stability. Runs sample-by-sample in recurrent mode at ADC rate.

Tier 2 — Energy Head (EnergyHead): Takes the full sequence of backbone features, applies learned attention pooling over time, and regresses photon energy with a 3-layer GELU MLP. Outputs (μ̂, log σ̂²) for Gaussian NLL training. Runs at photon rate (orders of magnitude slower than ADC rate).

Dual execution modes

• Parallel (__call__): Segmented Hillis-Steele associative scan over the full sequence, used during training. Accumulated in float64 to limit rounding error at long sequence lengths.
• Recurrent (step): O(1) per-sample with carried (h, conv_state), matches the eventual FPGA datapath.

Numerical agreement between the two paths is characterized across the full validation set by scripts/parallel_recurrent_distribution.py. Worst-case |e_par − e_rec| is 0.6 µeV on PtSi (sequence length 120) and 1.9 µeV on InHf (sequence length 426), about 10⁻⁵ of the per-pulse σ. Per-wavelength R values shift by ≤ 0.03 between the two modes.

Supporting modules

• backend.py — thin PyTorch layer that exports the mx, nn, optim namespace the model code is written against, plus DEVICE (CUDA, else CPU; never MPS).
• synthetic_pca.py — PCA-based synthetic pulse generator. Streaming Gram eigendecomposition → K=50 principal components, per-wavelength mean and covariance via shape-preserving PCHIP interpolation, energies drawn up to 10% beyond the outer lines, Cholesky sampling plus orthogonal CSD reconstruction residuals. Produces 120k synthetic pulses/epoch.
• synthetic.py — legacy template-interpolation synthetic generator (kept for reference; superseded by synthetic_pca.py).
• ssvkernel.py — vectorized locally-adaptive KDE from Shimazaki & Shinomoto (2010); used by compute_resolving_power for the KDE FWHM.

Winner configuration

Parameter	Value
d_model	16
d_state	8
n_layers	1
d_conv	4
expand	1
d_head_hidden	32
Total params	3,314

Installation

Python 3.11+.

# Core dependencies
pip install numpy scipy tqdm safetensors

pip install torch        # CUDA build for NVIDIA GPUs; CPU build otherwise

# Reports and figures
pip install matplotlib reportlab

VENOM_BACKEND may be unset, torch or auto; all mean PyTorch.

Paper checkpoints

opt/paper/ holds the checkpoint and metadata of every training run behind the paper, as written by python scripts/paper_campaign.py publish:

Path	Contents
opt/paper/headline/venom_{inhf,ptsi}_s{42,43,44}	Production checkpoints, three seeds per detector
opt/paper/baselines/, opt/paper/baseline_seeds/	Non-selective S4D, causal TCN and hand-feature MLP, three seeds each
opt/paper/loss/, opt/paper/clip/	Loss ablation and variance-floor sweep (PtSi)
opt/paper/arch/	One-knob architecture sweep
opt/paper/margin/	Synthetic energy margin check
opt/paper/pcak/	PCA component sweep, K = 1 to 20, three seeds per detector (Fig. 5)
opt/paper/index.json	Per-line validation R, command line and environment of every run
opt/paper/analysis/	PtSi optimal filter reanalysis, PTQ, variance head audit, parallel vs recurrent, InHf 313 nm check, PCA explained variance

Each _meta.json records the model, load and training configuration, the validation split and the per-line results. mkid_ssm.load_run(weights) rebuilds the model and reloads exactly the pulses it was scored on:

from backend import mx
from mkid_ssm import load_run, predict_batched

run = load_run("opt/paper/headline/venom_inhf_s42.safetensors")
mu = predict_batched(run.model, run.iq[mx.array(run.val_indices)])[:, 0]

The TCN and MLP baselines are other model classes (scripts/baseline_tcn.py, scripts/baseline_mlp.py); their metas hold their own configuration and results. Checkpoints are written in the true safetensors format. Legacy checkpoints that used a .safetensors suffix for PyTorch zip archives still load through model.load_weights().

Reproducing the winner runs

# InHf (10 wavelengths, 254-1310 nm); InHf needs 100 epochs to converge, PtSi 25
python mkid_ssm.py --downsample 4 --n-components 50 \
    --loss gaussian_nll --epochs 100 --n-synthetic 120000 \
    --n-layers 1 --d-model 16 --d-state 8 --expand 1 \
    --peak-sigma 2.0 --report-every 25

# PtSi (5 wavelengths, 808-1310 nm)
python mkid_ssm.py --ptsi --downsample 4 --n-components 50 \
    --loss gaussian_nll --epochs 25 --n-synthetic 120000 \
    --n-layers 1 --d-model 16 --d-state 8 --expand 1 \
    --peak-sigma 2.0 --report-every 25

The two commands differ only in --ptsi. Both use the same model, training recipe, and PCA synthetic generator. The trained checkpoints are in opt/paper/headline/.

Remote training via Ray

scripts/run_remote_train.py is a Ray-Client launcher that runs the standard PyTorch training path on a remote CUDA host without modifying any training code. The driver only needs ray==2.55.1 matched to the cluster's Python minor version (3.12 for our reference setup running an NGC PyTorch container with PyTorch 2.11, CUDA 13, RTX 5090); all heavy imports (torch, mkid_ssm) happen on the worker.

# One-time: driver env matched to the cluster's Python minor version
conda create -n pyray python=3.12 -y
/opt/anaconda3/envs/pyray/bin/pip install "ray[client]==2.55.1"

# Site-specific defaults: copy .env.example to .env and edit, or
# export the variables in your shell, or pass --ray-address /
# --data-root explicitly on the CLI.
cp .env.example .env   # then edit .env with your cluster URL and data path

# Submit a run; flags after `--` are forwarded verbatim to mkid_ssm.main()
/opt/anaconda3/envs/pyray/bin/python scripts/run_remote_train.py \
    --run-name ptsi_winner \
    -- --ptsi --downsample 4 --n-components 50 --loss gaussian_nll \
       --epochs 25 --n-synthetic 120000 --n-layers 1 --d-model 16 \
       --d-state 8 --expand 1 --peak-sigma 2.0 --report-every 25

The launcher uploads the repo source as runtime_env.working_dir (data, weights, and zips are excluded), submits a single @ray.remote(num_gpus=1) task that symlinks <data-root>/{InHf,PtSi}Data into the working dir, sets VENOM_BACKEND=torch, and pulls weights/, training_reports_*/, and VENOM_*.pdf back into runs/<run-name>/ after the run completes.

Common CLI flags

Flag	Meaning
--ptsi	Use PtSi dataset instead of InHf
--data-dir DIR	Any dataset directory with a metadata.json (see New datasets); exclusive with --ptsi
--downsample N	Anti-aliased FIR decimation factor (code default 1 for InHf, 2 for PtSi; the recipe uses 4)
--n-tau T	Record length after the trigger, T × τ_qp (default 5)
--pre-trigger-frac F	Pre-trigger context as a fraction of the post-trigger span (default 0.2)
--peak-sigma S	MAD peak-height clip in σ, thresholds fitted on the training side (code default 3.0; the recipe uses 2.0)
--loss {mse,gaussian_nll,...}	Code default mse; the recipe uses gaussian_nll
--n-synthetic N	Synthetic pulses per epoch (code default 41000; the recipe uses 120000)
--n-components K	PCA components for the synthetic generator (code default 10; the recipe uses 50)
--epochs N	Training epochs (default 20)
--micro-batch N	Pulses per forward/backward pass within a 256-pulse step (default 256)
--seed N	Seed for initialization and data order (default 42)
--report-every N	Emit multi-page PDF training report every N epochs
--full	Use full trace, no pulse windowing
--jitter	Per-wavelength label jitter
--plot-features	Plot all backbone feature channels after training
--synthetic	Use fully synthetic (no real) calibration for quick tests
--no-baseline-subtract	Keep each pulse's absolute IQ baseline (default: subtract it, so run-to-run baseline drift cannot identify a line)

Unknown flags are rejected; python mkid_ssm.py --help lists every flag.

Data loading

load_mkid_data() loads per-wavelength .npz files from a data directory. Each file must contain i_trace, q_trace, mask, peak_heights. Energy labels are computed from $E = hc/\lambda$.

When a metadata.json with an opt_filt_R dictionary is present, those per-wavelength OF resolving powers are used in preference to recomputing from peak_heights. This is how the published per-wavelength-template OF values from Zobrist 2022 are restored for InHf (shipped with InHfData/metadata.json). PtSi has no metadata.json and uses a post-clip KDE of the stored shared-template peak heights, matching Zobrist 2019.

from mkid_ssm import load_mkid_data

iq, energies, wl_ids, wl_labels, opt_R, peak_heights = load_mkid_data(
    data_dir='InHfData',
    downsample=4,
    n_tau=5.0,
    peak_height_sigma=2.0,
)

Supported datasets

Dataset	Dir	Wavelengths	Raw rate	Source
InHf bilayer	InHfData/	254–1310 nm (10)	0.8 MHz	Zobrist 2022
PtSi	PtSiData/	808–1310 nm (5)	2.0 MHz	Zobrist 2019

Obtaining the data

The raw IQ calibration data (InHfData.zip ≈ 1.1 GB, PtSiData.zip ≈ 580 MB) is too large to ship through GitHub. It is available on request — contact the Mazin Lab and we will provide a download link. Extract each archive in the repo root so InHfData/ and PtSiData/ sit next to mkid_ssm.py.

New datasets

A dataset is a directory of <λ>nm.npz files (i_trace, q_trace, mask, peak_heights) plus metadata.json. metadata.json gives sample_rate_mhz (raw rate), a tag that names outputs, and source_nm (line label to calibrated wavelength), which sets which files are lines, their order and their energy labels.

scripts/mkidcalc_baseline.py builds such a directory from raw mkidplotter data with Zobrist's mkidcalculator pipeline (loop fit, phase/dissipation, cuts, per-wavelength optimal filter templates, energy calibration). Its PEP 723 header pins the analysis environment, so uv run needs no setup:

uv run scripts/mkidcalc_baseline.py path/to/config.yaml --stage fit      # then process, metrics, filter, calcheck, export
python mkid_ssm.py --data-dir NewData --downsample 1 --n-tau 5 ...        # same flags as the recipe above

It writes calibrated optimal filter energies (eV) as peak_heights, so VENOM scores the optimal filter with its own KDE estimator on the same validation pulses; mkidcalculator's own resolving powers are kept in metadata.json under opt_filt_R_mkidcalc*. --n-tau and --pre-trigger-frac set the record length (defaults 5 and 0.2). Each checkpoint's _meta.json records the model, load and training configuration; mkid_ssm.load_run(weights) rebuilds the model and reloads exactly the pulses it was scored on.

Energy resolution metrics

• compute_resolving_power(values, energy) — KDE-based $R = E / \mathrm{FWHM}$ using the locally adaptive ssvkernel (Shimazaki & Shinomoto 2010). Headline metric; matches mkidcalculator's convention.
• summarize_resolving_power(pred, true, wl_ids) — Gaussian-approximation $R = E / (2.355 \cdot \mathrm{RMSE})$. Stricter (penalizes per-wavelength bias). Used as a training monitor. Edge wavelengths excluded from the aggregate.

Training details

• Optimizer: AdamW, weight decay 1e-4
• LR schedule: cosine annealing with 5-epoch linear warmup
• Loss: the recipe uses Gaussian NLL, which outputs (μ, log σ²). MSE, Huber and β-NLL reach the same resolution on PtSi and are equally stable across seeds; NLL is kept because it trains the variance head. With the log-variance floor lifted (VENOM_LOG_VAR_CLIP_LOW=-14, --loss beta_nll) that head gives a per-pulse σ calibrated over the validation set.
• Validation: 15% stratified-by-wavelength hold-out (fixed seed)
• Gradient accumulation: configurable micro-batch size to cap peak memory
• Gradient checkpointing: backbone layers

Scripts

Auxiliary scripts live in scripts/:

Script	Purpose
train_single.py	Train one config with periodic PDF reports
run_remote_train.py	Ray-Client launcher: run the PyTorch training path on a remote CUDA host
generate_paper_figures.py	Regenerate all paper figures from winner weights
parallel_recurrent_distribution.py	Characterize parallel-vs-recurrent mode agreement on the full validation set
convert_inhf_data.py	Convert mkidcalculator pickles to the common npz format
make_report.py	PDF data summary (raw-trace overview)
explore_data.py	Quick data-exploration plots
sweep_hyperparams_v2.py	Grid hyperparameter sweep with per-config reports
paper_campaign.py	Every training behind the paper (seeds, baselines, ablations, sweeps), resumable, memory gated
paper_tables.py	The paper's tables from the campaign summary
ptsi_of/	PtSi optimal filter reanalysis from the traces (reproduces the published R within 2%)
variance_head_audit.py	Calibration of the variance head (pinned fraction, χ², coverage)
compare_of_venom.py	VENOM against the optimal filter on the same pulses, with validity flags and figures
mkidcalc_baseline.py	mkidcalculator optimal filter baseline for a new dataset
run_grid.py	Record length × downsampling grid

Tests

pytest tests/
# In headless/CI sessions:
VENOM_BACKEND=torch pytest tests/

Covers: model shapes, recurrent/parallel equivalence, KDE accuracy against Gaussian theory, edge cases, stratified splitting, micro-batch weights, IQ normalization.

References

• Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces," COLM 2024
• Shimazaki & Shinomoto, "Kernel Bandwidth Optimization in Spike Rate Estimation," J. Comput. Neurosci. 29, 171 (2010)
• Zobrist et al., "Wide-band parametric amplifier readout and resolution of optical microwave kinetic inductance detectors," Appl. Phys. Lett. 115, 042601 (2019)
• Zobrist et al., "Membraneless Phonon Trapping and Resolution Enhancement in Optical Microwave Kinetic Inductance Detectors," Phys. Rev. Lett. 129, 017701 (2022)
• Smith et al., "MKIDGen3: Energy-resolving single-photon-counting MKID readout on an RFSoC," Rev. Sci. Instrum. 95, 114705 (2024)
• Fritsch & Carlson, "Monotone Piecewise Cubic Interpolation," SIAM J. Numer. Anal. 17, 238 (1980)

Citation

If you use VENOM in published work, please cite this repository and the associated paper (forthcoming). A CITATION.cff file is included; GitHub will render a "Cite this repository" shortcut in the sidebar.

License

BSD 3-Clause. See LICENSE for the full text. Briefly: you may use, modify, and redistribute this software with or without modification, provided the copyright notice is retained and the Mazin Lab / UCSB name is not used to endorse derivative products without written permission.