PaperlayerView exact revision on arXiv

Source-linked AI summary

Universal Transformers for Circuit Computations: Perfect Length Generalization in Tiny Transformers

Takuya Ito, Ruchir Puri, Murray Campbell, Parikshit Ram

arXiv:2608.31067v1cs.LG

TL;DR

Transformers often struggle to execute exact symbolic computations that generalize to longer inputs. This paper embeds circuit reduction in a compact transformer with depth-tracking positional encodings, proving arbitrary-depth Boolean evaluation and demonstrating perfect generalization on several compositional tasks. Its results rely on fully parenthesized, well-formed inputs, where reducible structure is explicit.

  • Problem

    Transformers often struggle with symbolic and length generalization, and strong benchmark performance does not guarantee a generalizable algorithm.

  • Method

    The paper embeds iterative circuit reduction in a small transformer using positional encoding for circuit depth, masked attention, and local semantic evaluation.

  • Results

    The construction provably evaluates arbitrary-depth Boolean expressions, while the architecture learns and perfectly generalizes across Boolean algebra, modular arithmetic, and ListOps.

  • Takeaways & Limitations

    The results provide an interpretable, efficient, compact transformer construction that implements exact circuit computation and can be learned from shallow data.

  • Takeaways & Limitations

    The construction assumes fully parenthesized, well-formed expressions, so its theorems exclude cases where structure must be inferred implicitly.

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses symbolic and length generalization by embedding circuit algorithms in a small transformer, proving exact Boolean evaluation and demonstrating perfect generalization across compositional tasks.

  • Standard transformers can perform well on benchmarks without learning a generalizable algorithm for longer symbolic inputs.
  • The authors parameterize a 280-parameter transformer that provably evaluates arbitrary-depth Boolean expressions through iterative circuit reduction.
  • A positional encoding tracks circuit depth and identifies evaluable subexpressions so the model can reduce expressions iteratively.
  • Circuits represent expressions as structured graphs whose operator gates are updated from computed subexpressions until the output is obtained.
  • The same architecture learns and perfectly generalizes from shallow data to modular arithmetic and ListOps, in addition to Boolean algebra.

2 Preliminaries

The paper frames Boolean expressions as circuit computations and evaluates them through iterative local reductions of deepest reducible subexpressions.

  • Boolean expressions provide a compact setting for universal symbolic computation, with a seven-dimensional one-hot vocabulary encoding.
  • Each input token is represented by a one-hot vector in a vocabulary-sized embedding space.
  • The model state X^(t) denotes the token sequence after t model iterations.
  • Expression depth is recursively defined, and a reducible subexpression is an innermost parenthesized span containing only an operator and its operands.
  • Each iteration identifies deepest reducible spans, routes information within them, computes local semantics, and preserves unreduced tokens through gated residuals.

3 Model Construction and Formalization

The construction embeds circuit reduction in a small transformer: positional encoding identifies deepest reducible spans, attention routes within independent subexpressions, and an MLP evaluates them locally. Iterating this depth-1 reduction yields exact arbitrary-depth evaluation with autonomous halting and total complexity O(n·d(x)).

  • 3 Model Construction and Formalization: The model is an exact, hand-parameterized transformer construction for Boolean circuit reduction, with 280 parameters and no training required for correctness.The construction targets valid Boolean expressions and provides a white-box parameterization for arbitrary depth.
  • 3.1 Positional encoding: Identifying a reduction mask and parallel subexpressions: Positional encoding tracks parenthesis depth to identify tokens belonging to the deepest unreduced subexpressions.The gating matrix assigns open parentheses +1, closed parentheses −1, and ignores other tokens before cumulative depth computations identify maximal depth.
  • 3.1 Positional encoding: Identifying a reduction mask and parallel subexpressions: The reduction mask is decomposed into contiguous clusters, each corresponding to one locally reducible subexpression that can be evaluated independently.Cluster IDs are assigned from transitions in the mask, while disjoint spans prevent independent subexpressions from interfering.
  • 3.2 Attention Mechanism: Type-constrained token routing: Type-constrained attention restricts routing to one reducible span and maps operand information to its operator according to the Boolean grammar.The construction uses linear attention to avoid the quadratic cost of standard softmax attention.
  • 3.3 Feedforward Layer: Local semantic evaluation: The attention output is a bag-of-words representation of a local reducible expression, whose operator and operand values determine its Boolean semantic value.A single-hidden-layer MLP realizes the finite depth-1 truth-table mapping.
  • 3.5 Depth-1 circuit reduction and global circuit evaluation: One transformer iteration evaluates every reducible subexpression simultaneously, preserves inactive positions, and decreases active depth by one.Reduced interiors become zero vectors while the operator position carries the computed value forward through the gated residual.
  • 3.5 Depth-1 circuit reduction and global circuit evaluation: Repeated application of the iteration evaluates any valid Boolean expression of depth d(x), with the construction halting autonomously when no parentheses remain.The positional encoding computes the active depth used by the halting criterion.
  • 3.5 Depth-1 circuit reduction and global circuit evaluation: Each iteration has near-linear attention cost O(n), and d(x) iterations give total complexity O(n·d(x)).The typed-routing cost is O(kn) for k reducible subexpressions and is effectively O(n) for structured inputs when k ≪ n.

4 Numerical Experiments

The universal transformer learns compositional arithmetic, Boolean, and ListOps tasks from shallow examples and achieves perfect generalization to substantially deeper problems. Compared with a standard looped transformer, the proposed construction remains exact out of distribution.

  • Tasks: The experiments cover modular arithmetic, Boolean evaluation, and ListOps over fully parenthesized nested expressions.ListOps uses fan-in limited to 3 to focus on depth generalization.
  • Training: The learnable models train positional encoding, attention, and MLP components from random initialization, using a depth-1-to-depth-2 curriculum.Training uses end-to-end optimization with task-dependent parameter counts and within-epoch curriculum learning.
  • Generalization: 100% generalization accuracy was achieved on modular arithmetic, Boolean logic, and ListOps after training on depth-1 and depth-2 problems and testing depths 3–30.Table 1 evaluates 100 problems at each depth.
  • Optimization: The models converged for 10/10 modular-arithmetic seeds and 9/10 Boolean and ListOps seeds before reaching zero training loss.Evaluation was limited to depth 30 because depth-30 ListOps inputs exceeded 41k tokens.
  • Baselines: Standard looped-transformer baselines collapsed toward chance beyond depth 2, whereas the proposed construction remained exact out of distribution.The comparison uses held-out depths 2, 3, and 10 under matched training protocols, with ground-truth depth supplied to baselines.

5 Discussion and Conclusion

The paper argues that a tiny, interpretable transformer can implement circuit reduction and generalize perfectly across several compositional tasks. Its scope is bounded by explicit structural assumptions, especially fully parenthesized inputs, while its complexity remains near-linear per iteration.

  • Discussion: The proposed transformer uses hand-chosen architectural components to implement an interpretable structured algorithm for circuit reduction.The construction combines positional gating, linear attention, gated residuals, and straight-through estimators.
  • Discussion: The approach targets universal symbolic computation by representing Boolean formulas and related algorithmic tasks as circuits.Circuits expose gates, edges, and intermediate values while scaling to arbitrary sizes.
  • Discussion: The construction uses learned positional gating and structured attention routing instead of global attention with O(n^2) cost.This routing is designed to compute circuit structure on the fly rather than relying on predefined graph structure.
  • Limitations: The model assumes fully parenthesized, well-formed expressions, so it counts brackets rather than inferring latent structure.Removing brackets changes the premise of the theorems into a structure-induction problem.
  • Conclusion: The construction requires 924 arithmetic, 280 Boolean, and 4416 ListOps parameters while achieving perfect length generalization across the three tasks.Its positional encoding tracks problem depth and induces hard attention masking over reducible subexpressions.
  • Computational complexity: Each iteration has effective O(n) cost for structured inputs, and evaluation takes exactly d(x) iterations for an expression of depth d(x).Total complexity is O(n · d(x)); balanced expressions require O(log n) iterations, while maximally unbalanced expressions can reach O(n^2).

A.2 Additional proofs for model construction

The construction reduces Boolean circuits one depth level per transformer iteration, evaluating all deepest reducible subexpressions while preserving inactive positions. Repeated application therefore reaches a Boolean atom after the expression’s depth number of iterations.

  • Feedforward lookup table: A finite feedforward lookup table maps each valid reducible subexpression’s bag-of-words representation to its evaluated output embedding.The table enumerates 8 valid fan-in-2 Boolean cases; invalid inputs map to the zero vector.
  • Gated residual connection: The gated residual update preserves positions outside the reduction mask and updates only masked reducible positions.The update combines the unchanged input with the feedforward output using the reduction mask and its complement.
  • Depth-1 circuit reduction: Each transformer iteration simultaneously evaluates every deepest reducible subexpression, replaces it with its correct value, and preserves inactive positions.The proof uses depth masking, parallel decomposition, attention-based operand aggregation, and the lookup-table MLP.
  • Depth-1 circuit reduction: One iteration reduces active depth by exactly one while preserving semantic equivalence.Formally, X(t+1) ≡ X(t) and δ(X(t+1)) = δ(X(t)) −1.
  • Global circuit evaluation: After d(x) iterations, no reducible subexpression remains and the result is a Boolean atom in {0,1}.The global evaluation theorem applies the one-step reduction inductively while preserving semantic equivalence.

B.1 Task details

The experiments evaluate fully parenthesized symbolic expressions across arithmetic, Boolean logic, and ListOps. The task definitions specify recursive composition, operator sets, and output semantics, with an implementation caveat for adjacent prefix subexpressions.

  • Arithmetic: Arithmetic expressions recursively combine digits 0–9 with + and *, and outputs are numerical values modulo 10.Depth increases through nesting, causing the average token count to grow exponentially with problem depth.
  • Boolean: Boolean expressions recursively combine TRUE and FALSE leaves with unary NOT and binary AND/OR to produce the full formula’s truth value.The task uses fully parenthesized recursive formulas, such as a depth-3 expression evaluating to TRUE.
  • ListOps: ListOps uses nested operators MAX, MIN, SM, and MED over digit leaves, with each operator returning its task-defined aggregate.Examples include maximum, minimum, sum modulo 10, and median operations on nested lists.
  • Tokenization caveat: Prefix ListOps subexpressions are adjacent, so closing parentheses are tokenized as a one-hot vector followed by an adjacent zero vector.This avoids doubling tokens with whitespace tokenization but increases input-token count in proportion to closed parentheses.

B.2 Model Training

The models are trained end-to-end on symbolic tasks using iterative reduction, discrete straight-through components, and shallow-depth curricula. Training is CPU-feasible but convergence varies across random seeds and relaxation choices.

  • Training setup: Training optimizes the full multi-step reduction procedure end-to-end, with iteration counts specified by problem depth.The objective is not a surrogate loss on intermediate states.
  • Domains: The experiments cover Boolean logic, modular arithmetic, and ListOps as three symbolic evaluation domains.Each domain uses its corresponding vocabulary and compositional task definition.
  • Model components: The model uses one-hot token inputs, ternary positional-encoding gates, binary straight-through attention routing, and a single dense MLP component.The value matrix is fixed to identity, while the MLP width and activation vary by task.
  • Discrete optimization: Straight-through estimators encourage discrete symbolic states, whereas Gumbel-Softmax training was considerably less stable and did not reliably reach exact solutions.The discrete forward pass is paired with continuous gradients during backpropagation.
  • Data and optimization: Training data comprise fully enumerated depth-1 and depth-2 expressions, with per-epoch random depth-2 sampling for ListOps.A within-epoch curriculum trains depth 1 first, then combines depths 1 and 2; a two-phase learning-rate schedule addresses shallow-expression vanishing gradients.
  • Reproducibility: 10 independent seeds were trained per domain, with Boolean seed 1936 and ListOps seed 7439 failing to converge.The reported seed failures concern convergence within the fixed training setup.

B.3 Evaluation Protocol

After training on depths 1–2, models are evaluated on held-out expressions at depths 3–30 to measure depth and length generalization.

  • Generalization testing: Held-out test sets at depths 3–30 evaluate generalization beyond the depths used during training.The deeper expressions were never seen during training.

B.4 Code & Implementation Details

The experiments use supplied code and a fixed command spanning Boolean, arithmetic, and ListOps models, trained on shallow depths with specified optimization settings.

  • Implementation resources: The supplementary zip includes experiment code, a README, Python notebooks, and the trained visualization model.The visualization notebook is named viz_boolean_params.ipynb.
  • Experiment configuration: The training command supports Boolean, arithmetic, and ListOps model types.The command sets model_type to boolean, arithmetic, or listops for the respective experiments.
  • Experiment configuration: Models are trained on depths 1 and 2 using AdamW, a 0.01 learning rate, no weight decay, CPU execution, balanced batches, and binary cross-entropy loss.The command uses 500-example training batches, ten listed seeds, and learnable positional encodings, attention, and MLP components.

C Comparison to standard architectures: details

The comparisons evaluate matched looped-transformer and CRvNN baselines against the exact construction, showing that fitted higher-capacity baselines degrade out of distribution while the construction remains exact.

  • Looped-transformer baseline: The looped-transformer baseline shares weights, uses full softmax attention, and is trained on the same inputs, curriculum, and task-specific epoch budgets.Because it lacks autonomous halting, testing supplies the ground-truth depth and runs that many iterations.
  • Evaluation protocol: Baseline evaluation uses five seeds per positional encoding, selecting the best checkpoint by training accuracy.Accuracy is reported only for seeds reaching at least 0.95 on both training depths.
  • Looped-transformer results: Among fitted baseline seeds, every looped-transformer model collapses toward chance out of distribution rather than failing to fit the training depths.The authors attribute non-fitting seeds to incomplete convergence within the fixed epoch budget and avoid claiming the architecture cannot fit shallow training data.
  • CRvNN results: CRvNN fits ListOps training depths but degrades out of distribution despite using 521,611 parameters, roughly 500 times the construction's capacity.The comparison uses the original CRvNN architecture, five seeds, and the same fan-in and depth-1/depth-2 protocol.
  • Overall comparison: The construction attains 1.0 accuracy at every depth on Boolean, arithmetic, and ListOps, contrasting with the baselines' out-of-distribution degradation.Table 4 states chance levels of 0.10 for arithmetic and ListOps and 0.50 for Boolean; Table 5 reports 1.000 for the construction on ListOps.
Loading 2608.31067v1…