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 develops an analogy between Transformer networks and dynamical systems of multiple interacting particles. They follow the work of Lu et al. 2019 using Lie-Trotter splitting scheme to derive differential equations approximating terms corresponding to operations that compose the Transformer network internals, namely multi-head self-attention and feed-forward transformation.   This paper develops the described technique beyond previous work by noting their analogous dynamical system motivates architecture choices that would require many fewer parameters than the associated Transformer. The authors investigate performance and memory gains by exploiting time-dependence in the dynamical transformation, achieving similar results to analogous weight-sharing between layers.