AF3's diffusion module never runs full attention over every atom, that would cost O(atoms squared). Instead it encodes atoms with the sequence-local attention from the previous screen, compresses them down to one vector per token (machine-learning terms; a token is one amino acid residue, or one heavy atom of anything else, in structural-biology terms), runs full attention only on that much smaller set, then broadcasts the result back down to atoms.
Full attention only ever touches the narrow waist, tokens, never the wide ends, atoms, that is what makes global reasoning affordable at all.
Ala-Gly-Val-Leu-Ser-Lys peptide + ATP + Mg²⁺. Click a level to see its breakdown.
Every heavy atom in the complex, before any compression.
Comparisons in pairwise attention computations (N^2 for full attention over N items, N x 128 for local attention with the previous screen's 128-atom window).
At our own small scale, the 128-atom window from the previous screen buys nothing: with only 71 atoms, a "local" 128-atom window already covers the whole complex, so local and full attention cost the same here. The real saving from windowing only appears once a complex exceeds 128 atoms, and grows without bound from there, which is why AF3's own paper frames the motivation at "tens of thousands of atoms for a large complex, vs. hundreds of tokens" rather than at our example's scale.