mrothroc/dna-enhancers-attention-mixlab

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 plain causal attention layers, modeldim 384, seqlen 512, nucleotide vocab 9 (A/C/G/T/N + BOS/EOS/PAD/MASK). A standard GPT-style stack, exported as a native Hugging Face GPT2LMHeadModel.
• —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