Stage 3: in-plugin autoencoder training + dataset building

[gustavokch]

What

Ports the upstream Python training pipeline into the live Rust plugin. Previously only the frozen decoder shipped (inference-only); training was offline-only in Python. Now the user builds a dataset and trains the autoencoder in-session, then plays their own model.

Highlights

• Full autoencoder in Rust (autoencoder.rs) — encoder + decoder, manual backprop, Adam, BatchNorm. Fixes two upstream bugs: broken Adam bias correction (per-param step counter t) and no batch shuffle. Converges better; not a numeric match to Python (by design).
• Dataset building — Capture pattern (snapshot the live grid) and Add audio… (decode .wav/.flac via symphonia → spectral-flux onset detection via rustfft → 32-dim sample). Both sources accumulate.
• Background training — runs on nih-plug's background thread, never the audio thread. Live epoch/loss progress bar + cancel.
• Lock-free hot-swap — trained decoder swaps into the audio path via ArcSwap (process() does a wait-free load(), never locks a Mutex). model_generation invalidates the regen cache on swap.
• Encoder at runtime — Encode pattern → latent sets the 4 latent sliders to the current grid's latent code.
• Persistence — trained model saved with the DAW session via #[persist]; baked decoder.json stays the default/fallback.

Design doc: docs/plans/2026-06-16-runtime-training-design.md.

New deps

arc-swap, rustfft, rfd, symphonia (wav+flac).

Testing

• 34 cargo tests — AE overfit-to-zero (gradient-correctness gate), Adam-fix lock-in, shuffle determinism, export↔decoder parity (1e-9), corpus_encode parity with the Python test_corpus_encode.py cases, onset sanity, and an end-to-end train→hot-swap→persist→encode integration test.
• cargo clippy --all-targets -- -D warnings clean.
• CLAP + VST3 bundle builds; both headless host scale-tests pass (audio path intact through real hosts).
• pluginval runs in CI.

Manual verification still needed

GUI check in Carla (VST3): train a model, watch the loss bar, confirm the pattern swaps, test Encode → latent, save/reload the project to verify persistence.

[gustavokch]

Code review

Overview

Solid, well-tested stage. Training is correctly isolated from the audio thread, the export format reuses the existing Decoder inference path (single forward impl, 1e-9 parity test), and the atomics carry clear ordering rationale. Test coverage is genuinely good: gradient-correctness gate (overfit-to-zero), Adam-fix lock-in, shuffle determinism, banker's-rounding parity with the Python encoder, synthetic onset sanity, and an end-to-end train→hot-swap→persist→encode test.

Findings

M1 — RT-safety: model hot-swap can deallocate the old decoder on the audio thread. (training.rs / lib.rs)
The audio thread does self.train.model.load() in maybe_regen. ArcSwap::load() is wait-free, but wait-free ≠ allocation-free: after the GUI/background thread calls model.store(new), the previous Arc<Decoder> is freed when the last outstanding guard is released — and that last release can happen on the audio thread, dropping a heap Vec of ops inside process(). The PR text ("wait-free load(), never locks a Mutex") is true but doesn't cover this. Impact is a brief, user-initiated (rare) glitch, not a crash — hence medium, not blocking.
Fix: retire old models off the audio thread, e.g. push the prior Arc into a "graveyard" Mutex<Vec<Arc<Decoder>>> drained by the GUI/background thread, or use arc_swap::cache::Cache so the audio thread holds a stable clone and never performs the final drop.

M2 — run_ingest status restore contradicts its own comment. (training.rs)
The comment says "keep showing Done", but the code restores Running only and otherwise forces Idle — so a finished training's Done status collapses to Idle after an audio ingest. Either preserve the captured prev status verbatim, or fix the comment.

L1 — encode→latent round-trip is lossy. The encoder ends at Dense(8→4) with no final activation, so its latent output is unbounded, but set_latents clamps to the params' [0,1]. "Encode pattern → latent" can therefore saturate and not reproduce the pattern. Acceptable, but note it in the UI/docs or constrain the encoder output range.

L2 — loss gradient isn't normalized by batch size. autoencoder.rs::fit uses grad = -d per element (summed over rows in the dense backward), so weight-gradient magnitude scales with batch size. Adam partly absorbs this, but changing the Batch field silently changes learning dynamics. Worth a doc note (matches the "not a numeric match" caveat, but users won't expect batch to act like a 2nd LR knob).

L3 — long audio files collapse to one bar. file_to_sample treats the whole file as bar_length = 1, so a multi-minute file's onsets quantize into 16 steps and most onset information is discarded. Documented as "approximate", but the UX may surprise; consider per-bar windowing later.

L4 — rfd::FileDialog::pick_files() blocks the editor thread. Called inside the egui frame closure; the native modal dialog stalls GUI repaint until dismissed. Audio thread is unaffected (good), but worth a comment.

L5 — minor Train/Ingest enable race. TrainStatus is set to Ingesting/Running only once the background task starts, so there's a small window after dispatch where the GUI still treats the plugin as idle. Low risk; tighten by setting status optimistically at dispatch if desired.

L6 — initialize() bumps model_generation on every call. Harmless (forces a regen), but a host that re-initializes repeatedly re-stores the model each time. Could guard on "model actually present and changed".

L7 — release.yml lives on both runtime-training and release-workflow. Potential merge conflict / duplicated CI; reconcile before merging both.

Positives

• No locks on the audio thread; atomics use Relaxed with a documented "no cross-variable invariant" justification.
• Deterministic SplitMix64 PRNG → reproducible init/shuffle without a rand dep.
• Export reuses Decoder::from_json_str, so there is exactly one inference impl, pinned by a 1e-9 parity test.
• UntrainedBatchNorm export guard mirrors the Python and fails loudly.

Verdict

Approve with nits. M1 (audio-thread dealloc) is the one I'd fix before relying on this in a live set; M2 and the L-items are cleanups.

_{🤖 Generated with Claude Code}

[gustavokch]

Review addressed — d6e3939

Thanks for the review. All actionable items are fixed; the two deferred ones are explained below. cargo test → 35 pass (added a graveyard test), clippy --all-targets -D warnings clean.

M1 — audio-thread dealloc on hot-swap. Fixed. Hot-swaps now go through TrainShared::swap_model, which uses ArcSwap::swap (not store) and parks the displaced Arc<Decoder> in a generation-tagged graveyard instead of letting it drop. The audio thread publishes gen_acked at the end of maybe_regen — after its load() guard is dropped, Release-ordered — and collect_garbage (GUI/background thread, Acquire) drops only the retired decoders the audio thread has provably moved past. So the heap free always runs off the audio thread; process() never deallocates. The graveyard is drained from the editor frame, the background executor, and initialize (all non-audio). New swap_model_retires_and_collects test covers retire → ack → collect and confirms the live decoder stays intact.

M2 — run_ingest status. Fixed. It now restores the captured prev status verbatim (a prior Done survives an audio ingest); only a transient Ingesting falls back to Idle. Code and comment now agree.

L1 — lossy encode→latent. The "Encode pattern → latent" button gained a hover note that the latents are clamped to 0..1, so re-decoding is approximate. Comment added at the clamp site.

L2 — gradient not batch-normalized. Documented on fit: the per-element loss gradient is summed (not averaged) over a batch's rows, so batch acts as a secondary LR knob (Adam absorbs most of it).

L4 — blocking pick_files(). Commented: it blocks only the editor thread, never audio; the decode/onset work is dispatched to the background thread.

L6 — initialize re-bumps generation. initialize now restores via swap_model and calls collect_garbage. We deliberately keep restoring on every initialize (preset/project reload depends on it), but the cost is now bounded by the graveyard and never touches the audio thread, which was the underlying concern.

L3 — long files collapse to one bar. Deferred by design. corpus_encode's bar_length = 1 is the upstream contract this port matches, and the audio path is already documented as approximate. Per-bar windowing is a feature, not a fix — tracked for later.

L5 — Train/Ingest enable race. Deferred. The window is sub-frame and the worst case is a redundant dispatch; not worth the extra atomic for now.

L7 — release.yml on two branches. Not a conflict: the release.yml on runtime-training and release-workflow are byte-identical, so git merges the additions cleanly. No action needed.

_{🤖 Generated with Claude Code}