Redesigning the Transformer Architecture with Insights from Multi-particle Dynamical Systems

1 Introduction . Neural networks have evolved from early feed-forward and convolutional networks , to recurrent networks , to very deep and wide ‘ Transformer ’ networks based on attention [ Vaswani et al. , 2017 ] . Transformers and their enhancements , such as BERT [ Devlin et al. , 2019 ] , T5 [ Raffel et al. , 2020 ] and GPT [ Brown et al. , 2020 ] are , by now , the default choice in many language applications . Both their training data and model sizes are massive . BERT-base has 110 million parameters . BERT-large , which often leads to better task performance , has 345 million parameters . GPT-3 has 175 billion trained parameters . Larger BERT models already approach the limits of smaller GPUs . GPT-3 is outside the resource capabilities of most research groups . Training these gargantuan models is even more challenging , with significant energy requirements and carbon emissions [ Strubell et al. , 2020 ] . In response , a growing community of researchers is focusing on post-facto reduction of model sizes , which can help with the deployment of pre-trained models in low-resource environments . However , training complexity is also critically important . A promising recent approach to faster training uses a 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) way of viewing layers of attention as solving ordinary differential equations ( ODEs ) defined over a dynamical system of interacting particles [ Lu et al. , 2019 , Vuckovic et al. , 2020 ] . We pursue that line of work . Simulating particle interactions over time has a correspondence to ‘ executing ’ successive layers of the Transformer network . In the forward pass at successive layers , the self-attention and position-wise feed-forward operations of Transformer correspond to computing the new particle states from the previous ones . However , the numeric function learned by the i-th attention layer has zero knowledge regarding the one learned by the ( i − 1 ) -th layer . This is counter-intuitive due to the fact that the whole evolution is temporal in nature , and this independence leads to growing numbers of trainable parameters and computing steps . We seek to develop time-evolution functionals from the initial condition alone . Such maps can then approximate the underlying ODE from parametric functions of time ( the analog of network depth ) and do not require computing self-attention over and over . We propose such a scheme , leading to a network/method we call TransEvolve . It can be used for both encoder-decoder and encoder-only applications . We experiment on several tasks : neural machine translation , whole-sequence classification , and long sequence analysis with different degrees of timeevolution . TransEvolve outperforms Transformer base model on WMT 2014 English-to-French translation by 1.4 BLEU score while using 10 % fewer trainable parameters . On all the encoder-only tasks , TransEvolve outperforms Transformer , as well as several strong baselines , with 50 % fewer trainable parameters and more than 3× training speedup . 2 Related Work . Our work focuses on two primary areas of machine learning — understanding neural networks as dynamical systems and bringing down the overhead of Transformer-based models in terms of training computation and parameters . Neural networks and dynamical systems . Weinan [ 2017 ] first proposed that machine learning systems can be viewed as modeling ordinary differential equations describing dynamical systems . Chang et al . [ 2018 ] explored this perspective to analyze deep residual networks . Ruthotto and Haber [ 2019 ] later extended this idea with partial differential equations . Lu et al . [ 2018 ] showed that any parametric ODE solver can be conceptualized as a deep learning framework with infinite depth . Chen et al . [ 2018 ] achieved ResNet-comparable results with a drastically lower number of parameters and memory complexity by parameterizing hidden layer derivatives and using ODE solvers . Many previous approaches applied sophisticated numerical methods for ODE approximation to build better neural networks [ Haber and Ruthotto , 2017 , Zhu and Fu , 2018 ] . Vuckovic et al . [ 2020 ] developed a mathematical formulation of self-attention as multiple interacting particles using a measure-theoretic perspective . The very first attempt to draw analogies between Transformers and dynamical systems was made by Lu et al . [ 2019 ] . They conceptualized Transformer as a numerical approximation of dynamical systems of interacting particles . However , they focused on a better approximation of the ODE with a more robust splitting scheme ( with the same model size as Transformer and the dot-product attention kept intact ) . We seek to parameterize the temporal dynamics to bypass attention up to a certain degree . Efficient variations of Transformer . Multiple approaches have been put forward to overcome the quadratic complexity of Transformers [ Wang et al. , 2020a , Choromanski et al. , 2020 , Peng et al. , 2021 , Xiong et al. , 2021 , Liu et al. , 2018 ] . Kitaev et al . [ 2020 ] sought to use locality-sensitive hashing and reversible residual connections to deal with long range inputs . Some studies explored the sparsified attention operations to decrease computation cost [ Liu et al. , 2018 , Ho et al. , 2019 , Roy et al. , 2021 ] . These sparsification tricks can be done based on the data relevant to the task [ Roy et al. , 2021 , Sukhbaatar et al. , 2019 ] or in a generic manner [ Liu et al. , 2018 , Ho et al. , 2019 ] . Wang et al . [ 2020a ] observed self-attention to be low-rank , and approximated an SVD decomposition of attention to linearize it . Peng et al . [ 2021 ] used random feature kernels to approximate the softmax operation on the attention matrix . Choromanski et al . [ 2020 ] sought to linearize Transformers without any prior assumption of low-rank or sparse distribution , using positive orthogonal random features . Lee-Thorp et al . [ 2021 ] achieved remarkable speedup over Transformers by substituting the attention operation with an unparameterized Fourier transform . A bulk of these works are suitable for encoder-only tasks and often incurs slower training ( e.g , Peng et al . [ 2021 ] observed 15 % slower training compared to Transformer ) . Our method overcomes both of these drawbacks . Transformer pruning and compression . Knowledge distillation has been used to build light-weight student model from large , trained Transformer models [ Behnke and Heafield , 2020 , Sanh et al. , 2019 , Wang et al. , 2020b ] . Michel et al . [ 2019 ] experimented with pruning different attention heads of BERT to observe redundancy in computation . However , these methods still require a trained , parameter-heavy model to start with . Tensorization approach has shown efficient compression of Transformer-based language models [ Ma et al. , 2019 , Khrulkov et al. , 2019 ] . 3 Transformers as Dynamical Systems . A single block of the Transformer encoder [ Vaswani et al. , 2017 ] is defined as a multi-head selfattention layer followed by two feed-forward layers , along with residual connections . The j-th head of self-attention operation in the l-th encoder block , with j ∈ { 1 , . . . , m } , on a given length-n sequence of d dimensional input vectors Xl : = { X li |X li ∈ Rd } ni=1 can be defined as : H lj = Softmaxi ( ( X lW lq ) ( X lW lk ) > / √ dk ) ( X lW lv ) ( 1 ) where W lq , W l k ∈ Rd×dk , and W lv ∈ Rd×dv are linear projection layers and dk = d/m . Conventionally , dv = dk . Each H lj is aggregated to produce the output of the multi-head self-attention as follows : Hl = Concat ( { H lj } mj=1 ) W lo + Xl ( 2 ) where Wo ∈ Rd×d is a linear projection . The subsequent feed-forward transformation can be defined as : Xl+1 = ( σ ( HlW lff1 +B l ff1 ) ) W l ff2 +B l ff2 + H l ( 3 ) where W lff1 ∈ Rd×dff , W lff2 ∈ Rdff×d , Blff1 ∈ Rdff , Blff2 ∈ Rd , and σ ( ) is a non-linearity ( Relu in [ Vaswani et al. , 2017 ] or Gelu in [ Devlin et al. , 2019 ] ) . As Lu et al . [ 2019 ] argued , Equations 1-3 bear a striking resemblance with systems of interacting particles . Given the positions of a set of interacting particles as x ( t ) = { xi ( t ) } ni=1 , the temporal evolution of such a system is denoted by the following ODE : d dt xi ( t ) = F ( xi ( t ) , x ( t ) , t ) +G ( xi ( t ) , t ) ( 4 ) with the initial condition xi ( t0 ) = si ∈ Rd . The functions F and G are often called the diffusion and convection functions , respectively — the former models the inter-dependencies between the particles at time t , while the latter models the independent dynamics of each particle . Analytical solution of such an ODE is often impossible to find , and the most common approach is to use numerical approximation over discrete time intervals [ t0 , t0 + δt , . . . , t0 + Lδt ] . Following Euler ’ s method of first order approximation and the Lie-Trotter splitting scheme , one can approach the numerical solution of Equation 4 from time t to t+ δt as : x̃i ( t ) = xi ( t ) + δtF ( xi ( t ) , x ( t ) , t ) = xi ( t ) + F ( xi ( t ) , x ( t ) , t ) xi ( t+ δt ) = x̃i ( t ) + δtG ( x̃i ( t ) , t ) = x̃i ( t ) + G ( x̃i ( t ) , t ) ( 5 ) Equation 5 can be directly mapped to the operations of the Transformer encoder given in Equations 1-3 , with Xl ≡ x ( t ) , Hl ≡ { x̃i ( t ) } ni=1 and Xl+1 ≡ x ( t + δt ) . F ( · , · , t ) is instantiated by the projections W lq , W l k , W l v , W l o and Softmax ( · ) operations in the multi-head self-attention . G ( · , t ) corresponds to the projections W lff1 , W l ff2 , B l ff1 , B l ff1 and the σ ( · ) non-linearity in Equation 3 . From the described analogies , it quickly follows that successive multi-head self-attention and the point-wise feed-forward operations follow a temporal evolution ( here time is equivalent to the depth of the encoder ) . However , Transformer and its subsequent variants parameterize these two functions in each layer separately . This leads to a large number of parameters to be trained ( ranging from 7 × 107 in neural machine translation tasks to 175 × 109 in language models like GPT-3 ) . We proceed to investigate how one can leverage the temporal evolution of the diffusion and convection components to bypass this computational bottleneck . 4 Time-evolving Attention . As Equation 5 computes x ( t ) iteratively from a given initial condition x ( t0 ) = s = { si } ni=1 , one can reformulate the diffusion map F ( · , x ( t ) , t ) as F̃ ( · , f ( s , t ) ) , i.e. , as a functional of the initial condition and time only . When translated to the case of Transformers , this means one can avoid computing pairwise dot-product between n input vectors at each layer by computing a functional form at the beginning and evolving it in a temporal ( depth-wise ) manner . We derive this by applying dot-product self-attention on hypothetical input vectors with augmented depth information , as follows . Let X′ = { X ′i|X ′i ∈ Rd+d ′ } be a set of vectors such that X ′i = Concat ( X0i , T l ) , where X0i = { xi1 , . . . , xid } and T l = { τ1 ( l ) , . . . , τd′ ( l ) } . W ′q , W ′k ∈ R ( d+d ′ ) × ( d+d′ ) are the augmented query and key projections , respectively , such that W ′q = [ ωij ] d+d′ , d+d′ i , j=1,1 , W ′ k = [ θij ] d+d′ , d+d′ i , j=1,1 . The pre-softmax query-key dot product between X ′i and X ′ j is given by a ′ ij = ( X ′ iW ′ q ) ( X ′ jW ′ k ) > . We can decompose W ′q as concatenation of two matrices Wq , W̃q such that Wq = [ ωij ] d , d+d′ i , j=1,1 and W̃q = [ ωij ] d+d′ , d+d′ i , j=d+1,1 . Similarly , we decompose W ′k into Wk and W̃k . Then a ′ ij can be re-written as : a′ij = ( X 0 iWq ) ( X 0 jWk ) > + ( X0iWq ) ( T lW̃k ) > + ( T lW̃q ) ( X 0 jWk ) > + ( T lW̃q ) ( T lW̃k ) > = a0ij +A1iT l > + T lA2j +A3 ( T l T l ) ( 6 ) where A1i = X0iWqW̃ > k , A2j = W̃qW > k X 0 > j , A3 = W̃qW̃ > k and signify hadamard product . Detailed derivation is provided in Appendix A . It is evident that a0ij is the usual dot-product pre-softmax attention between vector elements X 0 i , X 0 j . For the complete sequence of vector elements X , we write A0 = [ aij ] ij , A1 = { A1i } i and A2 = { A2j } j . By definition , T l is a vector function of the depth l. To construct T l as a vector function of l , we formulate T l = [ wl1 sin ( l P ) , . . . , wld′ 2 sin ( d′l 2P ) , wld′ 2 +1 cos ( l P ) , . . . , wld′ cos ( d′l 2P ) ] ( 7 ) where W lt = { wli } d ′ i=1 are learnable parameters at depth l , and P = d′L 2π . Such a choice of T l is intuitive in the following sense : let C = AT l + B for some arbitrary A = [ aij ] ∈ Rp×d , B = { bi } ∈ Rp , and C = { ci } ∈ Rp . Then we get ci = bi + d′ 2∑ j=1 aijw l j sin ( jl P ) + d′ 2∑ j=1 aijw l j+ d ′ 2 cos ( jl P ) ( 8 ) In other words , any feed-forward transformation on T l gives us a vector in which each component is a Fourier series ( thereby , approximating arbitrary periodic functions of l with period P ) . This enables us to encode the nonlinear transformations that X undergoes in successive blocks . With this , A1 , A2 , A3 constitute the time-evolution operators to map depth information to the attention , computed only from the initial conditions X0i , X 0 j . Figures 1 ( a ) and 1 ( b ) summarize the procedure . Given two input sequences X and Y ( in case of selfattention , they are the same ) , we first compute A0 , A1 , A2 , A3 using the linear projection matrices Wq , Wk , W̃q , W̃k , as described above . Additionally , we normalize A0 by a factor 1√ d/m , where m is the number of heads , similar to Transformer . Then , at any subsequent layer l ∈ { 1 , . . . , L } , instead of computing the query , key and value projections from the input Xl , we use the previously computed A0 , A1 , A2 , A3 along with T l. Then the time-evolved attention operation becomes Hl = Softmax ( A0 + A1T l > + T lA2 +A3 ( T l T l ) ) XlW lo + Xl ( 9 ) For the sake of brevity , we have not shown the concatenation of multi-headed splits as shown in Equation 2 . Therefore , one should identify W lo in Equation 9 with Equation 2 and not as value projection as in Equation 1 . We do not use value projection in our method . Also , in the multi-headed case , all the matrix-vector dot products shown in Equation 9 are computed on head-splitted dimension of size d/m . Complexity and parameter reduction . Given a d-dimensional input sequence of length n , computing the pre-softmax attention weight matrix in a single Transformer encoder requires O ( n2d ) multiplications . In our proposed method , the complexity of calculating dot-product attention ( corresponding to a0ij in Equation 6 ) invokes computations of similar order – O ( n2 ( d + d′ ) ) . However , this is needed only once at the beginning . The subsequent attention matrices are calculated using the components with T in Equation 6 . Both A1T and TA2 require O ( nd′ ) computation , and A3T requires only O ( d′ ) . Therefore , if one tends to use L successive attention operations , our proposed method is computationally equivalent to a single original dot-product self-attention followed by multiple cheaper stages . In addition to this , attention weight computation using Equation 6 eliminates the need for query , key , and value projections at each self-attention block . Thereby , it ensures a parameter reduction of O ( Ld2 ) for a total stacking depth of L .

Summary: This paper introduces TransEvolve, a new efficient transformer architecture that approximates the multi-headed self-attention and point-wise feed-forward transformations in the standard transformers.  Drawing from recent work by Lu et al., 2019 which established the connection between transformers and modeling the dynamics of a multi-particle dynamic systems, TransEvolve approximate multiple layers of the transformer with a temporal evolution scheme that bypasses costly dot product operations.  Experiments on encoder-decoder tasks show TransEvolve to be competitive with vanilla transformers while being slightly more efficient.  On encoder-only tasks, TransEvolve yields better predictive performance with half as many parameters.