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: In this work, the authors show that the sequence of self-attention and feed-forward layers within a Transformer can be interpreted as an approximate numerical solution to a set of coupled ODEs. Based on this insight, the authors propose to replace the first-order Lie-Trotter splitting scheme by the more accurate, second-order Strang splitting scheme. They then present experimental results that indicate an improved performance of their Macaron Net compared to the Transformer and argue that this is due to the former being a more accurate numerical solution to the underlying set of ODEs.