SNPBERT

Self-supervised transformer that learns individual-level representations from unlabeled genotype data. Each person is encoded as an ordered sequence of rsID tokens — one per locus carrying at least one alternate allele — and the model is pretrained by masked language modeling (masked rsID prediction) on 485,923 UK Biobank participants across 13,615 SNPs. Built on the ModernBERT architecture; its alternating local/global attention captures both short-range linkage disequilibrium and genome-wide structure in a single forward pass.

Code / reproduction: https://github.com/BioLab-D/SNPBERT

Two backbones in this repo

This repo ships two pretrained backbones as subfolders:

Subfolder	Model	Genotype encoding	Load with	Masked-SNP acc.
dominant/	ModernBertForMaskedLM	presence/absence (≥1 alt allele)	AutoModel + subfolder="dominant"	0.785
dosage/	ModernBertWithDosage (base + dosage embedding)	het (1) vs. hom-alt (2)	custom class + subfolder="dosage"	0.787

The dosage backbone adds an additive per-token dosage embedding on top of the dominant architecture, so it distinguishes heterozygous from homozygous-alternate genotypes.

• Architecture: ModernBERT — 22 layers, hidden 768, 12 heads, intermediate 1152, global attention every 3rd layer, max position 9936, vocab 13,620.
• Precision: bfloat16.

Precision note: weights are bfloat16 (trained on GPU). On CPU, cast to float32 (.float()) — CPU LayerNorm does not support bfloat16. On GPU, use torch_dtype=torch.bfloat16.

Tokenizer conventions (shared by both)

The vocabulary reserves the first 5 ids for special tokens, then SNP column i maps to token id i + 5:

0=[UNK]  1=[SEP]  2=[PAD]  3=[CLS]  4=[MASK]   5..=SNP tokens (e.g. rs11260596_T)

Each SNP token is rsID_allele. An individual is a sequence [CLS] <SNP tokens for loci with ≥1 alt allele> [SEP].

---

Usage — dominant backbone

Standard ModernBertForMaskedLM; loads directly with AutoModel.

import torch
from transformers import AutoModel, AutoTokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained("ccjsj/SNPBERT", subfolder="dominant")
model = AutoModel.from_pretrained("ccjsj/SNPBERT", subfolder="dominant",
                                  attn_implementation="sdpa").to(device)
model = model.bfloat16() if device == "cuda" else model.float()
model.eval()

# One individual: [CLS] + SNP tokens (loci with >=1 alt allele) + [SEP]
snps = ["rs11260596_T", "rs6603782_T", "rs11523819_A"]
ids = [tok.cls_token_id] + tok.convert_tokens_to_ids(snps) + [tok.sep_token_id]
input_ids = torch.tensor([ids], device=device)
attention_mask = torch.ones_like(input_ids)

with torch.no_grad():
    out = model(input_ids=input_ids, attention_mask=attention_mask)
cls_embedding = out.last_hidden_state[:, 0]   # (1, 768)

These [CLS] embeddings feed the paper's downstream tasks (population-structure clustering, disease-risk AMIL, variant identification).

---

Usage — dosage backbone

⚠️ Not a plain AutoModel — custom class required

The dosage/ checkpoint is not a standard ModernBertForMaskedLM. Its flat state dict is:

base.*              # a ModernBertForMaskedLM  (138 tensors)
dosage_emb.weight   # nn.Embedding(3, 768)     (the dosage embedding)

config.json still declares architectures: ["ModernBertForMaskedLM"], so loading with AutoModel.from_pretrained(...) will silently drop dosage_emb.weight. You must rebuild the ModernBertWithDosage wrapper. Its forward adds the dosage embedding to the token embeddings before the encoder:

inputs_embeds = base.get_input_embeddings()(input_ids) + dosage_emb(dosage_ids)

dosage_ids is per-token in {0, 1, 2} — 0 at special positions, 1 (heterozygous) / 2 (homozygous-alternate) at SNP positions.

Option A — install the SNPBERT package

pip install "git+https://github.com/BioLab-D/SNPBERT.git"

import torch
from huggingface_hub import snapshot_download
from snpbert.model import load_dosage_model

root = snapshot_download("ccjsj/SNPBERT")     # downloads both subfolders
ckpt = f"{root}/dosage"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = load_dosage_model(ckpt, ckpt, num_dosage=3, attn_impl="sdpa").to(device)
model = model.bfloat16() if device == "cuda" else model.float()  # CPU needs fp32
model.eval()

Option B — self-contained (no repo install)

import torch, torch.nn as nn
from transformers import AutoConfig, AutoModelForMaskedLM, AutoTokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

class ModernBertWithDosage(nn.Module):
    def __init__(self, base, num_dosage=3):
        super().__init__()
        self.base = base
        self.dosage_emb = nn.Embedding(num_dosage, base.config.hidden_size)
    def forward(self, input_ids, dosage_ids=None, attention_mask=None, labels=None):
        emb = self.base.get_input_embeddings()(input_ids)
        if dosage_ids is not None:
            emb = emb + self.dosage_emb(dosage_ids)
        return self.base(inputs_embeds=emb, attention_mask=attention_mask, labels=labels)

repo = "ccjsj/SNPBERT"
cfg = AutoConfig.from_pretrained(repo, subfolder="dosage")
cfg._attn_implementation = "sdpa"
model = ModernBertWithDosage(AutoModelForMaskedLM.from_config(cfg), num_dosage=3)
sd = load_file(hf_hub_download(repo, "model.safetensors", subfolder="dosage"))
model.load_state_dict(sd, strict=True)        # keys: base.* + dosage_emb.weight
model.float().eval()

Forward with dosage

tok = AutoTokenizer.from_pretrained("ccjsj/SNPBERT", subfolder="dosage")
snps = ["rs11260596_T", "rs6603782_T", "rs11523819_A"]
ids  = [tok.cls_token_id] + tok.convert_tokens_to_ids(snps) + [tok.sep_token_id]
input_ids  = torch.tensor([ids])
attn       = torch.ones_like(input_ids)
# per-token dosage: 0 at [CLS]/[SEP]; 1=het, 2=hom-alt at SNP positions
dosage_ids = torch.tensor([[0, 1, 2, 1, 0]])

with torch.no_grad():
    out = model(input_ids=input_ids, dosage_ids=dosage_ids, attention_mask=attn)
logits = out.logits              # (1, seq_len, 13620)

---

Downstream use

The frozen backbone feeds three analyses in the paper: population-structure clustering (matching PCA), disease-risk modeling with a sparse-gated attention-MIL head, and attention-based variant prioritization. Pipelines and a runnable synthetic-data demo are in the GitHub repo.

License

MIT. Trained on UK Biobank data under approved access; use of derived representations must comply with UK Biobank terms.