Thinking Like Transformers

TL;DR

The paper introduces RASP, a Restricted Access Sequence Processing Language, as a high-level computational model for transformer-encoder architectures. By mapping attention and feed-forward operations to s-ops and selectors, RASP enables symbolic programming of transformer-like computations and demonstrates tasks such as double-histograms and Dyck-k languages. It shows how RASP programs compile to transformer configurations, allowing explicit upper bounds on the required number of layers and heads and validating these bounds with attention-regularised training that aligns learned patterns with the compiled solutions. The work provides insights into restricted attention, sandwich-transformer effects, and symbolic reasoning in transformers, offering a practical framework and public code for further exploration.

Abstract

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.

Figures (19)

• Figure 1: We consider double-histogram, the task of counting for each input token how many unique input tokens have the same frequency as itself (e.g.: §hist2("§ aaabbccdef")=[§,1,1,1,2,2,2,2,3,3,3]§). (a) shows a RASP program for this task, (b) shows the selection patterns of that same program, compiled to a transformer architecture and applied to the input sequence §§ aaabbccdef§, (c) shows the corresponding attention heatmaps, for the same input sequence, in a 2-layer 2-head transformer trained on double-histogram. This particular transformer was trained using both target and attention supervision, i.e.: in addition to the standard cross entropy loss on the target output, the model was given an MSE-loss on the difference between its attention heatmaps and those expected by the RASP solution. The transformer reached test accuracy of $99.9\%$ on the task, and comparing the selection patterns in (b) with the heatmaps in (c) suggests that it has also successfully learned to replicate the solution described in (a).
• Figure 2: Visualizing the select and aggregate operations. On the left, a selection matrix s is computed by select, which marks for each query position all of the key positions with matching values according to the given comparison operator §==§. On the right, aggregate uses s as a filter over its input values, averaging only the selected values at each position in order to create its output, res. Where no values have been selected, aggregate substitutes $0$ in its output.
• Figure 3: RASP program for the task shuffle-dyck-2 (balance 2 parenthesis pairs, independently of each other), capturing a higher level representation of the hand-crafted transformer presented by bhattamishra2020abilitybhattamishra2020ability.
• Figure 4: Top: RASP code for computing reverse (e.g., §reverse("abc")="cba"§). Below, its compilation to a transformer architecture (left, obtained through §draw(reverse,"abcde")§ in the RASP REPL), and the attention heatmaps of a transformer trained on the same task (right), both visualised on the same input. Visually, the attention head in the second layer of this transformer corresponds perfectly to the behavior of the §flip§ selector described in the program. The head in the first layer, however, appears to have learned a different solution from our own: instead of focusing uniformly on the entire sequence (as is done in the computation of §length§ in RASP), this head shows a preference for the last position in the sequence.
• Figure 5: The RASP program for computing with-BOS histograms (left), alongside its compilation to a transformer architecture (cream boxes) and the attention head (center bottom) of a transformer trained on the same task, without attention supervision. The compiled architecture and the trained head are both presented on the same input sequence, §"§ aabbaabb"§. The transformer architecture was generated in the RASP REPL using §draw(hist,"§ aabbaabb")§.