feat(alpha_train): --multi-resolution CLI flag

Replaces --seq-len. Default '1:10,30:10,100:12' = 32 positions
covering 1510 ticks of context, the Phase 1 fix for temporal
receptive field mismatch. Logs the active config at preload start
so Argo logs surface the per-run choice.
This commit is contained in:
jgrusewski
2026-05-22 20:42:40 +02:00
parent 8d72c14a8b
commit 7e438aeba0

View File

@@ -2,8 +2,9 @@
//!
//! Reads predecoded MBP-10 sidecars via `MultiHorizonLoader`, drives
//! `PerceptionTrainer` end-to-end. Each training step consumes ONE
//! seq_len-snapshot window and emits one BCE loss against the labels
//! at the window's last position (per-horizon).
//! multi-resolution window (total positions = sum of per-scale counts)
//! and emits one BCE loss against the labels at the window's last
//! position (per-horizon).
//!
//! Emits `alpha_train_summary.json` on exit for downstream gate
//! consumption.
@@ -13,7 +14,7 @@
//! --mbp10-data-dir <path> \
//! --predecoded-dir <path> \
//! --epochs 5 \
//! --seq-len 32 \
//! --multi-resolution 1:10,30:10,100:12 \
//! --mamba2-state-dim 16 \
//! --lr-cfc 3e-3 \
//! --lr-mamba2 1e-3 \
@@ -54,15 +55,14 @@ struct Cli {
#[arg(long, default_value_t = 5)]
epochs: usize,
/// Snapshots per training sequence (Mamba2 + CfC in-window context).
/// Default 64 — empirically the 3-fold ISV CV ran K=32 and observed
/// the longest-horizon AUC saturating around 0.70-0.74. K=64 doubles
/// in-window context relative to K=32 and gives Mamba2's SSM state
/// more material to build long-horizon predictions from. Kernel
/// maximum is 96 (MAMBA2_KERNEL_SEQ_MAX). Per-epoch wall scales ~K
/// (more K-loop launches in the captured graph).
#[arg(long, default_value_t = 64)]
seq_len: usize,
/// Multi-resolution input window: `<scale>:<count>[,<scale>:<count>]...`.
/// Default `1:10,30:10,100:12` gives 32 positions covering 1510 ticks
/// (10 raw + 10 aggregated@30 + 12 aggregated@100). Use `1:32` for the
/// degenerate single-scale parity-check config.
///
/// See `docs/superpowers/plans/2026-05-22-multi-resolution-input.md`.
#[arg(long, default_value = "1:10,30:10,100:12")]
multi_resolution: ml_alpha::data::aggregation::MultiResolutionConfig,
#[arg(long, default_value_t = 16)]
mamba2_state_dim: usize,
@@ -341,7 +341,7 @@ fn main() -> Result<()> {
w.copy_from_slice(&parts);
w
} else if cli.auto_horizon_weights {
auto_horizon_weights(cli.seq_len, &HORIZONS)
auto_horizon_weights(cli.multi_resolution.total_positions(), &HORIZONS)
} else {
[1.0; N_HORIZONS]
};
@@ -354,7 +354,7 @@ fn main() -> Result<()> {
anyhow::ensure!(cli.batch_size >= 1, "batch_size must be >= 1");
let trainer_cfg = PerceptionTrainerConfig {
seq_len: cli.seq_len,
seq_len: cli.multi_resolution.total_positions(),
mamba2_state_dim: cli.mamba2_state_dim,
lr_cfc: cli.lr_cfc,
lr_mamba2: cli.lr_mamba2,
@@ -487,12 +487,18 @@ fn main() -> Result<()> {
// PRELOAD: construct loaders ONCE (each preloads its files into RAM).
// Per-epoch we call reset(seed) to re-shuffle anchor sampling — no
// disk IO. Cuts wall time ~5× on a 6-epoch run, ~7× on 15 epochs.
tracing::info!(
multi_resolution = %cli.multi_resolution,
total_positions = cli.multi_resolution.total_positions(),
lookback_ticks = cli.multi_resolution.required_lookback_ticks(),
"multi-resolution config",
);
tracing::info!("preloading train + val MBP-10 files into RAM (slow startup, fast per-epoch)…");
let preload_start = std::time::Instant::now();
let mut train_loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig {
files: train_files,
predecoded_dir: cli.predecoded_dir.clone(),
seq_len: cli.seq_len,
multi_resolution: cli.multi_resolution.clone(),
horizons: HORIZONS,
n_max_sequences: cli.n_train_seqs,
seed: cli.seed,
@@ -504,7 +510,7 @@ fn main() -> Result<()> {
let mut val_loader = MultiHorizonLoader::new(&MultiHorizonLoaderConfig {
files: val_files,
predecoded_dir: cli.predecoded_dir.clone(),
seq_len: cli.seq_len,
multi_resolution: cli.multi_resolution.clone(),
horizons: HORIZONS,
n_max_sequences: cli.n_val_seqs,
seed: cli.seed.wrapping_add(0xC0FFEE),
@@ -707,7 +713,7 @@ fn main() -> Result<()> {
val_steps += 1;
// probs_all is [K, B, N_HORIZONS] row-major. Score AUC from
// the LAST-position predictions for each sample in the batch.
let last = cli.seq_len - 1;
let last = cli.multi_resolution.total_positions() - 1;
for (b_idx, last_lbl) in val_last_label_batch.iter().enumerate() {
let off = (last * cli.batch_size + b_idx) * N_HORIZONS;
let last_probs = &probs_all[off..off + N_HORIZONS];
@@ -852,7 +858,7 @@ fn main() -> Result<()> {
let summary = AlphaTrainSummary {
epochs: epochs_completed,
seq_len: cli.seq_len,
seq_len: cli.multi_resolution.total_positions(),
horizons: HORIZONS,
final_train_loss: if train_steps > 0 { train_loss_running / train_steps as f32 } else { 0.0 },
final_val_loss,