mlm-baseline-mm39
A masked nucleotide language model for the Mus musculus genome (mm39).
What it predicts
The model is a masked language model over DNA: it learns the distribution of each nucleotide given its surrounding sequence context. At every position it outputs a probability distribution over the four nucleotides (A, C, G, T) — what base it expects there from the context alone. It is trained on the Mus musculus genome, and is useful for scoring how unexpected a variant is (its likelihood under the model) and as a sequence feature extractor.
Inputs and outputs
The model takes a LongTensor of shape (batch, length) that encodes DNA over a six-token alphabet: A, C, G, T, N, and a padding/mask symbol, mapped to indices 0 through 5. It returns a FloatTensor of shape (batch, length - 1548, 4): at each in-bounds position, four logits over [A, C, G, T]. A softmax over them is the predicted nucleotide distribution. Because it is a masked model, to read the prediction for a position you first replace that position's input token with the mask token (index 5).
Status
This revision is a partial run: it was released at step 84,500, roughly half of the configured cosine cycle, and the learning rate had not annealed. It is published so downstream benchmarks can run; a completed checkpoint is expected to replace it in a later revision. Pin a revision if you need reproducibility.
Training
Trained by masked language modeling: about 15% of input tokens are replaced with the mask token (index 5), and the loss is the cross-entropy between the model's per-position distribution and the true nucleotide at the masked positions. The network is reverse-complement equivariant by construction, so a sequence and its reverse complement receive consistent predictions.
How to use
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import importlib.util, torch
repo_id = "songlab/mlm-baseline-mm39"
model_path = hf_hub_download(repo_id=repo_id, filename="model.py")
weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
spec = importlib.util.spec_from_file_location("released_model", model_path)
module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
model = module.CNN() # defaults match the released checkpoint
model.load_state_dict(load_file(weights_path))
model.eval()
# Input: a (batch, length) LongTensor over [A, C, G, T, N, pad], length >= 1549.
x = torch.randint(0, 4, (1, 2000))
with torch.no_grad():
output = model(x) # (1, length - 1548, 4): per-position logits
embeddings = model.encode(x) # (1, length - 1548, 512): per-position features
This is a masked language model: to read the prediction for a position, replace that position's input token with the mask token (index 5) before the forward pass, then take a softmax over the four output logits. The output window is centered, so input position p maps to output index p - 774:
masked = x.clone()
masked[0, 874] = 5 # mask one position
with torch.no_grad():
probs = model(masked).softmax(-1)
prediction = probs[0, 100] # distribution over [A, C, G, T] at the masked position
Outputs cover only positions whose full 1549 bp context is in bounds. encode returns the per-position features (use it for embeddings).
Files
model.py: a self-containedCNNclass with no dependencies beyond PyTorch.model.safetensors: the released weights (~1.06 GB).config.json: architecture hyperparameters, tokenization, and training metadata.
License
MIT