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