Thinking Like Transformers

1 INTRODUCTION . While Yun et al . ( 2019 ) show that sufficiently large transformers can approximate any constantlength sequence-to-sequence function , and Hahn ( 2019 ) provides theoretical limitations on their ability to compute functions on unbounded input length , neither of these provide insight on how a transformer may achieve a specific task . Orthogonally , Bhattamishra et al . ( 2020 ) provide transformer constructions for several counting languages , but this also does not direct us towards a general model . This is in stark contrast to other neural network architectures , which do have clear computational models . For example , convolution networks are seen as as a sequence of filters ( Zhang et al. , 2018 ) , and finite-state automata and their variants have been extensively used both for extraction from and theoretical analysis of recurrent neural networks ( RNNs ) ( Omlin & Giles , 1996 ; Weiss et al. , 2018 ; Rabusseau et al. , 2018 ; Merrill et al. , 2020 ) , even inspiring new RNN variants ( Joulin & Mikolov , 2015 ) . In this work we propose a computational model for the transformer-encoder , in the form of a simple sequence-processing language which we dub RASP ( Restricted Access Sequence Processing Language ) . Much like how automata describe the token-by-token processing behavior of an RNN , our language captures the unique information flow constraints under which a transformer ( Vaswani et al. , 2017 ) operates as it processes input sequences . Considering computation problems and their implementation in the RASP language allows us to “ think like a transformer ” while abstracting away the technical details of a neural network in favor of symbolic programs . A RASP program operates on sequences of values from uniform atomic types , and transforms them by composing a restricted set of sequence processors . One pair of processors is used to select inputs for aggregation , and then aggregate the selected items . Another processor performs arbitrary but local computation over its ( localized ) input . However , access to the complete sequence is available only through aggregate operations that reduce a stream of numbers to a scalar . The key to performing complex global computations under this model is to compose the aggregations such that they gather the correct information , that can then be locally processed for a final output . Given a RASP program , we can analyze it to infer the minimal number of layers and maximum number of heads that is required to implement it as a transformer . We show several examples of expressive programs written in the RASP language , showing how complex operations can be implemented by a transformer . Thinking in terms of the RASP model also allows us to shed light on recent empirical observation of transformer variants ( Press et al. , 2020 ) and find concrete limitations of “ efficient transformers ” with restricted-attention ( Tay et al. , 2020 ) . 2 THE RESTRICTED ACCESS SEQUENCE PROCESSING LANGUAGE . In this section , we present the the Restricted Access Sequence Processing Language ( RASP ) . RASP assumes a machine composed of several Turing-complete processors , each of which can only run functions taking and returning a fixed number of primitive arguments , and a simple memory accessor that is controlled by these processors . The select , aggregate , and zipmap operations which we present will define and constrain how the processors work together to process an input sequence . We will focus here only on the language itself , leaving the discussion of its exact relation to transformers to Section 3 . Overview A RASP program works by manipulating sequences , occasionally with the help of selectors . Sequences contain values of uniform atomic type , such as booleans , integers , floats , or strings . They are functions used for selecting elements from sequences , and are used ( together with the appropriate operations ) only in the process of creating new sequences . All sequences in RASP are lazily evaluated , meaning that their length and contents are not populated until passed an input . The Base Sequences Every program in RASP begins from the same set of base sequences , and then creates new ones using a small number of core operations . These base sequences are indices , length , and tokens , evaluated on input x1 , x2 , ... , xn as their names suggest : ( 0 , 1 , ... , n − 1 ) , ( n , n , ... , n ) ( of length n ) , and ( x1 , x2 , ... , xn ) , respectively . Combining Sequences Sequences can be combined in an ‘ elementwise ’ manner , such that the value of the resulting sequence at each position i is a function of the values in the combined sequences at position i ( similar to a map operation ) , or have positions ‘ mixed ’ in more complicated ways using selectors , which are functions f : N × N → { True , False } whose sole purpose is to guide the combination of existing sequences into new ones . We present the basic ingredients of RASP using an example . Figure 1 shows a simple RASP function for sorting a sequence of values according to a sequence of keys . It accepts an input sequence vals and uses the base sequence indices , that is available to any RASP program , to compute its output in three operations as follows : 1. count_conditioned of Line 2 creates a new sequence that counts for each element of keys the number of “ previous items ” it has in keys , where the “ previous items ” are defined to be all items that have a lesser value , or equal value and lower index . Thus , num_prevs creates a sequence of numbers , representing the target sorted position of each item . 2. select of line 7 creates a new selector which will focus each position i on the corresponding position j for which indices [ i ] is equal to num_prevs [ j ] . Effectively , it will direct the elements in each position j towards their target location i . 3 . Finally , aggregate of line 8 applies select_sorted_val to vals , moving each i-th element of vals to its calculated sorted position num_prevs [ i ] . We now describe the base operations of RASP in-depth , occasionally presenting an example on the hypothetical input sequence x of length n. • zipmap The zipmap operation takes a tuple of sequences and an element-processing function f , and applies f per-index to the values in those sequences to create a new sequence . For a simple example , y1=zipmap ( ( indices , indices ) , lambda i , j : i+j ) creates a sequence that always evaluates to ( 0 , 2 , ... , 2n− 2 ) . • aggregate The aggregate operation takes a selector s , a sequence x , and an optional parameter default , and averages subsets of the values of x into a new sequence y as follows : for every index i , y [ i ] is the average of x [ j0 ] , x [ j1 ] , ... x [ jk ] , where j0 , j1 , ... , jk are the indices j ∈ [ n ] for which s ( i , j ) is True . We say k is the focus width of s at i . If k = 0 , then y [ j0 ] is assigned the value in default . For example : if s ( i , j ) returns True iff i is odd and j=0 , and the value in default is d , then y will evaluate to ( d , x [ 0 ] , d , x [ 0 ] , ... , y [ n− 1 ] ) where y [ n−1 ] is either d or x [ 0 ] depending on the parity of n. • select The select operation takes two sequences-tuples of lengths k and l , me= ( m1 , m2 , ... , mk ) and other= ( ot1 , ot2 , ... , otl ) , and a function f expecting k + 1 atomic values and giving boolean output . It composes these to create a selector s as follows : for every two indices i , j , s ( i , j ) is the output of f on the i-th and j-th slice of me and other respectively , i.e. , s ( i , j ) =f ( m1 [ i ] , ... , mk [ i ] , ot1 [ j ] ... otl [ j ] ) . For a simple example , in s=select ( ( indices , ) , ( indices , ) , lambda mi , oti : mi % 2==1 and oti==0 ) , then m1=indices , ot1=indices , and s is the same selector we used for our example in aggregate above . • count_conditioned This operation takes the same parameters me , other and f as select , but this time returns a sequence y describing the number of selected influencing positions j for each output position i that s=select ( me , other , f ) would have created . In other words , for each i , y [ i ] = k where j1 , ... , jk is the set of positions j for which s ( i , j ) =True . For example , h=count_conditioned ( ( tokens , ) , ( tokens , ) , lambda a , b : a==b ) returns an in-place histogram for the tokens in the input sequence : h ( “ abaa ” ) = ( 3 , 1 , 3 , 3 ) . This concludes the base operations of RASP – all other operations are shortcuts for combinations of the above 4 , occasionally with the base sequences . Sugar We implement RASP with a variety of syntactic sugar , presented fully in appendix E. Briefly : 1 . When applying zipmap to a single sequence , it may be passed directly without using a tuple , e.g . : zipmap ( indices , f ) is equivalent to zipmap ( ( indices , ) , f ) . 2. zipmap has sugar for most of the binary operators , e.g . : for two sequences x , y , then x+y is sugar for zipmap ( ( x , y ) , lambda a , b : a+b ) . 3 . Whenever the focus width of s at some index is ≤ 1 ( “ up-to-one selection ” ) , aggregate ( s , x , default=d ) does not explicitly compute the division . In this case the values of x do not have to be numbers . 4. aggregate accepts one additional optional parameter elementwise_function . The full order of parameters is s , x , elementwise_function , default , and the use of elementwise_function is as follows : aggregate ( s , x , f , d ) is equivalent to aggregate ( s , zipmap ( x , f ) , default=d ) . 2.1 EXAMPLES . We now present some more example RASP programs , by increasing order of complexity . Simple Examples The first and simplest example is to compute an in-place histogram for some sequence vals . This is achieved with a single application of count_conditioned : histogram=count_conditioned ( vals , vals , lambda a , b : a==b ) . From Length to Parity While length is provided as a primitive in the language , it can actually be achieved as a composition of the other base operations and sequences . This is done by computing full_s=select ( ( ) , ( ) , lambda : True ) followed by 1/aggregate ( full_s , indices , lambda i : int ( i==0 ) ) ( the fraction of elements equal to 0 in indices , inverted ) . From length and that same full_s we can then define count ( vals , v ) , a function taking any sequence vals and value v and returning a new sequence counting the number of appearances of v in vals . The implementation of count is simply length * aggregate ( full_s , vals , lambda e : e==v ) . count in turn enables us to write programs like parity simply as count ( tokens,1 ) % 2==01 . Reverse We can reverse a sequence seq with the help of an up-to-one selector mapping each position to its opposite : flip_s = select ( indices , length-1-indices , lambda m , oth : m==oth ) . We use flip_s to re-order the tokens of seq : reverse=aggregate ( flip_s , seq ) . Balanced Parentheses For balanced parentheses we use count_conditioned twice , storing in prev_opens and prev_closes the number of previous “ ( ” or “ ) ” ( respectively ) tokens each position has , including itself . The sequence is balanced if prev_opens-prev_closes has no negative values , and is 0 at position length-1 . These two qualities can be easily computed using two final select-aggregate pairs , and then combined with a zipmap . Most Frequent Tokens In fig . 2 we show how RASP can be used to arrange for any input sequence s the most frequent tokens in s , without repetition . The solution has 2 parts : first , we compute the histogram for all the tokens , and mask it such that all but the first of each token is given a negative value . Then , we sort the tokens according to the masked histogram . The solution uses the sort function from Figure 1 . 1This does not contradict with the findings of Hahn ( 2019 ) , who showed that parity is not computable in transformers when each selector is restricted to width 1 ( “ hard attention ” ) . Count Conditioned The operation count_conditioned is a powerful part of RASP , appearing in many other programs . Surprisingly , it is realisable as a composition of the other operations ( and base sequence indices ) . Understanding its implementation is interesting for learning how to “ truly ” think like a transformer , and we present the code in Figure 3 . The intuition is as follows : we compute the select whose width we want to calculate twice , once such that it also selects the position 0 , and once such that it only selects this position . We then aggregate both these values , broadcasting 1 from position 0 and 0 from everywhere else , and using default value 0 . The first aggregate computes for each position the inverse of the number of selected positions ( excluding 0 ) plus one , and the second computes whether that position would also focus on 0 . A straightforward zipmaps then gives us the result . To further help the intuition , we present in fig . 4 the computation flow for a histogram calculation , when count_conditioned is implemented as described here . Note Many of the functions provided with RASP can be expressed in terms of count_conditioned – such as count ( which counts how many elements in a sequence are equal to a given value ) and contains ( which checks if count is greater than 0 ) – but this is not necessarily an optimal implementation with respect to number of heads it uses . The RASPlibrary provides optimal implementations .

Summary: The authors introduce a DSL, the Restricted Access Sequence Processing (RASP) language, that they claim can serve as a computational model for the transformer-encoder.  They develop the reader's intuition for RASP by providing RASP implementations of many basic operations such as computing histograms, sorting, and reversing.  They also show how, for a given RASP program, to determine the minimum number of layers required and to upper-bound the number of heads required to implement it as a transformer.  Lastly, they analyze two transformer variants, restricted-attention transformers and sandwich transformers.  For the former, they use the RASP perspective to claim a theoretical limitation, and for the latter, they comment that a known empirical finding is intuitive in light of the RASP perspective.