Attention Is Only as Good as Its Witnesses: When Higher-Order Layers Help, and When They Do Not

Each layer we study computes the partition function of a family of witnesses, the objects a task’s answer sums over (tokens, pairs, walks, subsets), and reads it at a temperature.

Mostafa ‘Neo’ Mohsenvand · MIT Media Lab

A 250-vertex test question of the search benchmark (lookahead 31), lit by the walks to the goal of one propagator layer trained on graphs of 49 vertices: the walks are added one length at a time, then all at once by the solve, and the layer names the road from the start. Run it yourself.

Abstract

Transformers stop learning graph search as graphs grow, and learn the card game SET only after many distinct questions, although attention can represent both. We read these failures, and the uneven record of higher-order attention, from what a layer sums over: each layer we study computes the partition function of a family of witnesses and reads it at a temperature. We prove two laws. By the size law, existence is read correctly at every size, in the worst case, if and only if inverse temperature times margin is at least the logarithm of the number of candidates. By the depth law, $L$ layers over two-step walks reach exactly $2^{L} - 1$ steps: a construction reaches that far, and from anonymous initial states no such network reaches further, for every $L$ . Layers built on these laws, and on two principles that we test, learn where Transformers fail. One propagator layer trained on 49-vertex graphs answers every question at every search depth on 250-vertex graphs. Given the modulus 3, exponential Newton computes the softmax of 2-simplicial attention over pairs exactly in $3^{d} n$ operations for $d$ clocks; with six clocks it learns SET from as few deals as 2-simplicial attention and transfers from 8- to 20-card tables. On CLRS-30 without hints, a 186,000-parameter block of OR-join and power-mean channels matches (within half a point) or beats a model of at least 73 million parameters on 9 of 12 graph tasks of the official test set, and on 8 of 12 with freshly drawn test graphs. The same account explains where pair layers did not help and where parallel layers must loop.

Fig. 1 | **What a layer sums over.** Hand-set layers, nothing trained, on one 49-vertex search question (top; towns are vertices, roads directed edges) and one eight-card SET deal (bottom). a, A lookup sums over single witnesses: the start weighs every town by its own score (circle area, softmax weight: 0.100 for each of the eight children, 0.005 for every other town), and the card in hand weighs each table card 1/8. b, A join sums over pairs: the two-step walks start → $k$ → goal, each step up to four roads, through the three middle towns $k$ within four roads of both ends, and the 28 pairs of table cards, the four set pairs shaded. c, A closed form sums over every witness at once: the propagator $X = S (I - G)^{- 1}$ from the goal ( $t_{0} = 0.8$ ) lights only the route that reaches it, with weight $t_{0}^{D}$ at $D$ roads; Newton attention sums the cards’ phases and their squares, and $e_{2} = (p_{1}^{2} - p_{2}) /2$ gives the sum of the 28 pair terms.

What we found

Search

One propagator layer trained on 49-vertex graphs answers every question at every search depth on 250-vertex graphs.

Fig. 2 in the paper

SET

Given the modulus 3, exponential Newton computes the softmax of 2-simplicial attention over pairs exactly in $3^{d} n$ operations for $d$ clocks; with six clocks it learns SET from as few deals as 2-simplicial attention and transfers from 8- to 20-card tables.

Fig. 3 in the paper

The two laws

By the size law, existence is read correctly at every size, in the worst case, if and only if inverse temperature times margin is at least the logarithm of the number of candidates. By the depth law, $L$ layers over two-step walks reach exactly $2^{L} - 1$ steps: a construction reaches that far, and from anonymous initial states no such network reaches further, for every $L$ .

Fig. 4 in the paper

CLRS-30 and elsewhere

On CLRS-30 without hints, a 186,000-parameter block of OR-join and power-mean channels matches (within half a point) or beats a model of at least 73 million parameters on 9 of 12 graph tasks of the official test set, and on 8 of 12 with freshly drawn test graphs.

Fig. 5 in the paper