Restore sprint history on main#9
Merged
Merged
Conversation
Observation-level MAE tokenization produces up to 5376 tokens per sample (48 timesteps × 112 features), causing OOM on L4 (22GB) with batch_size=256 in fp32 due to quadratic attention memory. - batch_size: 256 → 64 (4× less peak memory) - precision: 32 → bf16-mixed (2× less memory, native L4 tensor core support) - accumulate_grad_batches: 1 → 4 (effective batch size stays 256)
…vel tokenization Unify all three SSL paradigms (MAE, JEPA, Contrastive) to use TransformerEncoder with obs_aware mode instead of ObservationTransformerEncoder. Each token now represents a full timestep (all D features) rather than individual observations, simplifying masking and decoding while preserving observation-aware input projection. Key changes: - Add obs_aware flag to TransformerEncoder with tokenize()/encode() API - Replace observation masking with timestep masking throughout - MAE decoder now predicts D features per timestep (not scalar per obs) - Extract build_sinusoidal_pe() to common utils - Update configs, finetune module, and all SSL tests
Split large files into smaller, single-responsibility modules: - dataset.py → tensor_preprocessing.py, tensor_cache.py - datamodule.py → splits.py - finetune_module.py → checkpoint_loading.py Also move EXPERIMENT_PLAN.md to docs/.
Revert pretraining precision to fp32 (bf16-mixed was a temporary OOM workaround, no longer needed with accumulate_grad_batches=4). Free timeseries_df (~3-5 GB) after tensor loading, use mmap for tensor cache deserialization, and delete temporary DataModule attributes after setup.
Replace instance-level (mean-pool) contrastive loss with temporal NT-Xent as the default mode. With two random masks at mask_ratio=0.5, ~25% of timesteps are visible in both views. These overlapping timesteps have two independently contextualized representations (different self-attention contexts), forming natural positive pairs. All other tokens across the batch are negatives. This preserves temporal granularity matching MAE/JEPA, ensuring fair paradigm comparison. The original instance mode is retained for ablation.
Python's `or` calls bool() on the first operand, which raises RuntimeError for multi-element tensors. The except clause silently swallowed this, making every cache lookup return None.
PyTorch 2.6 changed torch.load to use weights_only=True by default, which rejects OmegaConf objects stored in Lightning checkpoints.
Three targeted fixes for JEPA collapse observed in run ho5cwh3n: - Remove gradient accumulation (batch_size 64×4 → 256×1) to fix EMA updating 4× per optimizer step instead of 1× - Add target layer normalization in JEPA loss to prevent scale/shift collapse (standard in I-JEPA/BYOL) - Bottleneck predictor capacity (d_model 128→64, d_ff 512→256) to force information routing through the online encoder
The ModelCheckpoint filename templates used underscores (val_loss, val_auprc) but Lightning logs metrics with slashes (val/loss, val/auprc). Unmatched placeholders render as 0.0000 in filenames.
…onitoring - Increase EMA momentum_base from 0.996 to 0.999 to slow target encoder updates (~12%/epoch vs ~50%/epoch at 125 steps/epoch) - Add vectorized block masking (create_block_timestep_mask) that masks contiguous segments instead of random timesteps, forcing long-range prediction over neighbor interpolation - Add collapse monitoring metrics: jepa_target_repr_std (feature std across positions, collapse → 0) and jepa_target_repr_mean_cos (mean pairwise cosine similarity, collapse → 1.0)
Reduced max_epochs 500→100 and warmup 50→10 (same 10% ratio). Added min_epochs=10 so early stopping cannot trigger during LR warmup, which was killing JEPA runs before they reached peak learning rate.
- launch_sprint1.sh: runs all 4 paradigms (MAE, JEPA, Contrastive, Supervised) in parallel on MIMIC-IV with automatic pretrain→finetune chaining - auto_shutdown.sh: powers off GCP VM after 59min of no training activity (cron-based) - Update experiment plan with LR sensitivity ablation (Sprint 1b)
Disable early stopping for pretraining — JEPA's val loss rises as representations improve, making it unreliable as a stopping criterion. Follow I-JEPA/DINO/BYOL convention of fixed-schedule training. Save both last-epoch encoder (primary) and best-val-loss encoder (robustness check). Update experiment plan to reflect shared obs-aware timestep-level tokenization across all SSL paradigms.
All paradigms (MAE, JEPA, Contrastive, Supervised) now share the same encoder: d=64, L=2, H=4, obs_aware=True (~174K params). This eliminates model capacity as a confounding variable in the SSL comparison. - Enable obs_aware=True as default in base transformer config - Remove encoder size overrides from pretrain.yaml and supervised.yaml - Scale down SSL heads to match smaller encoder: MAE decoder: d=64, JEPA predictor: d=32, Contrastive proj: 256→64 - Update experiment plan to reflect shared architecture
Switch contrastive objective from temporal to instance-level (SimCLR-style) with temperature 0.07. Add alignment, uniformity, and effective rank metrics for collapse detection. Update experiment plan with dual evaluation protocols (linear probe + full finetune) and expanded Sprint 1 validation scope.
The obs_proj MLP was the only component in the encoder without dropout, causing severe overfitting when obs_aware=True (train auroc 0.993, val auroc 0.807). Added nn.Dropout to obs_proj for all paradigms. Also increased supervised baseline regularization: encoder dropout 0.2, label smoothing 0.1, weight decay 0.05. Added label_smoothing support to CrossEntropyLoss in FineTuneModule.
Dependency-aware scheduler that generates 915 runs across sprints 1-8, deduplicates shared pretrains, and executes jobs in parallel with crash recovery and atomic state persistence.
… weights Revert encoder.dropout (0.2→0.1), weight_decay (0.05→0.01), and label_smoothing (0.1→removed) overrides in supervised.yaml to maintain controlled variable parity with SSL paradigms per experiment plan. Enable class_weight: balanced in both supervised.yaml and finetune.yaml to match experiment plan specification of inverse frequency weighting for imbalanced classification tasks.
Detection window was filtering for hour >= obs_hours (48), but timeseries data only covers hours 0-47, producing zero positive labels. Changed to detect AKI after baseline period (hour >= baseline_hours, default 24). Also fixed absolute criterion to use all creatinine data with detection window filter, catching rises that span the baseline/detection boundary. This increased AKI positive rate from 6.4% to 12.1%. Added ValueError guard in supervised.py and finetune.py for zero-count classes when computing balanced weights.
- Head dropout 0.1 → 0.3 (all task configs) - Class weights: balanced → sqrt(balanced) to reduce gradient amplification - Downstream LR 1e-3 → 3e-4 (supervised config + PROTO_B) - Weight decay 0.01 → 0.05 (supervised + finetune configs) - Early stopping patience 20 → 10 (supervised config + PROTO_B) - Label smoothing 0.0 → 0.1 (supervised + finetune configs)
…l runs Adds a `warmup` command that sequentially builds tensor caches for all unique (dataset, task, seed, label_fraction) combinations needed by the requested sprints. This prevents OOM kills when running parallel experiments on the first run, where each process would otherwise independently load and preprocess raw parquet files.
Switch default contrastive masking from two independent random masks (~25% overlap) to complementary masks (view 2 = ~view 1, zero overlap). This forces the encoder to learn abstract temporal semantics from non-overlapping windows instead of exploiting shared timesteps. - Add complementary_masks config flag (default True) - Validate incompatibility with temporal mode - Update docstrings and YAML config - Add tests for zero overlap, full coverage, and validation
- Add Sprint 1c: mask ratio sensitivity (0.3, 0.5, 0.75) on MIMIC-IV, mortality_24h, seed=42 — validates shared mask_ratio=0.5 before committing to full experiment matrix - Add build_sprint1c() to runner script - Update Section 5.3 LR ablation to include 2e-4 (matching runner) - Fix Sprint 1b run count: 18 runs (9 pretrain + 9 finetune), not 12 - Update Sprint 8 to 60 runs (extra seeds for both LR and mask ratio) - Add sprint execution summary table to Section 6 - Update risk mitigation to reference both sensitivity studies
Refactor AKI label builder to use a proper forward-looking detection window (hours 48-72) instead of evaluating within the observation period. This prevents data leakage by ensuring the model never sees creatinine values used for label construction. Key changes: - AKI detection now uses [obs_window, obs_window + prediction_window) - Baseline window extended to full 48h observation period - RICU extraction horizon increased to 72h (LABEL_HORIZON_HOURS) - Stays with no creatinine in prediction window get null labels - prediction_window_hours is now required for AKI task - Tests updated to reflect forward-looking label semantics
Allow rerunning sprint experiments with changed settings via a post-processing transform that rewrites run IDs, output dirs, and injects W&B tags. Includes validation guards and revision-aware status display.
Add query_runs.py to wandb-analysis skill for filtering runs by tags, config, state, and group with aggregation support. Add sprint prefix to run names, protocol/mask_ratio/label_fraction/lr tags, and dynamic probe vs finetune run name distinction.
Skip unused tables (emar_detail, pharmacy, poe, prescriptions, etc.) that take 20+ min to import but aren't needed for concept extraction. Also remove unnecessary intermediate gc() calls in batch merging.
Fix generated downstream protocol overrides so SSL finetune and probe jobs compose while baseline configs keep protocol metadata without loading protocol groups. Add publication export safeguards for expected matrix coverage, duplicate fingerprints, W&B tag filters, and validation-before-write behavior. Adjust contrastive mask sweeps and AKI missingness documentation to preserve benchmark interpretation. Implications: generated finetune/probe commands now launch with protocol=A/B, XGBoost exports report xgboost.learning_rate, and publication exports fail closed on absent or duplicate matrix cells unless explicitly allowed.
Quarantine stale run outputs when launch commits change, require complete fairness reports before writing W&B metrics, and skip undefined single-class bootstrap samples for AUROC/AUPRC metrics. This prevents cross-commit checkpoint reuse and makes fairness/statistical analysis fail closed; no migration is required.
Bind runner state and checkpoint resume decisions to revision, W&B target, and launch commit so stale state cannot silently skip or resume a different corpus. Reject ambiguous finetune checkpoint inputs and emit per-task statistical comparisons alongside the omnibus tests. Behavioral implications: scoped reruns quarantine mismatched output directories and only resume checkpoints with matching launch markers; exported statistical tests now include comparison_scope rows for omnibus and per-task claims.
Fail closed when final experiment runs or retries omit revision and W&B project metadata, and make status reporting revision-aware for tmux monitoring. Tighten data and analysis safeguards by validating benchmark metadata, disambiguating W&B run names, logging label-support summaries, filtering fairness/export tasks to the fixed thesis set, and treating written undefined fairness metrics as complete. Behavioral implication: final launches must pass --revision and --project or set WANDB_PROJECT; exports and fairness now default to slices-thesis and exclude non-thesis task rows unless explicitly requested.
Update README, experiment plan, and SSL config comments to distinguish the shared mask budget and interface from objective-specific mask geometry. Document JEPA block masking and contrastive two-view masking so the benchmark contract matches the current experiment design. No behavioral, benchmark, or migration impact.
Fix JEPA block masking to honor the requested valid-timestep budget, stabilize SMART external-reference model size identity, and fail closed on missing fairness summaries or accidental multi-revision exports. Add subgroup comparability accounting for one-class fairness groups, document the export semantics, and add a root experiment-plan pointer. Behavioral implications: JEPA high-mask robustness now matches configured masking pressure, SMART appendix rows export as model_size=default, publication export now requires fairness metrics, and multi-revision exports require an explicit exploratory flag.
Relabel the contrastive robustness slice, expand default binary metrics, add aggregate confidence intervals, and fix LOS label and FLOPs reporting helpers. Benchmark implications: final exports now include threshold metrics and seed-level CIs, contrastive HP rows are no longer described as pure mask-ratio sweeps, and compute/label helper outputs are aligned with the 84-feature thesis corpus.
Preserve nullable hospital mortality outcomes, treat death-coded discharge locations as death evidence, and use discharge timing when it can place deaths inside mortality windows. Bumps the mortality label semantic version to invalidate stale processed labels; regenerated worktree artifacts update eICU, miiv, and combined manifests to 2.2.0.
Require launch commits and clean tracked worktrees for direct final run and retry commands while keeping dry-run audits convenient. Validate exported revision launch-commit homogeneity, surface skipped retry dependents, use locked tmux dependency sync, keep fairness/export jobs active for shutdown, and initialize XGBoost W&B runs before expensive setup. Final launch behavior now fails closed for missing or ambiguous provenance outside dry runs; retry --failed reports skipped downstream work and final exports warn on missing or mixed launch commits.
Log the full binary publication metric set for XGBoost baselines and add tests that keep the documented metric surface in sync. Stamp post-run fairness summaries with artifact, checkpoint, script, and subgroup metadata, and make export validation fail closed when summaries are missing or stale. Update fairness documentation with the class/revision-based corpus and final publication command sequence.
Update the thesis run matrix to remove unreliable 1% mortality_24h cells while adding the full hospital-mortality label-efficiency curve. Correct hospital mortality labeling, make result and fairness filtering robust to incomplete W&B tags, force fairness refreshes, and cover SMART imputation diagnostics and unknown gender handling with tests. Behavioral implications: the final matrix is now 2590 launched runs including SMART and 2455 excluding SMART. Processed data was regenerated locally after the mortality label semantic change, but ignored data artifacts are not part of this commit.
Emit lowercase Hydra protocol selectors for generated finetune runs, revalidate completed pretrain encoder artifacts, and require explicit W&B entity ownership for direct final launches. Add thesis task defaults and fail-closed combined rebuild checks, checkpoint/model SHA256 provenance for fairness/export linkage, and durable imputation result output. Behavioral implications: final reruns now fail closed on missing pretrain encoders, stale fairness artifacts, missing artifact digests, missing direct-launch W&B entity, and source datasets missing required thesis tasks.
Add fairness matrix coverage and artifact digest validation before post-hoc fairness inference. Enforce online W&B mode for final thesis launches, align timestep masking budgets across SSL objectives, and harden pooling/checkpoint edge cases. Behavioral implications: final launches now reject offline W&B mode, fairness evaluation fails on incomplete scoped matrices unless explicitly allowed, random timestep SSL masks use an exact per-sample budget, and CLS pooling checkpoint overrides are rejected instead of silently rebuilding incompatible modules.
Replace opaque A/B downstream protocol selectors with linear_probe and full_finetune across configs, runner commands, logging, export, fairness coverage, and regression tests. Legacy A/B protocol values are canonicalized during export and fairness matrix reconstruction so older W&B runs remain comparable.
Make W&B entity setup failures exit nonzero so the experiment scheduler cannot mark failed launches as completed. Propagate the explicit launch git-check skip into generated tmux runners and teach the Python runner to honor it for unchecked provenance runs. Behavioral implication: final thesis launches now fail closed on missing W&B entities, while SKIP_LAUNCH_GIT_CHECK=1 remains an explicit audited escape hatch.
Apply revision metadata to the full generated retry graph before class filtering, then include skipped downstream dependents caused by selected retry roots. This prevents class-scoped recovery from leaving label-efficiency or transfer runs skipped after an upstream SSL pretrain failure. Adds regression coverage for cross-class skipped dependents in revision-scoped retry. No experiment matrix definitions or benchmark semantics change.
Resolve symbolic launch commits to immutable hashes before recording run metadata, fail closed when the tmux thesis launcher selects a partial final corpus, and stop exposing los_days through model-facing static batch payloads while preserving static_df for analysis. This changes launch behavior for final runs: partial corpus tmux launches now require VALIDATE_FINAL_CORPUS=0, and direct launches record the resolved commit hash. Verified with uv run pytest -q and a 2590-run dry run.
Centralize the controlled SSL tokenization contract, move intrinsic missingness handling behind an encoder capability, and let the experiment runner rely on protocol Hydra configs instead of duplicating overrides. Also tighten data and pretraining config schemas, load FLOPs defaults from YAML, validate imputation data prerequisites, align the ruff pre-commit hook, sort imports, and make setup dependency sync lockfile-frozen. No script names or CLI arguments changed; setup now fails if the lockfile is stale. Verification: uv run ruff check src scripts tests; uv run mypy src; uv run pytest -q.
Replace existing W&B fairness summary keys during recomputation so stale subgroup metrics cannot leak into publication exports. Validate provided launch commits during dry runs, recover stale running entries from status, and update the experiment plan protocol metadata. Behavioral implications: fairness recomputation now clears old fairness summary keys by default, and status may update runner state when it detects dead or mismatched running processes.
Update the pretraining README to distinguish the core controlled SSL objectives from TS2Vec as a temporal-contrastive extension. This keeps the documentation aligned with the thesis benchmark contract without changing runtime behavior.
Keep VM scheduler slot costs while ensuring ready pretraining jobs get first claim on a full slot budget. This preserves existing launch guardrails and keeps tests aligned with the final runner policy.
hannesill
added a commit
that referenced
this pull request
Jun 7, 2026
Restore sprint history on main
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Verification
uv run pytest tests/test_fixes.pyNotes