Understanding and Improving Transformer From a Multi-Particle Dynamic System Point of View

1 INTRODUCTION . The Transformer is one of the most commonly used neural network architectures in natural language processing . Variants of the Transformer have achieved state-of-the-art performance in many tasks including language modeling ( Dai et al. , 2019 ; Al-Rfou et al. , 2018 ) and machine translation ( Vaswani et al. , 2017 ; Dehghani et al. , 2018 ; Edunov et al. , 2018 ) . Unsupervised pre-trained models based on the Transformer architecture also show impressive performance in many downstream tasks ( Radford et al. , 2019 ; Devlin et al. , 2018 ) . The Transformer architecture is mainly built by stacking layers , each of which consists of two sub-layers with residual connections : the self-attention sub-layer and the position-wise feed-forward network ( FFN ) sub-layer . For a given sentence , the self-attention sub-layer considers the semantics and dependencies of words at different positions and uses that information to capture the internal structure and representations of the sentence . The position-wise FFN sub-layer is applied to each position separately and identically to encode context at each position into higher-level representations . Although the Transformer architecture has demonstrated promising results in many tasks , its design principle is not fully understood , and thus the strength of the architecture is not fully exploited . As far as we know , there is little work studying the foundation of the Transformer or different design choices . In this paper , we provide a novel perspective towards understanding the architecture . In particular , we are the first to show that the Transformer architecture is inherently related to the Multi-Particle Dynamic System ( MPDS ) in physics . MPDS is a well-established research field which aims at modeling how a collection of particles move in the space using differential equations ( Moulton , 2012 ) . In MPDS , the behavior of each particle is usually modeled by two factors separately . The first factor is the convection which concerns the mechanism of each particle regardless of other particles in the system , and the second factor is the diffusion which models the movement of the particle resulting from other particles in the system . Inspired by the relationship between the ODE and neural networks ( Lu et al. , 2017 ; Chen et al. , 2018a ) , we first show that the Transformer layers can be naturally interpreted as a numerical ODE solver for a first-order convection-diffusion equation in MPDS . To be more specific , the selfattention sub-layer , which transforms the semantics at one position by attending over all other positions , corresponds to the diffusion term ; The position-wise FFN sub-layer , which is applied to each position separately and identically , corresponds to the convection term . The number of stacked layers in the Transformer corresponds to the time dimension in ODE . In this way , the stack of self-attention sub-layers and position-wise FFN sub-layers with residual connections can be viewed as solving the ODE problem numerically using the Lie-Trotter splitting scheme ( Geiser , 2009 ) and the Euler ’ s method ( Ascher & Petzold , 1998 ) . By this interpretation , we have a novel understanding of learning contextual representations of a sentence using the Transformer : the feature ( a.k.a , embedding ) of words in a sequence can be considered as the initial positions of a collection of particles , and the latent representations abstracted in stacked Transformer layers can be viewed as the location of particles moving in a high-dimensional space at different time points . Such an interpretation not only provides a new perspective on the Transformer but also inspires us to design new structures by leveraging the rich literature of numerical analysis . The Lie-Trotter splitting scheme is simple but not accurate and often leads to high approximation error ( Geiser , 2009 ) . The Strang-Marchuk splitting scheme ( Strang , 1968 ) is developed to reduce the approximation error by a simple modification to the Lie-Trotter splitting scheme and is theoretically more accurate . Mapped to neural network design , the Strang-Marchuk splitting scheme suggests that there should be three sub-layers : two position-wise feed-forward sub-layers with half-step residual connections and one self-attention sub-layer placed in between with a full-step residual connection . By doing so , the stacked layers will be more accurate from the ODE ’ s perspective and will lead to better performance in deep learning . As the FFN-attention-FFN layer is “ Macaron-like ” , we call it Macaron layer and call the network composed of Macaron layers the Macaron Net . We conduct extensive experiments on both supervised and unsupervised learning tasks . For each task , we replace Transformer layers by Macaron layers and keep the number of parameters to be the same . Experiments show that the Macaron Net can achieve higher accuracy than the Transformer on all tasks which , in a way , is consistent with the ODE theory . 2 BACKGROUND . 2.1 RELATIONSHIP BETWEEN NEURAL NETWORKS AND ODE . Recently , there are extensive studies to bridge deep neural networks with ordinary differential equations ( Weinan , 2017 ; Lu et al. , 2017 ; Haber & Ruthotto , 2017 ; Chen et al. , 2018a ; Zhang et al. , 2019b ; Sonoda & Murata , 2019 ; Thorpe & van Gennip , 2018 ) . We here present a brief introduction to such a relationship and discuss how previous works borrow powerful tools from numerical analysis to help deep neural network design . A first-order ODE problem is usually defined as to solve the equation ( i.e. , calculate x ( t ) for any t ) which satisfies the following first-order derivative and the initial condition : dx ( t ) dt = f ( x , t ) , x ( t0 ) = w , ( 1 ) in which x ( t ) ∈ Rd for all t ≥ t0 . ODEs usually have physical interpretations . For example , x ( t ) can be considered as the location of a particle moving in the d-dimensional space and the first order time derivative can be considered as the velocity of the particle . Usually there is no analytic solution to Eqn ( 1 ) and the problem has to be solved numerically . The simplest numerical ODE solver is the Euler ’ s method ( Ascher & Petzold , 1998 ) . The Euler ’ s method discretizes the time derivative dx ( t ) dt by its first-order approximation x ( t2 ) −x ( t1 ) t2−t1 ≈ f ( x ( t1 ) , t1 ) . By doing so , for the fixed time horizon T = t0 + γL , we can estimate x ( T ) from x0 . = x ( t0 ) by sequentially estimating xl+1 . = x ( tl+1 ) using xl+1 = xl + γf ( xl , tl ) ( 2 ) where l = 0 , · · · , L − 1 , tl = t0 + γl is the time point corresponds to xl , and γ = ( T − t0 ) /L is the step size . As we can see , this is mathematically equivalent to the ResNet architecture ( Lu et al. , 2017 ; Chen et al. , 2018a ) : The function γf ( xl , tl ) can be considered as a neural-network block , and the second argument tl in the function indicates the set of parameters in the l-th layer . The simple temporal discretization by Euler ’ s method naturally leads to the residual connection . Observing such a strong relationship , researchers use ODE theory to explain and improve the neural network architectures mainly designed for computer vision tasks . Lu et al . ( 2017 ) ; Chen et al . ( 2018a ) show any parametric ODE solver can be viewed as a deep residual network ( probably with infinite layers ) , and the parameters in the ODE can be optimized through backpropagation . Recent works discover that new neural networks inspired by sophisticated numerical ODE solvers can lead to better performance . For example , Zhu et al . ( 2018 ) uses a high-precision Runge-Kutta method to design a neural network , and the new architecture achieves higher accuracy . Haber & Ruthotto ( 2017 ) uses a leap-frog method to construct a reversible neural network . Liao & Poggio ( 2016 ) ; Chang et al . ( 2019 ) try to understand recurrent neural networks from the ODE ’ s perspective , and Tao et al . ( 2018 ) uses non-local differential equations to model non-local neural networks . 2.2 TRANSFORMER . The Transformer architecture is usually developed by stacking Transformer layers ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ) . A Transformer layer operates on a sequence of vectors and outputs a new sequence of the same shape . The computation inside a layer is decomposed into two steps : the vectors first pass through a ( multi-head ) self-attention sub-layer and the output will be further put into a position-wise feed-forward network sub-layer . Residual connection ( He et al. , 2016 ) and layer normalization ( Lei Ba et al. , 2016 ) are employed for both sub-layers . The visualization of a Transformer layer is shown in Figure 2 ( a ) and the two sub-layers are defined as below . Self-attention sub-layer The attention mechanism can be formulated as querying a dictionary with key-value pairs ( Vaswani et al. , 2017 ) , e.g. , Attention ( Q , K , V ) = softmax ( QKT / √ dmodel ) · V , where dmodel is the dimensionality of the hidden representations and Q ( Query ) , K ( Key ) , V ( Value ) are specified as the hidden representations of the previous layer in the so-called self-attention sublayers in the Transformer architecture . The multi-head variant of attention allows the model to jointly attend to information from different representation subspaces , and is defined as Multi-head ( Q , K , V ) = Concat ( head1 , · · · , headH ) WO , ( 3 ) headk = Attention ( QW Q k , KW K k , V W V k ) , ( 4 ) where WQk ∈ Rdmodel×dK , WKk ∈ Rdmodel×dK , WVk ∈ Rdmodel×dV , and WO ∈ RHdV ×dmodel are project parameter matrices , H is the number of heads , and dK and dV are the dimensionalities of Key and Value . Position-wise FFN sub-layer In addition to the self-attention sub-layer , each Transformer layer also contains a fully connected feed-forward network , which is applied to each position separately and identically . This feed-forward network consists of two linear transformations with an activation function σ in between . Specially , given vectors h1 , . . . , hn , a position-wise FFN sub-layer transforms each hi as FFN ( hi ) = σ ( hiW1 + b1 ) W2 + b2 , where W1 , W2 , b1 and b2 are parameters . In this paper , we take the first attempt to provide an understanding of the feature extraction process in natural language processing from the ODE ’ s viewpoint . As discussed in Section 2.1 , several works interpret the standard ResNet using the ODE theory . However , we found this interpretation can not be directly applied to the Transformer architecture . First , different from vision applications whose size of the input ( e.g. , an image ) is usually predefined and fixed , the input ( e.g. , a sentence ) in natural language processing is always of variable length , which makes the single-particle ODE formulation used in previous works not applicable . Second , the Transformer layer contains very distinct sub-layers . The self-attention sub-layer takes the information from all positions as input while the position-wise feed-forward layer is applied to each position separately . How to interpret these heterogeneous components by ODE is also not covered by previous works ( Tao et al. , 2018 ; Chen et al. , 2018a ) .

Summary: The paper points out a formal analogy between transformers and an ODE modelling multi-particle convection (the feed-forward network) and diffusion (the self-attention head). The paper then adapts the Strang-Marchuk splitting scheme for solving ODEs to construct a slightly different transformer architecture: “FFN of Attention of FFN”, instead of “FFN of Attention”. The new architecture, refered to as a Macaron-Net, yields better performance in a variety of experiments.