Attention Mechanisms

Why Attention Mechanisms Matter in Biology

Traditional neural networks process all inputs equally, treating every gene or nucleotide with the same importance. But biological systems are inherently selective—regulatory elements can affect genes millions of base pairs away, protein function depends on specific amino acid interactions, and cell identity is determined by a subset of marker genes. Attention mechanisms capture this selectivity.

How the Factory Works: Five Types of Attention

Key Applications in Biology

Single-Cell Genomics

Attention mechanisms enable models to learn which genes define cell types, predict how cells respond to perturbations, and integrate data across batches and technologies.

Model	Task	Attention Type	Key Achievement
scGPT	Cell annotation, perturbation prediction	Multi-head self-attention	33M cells, outperforms task-specific models
scBERT	Cell type classification	Performer (linear complexity approximation)	Handles whole transcriptome (16,000+ genes) with linear complexity
GeneCompass	Cross-species gene regulation	Multi-head with knowledge embedding	101.7M cells (53.5M human + 48.2M mouse) with cross-species transfer learning
CellPLM	Spatial transcriptomics	Spatial graph attention	Cell-level tokens capture spatial context
State (Adduri et al., Arc Institute, bioRxiv 2025)	Perturbation response prediction across cellular contexts	Multi-head self-attention, two-module design: State Embedding (SE) + State Transition (ST)	SE trained on 167M cells, ST on 100M+ perturbed cells; >30% improvement in discriminating perturbation effects; identifies strong perturbations in cellular contexts unseen during training
STACK (Brixi et al., Arc Institute, 2026)	Zero-shot perturbation prediction via in-context learning	Tabular attention — within-cell (across genes) + cross-cell (across rows); no fine-tuning at inference	Pretrained on 149M cells (Perturb Sapiens atlas); generalizes to unseen perturbations, donors, and tissues without weight updates by reading labeled context cells as table rows
Sigmoid Attention for scFMs (2026)	Single-cell foundation model pre-training	Sigmoid replaces softmax in attention scoring	25% higher cell-type separation, up to 10% faster training, and more stable training (bounded derivatives avoid the divergence softmax shows on long sequences without gradient clipping)
SAVE (Fudan Univ., ICLR 2026)	Multi-condition single-cell generation	Gene Block Attention — genes grouped into pathway-based blocks rather than treated as independent tokens	Generalizes across biological/technical conditions by modeling higher-order dependencies among gene modules

Genomic Sequence Analysis

Attention allows models to capture long-range regulatory interactions and learn sequence patterns across entire genomes.

Model	Task	Context Length	Key Innovation
Nucleotide Transformer	Variant effect prediction	12kb (2,000 6-mers)	Multi-species training (850 genomes)
GET	Expression from chromatin	200 genomic regions (~2-4 Mbp span)	Predicts expression from distal enhancers >1 Mbp away (r=0.94, R²=0.88)
AlphaGenome	Regulatory element discovery	~1 Mb window (1,048,576 bp, single-base resolution)	Multi-scale attention for different genomic features
HyenaDNA (Nguyen et al., NeurIPS 2023)	Genomic function prediction	Up to 1M nucleotides, single-nucleotide resolution	Hyena state-space operator (not attention) — no k-mer tokenization, full global context per layer at O(N) cost; SOTA on 12/18 Nucleotide Transformer benchmarks with far fewer parameters
Evo 2 (Brixi et al., Arc Institute, 2025)	Sequence generation + variant effect	Up to 1M tokens, single-nucleotide resolution	StripedHyena 2 (SSM); trained on 9.3T nucleotides across 128,000+ species; zero-shot variant-effect prediction

Protein Design and Structure

Attention mechanisms learn which amino acids interact in 3D space and generate functional proteins with specific properties.

Model	Application	Attention Approach	Experimental Validation
Structured Transformer	Inverse folding	k-NN graph attention (k=30)	27.6% native sequence recovery; 21,000× faster on GPU, 455× on CPU vs Rosetta
ProGen (Madani et al., Nat Biotechnol 2023)	Protein generation	Causal self-attention (1.2B params)	Functional lysozymes down to 31.4% identity (extreme low-identity case, ~200× lower efficiency). Not to be confused with the later-scaled ProGen2 (Nijkamp et al., Cell Systems 2023; 151M–6.4B params) — a different paper, different param range.
Chroma	Complex design	Random graph networks (O(N) edges)	High expression rates; crystal structures ~1Å RMSD to predictions
ProteinMPNN	Sequence design	Message passing with attention	State-of-the-art at publication (Dauparas et al., Science 2022) for fixed backbone design
LigandMPNN (Dauparas et al., preprint 2023 / Nat Methods 2026)	Sequence design for ligand/nucleotide/metal-binding contexts	Message passing extended with explicit nonprotein-atom context	ProteinMPNN successor: outperforms it on native sequence recovery at small-molecule (63.3% vs 50.5%), nucleotide (50.5% vs 34.0%), and metal (77.5% vs 40.6%) interfaces; 100+ experimentally validated binders, 4 X-ray-confirmed structures
AlphaFold3 (Abramson et al., Nature 2024)	Protein–nucleic acid–ligand complex structure prediction	Pair representation + triangle attention (PairFormer)	DOI 10.1038/s41586-024-07487-w; extends structure prediction beyond proteins alone to DNA/RNA/ligand complexes

Transformer Architectures in Biology

Most attention-based models in biology use the Transformer architecture, introduced by Vaswani et al. (2017). The core innovation is replacing recurrence with attention, allowing parallel processing of sequences while maintaining the ability to capture long-range dependencies.

Comparing Attention to Traditional Architectures

Feature	CNN	RNN/LSTM	Transformer (Attention)	SSM (Mamba/Hyena)
Long-range dependencies	Limited by receptive field	Degrades with distance (vanishing gradients)	Direct connections between any positions	Good, but compressed into a fixed-size hidden state (not exact all-pairs)
Computational complexity	O(N)	O(N) but sequential	O(N²) for self-attention, O(E) for graph attention (E = edges, sparse when avg. degree ≪ N)	O(N) time and memory
Parallelization	High	Low (sequential processing)	Very high (all positions processed together)	High during training (parallel scan); sequential-style recurrence at inference
Interpretability	Filter visualization	Hidden states (opaque)	Attention weights show relationships	State transitions (more opaque than attention weights)
Variable-length sequences	Requires padding	Natural support	Natural support	Natural support
Best biological applications	Local motifs, images	Short sequences, time-series	Long sequences, relationships, foundation models	Ultra-long sequences — whole chromosomes, 1M+ bp genomic context

Recent Innovations in Attention for Biology

Practical Benefits for Biologists

For Experimentalists

• In Silico Screening: scGPT predicts perturbation outcomes (r=0.94) before running CRISPR experiments, saving time and resources
• Variant Interpretation: Nucleotide Transformer scores clinical variants without functional assays
• Protein Design: ProGen and Chroma generate functional proteins in days vs. years of directed evolution
• Cell Type Discovery: Automated annotation with scGPT and scBERT reduces manual curation effort

For Computational Biologists

• Pre-trained Embeddings: Use gene/cell representations from foundation models as features (GeneCompass improved GEARS by 15%)
• Zero-shot Prediction: Apply models to new cell types/species without retraining
• Interpretable Models: Attention weights provide biological insights beyond predictions
• Transfer Learning: Fine-tune on small datasets leveraging knowledge from millions of cells/sequences

Computational Efficiency

Model	Task	Speed / Scale	Hardware
scGPT	Cell annotation	Pretraining scale: millions of cells (not an inference-speed figure, unlike the rows below)	8× A100 GPUs
Nucleotide Transformer	Variant scoring	1,000+ sequences/second	Single GPU
Structured Transformer	Protein inverse folding	GPU: 222 AA/s (21,000× faster); CPU: 0.488 AA/s (455× faster than Rosetta)	Single GPU or CPU
GET	Expression prediction	Minutes per cell type	8× A100 GPUs

⚖️ Side-by-Side: Who Attends to Whom — Global Self-Attention vs Graph-Based Attention

Two of the four attention types above sit at opposite ends of one axis: how many other tokens each token is allowed to look at. Global self-attention lets every token attend to every other token; graph-based attention restricts each token to its neighbors in a biological graph — a protein–protein network, a spatial neighborhood, a gene-regulatory graph.

🛠️ Hands-On Practice

The steps below implement the engine inside every one of the four types above — scaled dot-product attention — from scratch in NumPy, then show the PyTorch multi-head and graph-attention equivalents. Seeing the raw matrix math makes the "who attends to whom" idea concrete.

Environment & packages

numpy is enough for the core mechanism; torch provides the production multi-head module, and torch-geometric the graph-restricted variant.

conda create -n attn python=3.10 -y
conda activate attn
pip install numpy torch
# graph-restricted attention variant:
# pip install torch-geometric

Hardware. The toy example runs anywhere; training a real transformer wants a GPU, and long inputs (whole genomes, thousands of cells) need FlashAttention or an efficient-attention variant to fit the N×N matrix in memory.

Data structures & formats

• Input X — shape (n_tokens, d_model); tokens are genes, cells, nucleotides, or residues depending on the model
• Q, K, V — query / key / value projections of X, each (n_tokens, d_k)
• Attention weights — (n_tokens, n_tokens); each row is a softmax over all keys and sums to 1
• Mask — boolean (n_tokens, n_tokens) that blocks padding or future positions before the softmax
• Heads — d_model split into num_heads independent subspaces, concatenated after attention
• edge_index (graph attention) — (2, n_edges) list of connected token pairs

Minimal code walkthrough

Scaled dot-product attention in NumPy, then the PyTorch multi-head and graph equivalents.

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)      # subtract max: numerical stability
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-1, -2) / np.sqrt(d_k)   # 1/sqrt(d_k) scaling MATTERS
    if mask is not None:
        scores = np.where(mask, scores, -1e9)         # block disallowed pairs BEFORE softmax
    weights = softmax(scores, axis=-1)                # each row sums to 1
    return weights @ V, weights

# toy: 4 tokens (genes / cells / residues), model dim 8
rng = np.random.default_rng(0)
X = rng.standard_normal((4, 8))
Wq, Wk, Wv = (rng.standard_normal((8, 8)) for _ in range(3))
out, attn = attention(X @ Wq, X @ Wk, X @ Wv)
print("attention weights (rows sum to 1):\n", attn.round(2))
print("output shape:", out.shape)                    # (4, 8)

In practice you use the batched, multi-head module — and swap in graph attention when you have a biological graph:

import torch, torch.nn as nn

# Global multi-head self-attention (query = key = value = x)
mha = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
x = torch.randn(1, 4, 8)              # (batch, tokens, dim)
out, w = mha(x, x, x)                 # w: per-head attention weights

# Graph-restricted attention: attend ONLY along edges, not all-pairs
# from torch_geometric.nn import GATConv
# gat = GATConv(in_channels=8, out_channels=8, heads=2)
# out = gat(node_features, edge_index)   # edge_index = (2, n_edges)

Common pitfalls & tips

• Never drop the 1/√d_k scaling. Without it, dot products grow with dimension, push softmax into saturation, and gradients vanish — always divide scores by √d_k.
• Stabilize the softmax. Subtract the per-row max before exp(), or large scores overflow to inf.
• Mask before, not after, softmax. Set disallowed scores to a large negative value before softmax; masking afterward leaves leaked probability mass on padding or future tokens.
• Watch the O(N²) memory. The N×N matrix blows up for long inputs (whole genomes, thousands of cells) — use FlashAttention or a linear/efficient-attention variant.
• num_heads must divide d_model. Each head gets d_model / num_heads dimensions; a mismatch errors or silently mis-shapes the tensors.
• Graph attention is only as good as its edges. GATConv attends solely along edge_index; a wrong or too-sparse graph silently starves the model of context.

Scaling up: FlashAttention

The pitfall above (O(N²) memory) is what FlashAttention solves in practice — same exact attention output, computed via tiling so the full N×N matrix is never materialized in memory:

# pip install flash-attn  (requires a CUDA-capable GPU)
from flash_attn import flash_attn_qkvpacked_func

# qkv: (batch, seqlen, 3, num_heads, head_dim) — same math as nn.MultiheadAttention,
# but tiled so the N x N score matrix is never fully materialized in memory.
# out = flash_attn_qkvpacked_func(qkv, dropout_p=0.0, causal=False)
#
# Exact (not approximate) attention, typically ~3-7x faster and using an order
# of magnitude less memory than the naive O(N^2) implementation above — the
# difference between truncating to top-K genes and attending over the full
# ~20K-gene transcriptome in one pass.

Key Takeaways