Universal and Efficient Computation with 2D Attention

Transformers are powerful models but are inefficient with large context due to the quadratic complexity of the attention mechanism. Limiting their context window addresses the efficiency concern but at a cost: they cannot remember everything from the past, which severely limits their computational power. We show that we can simultaneously achieve efficiency and universal computation if we restrict attention heads to be 2-dimensional. First, we give the first algorithm for standard 2D attention that serves every insertion and every query in amortized polylogarithmic time per token. Prior work either restricted key/query magnitudes or paid polynomial overhead per step. Second, we show that the class of 2D-attention models is universal and efficient: a vanilla decoder transformer efficiently simulates an arbitrary RAM machine via chain-of-thought decoding so that, combined with our first result, the simulation runs in $\mathrm{polylog}(N)$ time per simulated RAM step end-to-end. In contrast, earlier universality results addressed only computability, showing Turing completeness while incurring large polynomial factors in their execution time. Third, as an application, we turn a transformer into a computer. We provide a compiler that supports the full $\texttt{i32}$ WebAssembly subset with integer computations (reachable from C via $\texttt{clang}$) and turns such programs into the weights of a transformer, which then executes them over millions of steps autoregressively. We illustrate how this can be used to solve combinatorial problems, tasks at which even large reasoning LLMs struggle.