You can deploy this model and start earning money today!

msparsa/
splicemamba

Name: splicemamba
Availability: InStock

splice-site-prediction

mamba

dna

mit

Views

No views yet

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, 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}.

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.

msparsa/splicemamba