Learning to compose

What should one layer be able to do?

Attention compares tokens in pairs. Some relations only appear when several tokens come together. Higher-order attention gives a layer that operation directly—and wins on puzzles built around it.

Language did not reward more order at our scale. The useful turn was composition: learning how small operations combine. Clock hands add without caring about order. Rotations can remember it.

Given only answers to a relational task, a layer learns rotations that nearly obey the symmetries of an icosahedron. The route from token pairs to those rotations is the story here.

One article with live figures, and two labs where you can open up and train the models.

Read the paper Open a trained layer

Scores for whole triples

A standard head fills a matrix, one score per query and key. A higher-order head fills a tensor: every query against every pair of keys, so one step can weigh how two tokens bear on a third.

A_{ij} = ⟨ Q_{i}, K_{j} ⟩ ⟶ A_{ij k} = ⟨ Q_{i}, K_{j}^{1} ⊙ K_{k}^{2} ⟩

Attention on hyperedges

The same scores as a graph. Weight lands on triangles, not edges: the query reads what a group of tokens holds together, so a relation that exists only in the group takes one step.

y_{i} = \sum_{j, k} softmax_{j k} (A_{ij k}) V_{j}^{1} ⊙ V_{k}^{2}

Relations as composition

A token becomes a rotation. A layer multiplies rotations in sequence order, then reads relations among subsets. The learned heads approximate the symmetries shown here.

More per layer

On these three-way relations, cubic attention reaches 100.0% held-out accuracy. The softmax baselines stay behind within this training budget. Held-out accuracy, trained in the browser.