Source-linked AI summary
Thinking Like Transformers
Gail Weiss, Yoav Goldberg, Eran Yahav
TL;DR
Transformers lack a familiar computational abstraction comparable to finite-state machines for recurrent networks. The paper introduces RASP, a symbolic language that maps attention and feed-forward computation to sequence operations, and shows that its programs describe tasks and predict transformer architectures. Experiments found that trained transformers often achieved at least 99.5% accuracy and sometimes learned attention patterns similar to compiled RASP solutions.
Problem
Transformers lack a familiar computational model that supports clear reasoning about their architectures, trained models, and information-flow constraints.
Method
The paper introduces RASP, a symbolic sequence-processing language whose primitives represent transformer attention, aggregation, and feed-forward computation, and whose programs compile into transformer architectures.
Results
Most trained transformers reached accuracies of 99.5% and over, while some learned attention patterns strongly similar to compiled RASP solutions.
Takeaways & Limitations
RASP provides a way to reason about transformer computations symbolically and infer the minimum layers and maximum heads needed to realize a program.
Takeaways & Limitations
The approach is restricted to tasks for which a human can encode a solution, rather than systems such as strong language models or machine translation.
Abstract
from arXiv · showhide
What is the computational model behind a Transformer? Where recurrent neural networks have direct parallels in finite state machines, allowing clear discussion and thought around architecture variants or trained models, Transformers have no such familiar parallel. In this paper we aim to change that, proposing a computational model for the transformer-encoder in the form of a programming language. We map the basic components of a transformer-encoder -- attention and feed-forward computation -- into simple primitives, around which we form a programming language: the Restricted Access Sequence Processing Language (RASP). We show how RASP can be used to program solutions to tasks that could conceivably be learned by a Transformer, and how a Transformer can be trained to mimic a RASP solution. In particular, we provide RASP programs for histograms, sorting, and Dyck-languages. We further use our model to relate their difficulty in terms of the number of required layers and attention heads: analyzing a RASP program implies a maximum number of heads and layers necessary to encode a task in a transformer. Finally, we see how insights gained from our abstraction might be used to explain phenomena seen in recent works.
1. Introduction
The paper introduces RASP, a programming-language abstraction for transformer computation that makes information flow explicit through symbolic sequence operations. It uses RASP to express transformer solutions, compile them into architectures, and relate task difficulty to layers and attention heads.
- Computational model: RASP models transformer computation as a simple language capturing the architecture’s information-flow constraints.It abstracts away neural-network details while preserving the constraints under which transformer encoders process sequences.
- Computational results: RASP fully recognizes Dyck-k languages for all k, extending an earlier transformer construction for a simplified Dyck-k variant.This is presented as a computational result obtained by expressing and improving the earlier construction in RASP.
- Computational model: RASP expresses transformer solutions as symbolic programs rather than hand-crafted neural-network constructions.The paper presents this as a way to communicate solutions in a few lines of code and reason at a higher level of abstraction.
- Empirical realization: A 2-layer, 2-head transformer reached 99.9% test accuracy on double-histogram while receiving target and attention supervision.Its learned attention heatmaps suggested replication of the RASP solution’s selection patterns.
- Compilation and analysis: RASP programs compile into transformer architectures whose layers and attention heads can be inferred from the program’s computation flow.The paper demonstrates compilation for tasks including double-histogram and other sequence-processing problems.
2. Overview
RASP represents transformer computations over fixed-length sequences and n × n selection matrices using elementwise, selection, and aggregation operations. Its examples show how symbolic programs describe histogram and Shuffle-Dyck solutions while compiling into bounded transformer architectures.
- Computational model: A RASP computation over length-n input manipulates only sequences of length n and matrices of size n × n.Its inputs are tokens and indices, and its output is a sequence whose consumer may inspect selected locations.
- Core operations: Elementwise operations transform sequence values, while select and aggregate combine information across positions.Select creates a matrix from pairwise predicates; aggregate combines values selected in each row.
- Core operations: Aggregate operations are the only RASP mechanism for moving or combining values from different sequence positions.They can filter values through a selector and broadcast a selected value across a sequence.
- The RASP language: RASP programs are lazy functional compositions of sequence operators and selectors, including select, aggregate, selector_width, and elementwise operators.These primitives correspond to attention and feed-forward computation in a transformer.
- Examples: The double-histogram program counts each token’s frequency and then counts how many unique tokens share that frequency.The Shuffle-Dyck example independently balances multiple parenthesis types without requiring their opening and closing symbols to be ordered relative to one another.
- Examples: The Shuffle-Dyck-2 RASP solution compiles to 2 layers and 3 total heads, matching the cited hand-crafted transformer construction.RASP infers these architectural requirements without requiring the programmer to specify such implementation details.
- Scope and constraints: RASP excludes input-dependent loops and restricts selector decisions to pairwise position comparisons, reflecting transformer information-flow constraints.The model therefore prevents programs from expressing computations beyond the supported transformer setting.
3. The RASP language
RASP represents transformer computation as compositions of sequence operators, with primitives corresponding to attention-based selection and aggregation and feed-forward elementwise operations. Its compilation traces these dependencies into transformer heads and layers, while abstracting away concrete neural-network details.
- Sequence operators: RASP programs compose same-length sequence operators, constants, and elementwise operations to compute target tasks.Built-in operators include tokens, indices, and length; elementwise combinations can include arithmetic, comparisons, and conditional expressions.
- Selection and aggregation: Selectors compare keys and queries to form selection matrices, while aggregate averages selected values for each output position.Together, select and aggregate combine information across sequence positions in a two-dimensional map-reduce pattern.
- Selector operations: RASP exposes selector_width to count how many input values each selector chooses, enabling operations such as sequence histograms.Selectors can also be reused across multiple aggregations, making them first-class objects in the language.
- Relation to transformers: Elementwise operations correspond to transformer feed-forward sublayers, while selection and aggregation correspond to attention matrices and weighted-value mixing.The feed-forward mapping is intentionally broad because sufficiently large MLPs can approximate arbitrary Borel-measurable functions with arbitrary accuracy.
- Compilation: Compilation places each aggregation in a later layer than its inputs and places elementwise operations in the earliest layer containing all dependencies.Aggregations with the same selector at the same layer may be merged into one attention head; deriving concrete MLP weights is beyond the work’s scope.
4. Implications and insights
RASP is used to analyze restricted-attention and reordered-sublayer transformers, as well as symbolic reasoning and Dyck-k recognition. These analyses connect architectural restrictions to computable operations and establish supported task-specific scope boundaries.
- Restricted-Attention Transformers: Restricted-attention variants permanently force some selector index pairs to False, allowing RASP to characterize their computational restrictions.The paper frames the question as whether reducing possible interactions necessarily weakens the transformer.
- Restricted-Attention Transformers: Ω(n log(n)) operations are required for sorting, so variants restricting attention to o(n log(n)) operations incur a loss in expressive power.The RASP sorting algorithm applies to arbitrary input lengths and alphabet sizes.
- Sandwich Transformers: Repeated feed-forward operations before the first aggregation add little value because each position has only its initial input and cannot generate new information.Architectures that begin with attention can gather information early, after which elementwise operations can process the aggregated information.
- Scope and limitations: RASP solutions require sufficiently stable word and positional embeddings when realized in real transformers.This is presented as a practical limitation applying to transformer variants generally.
- Recognising Dyck-k Languages: RASP provides a fixed-head-and-layer solution for Dyck-k-PTF for any k, classifying each prefix as legal-unbalanced, balanced, or illegal.The task is defined over every prefix rather than only the complete sequence.
- Symbolic Reasoning in Transformers: RASP programs express simplified symbolic containment and inference computations, offering a high-level account of transformer-style logical operations.The example stores set-membership information at set and element positions before performing logical inferences.
5. Experiments
The experiments test whether RASP-derived transformer sizes are sufficient and whether trained transformers learn similar mechanisms. Across tasks, compiled bounds generally support high accuracy, while reducing heads or layers usually causes accuracy drops, with sorting an exception.
- Evaluation setup: The experiments evaluate RASP’s ability to upper-bound required heads and layers, the tightness of those bounds, and feasibility in neural transformers.Tasks include reverse, histograms, double-histograms, sorting, most-frequent tokens, and Dyck-i.
- Upper bounding the difficulty of a task: Most transformers trained at RASP-predicted sizes reached 99.5% accuracy or higher, suggesting the compiled bounds are sufficient for these tasks.The comparison uses the best of four trained transformers per task.
- Feasibility of a RASP program: Reverse and histogram-BOS models showed strong compiled–learned attention similarity, although the reverse model used a different apparent mechanism for computing length.In reverse, the second-layer head matched the flip selector, while the first-layer head preferred the final position rather than attending uniformly.
- Tightness of the bound: For most tasks, reducing heads or layers below the compiled size caused a clear accuracy drop, and several reduced transformers failed to learn their target languages.Table 2 reports average test accuracy for reduced architectures.
- Tightness of the bound: Sort was an exception: removing one layer did not hurt accuracy, apparently because the transformer exploited the bounded alphabet to implement bucket sort.A single full-attention head can count each token’s appearances, after which outputs are computed locally.
- Feasibility of a RASP program: With target and attention supervision, double-histogram, sort, and most-frequent transformers reached over 99% test accuracy and learned patterns similar to compiled solutions.The models were trained for 250 epochs with both losses.
6. Conclusions
The conclusion presents RASP as an abstraction of transformer-encoder information flow that enables symbolic reasoning about transformer computation. It reports that RASP programs predict required heads and layers, correspond to realizable transformers, and expose limitations of some efficient-transformer variants.
- Conclusions: RASP captures transformer-encoder information-flow constraints while replacing neural-network details with symbolic sequence-processing programs.The language provides a higher-level computational model for transformer encoders.
- Conclusions: Analyzing a RASP program infers the minimum number of layers and maximum number of heads needed to realize it in a transformer.The paper trains transformers on these tasks and finds that RASP helps predict their required architectural resources.
- Conclusions: The paper uses RASP to explain an empirical observation about transformer variants and identify concrete limitations of some efficient transformers.These conclusions extend the abstraction beyond the paper’s constructed task programs.
Appendices
The appendices provide experimental training details, additional results, complete RASP solutions, and implementation and compilation-flow details.
- Appendix A: Appendix A gives training details and additional results for transformers trained to mimic RASP-predicted attention patterns.It covers the experiments reported in the paper.
- Appendix B: Appendix B presents exact RASP solutions for all considered tasks and implements selector_width using operations that directly translate to a transformer.It also presents computation flows for compiled architectures.
A.1. Results: Attention-regularised transformers
Attention-regularised transformers achieved at least 99% sequence-to-sequence accuracy on the reported tasks, with attention patterns resembling compiled RASP programs.
- Results: 99+% accuracy was reached across the double-histogram, sort, and most-freq sequence-to-sequence tasks.Accuracy is computed as the fraction of output tokens predicted correctly.
- Attention patterns: The transformers’ attention patterns showed clear similarity to the corresponding compiled RASP programs.The paper presents attention-pattern comparisons for double-histogram, sorting, and most-freq.
- Sorting: 99.0% test accuracy was achieved on sorting with an input alphabet of size 52.The metric is the percentage of output positions where the correct token had maximum probability.
- Most-freq: 99.9% test accuracy was achieved on most-freq with an input alphabet of size 26.The task returns unique input tokens by descending frequency and pads the remainder with the BOS token.
A.2. Training Details
The experiments trained transformers under fixed architectural, optimization, data, and attention-regularization settings to evaluate task performance and RASP-aligned attention.
- Upper-bound and tightness experiments: Upper-bound and tightness experiments used transformers with embedding dimension 256 and feed-forward dimension 512 for 100 epochs.Four transformers were trained for each task and layer/head specification.
- Optimization: Training used ADAM without dropout, learning rates 0.0003 and 0.0001, and decay γ = 0.98 and 0.99.The models were trained on sequences of length 0−−100 with separate train, validation, and test sets.
- Attention regularisation: Attention-regularised experiments trained one transformer per language for 250 epochs with learning rate 0.0003 and decay 0.98.The loss included an MSE component comparing each attention distribution with its expected pattern, scaled by 100.
B. RASP programs and computation flows for the tasks considered
The paper implements the powerful RASP operation selector_width using only select and aggregate, enabling its compilation into transformer operations.
- RASP implementation: selector_width is implemented with only the RASP primitives select and aggregate.The implementation computes selector width through an inverse-width construction and a correction based on whether the selector focuses on position 0.
B.2. RASP solutions for the paper tasks
The paper gives RASP constructions for histograms, sorting, frequency sorting, and Dyck languages, while relating alternative constructions to compiled transformer depth and head counts.
- RASP solutions: RASP solutions are presented for histograms, double-histograms, general sorting, most-freq, and Dyck-PTF languages.The task-specific programs are accompanied by descriptions in figure captions and computation-flow figures.
- Dyck-PTF: Dyck-3-PTF extends to Dyck-n by adding parenthesis pairs without introducing additional layers or heads.The pure RASP code handles the parentheses (,), {,} and [,].
- Dyck-2-PTF: The pure RASP construction computes structure-matched openers using depth numbering, requiring two steps for open_for_close.For (())(), the depths are [1,2,2,1,1,1] and the depth-index is [1,1,2,2,3,3].
- Dyck-2-PTF: The theoretical select_best operation reduces open_for_close by finding the last earlier opener at the closer’s depth, saving one layer and two heads.This avoids computing depth_index for that operation.
- select_best: select_best retains only the selected value with the highest score for each position.Its score operation multiplies numeric values to create scorers, unlike predicate-based select operations.
- Compilation: selector_width compiles into transformer architecture through its implementation in terms of other RASP operations.Figure 9 provides the corresponding implementation.
- Reverse: The reverse one-liner compiles to two layers because length requires an attention head and reverse uses a select-aggregate pair.The reverse operation uses the length s-op among its inputs.
- Sorting: The general sorting program focuses each position on smaller keys, uses input position as a tiebreaker, computes target positions, and moves values there.The construction uses selector_width and a further select-aggregate pair; sort_in is the experimental sorting task.