| --- |
| license: mit |
| tags: |
| - biology |
| - genomics |
| - splicing |
| - splice-site-prediction |
| - mamba |
| - dna |
| library_name: pytorch |
| --- |
| |
| # SpliceMamba — checkpoints |
|
|
| Splice site prediction from pre-mRNA sequence: every position in a 15 kb window is |
| classified as **donor**, **acceptor**, or **neither**. Same task and same |
| train/test split as [SpliceAI](https://github.com/Illumina/SpliceAI), with a |
| Mamba-based architecture instead of dilated convolutions. |
|
|
| This repository holds the **full training checkpoints** for every model version and |
| ablation. Each `.pt` retains `optimizer` and `scheduler` state, so any of them can be |
| resumed, not just used for inference. |
|
|
| ## Architecture (v3 onward) |
|
|
| ``` |
| Input (B, 4, 15000) one-hot DNA |
| -> DeepConvStem: 4 dilated residual blocks (dilations 1/4/10/25, ~641 bp receptive field) |
| -> Sinusoidal positional encoding |
| -> BiMamba2 encoder: 8 layers/direction, d_model=256, d_state=64, expand=2 |
| -> Coarse head (auxiliary loss, lambda=0.1) |
| -> Sliding-window attention: 4 layers + FFN, window_radius=400, 8 heads |
| -> Refined head (primary loss, lambda=1.0) |
| ``` |
|
|
| Dual-head design: the coarse head supplies auxiliary gradient to the encoder, the |
| refined head produces final predictions. Both emit `(B, L, 3)` logits. |
|
|
| Only positions `[5000:10000]` of each window carry labels; the outer 5 kb on each |
| side is context-only flanking. Classes are `{0: neither, 1: acceptor, 2: donor}`. |
|
|
| ## Contents |
|
|
| | Path | What it is | |
| |---|---| |
| | `v2-ce/`, `focal-loss/` | Early loss-function comparison (weighted CE vs focal) | |
| | `v3/` | First dual-head model | |
| | `v5.0/` | Main classification model | |
| | `v6/`, `v6.1/` | v5 + adversarial codon-phase removal (DANN guesser head) | |
| | `v7/`, `v7.1/` | Stabilised adversarial phase removal | |
| | `v8/` | Quantitative / PSI-style objective — different task, see below | |
| | `ablation/` | 12-config sweep over `d_model`, `n_mamba_layers`, `n_attn_layers`, `window_radius` | |
| | `ensemble/model_1..5/` | 5-model ensemble, different seeds and validation splits | |
| | `aug-*/` | Trained on CDS-shuffled / codon-masked augmentations | |
|
|
| ## Recorded validation AUPRC |
|
|
| Values below are the `best_auprc` field stored inside each `best.pt` at checkpoint |
| time (validation set, splice classes). They are **not** chr1 test-set numbers. |
|
|
| | Model | Val AUPRC | Epoch | |
| |---|---|---| |
| | `aug-codon_mask1` | 0.9848 | 5 | |
| | `ablation/d_model-512_n_attn-8_w200` | 0.9716 | 9 | |
| | `ablation/d_model-512` | 0.9707 | 9 | |
| | `v5.0` | 0.9705 | 23 | |
| | `v6.1` | 0.9679 | 5 | |
| | `v7` | 0.9651 | 2 | |
| | `ablation/n_attn-8` | 0.9647 | 9 | |
| | `v6` | 0.9636 | 10 | |
| | `v7.1` | 0.9629 | 1 | |
| | `ensemble/model_2` | 0.9584 | 14 | |
| | `ablation/window_radius-1600` | 0.9587 | 3 | |
| | `ablation/window_radius-200` | 0.9584 | 4 | |
| | `ablation/n_mamba-16` | 0.9576 | 3 | |
| | `ablation/baseline-10ep` | 0.9573 | 3 | |
| | `ensemble/model_3` | 0.9556 | 5 | |
| | `ensemble/model_4` | 0.9554 | 5 | |
| | `ensemble/model_1` | 0.9533 | 14 | |
| | `v3` | 0.9531 | 6 | |
| | `ensemble/model_5` | 0.9531 | 6 | |
| | `ablation/n_mamba-4` | 0.9513 | 4 | |
| | `ablation/d_model-128` | 0.9477 | 4 | |
| | `aug-codon_shuffle` | 0.9399 | 9 | |
| | `v2-ce` | 0.9287 | 19 | |
| | `aug-nt_shuffle` | 0.9121 | 5 | |
| | `focal-loss` | 0.9133 | 10 | |
| | `ablation/n_mamba-12` | 0.2647 | 1 | |
|
|
| `ablation/n_mamba-12` diverged and is kept only for completeness. |
|
|
| The `aug-*` models are trained *and* validated on shuffled-CDS data, so their AUPRC |
| is not comparable to the others — the shuffling changes the label distribution. |
|
|
| ## v8 is a different objective |
|
|
| `v8/` checkpoints optimise a quantitative (PSI-style) target rather than 3-class |
| site classification, so they store `step` and a `val` dict of correlation metrics |
| instead of `best_auprc`. Best recorded validation Spearman: |
|
|
| | Checkpoint | Spearman (mean) | Pearson (global) | Match rate | |
| |---|---|---|---| |
| | `protected_best.pt` | 0.347 | 0.670 | 0.657 | |
| | `complete_best.pt` | 0.343 | 0.649 | 0.568 | |
| | `phase2careful_best.pt` | 0.336 | 0.671 | 0.674 | |
| | `frozenlocal_best.pt` | 0.328 | 0.629 | 0.843 | |
| | `strengthen_best.pt` | 0.315 | 0.649 | 0.796 | |
|
|
| These also carry a different `model` state_dict shape and will not load into the |
| classification model class. |
| |
| ## Loading |
| |
| ```python |
| import torch |
| from huggingface_hub import hf_hub_download |
|
|
| path = hf_hub_download("msparsa/splicemamba", "v5.0/best.pt") |
| ck = torch.load(path, map_location="cpu", weights_only=False) |
|
|
| model.load_state_dict(ck["model"]) # inference |
| # ck["ema"], ck["optimizer"], ck["scheduler"], ck["config"] also present |
| ``` |
| |
| To resume training, pass the file to `train.py --resume`. Model code (`model.py`, |
| `train.py`, `dataset.py`) is not included here — see the project repository. |
| |
| ## Notes |
| |
| - Checkpoints from v5.0 onward include an `ema` state dict alongside `model`. |
| - `ensemble/model_1/best_copy.pt` and `last_copy.pt` are distinct checkpoints, not |
| duplicates of the adjacent files. |
| - Class imbalance is roughly 6200:1 neither-to-splice, handled with per-class alpha |
| weights `[0.1, 1.0, 1.0]` and weighted random sampling. |
| |
