Learnable mask token + Kronecker marker covariance GP - #22
Open
mzmyslowpro wants to merge 47 commits into
Open
Conversation
- Add train.sh SLURM batch script for standard and GP training - Add setup_venv.sh for uv-based venv setup on remote server - Fix Comet.ml empty tags error (skip add_tags when tags=[]) - Add val_standard_nll and val_gp_nll to log_validation_metrics - Fix train_masked_config.yaml and train_masked_gp_config.yaml Comet placeholders - Set gp_lengthscale 0.1→5.0, lambda_gp 0.0→0.1, batch_size 1→8, use_kronecker_gp=true in GP config Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…aining Save total_steps in checkpoint so scheduler can be reconstructed with identical warmup/annealing boundaries when resuming. Previously resuming with different epochs config would miscalculate LR schedule. Also bump epochs 10→100 in GP config. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
When training for another N epochs from a fully-converged checkpoint, the saved scheduler state has LR near zero. reset_lr_schedule=true ignores checkpoint optimizer/scheduler state and starts a fresh cosine cycle from the trained weights. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Two bugs when using reset_lr_schedule=True: 1. epochs:100 + start_epoch=100 → range(100,100) empty, no training Fix: epochs:200 so range(100,200) = 100 new epochs 2. total_steps calculated from config.epochs (200) but only 100 epochs will run → cosine schedule only half-completed at end of run Fix: use remaining_epochs = config.epochs - start_epoch for total_steps Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Triple Kronecker (K_x ⊗ K_y) ⊗ K_C + Woodbury approach for modeling inter-marker uncertainty from Hyperkernel embeddings. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
9 tasks covering: config, triple Kronecker solver, log_prob_joint, loss classes, training/validation integration, exports, and e2e test. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…arkerCovariance Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…n training script Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…g loop Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…s logging Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
… config Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
… marker cov diagnostics
Add check_untyped_defs = true so mypy inspects bodies of unannotated functions (train_masked_gp, test_masked_gp) and catches call-arg errors like wrong keyword names — the class of errors that caused the cluster crashes. Also enforce disallow_untyped_defs on multiplex_model.utils.* since those modules are already fully annotated. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…raining scripts - train_logging: fix NPE in get_next_version_number (re.match returns None when no experiment matches pattern); add int() cast for dict[int,str] lookups; suppress plt.cm.CMRmap attr-defined false positive; add return type annotations to plot_reconstructs_with_uncertainty and plot_reconstructs_with_masks - masking: rename loop accumulators to *_list so torch.cat reassignment doesn't confuse mypy about return tuple element types - optim: add type annotations to ClampWithGrad forward/backward; suppress FunctionCtx dynamic attribute errors with targeted type: ignore comments - configuration: switch to ValidationInfo (correct pydantic v2 type) for field validator info parameters - train_masked_model_gp / train_masked_model: rename all_latents/all_channel_* post-cat variables to avoid list[Tensor] -> Tensor reassignment confusion Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
K_C = E @ E.T + jitter * I was ill-conditioned (condition number 10^5-10^6) because embedding_projection output scale was O(1/sqrt(hyperkernel_model_dim)) ≈ 0.02, making E @ E.T entries much smaller than typical values. With many channels C, the max eigenvalue grew as O(C * D * scale^2) while the jitter stayed at 0.01, causing condition numbers far beyond numerical stability. Normalizing E rows to unit norm before K_C makes the matrix a proper cosine- similarity correlation matrix (diagonal = 1 + jitter, off-diagonals = cosine sim bounded by [-1, 1]). Condition number is now bounded by C, typically 10-100. This eliminates the nan GP NLL that poisoned training from epoch 0. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…zed E Mirror the module's row-normalization (nn.functional.normalize) when building the reference K_C for the dense log-prob comparison test. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…alues When C > marker_embed_dim (32), E_norm @ E_norm.T is rank-deficient with C-D eigenvalues exactly at marker_jitter. LAPACK's divide-and-conquer eigh fails with 'too many repeated eigenvalues' in float32. Casting K_C to float64 before eigh and back to float32 after resolves convergence reliably with negligible overhead (C is at most ~40, a tiny matrix). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Add use_mask_token and mask_token_init parameters to MultiplexImageEncoder to enable replacing spatially masked pixels with a learnable scalar token during encoding. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…hitecture config utilities - Add spatial_mask parameter to encode() and forward() methods to support learnable mask tokens - Add get_architecture_config() method to retrieve model configuration for checkpointing - Add load_from_checkpoint() classmethod to reconstruct model from checkpoint dict with config - Store architecture config in _architecture_config on init for later retrieval Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Adds optional mask_token parameter to log_training_metrics function to support logging learnable mask token values during training. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…mask utility Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Combines learnable spatial mask token (from learnmask script) with additive K_C marker covariance GP loss (Woodbury solver). Loss path runs outside autocast to avoid bfloat16/float32 mismatch in linalg.solve inside the Kronecker Woodbury step.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- losses.py: assert→ValueError for grid size check, trailing newline - gp_covariance.py: guard non-positive triple_eigs before log(), return float32-cast K_C from _compute_marker_eigen - configuration.py: marker_jitter gt=0, model_validator for use_marker_covariance→use_kronecker_gp - train scripts: NaN loss guard before backward, warning on missing gp_covariance_state_dict, math.isfinite guard on corrcoef logging, mask_token None guard - tests: C>marker_embed_dim float64 eigh path, learnmask+GP validation loop smoke test, mask token gradient flow test Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…cs/superpowers) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds learnable spatial mask-token training and Kronecker marker-covariance GP loss support, with new training entry points, configs, logging, and tests.
Changes:
- Adds mask-token propagation through
MultiplexAutoencoderand new learnmask/learnmask+GP training scripts. - Adds
KroneckerMarkerCovarianceand hybrid marker GP NLL loss, plus config fields and tests. - Updates training configs, logging metrics, scheduler/checkpoint handling, and utility scripts.
Reviewed changes
Copilot reviewed 27 out of 29 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
train.sh |
Adds dispatch for learnmask and learnmask+GP training modes. |
train_masked_model.py |
Adds variance/MSE correlation validation logging. |
train_masked_model_learnmask.py |
New learnable mask-token training script. |
train_masked_model_learnmask_gp.py |
New combined learnmask + marker GP training script. |
train_masked_model_gp.py |
Extends GP training for marker covariance and new diagnostics. |
train_masked_learnmask_gp_config.yaml |
New combined learnmask + marker GP config. |
train_masked_learnmask_config.yaml |
New learnmask-only config. |
train_masked_gp_marker_config.yaml |
New marker GP config. |
train_masked_gp_config.yaml |
Updates default GP training settings. |
train_masked_config.yaml |
Updates Comet logging defaults. |
tests/test_training_integration.py |
Adds integration tests for training, logging, and mask token flow. |
tests/test_kronecker_marker.py |
Adds numerical tests for marker Kronecker GP covariance. |
summary_pl.py |
Adds Polish PDF summary generator for the method. |
setup_venv.sh |
Adds environment setup helper. |
run_validation_leave_one_out.py |
Adds leave-one-out validation utility. |
run_embed.py |
Adds embedding extraction utility. |
pyproject.toml |
Tightens mypy checks for utils. |
multiplex_model/utils/train_logging.py |
Adds logging support for mask token, GP components, batch metrics, and image suffixes. |
multiplex_model/utils/optim.py |
Adds typing to optimizer utilities. |
multiplex_model/utils/masking.py |
Adds get_pixel_mask helper. |
multiplex_model/utils/configuration.py |
Adds mask-token, marker GP, and reset scheduler config fields. |
multiplex_model/utils/__init__.py |
Exports new masking/logging helpers. |
multiplex_model/modules/immuvis.py |
Adds mask-token support and architecture checkpoint helpers. |
multiplex_model/modules/gp_covariance.py |
Adds marker covariance Kronecker GP module. |
multiplex_model/modules/__init__.py |
Exports marker covariance module. |
multiplex_model/losses.py |
Adds marker Kronecker GP and hybrid loss. |
immuvis.py |
Adds a duplicate top-level autoencoder module. |
.gitignore |
Ignores additional local files/directories. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| fully_masked_channels_max_frac=0.25, | ||
| mask_patch_size=2, | ||
| use_gp_loss=True, | ||
| use_marker_covariance=True, |
|
|
||
| final_model_path = f"{checkpoints_path}/final_model-{run_name}.pth" | ||
| print(f"Training completed. Saving final model at {final_model_path}...") | ||
| final_checkpoint: dict[str, Any] = {"model_state_dict": model.state_dict()} |
Comment on lines
+8
to
+9
| from .base_modules import Block, Encoder, Identity, LayerNorm | ||
| from .registry import resolve_block_class, resolve_encoder_class |
| fully_masked_channels_max_frac: 0.5 | ||
| mask_patch_size: 8 | ||
| from_checkpoint: null | ||
| from_checkpoint: checkpoints/final_model-ImVs-12.pth |
- remove invalid use_marker_covariance param from test_masked_learnmask_gp call - add gp_covariance_state_dict to final checkpoint in learnmask_gp script - untrack top-level immuvis.py (broken relative imports, shadow of package module) - fix train_masked_gp_config.yaml: final_model → last_checkpoint for from_checkpoint Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Parallel jobs race on the Comet version query and get the same ImVs-N number, overwriting each other's checkpoints. Append SLURM_JOB_ID to the run name so concurrent runs get distinct checkpoint files. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds dump_marker_covariance.py to reconstruct and inspect the learned marker covariance K_C from an ImVs checkpoint (per-panel, shared vs residual structure), plus the ImVs-34 figures and a summary note. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KZDVdCQTc9ULyv73Qvmivo
…ests) Adds three LOO-based probes of the learned marker covariance K_C and records the findings in notes: (1) K_C does not predict correlated LOO errors, (2) the apparent redundancy->lower-MSE effect is a scale confound, gone under NMSE/R2. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KZDVdCQTc9ULyv73Qvmivo
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
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.
Co zostało zmienione?
Scalenie dwóch funkcjonalności:
use_mask_token: truei zastępuje zamaskowane patche uczonym skalarnym tokenem zamiast zerowania; propagacjaspatial_maskprzez cały pipelineMultiplexAutoencoder.KroneckerMarkerCovariancez(K_x ⊗ K_y) ⊗ K_C, gdzie K_C pochodzi z embeddingów hyperkernel (przeznn.Linear+ L2-normalize, float64 dla eigh).HybridKroneckerMarkerGPNLLLossz λ=0.1.train_masked_model_learnmask_gp.py+ configtrain_masked_learnmask_gp_config.yaml.train.shdla trybulearnmask_gp.Szczegóły numeryczne: row-normalizacja embeddingów K_C zapobiega warunkowi K_C → NaN; float64 dla
linalg.eighprzy C > marker_embed_dim (powtórzone wartości własne).Dlaczego?
Sama mask token trenuje rekonstrukcję bez informacji o maskowaniu (lepsza generalizacja). GP z Kronecker marker covariance modeluje zależności przestrzenne i między-markerowe w niepewności — dekoruje σ od μ (korelacja σ↔output ~93% zamiast ~100%), co daje lepiej skalibrowane predykcje.
Jak zweryfikować?
python -m mypy multiplex_model/ train_masked_model_learnmask_gp.pypytest tests/ -v— testy wtest_kronecker_marker.pyitest_training_integration.pysbatch train.sh train_masked_learnmask_gp_config.yaml learnmask_gp— po 1. epoce log powinien zawieraćmask_token,gp_nll,standard_nlli diagnostykę K_C bez NaN🤖 Generated with Claude Code