File size: 5,019 Bytes
71f6d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | ---
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.
|