Universal Transformers for Circuit Computations: Perfect Length Generalization in Tiny Transformers
2026-08-31 • Machine Learning
Machine Learning
AI summaryⓘ
The authors created a special type of transformer neural network that can perfectly solve Boolean algebra problems of any size by treating them like circuits and evaluating them step-by-step. They use a new way to represent the position of elements in these expressions so the model can know which parts to solve at each step. Training on simple, small problems helps the model learn exactly how to solve bigger, more complex ones without losing accuracy. They prove their method works perfectly for Boolean expressions and also show it can handle other tasks like modular arithmetic and ListOps with 100% accuracy.
transformerBoolean algebraalgorithmic generalizationpositional encodingmasked attentionlinear attentiondepth generalizationmodular arithmeticListOpssymbolic computation
Authors
Takuya Ito, Ruchir Puri, Murray Campbell, Parikshit Ram
Abstract
Learning generalizable algorithmic computations remains a challenge for neural networks, as reflected in persistent failures on compositional and length generalization benchmarks. We present a provably correct, transformer parameterization (with only 280 learnable parameters for Boolean algebra tasks) capable of learning and evaluating problems of any depth or length. We assume inputs are fully parenthesized, well-formed expressions. Our approach conceptualizes algorithmic tasks as circuit models embedded in transformers, enabling depth-1 circuit reduction in a single forward pass. To achieve depth generalization, we introduce a positional encoding that tracks each gate's depth within the circuit, enabling the model to identify evaluable subexpressions at each iteration via masked hard attention, with $O(n)$ per-iteration complexity via linear attention. Combined with an autonomous halting criterion, the model terminates after $d$ iterations for problems of depth $d$, yielding $O(n \cdot d)$ total complexity. We show that training on shallow problem instances (depth 1 and depth 2) effectively recovers interpretable parameters that {\em snap} into place, resulting in exact length generalization. Though we establish that our construction provably evaluates Boolean expressions -- a universal symbolic computation -- of arbitrary length perfectly, in other experiments we also demonstrate that our transformer variant can learn and generalize perfectly (100% accuracy) on other common length generalization benchmarks, including modular arithmetic and ListOps.