🧬 NeuroMamba

A from-scratch selective state-space (Mamba/S6) protein language model for de novo generation of novel PDZ-domain sequences.

Author: Dr. Sanjay Anbu

Flagship target: the PDZ domain — the synaptic scaffold (PSD-95) that organizes the post-synaptic density, a natural building block for neural-interface tooling.

Architecture · Quickstart · Results · How it fits the de-novo-design brief

---

Why this exists

De novo protein design needs a generative engine — a model that invents new sequences rather than mutating old ones. Most protein generators are either Transformers (quadratic attention, heavy) or structure-only diffusion models. State-space models (SSMs) — the Mamba/S6 family — offer linear-time sequence modelling with a learned, content-selective memory, and they are conspicuously under-explored for proteins.

NeuroMamba is a from-scratch selective-SSM protein language model that learns the grammar of a protein family and autoregressively samples novel members of it. It is deliberately small (~0.7 M parameters) so the whole thing trains on a 6 GB laptop GPU in resumable batches — stop and resume any time — and every generation is measured for validity, novelty, and diversity, then scored by a built-in in-silico oracle. No CUDA-kernel dependency, no cloud, no external weights.

In one sentence: a laptop-scale, from-scratch Mamba language model that invents brand-new PDZ-domain proteins and grades them — the design half of a closed-loop discovery platform.

What makes it different from a generic protein generator

	Typical baseline	NeuroMamba
Sequence model	Transformer (O(L²) attention)	From-scratch selective SSM (Mamba/S6), work-efficient parallel scan
Inference	O(L²) / KV-cache grows with L	O(1) per token, constant-memory recurrent decode
Selectivity	Fixed mixing	Input-dependent Δ, B, C — content-based memory (the S6 idea)
Dependencies	mamba-ssm CUDA kernel	Pure PyTorch — trains on CPU or a 6 GB GPU
Training	One long run	Step-based resumable checkpointing ("train in batches")
Generation quality	Loss only	Validity · uniqueness · novelty vs. train · diversity
Scoring	External	Built-in proxy oracle (stability · solubility · PDZ-groove)
Baseline	None	Torch-free order-k Markov baseline to beat

Architecture

flowchart TD
    D([PDZ family sequences<br/>UniProt Pfam PF00595 · synthetic fallback]) --> T

    subgraph GEN [NeuroMamba — from-scratch selective SSM]
        direction TB
        T["Tokenize<br/>20 AA + BOS/EOS/PAD"]
        M["Stacked Mamba/S6 blocks<br/>causal conv → selective scan (Δ,B,C from input)<br/>→ gated output · RMSNorm · weight-tied head"]
        S["Autoregressive sampling<br/>temperature · top-k · nucleus"]
        T --> M --> S
    end

    S --> O["Proxy oracle<br/>stability · solubility · PDZ-binding groove"]
    S --> E["Generation metrics<br/>validity · novelty vs. train · diversity"]
    O --> R([Scored library of novel PDZ designs])
    E --> R

Loading

Each Mamba block runs a strictly-causal depthwise conv, then a selective state-space scan whose step-size Δ and input/output matrices B, C are functions of the current residue — so the model chooses, position by position, what to remember. The recurrence only looks backward, which is exactly what makes left-to-right generation and next-token training valid. Because that recurrence is associative, it is computed as a work-efficient parallel prefix scan (Hillis-Steele, ⌈log₂ L⌉ vectorised passes) instead of a per-timestep Python loop — the GPU stays busy without a fused CUDA kernel, and a reference sequential scan is asserted numerically identical in the tests. See src/neuromamba/mamba.py for the annotated implementation.

Quickstart

Setup — macOS / Linux:

# Python 3.10–3.12 recommended (PyTorch has no CUDA wheels for 3.13/3.14 yet)
python3 -m venv .venv
source .venv/bin/activate

pip install --upgrade pip
pip install -e ".[torch]"          # CPU/Mac; for a CUDA box use the cu121 wheel:
# pip install torch --index-url https://download.pytorch.org/whl/cu121

Setup — Windows (see RUN.md for the full verified walkthrough):

py -3.10 -m venv .venv
.venv\Scripts\activate.bat
python -m pip install --upgrade pip
python -m pip install torch --index-url https://download.pytorch.org/whl/cu121
python -m pip install -e ".[dev]"

Two Windows gotchas: use Python 3.10 (no CUDA torch wheels for 3.13/3.14), and call pip as python -m pip ... (Device Guard can block the pip.exe shim).

Run (paths use /; on Windows swap to \):

# 1. Get real PDZ-domain sequences (UniProt); falls back to synthetic offline
python scripts/download_pdz_family.py --out data/pdz.jsonl

# 2. Train — parallel scan + gradient checkpointing fit 6 GB; a held-out
#    validation split + early stopping stop at the generalization sweet spot
#    (so it won't memorize), and the best-by-val model is saved to model.pt.
python -m neuromamba.train --data data/pdz.jsonl --max-steps 4000 --batch-size 64 --device cuda
#    ...stop any time, then continue (resumes from last.pt):
python -m neuromamba.train --data data/pdz.jsonl --max-steps 8000 --device cuda --resume
#    (hit CUDA OOM? lower --batch-size to 32/16 · macOS Apple-Silicon: --device mps)

# 3. Generate + score novel sequences (writes designs.jsonl + metrics.json)
python scripts/generate.py --ckpt outputs/neuromamba/model.pt --n 64

# 4. No GPU / no torch? The whole pipeline still runs on the Markov baseline:
python scripts/generate.py --markov --data data/pdz.jsonl --n 64

Everything degrades gracefully: no network → synthetic data; no checkpoint or no PyTorch → torch-free Markov baseline; no GPU → CPU. The pipeline always runs.

Project layout

neuromamba/
├── src/neuromamba/
│   ├── mamba.py        # from-scratch selective-SSM (Mamba/S6) language model
│   ├── tokenizer.py    # 20-AA + special-token tokenizer
│   ├── data.py         # PDZ-family data (real UniProt + synthetic fallback)
│   ├── train.py        # resumable, 6 GB-tuned training loop
│   ├── generate.py     # autoregressive sampling (temperature/top-k/nucleus)
│   ├── metrics.py      # validity · uniqueness · novelty · diversity
│   ├── oracle.py       # self-contained in-silico proxy scorers
│   ├── markov.py       # torch-free order-k baseline generator
│   └── utils.py        # logging · seeding · device
├── scripts/            # download_pdz_family · train · generate · plot_training
├── tests/              # torch-free core + PyTorch model/resume tests
├── docs/REPORT.md      # technical report
└── .github/workflows/  # CI (ruff + pytest on 3.10 / 3.11)

Results

Trained on 646 real PDZ-domain sequences (UniProt Pfam PF00595) on a 6 GB laptop GPU (RTX 3000). A held-out validation split + early stopping selected the best model at step 300 (val perplexity 3.37) — before it memorised.

Generation — 64 sampled sequences, NeuroMamba vs. the torch-free order-3 Markov baseline, scored on the same proxy oracle (higher is better, except max-train-identity where lower = more novel):

Metric	Markov (order-3 baseline)	NeuroMamba
Validity (canonical, sensible length)	0.81	0.86
Uniqueness	1.00	1.00
Novel fraction (max id to train < 0.8)	1.00	0.85
Mean max-train-identity (lower = novel)	0.39	0.58
Diversity (1 − mean pairwise id)	0.77	0.77
Proxy — stability	0.864	0.900
Proxy — solubility	0.715	0.727
Proxy — PDZ-binding groove	0.777	0.816

Read: NeuroMamba beats the memoryless k-gram on validity and all three proxy objectives (stability, solubility, PDZ-groove) while keeping 85% of samples novel. The Markov baseline scores "100% novel" only because a 3-gram wanders off the family manifold — that extra "novelty" is noise, which is exactly why its proxy scores are lower. The SSM learns the family's grammar, so its novel samples stay protein-like: it sits in the sweet spot (novel and functional).

Early stopping matters — the ablation:

The curve above is the tell: train loss keeps falling while validation loss bottoms at step 300 and then rises — the model is starting to memorise. Early stopping saves the step-300 model (green star) instead of the overfit one.

NeuroMamba checkpoint	novel_fraction	id-to-train	validity	PDZ-groove
Overfit (step 4000, memorised)	0.05	0.95	1.00	0.92
Early-stopped (step 300, best val)	0.85	0.58	0.86	0.82

Trained to convergence the model memorises — highest proxy scores but only 5% novel (near-copies of the training set). Early-stopping on held-out loss trades a little validity/score for an 18× jump in novelty — the difference between retrieving training data and generating new proteins.

Reproduce:

python scripts/generate.py --ckpt outputs/neuromamba/model.pt --n 64   # NeuroMamba
python scripts/generate.py --markov --data data/pdz.jsonl --n 64        # baseline

Scientific honesty

• The oracle scorers are cheap in-silico proxies, not validated wet-lab assays — transparent heuristics for stability, solubility, and the PDZ peptide-binding groove. They exist to demonstrate a generate → score → select loop end-to-end; swap any proxy for a learned model or an experimental readout and the interface is unchanged.
• The synthetic data path is clearly labelled synthetic and used only when the UniProt fetch is unavailable, so training never silently trains on made-up data believing it is real.
• The contribution is the from-scratch selective-SSM generative model and the measured, reproducible pipeline around it — not a claim of validated designs.

How it fits the de-novo-design brief

Requirement	Where it lives
State-space models (SSMs)	From-scratch Mamba/S6 blocks — mamba.py
Generative / LLM modelling	Autoregressive protein language model + sampling
Transfer / domain priors	Learns a real protein family; PDZ-groove prior in the oracle
Sparse, high-cost data	Trains on a small real PDZ set; resumable, laptop-scale
Production-grade code	Typed, tested (CI on 3.10/3.11), ruff-clean, graceful fallbacks
Democratization	One-command generate-and-score; runs with zero heavy deps
Neuroscience (nice-to-have)	PDZ / PSD-95 synaptic-scaffold target

License & data notes

• Code: MIT (see LICENSE).
• Sequence data: fetched from UniProt (freely available); the synthetic fallback is generated locally and clearly labelled.