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

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

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.