Attention DNA language model (trained in mixlab)
A small attention (GPT-style) DNA language model, pretrained on human genome in
mixlab on a single Apple-silicon Mac. It is a reusable backbone: load
it, then fine-tune a classifier on your own genomic task. On the human_enhancers_cohn benchmark, a
classifier fine-tuned on this backbone reaches 0.717 accuracy / 0.799 AUROC, above the benchmark's
baseline CNN (~0.69-0.70).
This is the attention baseline for a mixer comparison: its companion, mrothroc/dna-enhancers-mamba3-mixlab, is the same size and budget with the mixer swapped to canonical Mamba-3, and reaches 0.728. The point is how little changes between them: one block in a JSON config.
What it is
- 14.4M parameters: 8
plaincausal attention layers, model_dim 384, seq_len 512, nucleotide vocab 9 (A/C/G/T/N + BOS/EOS/PAD/MASK). A standard GPT-style stack, exported as a native Hugging FaceGPT2LMHeadModel. - Pretrained next-base on GRCh38 chromosomes 1-3 (a continuous nucleotide stream), 20,000 steps, with reverse-complement augmentation. Validation loss 1.08.
Load and generate
It loads with stock transformers as a GPT2LMHeadModel, no trust_remote_code. Token ids are BOS=1,
EOS=2, PAD=0, and A/C/G/T/N = 4/5/6/7/8.
import torch
from transformers import AutoModelForCausalLM
repo = "mrothroc/dna-enhancers-attention-mixlab"
model = AutoModelForCausalLM.from_pretrained(repo).eval() # GPT2LMHeadModel
ids2base = {4: "A", 5: "C", 6: "G", 7: "T", 8: "N"}
out = model.generate(torch.tensor([[1]]), do_sample=True, max_length=60, temperature=1.0, top_k=0)
print("".join(ids2base.get(t, "") for t in out[0].tolist() if t in ids2base)) # a DNA string
For classification, load with AutoModelForSequenceClassification and fine-tune. The full recipe is in the
cookbook below.
Reproduce / learn more
Full recipe (reproduce this baseline, then improve it by swapping the mixer to Mamba, all on a laptop): mixlab cookbook, genomics-mamba.
For the story behind it, read I Built an ML Architecture Lab in Go.
Citation
Please cite what this builds on:
- Genomic Benchmarks (the task): Gresova, K. et al. "Genomic benchmarks: a collection of datasets for genomic sequence classification." BMC Genomic Data 2023, 24, 25. doi:10.1186/s12863-023-01123-8
- GRCh38 human reference genome (Genome Reference Consortium), via the UCSC Genome Browser: Kent, W. J. et al. "The Human Genome Browser at UCSC." Genome Research 2002, 12 (6), 996-1006. doi:10.1101/gr.229102