cadd-distillation

A convolutional network distilled from CADD on the human genome (hg38).

What it predicts

CADD (Combined Annotation Dependent Depletion) scores how deleterious a variant is likely to be. This model is a fast convolutional network trained to reproduce CADD's score for every possible single-nucleotide substitution at once: at each position it emits one value per alternate allele, where larger values flag substitutions predicted to be more damaging.

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, one value per nucleotide in [A, C, G, T]. Each value is the model's CADD RawScore for substituting the reference base to that nucleotide, measured relative to the reference — so the reference allele reads about zero and is not itself a prediction. A larger value means a more deleterious substitution; applying a sigmoid recovers the 0-to-1 deleteriousness probability the teacher was trained on.

Training

Distilled from CADD with a per-position binary cross-entropy loss. The teacher's CADD RawScore for each candidate substitution is squashed through a sigmoid to form the target probability, and the student's output — shifted so the reference allele reads zero — is fit to it. 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/cadd-distillation"
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, 512): 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 (~1.06 GB).
• config.json: architecture hyperparameters, tokenization, and training metadata.

License

MIT