Architecture#

These diagrams take RFD4-Proteina apart one level at a time, from the whole generation loop down to a single attention head. The top level is a map of the system. Every level below it is a neural circuit diagram generated from an algebraic expression of the code, drawn at the dimensions of the reference model (pair width 384, triangle hidden width 256).

Hover a box or an operation to see its formula, a description and links to the lines of code it stands for. Click a box whose details offer “Open diagram” to open its own diagram. The breadcrumbs, the tree and the browser’s back button take you back up.

The latent explorer takes the autoencoder’s latents apart.

RFD4-Proteina at a glance An annotated CIF is featurized into design features that do not depend on the noisy state. A flow-matching sampler runs 60 steps from Gaussian noise; at every step the denoiser embeds the noisy state, runs its trunk R times in a recycling loop and predicts velocities, and an integration step moves the state. The decoder turns the final CA positions and latents into a sequence and atoms. The autoencoder is trained first and defines the latents that flow training regresses. lat = x₁ YOU SPECIFY ANNOTATED CIF structure and intent in one file DESIGN FEATURES 729 token, 77 pair channels, CATH labels, token types fixed during a generation, except 37 coordinate-group channels, relabelled at every network evaluation FLOW-MATCHING SAMPLER T = 60 STEPS from Gaussian noise at t = 0 to data at t = 1 x at t + Δt becomes the input of the next step DENOISER LocalLatentsTransformer · 178 M parameters Embeddings s₀ 768 · z₀ 384 Conditioning times, CATH → c RECYCLE × R R(i) = 1, 2, 3, 5, 6, 5, 3, 2, 1 Trunk 14 layers · 7 pair updates Layer Pair update Attention Triangles s, z feed the next cycle Heads v̂ CA, v̂ lat Step vf_tsr COMES BACK DECODER CA positions and latents → 33 residue classes → atom37 coordinates DESIGN sequence + all-atom CIF TRAINED FIRST AUTOENCODER 256 M parameters, frozen afterwards Encoder AE layer THEN, FOR THE FLOW TRAINING STEP encode → interpolate with noise → denoise with R ~ Poisson(3) → squared velocity error per modality + insertion loss

drag to pan · Ctrl + scroll to zoom · hover for details · click a box whose details offer "Open diagram" to go deeper · Tab into the diagram, arrow keys move between blocks, Enter opensdrag to pan · pinch to zoom · tap a box for details · tap "Open diagram" to go deeper

Reading the diagrams#

A neural circuit diagram draws arrays as bundles of wires, one wire per axis, and operations as shapes across them. A wire that passes an operation is an axis it is broadcast over; a wire that ends at an operation is an axis it consumes. Dashed lines separate the arrays of a product. A bold box is a block of the model: hover it for its formula, description and code. When its details offer “Open diagram”, click it to open the block as its own diagram; clicking any other box pins its details. A bracket with a count is a loop: T for the 60 sampling steps, R for the recycles, P for the 7 repetitions of [layer, pair update, layer] in the trunk and L for the 12 layers of each half of the autoencoder. Coloured labels at the top and bottom of a box are tape slots: values written once and read elsewhere, such as the conditioning c, which every recycling cycle reads and the trunk passes through its layers.

Axis letters and value symbols are read differently. An axis label is a letter with its size as a subscript, as p384 for the pair channels, except on the token axis x, whose size is the length of the design. The legend beside every diagram lists each axis with its identifier in the code where it has one and a descriptive name otherwise. Tape-slot labels are coloured and carry no size, and they and the formulas name values, not axes: s is the single representation, z the pair representation and c the conditioning. The two modalities are marked CA for the CA positions and lat for the latents, as in Step_CA and v̂ lat, so z always means the pair representation. The CA position of a token is atom37 slot 1: the Cα of an amino acid, C1′ of a nucleotide and the atom itself for an atomized token, which carries no latent. Diagrams Explained (Quickly) introduces the notation in full.

Sizes of the diagrams#

Axis

Size

Name in the code

What it indexes

x

design

tokens

residues, nucleotides and atomized ligand atoms of the design

m

768

token_dim

the single representation

p

384

pair_repr_dim

the pair representation

u

256

dim_cond

the conditioning vector

h, d

12, 64

nheads, dim_head

attention heads

f

3072

dim_inner

the SwiGLU of every trunk and autoencoder layer

e

256

tri_mult_c

the triangle multiplications

l

16

latent_dim

autoencoder latents per residue or nucleotide

T

60

nsteps

sampling steps

R

1 to 6

n_recycle

trunk passes per network evaluation, set per sampling step

P

7

nlayers / update_pair_repr_every_n

repetitions of [layer, pair update, layer] in the trunk

L

12

nlayers of the autoencoder

layers of the encoder and of the decoder