PhyloGPN2

A reverse-complement-equivariant CNN that predicts per-nucleotide fitness, trained on a 447-mammal genome alignment.

What it predicts

For each genomic position the model outputs a fitness value for each of the four nucleotides: how tolerated that base is at the site under purifying selection. Fitness is identified only up to a per-position constant, so the absolute values are not meaningful on their own — only differences between alleles are. A variant's effect is the difference between its alternate and reference fitness, F(alt) - F(ref), where a more negative value means a more deleterious substitution.

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 symbol, mapped to indices 0 through 5. It returns a FloatTensor of shape (batch, length - 1548, 4): at each in-bounds position, a fitness value for each nucleotide in [A, C, G, T]. Because fitness is defined up to a per-position constant, score a variant as F(alt) - F(ref) (the constant cancels). encode returns the per-position 256-dimensional features.

Training

Trained by maximum likelihood on a 447-mammal genome alignment under a Halpern-Bruno mutation-selection model: substitutions follow a neutral rate matrix reweighted by the predicted fitness, and the loss is the negative phylogenetic (continuous-time Markov chain) log-likelihood of the alignment along the tree, with a fixed neutral (Zoonomia) rate matrix. CpG hypermutability is de-contaminated so it does not leak into the fitness readout. 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/PhyloGPN2"
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 predictions
    embeddings = model.encode(x)  # (1, length - 1548, 256): per-position features

Call forward (i.e. model(x)) for predictions and encode for embeddings: encode returns the final hidden representation, and forward is just a linear head on top of it. Both collapse the spatial dimension to the positions whose full context window is in bounds.

Files

• model.py: a self-contained CNN class with no dependencies beyond PyTorch.
• model.safetensors: the released weights (~0.26 GB).
• config.json: architecture hyperparameters, tokenization, and training metadata.

License

MIT