Max Zhuang
Quick digests

Looped transformers and AlphaFold recycling

· 6 min read

Digests are written primarily by LLMs with only readability edits by me. If this one reads like slop, please shoot me an email at max.zhuang@berkeley.edu.

A looped transformer runs the same block of layers several times before producing an output. AlphaFold 2 recycling does something similar: it makes another structure prediction using information from the previous one. Both can spend more computation on a problem without adding more learned weights. Dehghani et al., Jumper et al.

The two models carry different information between passes. A language model can work entirely with internal numerical states. AlphaFold 2 also feeds back information from a predicted three-dimensional structure.

What a looped transformer repeats

A transformer represents its input as arrays of numbers called hidden states. Its layers use attention to combine information from different positions, then transform those states. In a conventional transformer, successive layers usually have different weights: the numerical parameters learned during training.

A looped transformer reuses a layer or block of layers. Four layers run six times give 24 layer applications while storing four distinct layers’ weights in that block. Each pass receives the hidden state produced by the previous one, so it can build on earlier computations. The weights stay fixed during inference, when the trained model is used to make a prediction. Saunshi et al.

This is called recurrent depth because the model repeatedly processes states at the same input positions. A conventional recurrent neural network instead carries state from one sequence position to the next. Universal Transformers is an early example, combining depth recurrence with attention.

A shared block updates the current state
A shared block updates the current stateThe current state enters a shared block, which updates it. The resulting state can return for another pass or proceed to the readout. Three illustrative passes use the same weights; the small bars represent changing state values, not measured results.Current stateShared blocksame weights each passUpdated stateReadout when finished
The readout converts the final state into an output. Some architectures also supply the original input to the block on each pass. Three passes are shown for illustration; further passes need not improve the result.

In a language model, several loops can run before the readout produces probabilities for the next text token. A token is a unit of text, often part of a word. When intermediate problem-solving stays in these internal states, it is called latent reasoning.

The model can still write a chain of thought: intermediate reasoning expressed as text. It can use depth recurrence before emitting each token of that text. Geiping and colleagues demonstrate this combination in a model with a repeated middle block. Geiping et al.

What AlphaFold 2 recycles

AlphaFold 2 predicts a protein’s structure from its amino acid sequence, aided by related sequences and, when available, structural templates. A multiple sequence alignment, or MSA, lines up those related sequences. The Evoformer exchanges information between sequence features and pair features, which describe relationships between pairs of amino acid residues. The structure module then predicts coordinates. Jumper et al.

Recycling feeds three things into the next pass:

  1. Target-sequence features, from the first row of the previous MSA representation.
  2. Pair features from the previous pass.
  3. Distances from the predicted structure, measured between representative points on the residues and encoded into distance bins.

AlphaFold combines this feedback with the input representations, then runs the Evoformer and structure module again with the same weights. A recycle repeats an entire predictor containing many layers. Official AlphaFold 2 implementation

The original setup used four total passes: the initial prediction plus three recycles. The structure module also performs its own iterations within each pass. Recycling is the outer loop, encompassing both the Evoformer and the structure module. Supplementary methods

What AlphaFold 3 recycles

AlphaFold 3 recycles features before generating coordinates. Its Pairformer updates features for individual molecular tokens and pairs of tokens. A token here represents a residue or an atom. The resulting features guide a separate diffusion module, which generates atomic coordinates through successive updates to noisy coordinates. Those coordinates are not fed back through the feature-recycling loop. Abramson et al., released inference code

In this comparison, AlphaFold 3’s feature recycling is closer to a looped transformer’s internal-state updates. Diffusion is a second iterative computation with its own state and training objective. More recycles, more diffusion steps, and more independent samples spend computation in different parts of the prediction process.

Recycling in AlphaFold 2 and 3
AlphaFold 2 predicts coordinates inside the recycle. AlphaFold 3 runs diffusion after feature recycling. Sources: AlphaFold 2, AlphaFold 3.

Why another pass can help

A later pass can combine relationships computed by an earlier one. Suppose the input says Alice is older than Bob, and Bob is older than Carol. An update could represent the two comparisons; another could combine them to establish that Alice is older than Carol.

For a protein, suppose two regions have plausible local shapes but are poorly placed relative to each other. The predicted distances describe that placement. Feeding them back gives the next pass information it can use to revise the structure.

These examples illustrate computations the feedback allows. They do not identify the steps a trained model actually takes. In particular, an AlphaFold recycling trajectory shows successive predictions, not the physical process by which a protein folds.

What the experiments show

Saunshi and colleagues compared looped models with deeper models using separate weights, matching the total number of layer applications. Looped models approached or exceeded the deeper models on several arithmetic and reasoning tests. In the language-model experiments, the deeper models did better on next-token prediction and questions testing memorized facts. Saunshi et al.

Geiping and colleagues trained a 3.5-billion-parameter recurrent-depth language model on 800 billion tokens. Running more loops improved several reasoning evaluations. Their paper establishes that useful depth recurrence can be learned at this scale; it leaves broader comparisons at equal training and inference cost for further work. Geiping et al.

Fewer parameters do not make a looped model faster. Six passes through a block require roughly six times that block’s arithmetic, and later passes depend on earlier ones. Total memory also includes intermediate states and cached attention data, not only the stored weights.

Choosing a loop count

A model can use a fixed number of passes or a stopping rule. The Universal Transformer’s adaptive version learns when to stop processing each position, allowing different positions to receive different amounts of computation. Dehghani et al.

Useful loop counts have to be measured for the trained model and task. Additional passes do not necessarily improve accuracy, and a stable output does not mean a correct one: the model can repeat the same wrong answer. Different hidden states can also produce the same answer, so an unchanged output need not imply an unchanged state.

For AlphaFold, more recycles, different input features, and additional sampled predictions are separate choices. The subset of aligned sequences can also be resampled between recycles. EMBL-EBI’s guide and the AlphaFold 2 code describe these options.

How the loops are trained

Training depths and gradients

The repeated block is trained on states produced by earlier passes, so it learns to process its own outputs. An arbitrary pretrained layer need not handle those states well.

Backpropagation computes how the prediction error depends on the weights. Tracing this calculation through every pass can require substantial memory.

AlphaFold 2 samples a training pass count and applies the loss at the final pass. It uses stop-gradient between recycles: earlier outputs become fixed inputs to the final pass, and the derivative calculation stops at that boundary. The shared weights still learn through the final pass, but their effect through the earlier passes is omitted from the gradient. Supplementary methods

Geiping’s model samples varying depths and limits backpropagation through the recurrent block to the last eight loops. This retains a training signal through several updates while limiting memory. These are training procedures; ordinary inference uses the resulting weights without updating them. Geiping et al.

Coconut and RFdiffusion

Coconut feeds a language model’s last hidden state directly into the next sequence position, bypassing the conversion to a text token. It adds latent positions to the sequence, whereas depth recurrence revisits existing positions. Both allow intermediate computation without writing it out. Hao et al.

RFdiffusion uses self-conditioning: the previous estimate of a clean protein structure helps condition the next denoising step. The authors explicitly connect this choice to AlphaFold recycling. Here, prediction feedback operates inside a diffusion process. Watson et al.

Sources