Update and Major fixes to the Model and Pathway Loss#5
Merged
BenjaminIsaac0111 merged 10 commits intomainfrom Mar 4, 2026
Merged
Update and Major fixes to the Model and Pathway Loss#5BenjaminIsaac0111 merged 10 commits intomainfrom
BenjaminIsaac0111 merged 10 commits intomainfrom
Conversation
…ored for better maintainbility.
…tion
- Refactor train.py by extracting logic into training submodules:
- arguments.py: CLI parameter definitions
- builder.py: Model and criterion setup
- checkpoint.py: Robust saving and loading logic
- Fix learning rate plateau by replacing disjoint schedulers with SequentialLR
to properly chain linear warmup and cosine decay phases.
- Simplify SpatialTranscriptFormer architecture:
- Remove redundant log_temperature parameter to reduce gradient variance.
- Implement L1-normalization for MSigDB pathway weight initialization to
prevent exponential prediction explosion at startup.
- Enhance load_checkpoint with robust error handling for EOFError (corrupted
files) and ValueError (architecture/optimizer mismatches) to ensure
graceful fallbacks.
- Fix test_checkpoint.py by adding required schedulers argument to save_checkpoint and load_checkpoint calls. - Fix test_checkpoints.py by updating pathway initialization assembly assertion to expect L1-normalized weights instead of raw binary values. - Fix test_spatial_interaction.py by removing obsolete test_temperature_scaling (log_temperature removed from model). - Fix test_pathways.py by providing mock args to _compute_pathway_truth to satisfy new visualization parameters. - Add test_pathway_stability.py to verify numerical stability and gradient flow of the final stabilized pathway-informed architecture.
Implemented spatial Z-score normalization and mean-aggregation for biological pathway ground-truth calculation. This ensures that every member gene in a pathway (even lowly-expressed transcription factors) contributes equally to the spatial activation signature, preventing high-count housekeeping genes from dominating the pathway patterns. Changes: - Updated AuxiliaryPathwayLoss to spatially standardize genes before projecting onto the pathway matrix. - Handled normalization across batch (patch-level) and spatial (whole-slide) dimensions with proper masking. - Switched from raw summation to mean-aggregation (averaging by pathway member counts). - Synchronized visualization.py ground-truth logic with the new objective. - Fixed mock tests in test_losses.py to match the normalized targets. Variance analysis on HEST data indicated raw gene variance ratios exceeding 300,000x, necessitating this standardization for biologically relevant pathway supervision.
…ation Deprecates the experimental data-driven pathway discovery in favor of strictly biologically-prior-driven interpretability. Updates the auxiliary pathway loss to use spatial Z-score normalization, ensuring lowly-expressed transcription factors contribute equally to the spatial objective. - Remove `--sparsity-lambda` and associated L1 regularization logic. - Implement spatial Z-score normalization in `AuxiliaryPathwayLoss`. - Synchronize visualization ground truth calculation with the new math. - Add `--plot-pathways-list` for dynamic user control over heatmaps. - Update plot labels to reflect Z-scored spatial patterns. - Cleanup: Delete LATENT_DISCOVERY.md and scrub legacy doc references. Ref: docs/PATHWAY_MAPPING.md, src/spatial_transcript_former/training/losses.py
…selection logic..
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.
Double-checking that the model and pathway loss are correctly ingesting and predicting the data correctly. I realised there were some things we did not calculate correctly, such as the auxiliary loss for the pathway.
I noted that every member gene in a pathway (even lowly expressed transcription factors) should contribute equally to the spatial activation signature, preventing high-count housekeeping genes from dominating pathway patterns.
I also removed the data-driven pathways learning module. I might get back to this at some point, as this might be interesting to implement as a feature, but for now, I am not 100% on this - at least without more contextual reading on how pathways might interact. Would be good to get an expert on this.
Updated the documentation to reflect these changes and made the preset model a little bit easier to run from CLI.