Visually Explaining AlphaFold 2 & 3

Tokenization

AF2 only handled proteins, so every token was one amino acid, one token, fixed. AF3 needs to represent proteins, nucleic acids, small molecules, and modified residues all in one framework, so the scheme is more involved.

The rule: standard amino acids and nucleotides get one token per residue, no matter how many atoms they contain. Everything else (modified residues, ligand atoms, ions) gets one token per atom. Hydrogens are excluded throughout; only heavy atoms count.

The reason the compression is valid for standard residues: the model has access to a reference conformer for each one: a canonical 3D geometry from a chemical lookup table (the CCD). Atom-level positions can be recovered downstream from the residue token using this reference. Modified residues and ligands have no such lookup, so each atom must carry its own token.

The compression matters beyond just cleanliness. The pairformer (the main attention stack) operates on an $N_\text{token} \times N_\text{token}$ pair representation. Its memory cost scales quadratically with token count, not atom count. Without this compression, a typical protein complex would be computationally intractable at the atom level.

Genetic Search: MSA Construction

Before any neural network runs, AF3 does something that resembles retrieval-augmented generation: search existing databases for related sequences, then hand them to the model as extra context. The genetic search handles the sequence side of this.

For each protein or RNA chain in the input, AF3 runs HMM-based tools (jackhmmer for proteins, nhmmer for RNA) against several large databases. For proteins: UniRef90, BFD, MGnify, UniClust30. For RNA: RNAcentral, Rfam, and the NCBI Nucleotide Database. These databases collectively span hundreds of millions of sequences: UniRef90 alone exceeds 300 million, and BFD runs into the billions. No neural networks are involved; these are classical bioinformatics tools, and they run before the model sees anything.

The hits get aligned into a multiple sequence alignment (MSA): a matrix where each row is a related sequence from another organism, and each column corresponds to a position in the input. The MSA is capped at $N_\text{MSA} < 2^{14}$ sequences, since complexity scales with $N_\text{MSA}$, so this is a hard ceiling even when far more hits exist.

One notable change from AF2: AF3 substantially de-emphasizes MSA processing. The MSA module is only 4 blocks vs. 48 in AF2's Evoformer, and after those 4 blocks the MSA representation is discarded entirely. Only the pair representation carries forward. The evolutionary signal still matters, but AF3 bets it can be absorbed quickly and that the Pairformer does the heavy lifting.

Template Search

Alongside the MSA, AF3 retrieves structural templates: known 3D structures from the PDB that resemble the input. The process uses hmmsearch against PDB sequences using the constructed MSA as a query profile. Up to 4 high-quality template structures are sampled. Only individual chains are used, not full complexes.

Each template is represented as a pairwise distance matrix. For every pair of tokens, the Euclidean distance between their center atoms is computed: Cα for amino acids, C1' for nucleotides, or the atom itself for single-atom tokens. Rather than storing raw distances, they're binned into a distogram: a discretized distribution over distance ranges. Each entry is augmented with per-token metadata. Inter-chain distances are masked out, so templates inform only intra-chain geometry.

Conformer Generation & Atom Representations

Every atom in the input needs a reference 3D position: a starting geometry the model can build on. For standard amino acids and nucleotides, this is a lookup: each has a canonical low-energy conformation in the Chemical Components Dictionary (CCD), retrieved directly. For small molecules there is no such lookup. AF3 uses RDKit's ETKDGv3 algorithm to generate a conformer from the SMILES string, combining experimental torsion angle preferences with distance geometry to produce chemically reasonable 3D coordinates.

These positions are not predictions. They are priors: rough geometry reflecting known chemistry, used to initialize the atom-level representations before any learned processing. The model doesn't commit to them; they're a starting point that the network updates and the diffusion module ultimately discards in favor of its own generated coordinates.

One practically important point: the conformer only needs to be reasonable, not perfect. Small errors in ring conformations or torsion angles don't catastrophically affect downstream predictions because the diffusion module re-optimizes from learned embeddings, not from the conformer directly. That said, gross failures such as invalid stereochemistry or highly strained geometries can degrade the initial atom pair representation in ways that are harder to recover from, so RDKit quality matters more than it might initially seem.

Building the atom representations

c: atom single representation. For every atom, concatenate its conformer position (relative to other atoms in the same token), charge, atomic number, and other identifiers. This gives c (N_atoms × C_atom): one feature vector per atom.

Distances → local distances → p. From the conformer positions, compute a full N_atoms × N_atoms pairwise distance matrix: d_l,m = |r_l − r_m|. This covers every atom pair. Then apply mask v, a block-diagonal matrix that keeps only within-token pairs and zeros out everything else, producing the local distances matrix. These local distances get embedded as 1/d² (more informative at short range), the feature vectors c_l and c_m are projected and added to each entry, and three linear layers with residual connections refine this into p (N_atoms × N_atoms × C_atompair). The 3 linear layers are what build p. This is separate from what produces q.

q: atom Transformer output. Separately, the Atom Transformer runs 3 blocks (adaptive LayerNorm + attention biased by p + conditional gating + SwiGLU) to update c using p as context. The result is copied to q, the atom-level single representation passed forward. The original c is saved and reused later in the diffusion module.

Input Embedder: Updating Atom Representations

With q (atom-level single) and p (atom-level pair) initialized from the conformer, the input embedder uses the Atom Transformer to update them based on what neighboring atoms look like. c, the original atom features before any learned processing, stays fixed and acts as a conditioning signal throughout. It's effectively a residual anchor to the initial representation.

Each of the 3 Atom Transformer blocks has 4 steps:

c and p condition each step  ·  residual connections (q = q + step output) throughout each block

Step 1: AdaNorm

Standard LayerNorm normalizes each token's activations by mean and standard deviation, then rescales with fixed learned parameters γ (scale) and β (shift) that are the same regardless of input. AdaNorm keeps the normalization step but makes γ and β adaptive: a linear projection of c generates them on the fly. The conditioning is one-directional (c → γ,β → applied to q); c itself is never updated. This lets the initial atom features dictate how each block's normalization behaves, rather than using a single fixed rescaling learned from all examples.

Step 2: Attention with Pair Bias

Self-attention on q, with three differences from standard multi-head attention:

• Pair bias: after computing QK dot products for query atom l, row l of p is linearly projected to a scalar per head and added as a bias to the attention logits before softmax. This is strictly one-directional: p shapes which atoms attend to each other, but the attention output never updates p. Dimensions: C_atom = 128, C_atompair = 16.
• Gating: q is also projected through a sigmoid gate ∈ [0,1]. The attention output is multiplied by this gate before heads are concatenated, controlling how much of each attention update enters the residual stream. The model learns to filter which information gets written into the running representation.
• Sparse (sequence-local) attention: because N_atoms ≫ N_tokens, full N_atoms × N_atoms attention is prohibitively expensive. AF3 uses sequence-local atom attention: groups of 32 query atoms each attend to 128 nearby key atoms.

Step 3: Conditioned Gating

After the attention step, a second gate filters q. This time the gate is generated from c, not from q itself. A linear projection of c through sigmoid gives per-channel values in [0, 1] that element-wise scale what gets absorbed from the attention step into the residual stream. The gate depends on the original atom features (c), not the current learned state of q, so the initial chemistry of each atom determines how strongly the attention update registers, independent of what q has learned so far.

Step 4: Conditioned Transition (SwiGLU MLP)

The MLP step, analogous to the FFN in a standard transformer. It's "conditioned" because it's wrapped in AdaNorm and Conditioned Gating, both of which depend on c. AF3 uses SwiGLU here instead of the ReLU used in AF2. ReLU projects up to 4×D, applies ReLU, and projects back down. SwiGLU takes two parallel up-projections (D → K×D), applies Swish (x ⋅ σ(x), a smooth non-linearity that does not hard-zero negative values) to one, element-wise multiplies the two, then projects down. Swish's smoothness helps gradient flow, particularly when the network is deep. SwiGLU has become standard in large models (LLaMA, PaLM) and generally outperforms ReLU without meaningful additional cost.

Aggregating to Token Level: s^init and z^init

Everything built so far (c, q, p) lives at the atom level. The representation learning section operates at the token level, so we aggregate here. This is the last step of the input embedder.

s^init (token-level single). q is linearly projected from C_atom (128) to C_token (384). For each token, all atom representations belonging to that token are averaged, compressing N_atoms rows down to N_tokens. The result is concatenated with token-level MSA features where available (residue type, MSA statistics at position i, etc.), growing the channel dimension to C_token + 65. A final linear projection brings it back to C_token, producing s^init. The pre-projection version, s^inputs, is saved separately for later use in structure prediction.

z^init (token-level pair). Starting from s^init, project each token's representation from C_token to C_z (128), giving vectors s_i and s_j. The pair entry z_i,j is initialized as s_i + s_j. Two things are then added: a relative positional encoding capturing how far apart tokens i and j are in sequence, and any user-specified bond information linearly embedded into the pair representation.

At this point all inputs are fully embedded. The atom-level representations (c, q, p) are set aside for the diffusion module. The representation learning section from here works entirely with the token-level s and z, updated using the MSA (m) and templates (t).

Template Module

The template module injects structural prior information from PDB homologues into the pair representation z. The templates retrieved earlier each contribute a pairwise distance matrix: all token-pair distances binned into 38 bins spanning 3.15–50.75 Å. Cross-chain distances are masked to zero because templates only inform intra-chain geometry, never inter-chain contacts.

The processing steps for each template:

1. Each template's pair representation is linearly projected into channel dimension C, giving one matrix v_t per template.
2. z is also linearly projected (C_z → C) and added to each v_t, giving every template context about the current state of the pair representation.
3. Each combined v_t passes through 2 blocks of the Pairformer stack (same operations as the main Pairformer, described in that section).
4. All template representations are averaged into a single N_tokens × N_tokens × C matrix.
5. A linear layer followed by ReLU produces u, notably one of only two uses of ReLU in AF3. The other is in the input embedder; every other non-linearity in the network uses SwiGLU.
6. u is added to z, injecting structural template knowledge into the pair representation before the main Pairformer runs.

The design is efficient: templates are processed in parallel, averaged rather than concatenated (keeping z's dimensionality fixed regardless of how many templates are found), and the lightweight 2-block Pairformer does the heavy lifting of extracting useful structural signal before averaging collapses the template dimension.

MSA Module

The MSA module is AF3's version of AF2's Evoformer. Its job is to simultaneously refine the MSA representation (m) and the pair representation (z), letting information flow between them. It runs 4 blocks, each containing three distinct steps: updating z from m (outer product mean), updating m from z (row-wise attention), and updating z through triangle operations.

Step 1: Outer Product Mean (m → z)

The first thing each block does is push MSA information into the pair representation. Comparing two columns of the MSA tells you how correlated those two positions are across evolution. For every token pair (i, j), take the outer product of a_s,i and b_s,j (two linearly projected MSA vectors) for each evolutionary sequence s, then average across s. This gives one N_tokens × N_tokens × C² tensor capturing cross-position correlations. Flatten the C² depth, project down to C_z, and add to z_i,j.

This is the only point in the entire model where information is shared across evolutionary sequences. Everything else in the MSA module runs independently per sequence row. This is a significant simplification from AF2's Evoformer, which was much more computationally intensive in how it mixed sequences.

Step 2: Row-wise Gated Attention (z → m)

Having updated z from m, the block then does the reverse: update m using z. This is attention with no queries or keys. Instead, row i of z directly provides the attention scores for sequence row s. Each z_i,j vector is linearly projected to a scalar per head, softmaxed along j, and used as attention weights to compute a weighted sum over column i of m. A sigmoid gate (also from m) filters the output before concatenating heads. This runs independently per sequence row, with no cross-sequence information shared.

The key intuition: z_i,j already encodes the relationship between tokens i and j, so projecting it to a scalar gives an attention map directly from that relationship. No learned Q/K projections are needed.

Step 3: Triangle Updates (z → z)

The rest of each block refines z using the same four triangle operations described in the Pairformer section: triangle multiplicative update (outgoing), triangle multiplicative update (incoming), triangle self-attention (start node), and triangle self-attention (end node). Each is followed by a residual add. The block closes with a SwiGLU transition. No MSA features are used in this step; it operates purely on z.

Pairformer

After the template and MSA modules have enriched z and s, the model sets them aside. From here, only the pair representation z and single representation s pass through the Pairformer: 48 blocks that refine them against each other. Each block has two streams: the pair stream (4 triangle operations + transition) and the single stream (single attention with pair bias + transition). The updated z feeds down into the single stream each block.

Why triangles?

The pair representation z_i,j encodes the relationship between token i and token j, a learned analog of their geometric relationship. The triangle inequality says: if you know the distance from i to k, and from j to k, you have a strong constraint on what the distance from i to j can be. The triangle operations bake this principle in by ensuring every z_i,j is updated by looking at all possible third tokens k simultaneously.

Because z encodes directional relationships, we need to consider two types of triangles. If we think of tokens as nodes in a directed graph with z as a weighted adjacency matrix, "outgoing edges" from i form one type of triangle and "incoming edges" to i form another.

Triangle Updates

Each triangle update creates three linear projections of z (called a, b, g). To update z_i,j, take row i from a and row j from b, multiply them element-wise for each k, then sum across all k. This is equivalent to asking: "for every possible third token k, how does the relationship i→k combine with the relationship j→k?" The result is gated by g_i,j and added back to z. The incoming update does the same but transposes the logic: use column i from a and column j from b, capturing k→i and k→j triangles instead.

Triangle Attention

Triangle attention extends axial attention with the triangle principle. For the starting node variant, to compute attention scores for z_i,j along row i, we compare query z_i,j with key z_i,k for all k, then add a bias from z_j,k projected to a scalar per head. The bias from z_j,k nudges the attention weights based on the third edge of the triangle, encoding how j relates to k. Values also come from row i.

The ending node variant transposes this: keys and values come from column j of z instead of row i, and the pair bias comes from column i (z_k,i). Same logic, opposite axis. Sigmoid gating also appears throughout.

Single Attention with Pair Bias

After the four triangle steps and a transition block refine z, the model updates s using z as an attention bias. Q/K/V projections from s, a bias from row i of z projected to a scalar per head, and a sigmoid gate. This runs at the token level with full attention (no sparse windowing, since N_tokens ≪ N_atoms). After 48 blocks the outputs are s_trunk and z_trunk, which feed into the diffusion module.

Recycling

Before the 48 blocks run, the entire trunk (input embedder, MSA module, template module, and Pairformer) runs multiple times end to end. This is called recycling (3 iterations by default). At the start of each new pass, the pair and single representations from the previous pass are added back into the input, giving the model a better initialization for the next round.

Each recycling pass lets the representations converge: early passes establish rough global geometry, later passes refine local structure. Only the final pass produces s_trunk and z_trunk that feed into the diffusion module.

Diffusion Module

The diffusion module is where 3D coordinates are actually generated. At each denoising step, the module takes noisy atom coordinates x and predicts how much noise to remove. To do this it conditions on everything the trunk computed (s_trunk, z_trunk) plus the original atom-level embeddings from the input embedder (s^inputs, c, p). The unprocessed versions matter because the trunk representations are token-level, while diffusion also needs raw atom-level geometry.

The process has 4 steps that repeatedly move between atom-level and token-level representations within each denoising step.

Step 1: Token-Level Conditioning

Two conditioning tensors are initialized. The token-level single s is formed by concatenating s^inputs and s_trunk, projecting back to C_token, then adding a Fourier embedding of the current diffusion timestep t. Including t ensures the model knows what noise scale to remove at this step. The token-level pair z is formed similarly by concatenating z_trunk and the relative positional encoding, projecting to C_z, and running through transition blocks.

Step 2: Atom-Level Conditioning, Coordinate Update, and Aggregation

The atom-level tensors c and p (from the input embedder) are updated to incorporate trunk information. s_trunk is broadcast from N_tokens to N_atoms (each token's representation is copied to all its atoms), projected to C_atom, and added to c. z_trunk is similarly broadcast and added to p.

The noisy coordinates x (N_atoms×3) are scaled to unit variance to create dimensionless r, then linearly projected to C_atom and added to q. This gives q awareness of each atom's current position. q passes through the Atom Transformer (conditioned on c and p), and the output is aggregated back to token-level to produce a.

Step 3: Token-Level Attention (Diffusion Transformer)

The aggregated a is updated using a token-level transformer: AdaNorm (conditioned on s), attention with pair bias (using z), conditioned gating, and a SwiGLU transition. This runs at token level with full attention and uses s and z from Step 1 as conditioning.

Step 4: Atom-Level Attention and Coordinate Prediction

a' (updated token-level) is broadcast back to atom-level, then used to update q via another Atom Transformer (conditioned on c' and p'). A final linear layer maps q back to ℝ³, giving the coordinate update. Because the prediction was made in unit-variance space, the update is rescaled and applied: x_new = x + x_update. That's a single diffusion step. The process repeats until coordinates converge.

Confidence Module & Loss Function

The confidence module predicts how accurate the model's own output is. Its goal is not to improve the structure but to give the user a signal about which parts of a prediction to trust. Its outputs (pLDDT, PAE, PDE, and resolved atom predictions) are trained via a dedicated confidence loss. The full training objective combines three terms:

$\mathcal{L}_\text{distogram}$: token-level structure accuracy

The model's predicted atom coordinates are converted to a token-level distogram by using the center atom of each token (Cα for amino acids, C1' for nucleotides). Pairwise distances between these center atoms are binned into distance categories, and the predicted distogram is compared to the true one via cross-entropy. This acts as a fast, coarse signal on overall structural accuracy.

$\mathcal{L}_\text{diffusion}$: atom-level structure accuracy

The diffusion loss operates at full atom resolution and is itself a weighted sum:

ℒ_MSE measures mean squared error between predicted and ground-truth atom positions across all atoms (with DNA, RNA, and ligand atoms upweighted). The full-atom MSE is scaled by a factor that depends on the diffusion noise level $\hat{t}$, giving more weight to predictions at low noise (fine-grained refinement steps). ℒ_bond adds an extra MSE penalty specifically for protein-ligand bond lengths.

ℒ_smooth-lDDT (smoothed local distance difference test) measures local accuracy. For each atom pair, the predicted and true distograms are compared. The difference is passed through a sigmoid at four thresholds (0.5, 1, 2, 4 Å), converting each into a smooth probability of passing that test. The average pass probability across all four thresholds is the smooth-lDDT score. Pairs where the true distance is large are excluded.

$\mathcal{L}_\text{confidence}$: self-predicted accuracy

The confidence loss trains the model to predict its own error metrics. The key insight: if the predicted structure is highly inaccurate but the model correctly predicts that it will be, the confidence loss is low. This teaches calibration, not structural accuracy. Gradients from this loss only update the confidence heads. They do not affect the rest of the network.

Confidence Head Architecture

At a selected diffusion timestep t, the predicted coordinates r_t are used to construct a distogram d (by linear projection and binning). This d is added to s_i + s_j to initialize z_i,j, combining sequence and structural information into the pair representation for the confidence Pairformer. s_trunk and s^inputs are concatenated, projected, and combined with z_trunk to initialize s and z for a lightweight Pairformer update. The updated s and z then feed the four confidence heads.

The confidence predictions are generated mid-run during diffusion. At a selected timestep t, the current noisy coordinates are used to update s_trunk and z_trunk, and the confidence heads predict errors from those updated representations. The actual error metrics are then computed on those same coordinates and used as targets. This "mini rollout" during training is what makes confidence calibration possible without running full diffusion at every training step.

Before AlphaFold3 unified structure prediction across proteins, nucleic acids, ligands, and complexes into a single framework, DeepMind took an intermediate step with AlphaFold-Multimer (2021), a version of AF2 retrained specifically to handle multi-chain assemblies. Rather than bolting complexes onto a single-chain model with tricks like residue gaps or flexible linkers, AlphaFold-Multimer was trained end-to-end on oligomeric structures, with native handling of cross-chain genetics (pairing MSAs across chains using species annotations), permutation symmetry (so that identical chains in a homomer aren't unfairly penalized for being "out of order"), and a new confidence metric called interface pTM (ipTM) that specifically scores how well predicted interfaces match reality. Visually, this showed up as much sharper predicted aligned error (PAE) maps at chain interfaces: a model could now say not just "this chain folds correctly" but "these two chains dock correctly," and flag low-confidence interfaces (like PDB 6QF7) even when each individual chain was predicted well. It's a useful reference point for understanding why visualizing complexes, not just single folds, became such a central challenge on the road to AlphaFold3.