Mnemos-DNA

Pretrained weights for Mnemos-DNA: Gated Delta-Rule Linear Attention with k-mer Fast-Weight Memory for Edge-Scale DNA Modeling (Akshat Balyan, 2026).

~2.3M-parameter character-level DNA language models (vocabulary A/C/G/T/N + specials), pretrained on the human reference genome (GRCh38/hg38, HyenaDNA's held-out-chromosome split) for 4.5B tokens at a 131,072 bp causal context — trained entirely on one 8 GB consumer GPU. The k-mer fast-weight memory arms support exact chunked inference to 1 Mbp in 5.6 GB VRAM.

Code, evaluation harness, and paper: https://github.com/B4K2/Mnemos-DNA

Checkpoints

folder	arm	params	step	description
mnemos-dna/	D06	2.32M	34,177	KDA + per-8-mer fast-weight memory (main model)
mnemos-dna-c/	D07	2.29M	33,900	Mnemos-DNA + local-attention block (w=256) + length curriculum
pure-kda-d01/	D01	2.25M	34,408	pure KDA, fully matched control (no k-mer memory)
kda-local-d02/	D02	2.23M	19,073	KDA + windowed local attention, NTK-RoPE, 262K ctx, 5.0B tokens (pre-freeze protocol)

Each folder contains model.pt — {"model_state_dict", "config", "step"} — and the training config.yaml.

Usage

These are research checkpoints for the custom architecture in the GitHub repo (not transformers-compatible). From a clone of the repo:

import torch
from huggingface_hub import hf_hub_download
from training.config import load_config
from models import build_model  # see repo for the exact builder entry point

ck = torch.load(hf_hub_download("B4K2xx/Mnemos-DNA", "mnemos-dna/model.pt"),
                map_location="cpu", weights_only=False)
model = build_model(ck["config"])
model.load_state_dict(ck["model_state_dict"])

Key results (see paper for full tables)

• Pure delta-rule linear attention has zero exact recall at any distance on genomic background at this state size; the k-mer fast-weight memory restores 0.53–0.87 nats of repeat-copy advantage out to 64 kbp, robust to SNP-level substitutions.
• Held-out hg38 cross-entropy ≈1.07 nats/base, flat from 8 K to 1 M bp of evaluation context.
• Fixed ultra-long-context pretraining damages sub-kilobase grammar (random init beats these checkpoints on splice-site fine-tuning) — reproduced on HyenaDNA-1M as well; see the paper.