Genomic LLM Embeddings

genomic-llm-embeddings

When to Use

• Turning raw DNA sequences into fixed-length numeric vectors for a downstream classifier or clustering task.
• Deciding between character-level, k-mer, or BPE tokenization before fine-tuning a genomic LM.
• Sanity-checking a new sequence embedding (from a pretrained model or your own encoder) against a cheap k-mer baseline before trusting it.
• Choosing between short-context transformers (DNABERT-2, Nucleotide Transformer) and long-context models (HyenaDNA, Evo) based on how far apart regulatory elements are.
• Building a synthetic probe task (e.g., promoter motif detection) to validate an embedding pipeline end-to-end.

Version Compatibility

• Python ≥3.10, NumPy ≥1.24
• transformers ≥4.40 (only if loading pretrained genomic LMs)
• Reference models: Nucleotide Transformer v2 (InstaDeepAI, 500M-multi-species), DNABERT-2 (117M), HyenaDNA (up to 1M bp context), Evo (prokaryotic/viral, StripedHyena backbone)

Prerequisites

• pip install numpy (add transformers torch only for loading real pretrained models)
• Familiarity with basic sequence handling (biopython skill) helps but is not required
• Related skill: genomic-foundation-models for fine-tuning details of the models named above

Tokenization Strategy Selection

Goal: Turn a DNA sequence into a fixed-length vector without training anything. Approach: Slide a window of length k across the sequence, count k-mer occurrences, and normalize into a frequency vector over the fixed 4**k vocabulary. This is the standard sanity-check baseline before reaching for a pretrained model.

import numpy as np
from collections import Counter

def kmers(seq: str, k: int = 6) -> list[str]:
    """Slide a window of length k across seq, returning overlapping k-mers."""
    seq = seq.upper()
    return [seq[i:i + k] for i in range(len(seq) - k + 1)]

def kmer_embedding(seq: str, vocab: list[str], k: int = 3) -> np.ndarray:
    """Normalized k-mer frequency vector over a fixed vocabulary (order-independent)."""
    tokens = kmers(seq, k)
    counts = Counter(tokens)
    vec = np.array([counts[v] for v in vocab], dtype=float)
    return vec / (vec.sum() + 1e-9)

alphabet = ["A", "C", "G", "T"]
vocab_3 = [a + b + c for a in alphabet for b in alphabet for c in alphabet]  # 64-dim

example_vec = kmer_embedding("ATGATGATGCCC", vocab_3, k=3)
assert example_vec.shape[0] == 64
assert abs(example_vec.sum() - 1.0) < 1e-6

Goal: Check whether an embedding (k-mer or model-derived) actually separates two classes before trusting it downstream. Approach: Train-free nearest-centroid probe — compute per-class centroids on train embeddings, classify test embeddings by nearest centroid. Cheap, has no hyperparameters, and exposes garbage embeddings immediately.

import numpy as np

def nearest_centroid_predict(X_train: np.ndarray, y_train: np.ndarray, X_test: np.ndarray) -> np.ndarray:
    """Binary nearest-centroid classifier: assign each test row to the closer class mean."""
    c0 = X_train[y_train == 0].mean(axis=0)
    c1 = X_train[y_train == 1].mean(axis=0)
    d0 = ((X_test - c0) ** 2).sum(axis=1)
    d1 = ((X_test - c1) ** 2).sum(axis=1)
    return (d1 < d0).astype(int)

def random_dna(n: int, rng: np.random.Generator) -> str:
    """Generate a random ACGT sequence of length n."""
    return "".join(rng.choice(list("ACGT"), size=n))

def inject_motif(seq: str, motif: str, pos: int) -> str:
    """Splice a motif into seq at position pos (overwrites in place, keeps length)."""
    return seq[:pos] + motif + seq[pos + len(motif):]

def demo():
    """Synthetic promoter-vs-background probe: TATA box injected at pos 20 in half the sequences."""
    rng = np.random.default_rng(7)
    n_samples, length, motif = 120, 80, "TATAAA"
    seqs, labels = [], []
    for _ in range(n_samples):
        s = random_dna(length, rng)
        if rng.random() < 0.5:
            s = inject_motif(s, motif, pos=20)
            labels.append(1)
        else:
            labels.append(0)
        seqs.append(s)

    X = np.stack([kmer_embedding(s, vocab_3, k=3) for s in seqs])
    y = np.array(labels)
    X_train, y_train = X[:90], y[:90]
    X_test, y_test = X[90:], y[90:]

    pred = nearest_centroid_predict(X_train, y_train, X_test)
    acc = (pred == y_test).mean()
    assert acc > 0.8, f"probe accuracy too low: {acc}"  # motif is easy; k-mer baseline should nail it
    print(f"nearest-centroid probe accuracy: {acc:.3f}")

if __name__ == "__main__":
    demo()

Goal: Get real embeddings from a pretrained genomic LM instead of a k-mer baseline. Approach: Use transformers to load a genomic foundation model, tokenize, and mean-pool the last hidden state into a single vector per sequence.

def embed_with_nucleotide_transformer(seqs: list[str]) -> "np.ndarray":
    """Mean-pooled embeddings from Nucleotide Transformer v2 (500M-multi-species).
    Requires: pip install transformers torch
    """
    import torch
    from transformers import AutoTokenizer, AutoModel

    model_name = "InstaDeepAI/nucleotide-transformer-v2-500m-multi-species"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).eval()

    tokens = tokenizer(seqs, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        out = model(**tokens)
    mask = tokens["attention_mask"].unsqueeze(-1)
    pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1)
    return pooled.numpy()

Pitfalls

• Mixing tokenization schemes between train and inference silently destroys performance — a model trained on 6-mer tokens fed character-level input at inference will not error, just degrade quietly.
• Always compare against a k-mer baseline first; without one, overfitting or a broken embedding pipeline is very hard to detect.
• k-mer vocab explodes with k: 4^6 = 4096, 4^8 = 65,536 — keep explicit-count k-mers to k≤6, use BPE or learned embeddings beyond that.
• Match sequence window length across models being compared — truncating to a short window can flip labels that depend on distal motifs (see the distal_interaction_label pattern: a motif near the end is silently dropped if you truncate to 200bp).
• Evo is trained on prokaryotic/viral genomes only — do not use it for human/mammalian regulatory tasks.
• Nearest-centroid probes only validate that classes separate at all; a low probe score means "fix the embedding," not "add a bigger downstream model."

See Also

• genomic-foundation-models — fine-tuning and inference details for NT, DNABERT-2, HyenaDNA, Evo
• protein-language-models — ESM2 embeddings for protein sequences (analogous workflow for proteins)
• ai-science-epigenomic-sequence-models — regulatory-activity prediction (Enformer/AlphaGenome) from DNA
• bio-sequence-manipulation-motif-search — motif scanning utilities used alongside embedding probes