PaperlayerView exact revision on arXiv

Source-linked AI summary

Looped Transformers as Programmable Computers

Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D. Lee, Dimitris Papailiopoulos

arXiv:2301.13196v1cs.LGcs.AI

TL;DR

The paper addresses whether shallow transformers can serve as programmable computers for general programs and iterative algorithms. It hardcodes computing primitives and instruction-set operations into a transformer whose output is looped back as input, then demonstrates broad algorithmic emulation. The resulting constructions include a basic computer, calculator, linear-algebra routines, and in-context backpropagation, while their width and approximation requirements impose scope boundaries.

  • Problem

    Existing constructions show transformer computation or specialized programs, but the paper investigates a programmable general-purpose computer that can execute iterative algorithms with shallow depth.

  • Method

    The paper hardcodes attention-based computing blocks and a flexible SUBLEQ-like instruction called FLEQ into fixed-weight transformers connected by a single recursive loop.

  • Results

    The constructions use constant-depth looped transformers to emulate a one-instruction computer, a calculator, numerical linear algebra, and in-context learning with backpropagation or SGD.

  • Takeaways & Limitations

    The results demonstrate that attention mechanisms and a single loop can support programmable transformer models for complex iterative algorithms and general programs.

  • Takeaways & Limitations

    The constructions’ width scales logarithmically with program and memory length, while arbitrary loss-function approximation can require polynomial width and specific derivative forms avoid that error.

Abstract

from arXiv · show

We present a framework for using transformer networks as universal computers by programming them with specific weights and placing them in a loop. Our input sequence acts as a punchcard, consisting of instructions and memory for data read/writes. We demonstrate that a constant number of encoder layers can emulate basic computing blocks, including embedding edit operations, non-linear functions, function calls, program counters, and conditional branches. Using these building blocks, we emulate a small instruction-set computer. This allows us to map iterative algorithms to programs that can be executed by a looped, 13-layer transformer. We show how this transformer, instructed by its input, can emulate a basic calculator, a basic linear algebra library, and in-context learning algorithms that employ backpropagation. Our work highlights the versatility of the attention mechanism, and demonstrates that even shallow transformers can execute full-fledged, general-purpose programs.

1 Introduction

The paper proposes programming transformers with fixed weights and a loop so they can execute general-purpose programs and iterative algorithms. Its framework uses instruction-set-like operations to emulate computing systems and demonstrates calculator, linear-algebra, and in-context-learning capabilities.

  • Transformers can simulate complex algorithms and programs when hardcoded with specific weights and placed in a loop.
  • The framework reverse-engineers attention to implement edit operations, nonlinear functions, function calls, program counters, and conditional branches.
  • SUBLEQ provides a universal one-instruction computer by subtracting one memory value from another and branching when the result is nonpositive.
  • FLEQ generalizes SUBLEQ by allowing hardcoded functions such as matrix multiplication, nonlinear functions, and polynomials, while loop depth depends on one instruction rather than program length.
  • The constructions emulate a basic calculator, numerical linear-algebra methods, and in-context learning with stochastic-gradient-descent algorithms using fewer than 13 layers.
  • The proposed framework highlights attention’s flexibility and the role of a single loop in executing complex iterative algorithms, while cautioning that its construction differs from trained language models.

2 Prior Work

Prior work establishes transformers’ computational expressivity, programmable abstractions, and ability to implement learning algorithms in context. The paper positions its approach as combining universal-computer functionality with iterative algorithms whose transformer depth need not grow with program length.

  • Earlier constructions show that transformers can simulate Turing machines, but typically rely on sufficient depth, recursive attention links, or high precision.
  • RASP maps selection and aggregation programs to transformer networks for tasks such as counting, sorting, histograms, and Dyck-k recognition.
  • RASP-based approaches typically scale in model size with program size, while their expressivity and Turing completeness remain discussed limitations.
  • The paper’s framework instead targets arbitrary nonlinear functions and iterative nonlinear algorithms without requiring transformer depth to scale with executed lines of code.
  • Other work constructs or trains transformers to emulate in-context learning algorithms, including gradient descent or SGD, mainly for restricted model and loss-function settings.

3 Preliminaries

The looped-transformer setup represents inputs as embedding matrices partitioned into commands, memory, and scratchpad regions. Repeated application of fixed transformer layers feeds outputs back as inputs for iterative computation.

  • The model processes d-dimensional embedding vectors arranged as columns rather than relying on token semantics.
  • A multilayer transformer is recursively applied by feeding its output sequence back as input, enabling iterative computation through a fixed-point-like loop.
  • The CPU analogy describes continual instruction processing and memory reads/writes, but the paper warns that implementation differences limit the analogy.
  • The input sequence is partitioned into scratchpad, memory, and command regions, with positional encodings acting as pointers to data and instructions.
  • The scratchpad records computation inputs and outputs as temporary workspace for copying, transforming, and manipulating data.
  • Memory stores scalars, vectors, and matrices that are copied to the scratchpad for updates before results are returned.

4 Building Transformer Blocks towards General Computation

The construction organizes transformer computation around interoperable blocks for editing sequences, tracking pointers and program counters, and moving data through a scratchpad. Binary positional encodings support addressing, while shallow transformer components implement read/write operations and counter updates.

  • Transformer layers provide interoperable building blocks for sequence edits, approximate functions, program counters, and conditional computation.These operations are assembled into more complex routines and algorithms.
  • Binary positional encodings represent column indices with logarithmic-dimensional ±1 vectors, allowing tokens to be addressed by location.The encoding dimension is logarithmic in sequence length, and distinct columns differ in at least one coordinate.
  • A program counter stores the next command location, while data pointers identify locations to read or write using the same positional encodings.A ReLU feedforward network can increment binary counters when the sum remains within range.
  • Read/write operations: The scratchpad acts as temporary computational workspace where commands and data are copied, manipulated, and prepared for subsequent operations.The input layout separates scratchpad, memory, and commands, with positional information organizing these regions.
  • Read/write operations: A one-layer, one-head transformer of width O(log n + d) can write a scratchpad vector to a designated input location, while the analogous read operation copies pointed-to data into the scratchpad.The read construction uses positional pointers and copies data or commands into scratchpad rows for further computation.
  • Conditional branching: The conditional-branching construction requires 2 transformer layers because the full operation uses 3 feedforward layers while the attention layer is prevented from modifying the input.The attention value matrix is set to zero to neutralize its direct input modification.

5 Emulating a Generalized One-instruction Set Computer

The paper constructs looped transformers that execute one-instruction computer programs, using attention-based read/write operations, scratchpad computation, conditional branching, and program counters. It generalizes SUBLEQ to FLEQ, supporting multiple hardcoded functions while keeping transformer depth independent of program length.

  • 5.1 A SUBLEQ Transformer: SUBLEQ accesses two memory locations, stores their difference in one location, and conditionally jumps based on whether the result is nonpositive.
  • 5.1 A SUBLEQ Transformer: The construction restricts SUBLEQ by separating memory or registers from instructions, while the paper states that this version remains Turing complete.
  • 5.1 A SUBLEQ Transformer: The SUBLEQ implementation uses a nine-layer, two-head architecture with width O(log(n)+N), where n scales with program and memory length and N stores integer bits.
  • 5.1 A SUBLEQ Transformer: The external loop is crucial for sequential execution and prevents transformer size from scaling with the number of program lines, although total runtime still scales with program length.
  • 5.1 A SUBLEQ Transformer: The looped transformer tracks instructions, memory, and a program counter while using scratchpad operations for pointers, intermediate values, reads, writes, and branching.
  • 5.1 A SUBLEQ Transformer: The input separates scratchpad, memory, and instruction columns, with positional encodings and binary pointers identifying data, commands, and the next instruction.
  • 5.1 A SUBLEQ Transformer: Executing SUBLEQ proceeds by reading the current instruction, reading operands, subtracting the addressed values, storing the result, and selecting a branch or the next instruction.
  • 5.2 FLEQ: A More Flexible Attention-based Computer: FLEQ generalizes SUBLEQ by allowing each instruction to apply one of M predefined functions, and a recurrent transformer can execute programs whose instruction count exceeds its fixed depth.

6 Functions in the Unified Template Form

The section develops transformer-based function blocks within a unified template, encoding nonlinear functions and matrix operations with shallow attention layers. These constructions provide constant-depth primitives for broader computation.

  • Unified template: Function blocks in the unified template use consistent input/output locations and can be combined to construct iterative algorithms.Each block may use positional encodings to transfer outputs between locations.
  • 6.1 Encoding Non-linear Functions within the Attention Mechanism: Attention encodes nonlinear functions by forcing softmax to act as a sigmoid and storing coefficients in query and value matrices.An input indicator vector selects among encoded linear combinations of sigmoids.
  • 6.1 Encoding Non-linear Functions within the Attention Mechanism: 3 layers, m heads, and dimensionality r = 2 log(d) + d + 1 = O(d) suffice for a function block encoding N selectable functions.The selector e = e_j ∈ R^N identifies the function to execute.
  • 6.1 Encoding Non-linear Functions within the Attention Mechanism: Arbitrary functions g ∈ ΓC,B can be well approximated by attention layers using the encoded sigmoid representation.The construction ignores inconsequential output values downstream and can be zero-padded when additional scratchpad space is required.
  • 6.2 Matrix Transposition and Multiplication by Linearizing the Softmax: Matrix representations are linearized into sequences, while fixed column permutations implement transposition in the unified template.Vectorizing a matrix supports transpose construction through a fixed permutation of columns.
  • 6.2 Matrix Transposition and Multiplication by Linearizing the Softmax: 4 layers, 1 head, and dimensionality r = 2d + 2 log d = O(d) implement matrix transposition with an approximation A′ = A⊤ + ϵM, where ∥M∥ ≤ 1.The error depends on the temperature used in the read/write operations.
  • 6.2 Matrix Transposition and Multiplication by Linearizing the Softmax: 2 layers, 1 head, and dimensionality r = O(d) implement A⊤B + ϵM for matrices A ∈ R^{k×m} and B ∈ R^{k×n}, with ∥M∥ ≤ 1.The softmax is linearized using a large constant, and the error ϵ is controlled by that constant.
  • 6.3 Advantage of attention over fully-connected networks: x^2 can be implemented in essentially 2 attention layers, whereas ReLU networks require depth proportional to log(1/ϵ) for comparable approximation quality.The paper supplies explicit hardcoded attention layers rather than only an existence argument.

7 A Basic Calculator

The paper assembles shallow transformer function blocks into a basic calculator. The resulting construction supports common arithmetic operations while giving explicit layer, head, range, and approximation guarantees.

  • Calculator construction: The calculator combines six function blocks for addition, subtraction, multiplication, inversion, square root, and percentage.These blocks are assembled into a program executed in the unified framework.
  • Calculator construction: 12 layers, m heads, and dimensionality O(log n) implement a calculator for addition, subtraction, multiplication, inverse, square root, and percentage.Here n denotes the number of operations.
  • Approximation conditions: The calculator’s error guarantee changes with the design choice between head-based and dimension-based approximation bounds.The paper notes that the relevant guarantees depend on the number of heads or transformer dimension.
  • Program execution: A sample program executes calculator operations sequentially by reading inputs from memory, writing intermediate results, and applying function calls.The example computes a composed expression using addition, subtraction, multiplication, inversion, and square root.
  • Arithmetic blocks: Addition and subtraction use 3-layer, 1-head, O(1)-dimensionality blocks.The constructions implement f(a,b) = a + b and f(a,b) = a − b.
  • Arithmetic blocks: Multiplication uses 2 layers, 1 head, and dimensionality O(d), while percentage uses 2 layers, 1 head, and dimensionality O(1).The corresponding functions are f(a,b) = ab and f(a) = a/100.
  • Approximate functions: Inversion uses a 3-layer block with error ˜O(1/√m) over a specified nonzero operand range.Division follows by combining multiplication with inversion.
  • Approximate functions: Square root uses a 3-layer, m-head block with error O(1/m) for inputs a ∈ [0, O(m^2)].The construction encodes square root through sigmoid approximations.

8 Linear Algebra

The framework maps matrix operations and iterative numerical algorithms into looped transformer programs. It provides constructions for matrix inversion and power iteration, with explicit shallow-network bounds and a matrix-transpose comparison.

  • Algorithmic applications: Matrix transpose and multiplication function blocks are used to execute Newton-Raphson matrix inversion and Power Iteration for dominant eigenvectors.These algorithms illustrate iterative numerical linear algebra in the unified framework.
  • Comparison: For 30 × 30 matrix transposition, the construction uses 1 layer, 1 head, and width 168, compared with Charton’s 1 layer, 8 heads, and width 256.The reported widths reflect the specific transposition setting and encoding construction.
  • Scope: The paper leaves implementation and evaluation of errors as future work, including dependence on proof constants and softmax temperature.The stated future evaluation concerns errors involved in the constructions.
  • Matrix inversion: The matrix-inversion program initializes X−T = ϵA and iterates X_{i+1} = X_i(2I − AX_i).The transformer program implements multiplication, subtraction, transpose, counters, and conditional looping.
  • Matrix inversion: 13 layers, 1 head, and dimensionality r = O(d) emulate Newton’s matrix-inversion algorithm for any ϵ > 0.The implementation translates each algorithmic step into unified-framework commands.
  • Power Iteration: Power Iteration repeatedly multiplies an initial vector by A, then normalizes the final vector using a Newton iteration for inverse square root.The pseudocode specifies b_{k+1} = Ab_k followed by normalization.
  • Scope: The framework is presented as a proof of concept for small linear algebra libraries, with additional iterative algorithms requiring incorporation of further functions.Examples named include QR decomposition, Gauss-Seidel, Arnoldi iteration, and Lanczos methods.

9 Emulating Learning Algorithms at Inference Time

The framework uses looped transformers to execute learning algorithms at inference time, including SGD and backpropagation, with constant-depth constructions whose repeated cycles implement iterative updates. It supports linear models and two-layer neural networks, while deeper networks and arbitrary loss functions introduce additional costs or approximation requirements.

  • Backpropagation and SGD: The framework calls backpropagation as a function at each iteration, allowing the transformer to update implicit model weights in context.Different data points are supplied through changing pointers and repeated program execution.
  • Stochastic Gradient Descent in linear models: The SGD program iterates over D feature-label pairs, updates the weight vector, and resets pointers after each pass.The step size is supplied as an input parameter, while program and data pointers control the repeated passes.
  • Stochastic Gradient Descent in linear models: Constant-depth looped transformers can implement T iterations of SGD on linear models with error up to ϵ.The construction uses 13 layers, one head, dimensionality O(log(D) + d), and takes the step size as input.
  • Loss functions: Arbitrary loss functions generally require poly(TD) attention heads to control approximation error, whereas sums of sigmoids avoid this additional error.The extra error arises from approximating the loss derivative through Barron’s theorem.
  • Backpropagation and SGD: Backpropagation computes hidden activations, outputs, error terms, and gradients for the two-layer network before applying SGD updates.The algorithm uses z = W1x + b1, a = σ(z), o = W2a + b2, and gradients for W2 and W1.
  • Backpropagation and SGD: 13-layer looped transformers can implement T iterations of SGD on two-layer neural networks with error up to ϵ.The construction uses one head and dimensionality O(log(D) + d) over D data points, with the step size given as a parameter.
  • Generalizing to arbitrary depth: Generalizing to arbitrary-depth networks requires loop cycles proportional to network depth, increasing training cost as depth grows.The code length and number of transformer cycles scale with the number of network layers.

10 Conclusion and Open Problems

The paper concludes that programmed, looped transformers can act as general-purpose computers and execute calculators, linear algebra, and backpropagation-based in-context learning. It identifies natural-language control, model fusion, distillation, smaller-network validation, and easier-to-train architectures as open directions.

  • Conclusion: Constant-depth transformer networks can emulate computing blocks and execute iterative algorithms through a loop.The paper constructs a one-instruction set computer whose programs run on the transformer network.
  • Conclusion: The resulting systems emulate a basic calculator, a basic linear algebra library, and a full backpropagation in-context learning algorithm.These results are presented as evidence that transformers can serve as programmable compute units.
  • Conclusion: A single loop is presented as sufficient to support models that mimic complex iterative algorithms and execute general programs.The conclusion links this capability to the flexibility of attention and mathematical and algorithmic task execution.
  • Open Problems: Open problems include fusing hardcoded models with pretrained transformers and tokenizing natural-language commands for input control.The paper also proposes distillation, experimental validation with smaller networks, and architecture changes that ease implementation and training.

A Ommited proofs

The omitted proofs establish transformer-based blocks for arithmetic, function selection, and routing data through scratchpad-like representations. These constructions support implementing nonlinear functions and moving selected values across structured input blocks.

  • Binary addition: A 1-hidden-layer ReLU network with 8d hidden activations computes the sum of two d-dimensional binary integers when the sum is below 2^(d+1).The construction computes each output bit from the lower-order input bits and uses six neurons per bit before accounting for residual subtraction.
  • Binary addition: The binary-addition proof works for both {0,1}^d and {±1}^d representations because the conversion is affine.The proof is presented in the {0,1}^d representation, while the paper notes that affine transformations transfer it to the alternative encoding.
  • Nonlinear functions: A transformer-based function block uses an indicator vector to select one function from a set represented by sums of sigmoids.The construction encodes the selected function through one-hot indicators and implements the sigmoid terms with attention and feedforward layers.
  • Routing and control: Attention mechanisms can control which input positions are read by using embedding dimension and encoded indicators.The passage explicitly identifies the attention mechanism’s size, controlled by embedding dimension d, as another way to control the input.
  • Nonlinear functions: The function-selection construction combines attention and feedforward layers to place the selected nonlinear-function output in the required representation.The proof moves inputs between row blocks, uses indicator bits to control column selection, and applies a final layer to produce the desired output.

A.3 Matrix Transposition

The matrix-transposition construction vectorizes a matrix, applies a fixed permutation to obtain the vectorized transpose, and matricizes the result. The paper gives a four-layer transformer block realizing this operation with linear dimensionality.

  • Construction: The construction first makes d copies of A and vectorizes the matrix into a row block using positional encodings.The first layer copies the matrix, and the second-layer feedforward component places entries into a vectorized representation.
  • Transposition: A fixed permutation of the vectorized columns converts the representation of A into the vectorized form of A^T.The proof describes transposition as a column permutation after vectorization, followed by reconstruction of the matrix format.
  • Transposition: Attention and feedforward layers then copy the permuted values back into the desired matrix layout.The stated intermediate outputs show the transposed entries being assembled in the second row block and subsequently moved using the positional encodings.

A.4 Matrix Multiplication by Linearizing the Softmax

The matrix-multiplication construction linearizes softmax attention to compute products through attention weights and residual/feedforward transformations. It yields several approximate matrix products with two-layer, one-head blocks and O(d) dimensionality.

  • Corollaries: The same construction gives B^T A + ϵM, B^T B + ϵM, and A^T A + ϵM under the corresponding input arrangements.The latter three results are stated as corollaries of the first multiplication result.
  • Construction: The construction pads matrices to fit a 2d × 2d representation and uses fixed identity and all-ones components in the input design.The proof describes these components as part of the fixed input structure, while encoding rows are reserved for copying outputs to target locations.
  • Error control: The attention mechanism forms weighted sums whose approximation error is controlled by softmax scaling and subsequent feedforward transformations.The proof selects constants controlling the attention concentration and then copies the computed product to its desired location.

B Error Analysis

The error analysis tracks approximation errors from bounded inputs, read/write operations, matrix multiplication, and function approximation. It shows how per-operation errors accumulate across iterative computation and can be controlled by choosing construction parameters.

  • Assumptions: The analysis assumes every input-matrix element is bounded by a constant G, with |v_i| ≤ G.This bounded-input condition is used when setting constants for the matrix-multiplication error.
  • Read/write error: Read/write errors arise from imperfect positional attention, whose off-target weights are bounded by e^(log n−λ).The positional encodings make the target position more aligned than other positions, while λ controls the temperature and concentration of softmax attention.
  • Matrix multiplication error: Choosing constants so that ϵ_2 = ϵ/4 and ϵ_1 = C_1 log(nϵ) makes the total matrix-multiplication error less than ϵ.The passage states this parameter choice as the final bound for the multiplication error.
  • Function approximation: Function-approximation error is proportional to 1/√m, where m is the number of attention heads used.The passage attributes this rate directly to the approximation result used for the function block.
  • Accumulation: After T operations, the tracked error can grow with iteration index as 0 < ϵ_t+1 < (t + 1)ϵ/T.The analysis represents the accumulated deviation with a matrix M_t+1 satisfying ∥M_t+1∥ ≤ 1.

C subleq is Turing Complete

The restricted SUBLEQ instruction is shown to be Turing complete, and a looped transformer implements its execution through structured input, scratchpad memory, instruction decoding, function blocks, and branching.

  • C subleq is Turing Complete: The restricted SUBLEQ instruction is Turing complete, established by translating Turing-complete Minsky-machine operations into SUBLEQ programs.The translation uses fixed memory locations containing −1, 0, and +1, plus an extra register.
  • C subleq is Turing Complete: Minsky-machine add(a) translates to one SUBLEQ instruction, while sub(a,n) translates to a sequence of five SUBLEQ instructions.The translated subtraction sequence branches according to whether mem[a] is positive, negative, or zero.
  • Transformer Instruction Execution: Each transformer instruction encodes operand, function, flag, and branch locations, then computes mem[c] = f_m(mem[a], mem[b]) with conditional control flow.The encoded pointers are binary ±1 vectors, while the function index selects among available function blocks.
  • Transformer Instruction Execution: Each execution iteration reads the next command, copies operands into scratchpad memory, dispatches computation to a function block, writes results back, and updates the program counter.The branch condition stores the target pointer or advances to the next instruction.
  • Transformer Instruction Execution: The input X is partitioned into Scratchpad, Memory, and Instructions blocks, providing storage for variables, current execution state, and encoded commands.Variables may be scalars, vectors, or matrices, and the scratchpad includes the program counter, current instruction, temporary memory, encodings, and function blocks.
Loading 2301.13196v1…