R2D2: Reuse & Reduce via Dynamic Weight Diffusion for Training Efficient NLP Models

We propose R2D2 layers , a new neural block for training efficient NLP models . Our proposed method is characterized by a dynamic weight diffusion mechanism which learns to reuse and reduce parameters in the conventional transformation layer , commonly found in popular Transformer/LSTMs models . Our method is inspired by recent Quaternion methods which share parameters via the Hamilton product . This can be interpreted as a neural and learned approximation of the Hamilton product which imbues our method with increased flexibility and expressiveness , i.e. , we are no longer restricted by the 4D nature of Quaternion weight sharing . We conduct extensive experiments in the NLP domain , showing that R2D2 ( i ) enables a parameter savings of up to 2× to 16× with minimal degradation of performance and ( ii ) outperforms other parameter savings alternative such as low-rank factorization and Quaternion methods . 1 INTRODUCTION . The transformation layer is one of the most ubiquitous and dominant component in existing and current deep learning literature ( Goodfellow et al. , 2016 ) . Its pervasiveness can not be understated , given its centrality to many core building blocks in neural network research . Given widespread adoptions of FC layers , e.g. , within Transformer ( Vaswani et al. , 2017 ) models and recurrent ( Hochreiter & Schmidhuber , 1997 ) networks , a further reduction in parameter complexity and memory footprint could be extremely useful to many real world applications . This paper proposes a new neural component for reusing and reducing parameter costs of the transformation layer ( i.e. , memory footprint ) . Concretely , our proposed R2D2 layer is characterized by a new dynamic weight diffusion mechanism . The central idea is that we start off with a core set of parameters ( weights ) and learn dynamic ‘ diffusion ’ of these weight partitions to construct the final transformation weight W . The key novelty lies in our method of incorporating soft weight reuse , i.e. , taking inspiration from multiplication in hypercomplex spaces . Recent work in Quaternion spaces and Hamilton products have demonstrated reasonable success ( Parcollet et al. , 2018 ; 2019 ; Tay et al. , 2019 ) . The Hamilton product , which multiples two Quaternions by fusing latent inter-component representations , enables a four time parameter saving as compared to the real-valued adaptation . Unfortunately , Hamilton products operate on 4D hypercomplex numbers , which limits its expressiveness and utility . To this end , our proposed R2D2 layer can be viewed as a neural approximation of the Hamilton product , learning the permutation of the inter-latent component interactions in a soft differentiable fashion . Moreover , our method can operate on any arbitrary n , aside from only on 4D Quaternion representations , facilitating up to n times parameter savings . To demonstrate applicability , we equip two well-established models ( LSTMs and Transformers ) with R2D2 layers . We conduct extensive experiments on flagship benchmarks , i.e. , neural machine translation for Transformers and natural language inference for LSTMs . Additionally , we include further validation on text style transfer and subject-verb agreement tasks . All in all , we find that R2D2 generally enables up parameter savings up to approximately 1/16 ∼ 1/2 with minimal degradation in performance . Moreover , our proposed method also enables slight speed-up in terms of inference ( decoding speed ) . 2 OUR PROPOSED METHOD . This section introduces our proposed R2D2 components . 2.1 VANILLA FC LAYERS . Before we delve into our proposed method , the standard FC transform layer is defined as : y = FC ( x ) = Wx + b . ( 2.1 ) The FC layer ( equation 2.1 ) is fundamental to many modern and traditional neural network architectures . Owing to its ubiquity and widespread usage , an efficient reduction of parameters in the FC layer may result in massive parameter savings , especially for models that heavily use transformation equation 2.1 . Remark 2.1 The degree of freedom for the weight parameters W in equation 2.1 is kd . Since W dominates parameterization , the parameter size of the FC layer equation 2.1 is O ( kd ) . In the following , we describe the proposed R2D2 layer and its relationships with matrix multiplication in real space and Hamilton product in hypercomplex space . 2.2 REUSE AND REDUCE WITH DYNAMIC WEIGHT DIFFUSION ( R2D2 ) . Preserving the same notation from the FC layer equation 2.1 , a R2D2 layer transforms an input x ∈ Rd into an output y ∈ Rk with H ∈ Rk×d and b ∈ Rk : y = R2D2 ( x ) = Hx + b . ( 2.2 ) The distinct difference in R2D2 is that we construct the parameter H in equation 2.2 with a reduced degree of freedom in order to reduce the parameter size . This is achieved by our Dynamic Weight Diffusion method described as follows : 2.2.1 DYNAMIC WEIGHT DIFFUSION . This section describes our proposed Dynamic Weight Diffusion mechanism . The central idea is to operate on partitioned1 weight blocks and learn a dynamic diffusion of weights . There are two key parameter blocks A and S which are central to our approach . Intuitively , A ∈ Rn×n×n controls the weight diffusion process and learns the soft interaction between S partitions . Here , n is a user defined hyperparameter . Suppose that both d and k are divisible by n ∈ Z > 0 . For i = 1 , . . . , n and j = 1 , . . . , dn , denote by each partitioned parameter block Sj ∈ Rn× k n . Ai ∈ Rn×n is the weight diffusion matrix assigned to each of n blocks . The parameter H in equation 2.2 is constructed by column-wise concatenation ( ; ) : H = [ s ( A1 ) ; s ( A2 ) ; . . . ; s ( An ) ] , ( 2.3 ) where each segment s ( Ai ) is also formed by column-wise concatenation : s ( Ai ) = [ ψ ( AiS1 ) ; ψ ( AiS2 ) ; . . . ; ψ ( AiS d n ) ] . ( 2.4 ) In equation 2.4 , function ψ : Rp×q → Rpq , where ψ ( X ) flattens the matrix X ∈ Rp×q by concatenating each row of X then transposes the concatenated row vector into a column vector of dimension pq . It is easy to see that , ψ ( AiSj ) ∈ Rk , s ( Ai ) ∈ Rk× d n , thus H ∈ Rk×d . 1Here , the partitioning ( or splitting ) of weights into different components is analogous to complexification , i.e. , splitting real vectors into real and imaginary components . As illustrated in Figure 1 , the reuse of partitioned parameter blocks S1 , . . . , S d n in the n segments s ( Ai ) ( i = 1 , . . . , n ) in equation 2.3 and the reuse of the partitioned parameter block Ai in each individual s ( Ai ) segment in equation 2.4 may reduce the degree of freedom for H. Remark 2.2 It is partitioned parameter blocks Ai ( i = 1 , . . . , n ) and Sj ( j = 1 , . . . , dn ) that determine the degree of freedom for H , which is kdn + n 3 . Since H dominates parameterization , the parameter size of the R2D2 in equation 2.2 is O ( kdn ) , where kd ' n 4 is assumed : this condition is mild for real-world problems , such as in our experiments ( e.g. , d = 512 , k = 2048 , n = 2 , 4 , 8 , 16 ) . Thus , for the same input and output sizes , the parameterization cost of a R2D2 layer is approximately 1 n of that of an FC layer under mild assumptions . 2.3 RELATIONSHIP WITH MATRIX MULTIPLICATION IN REAL SPACE . The R2D2 layer subsumes matrix multiplication in real space . Thus , it is a generalization of the FC layer via hyperparameter n. To explain , referring to equation 2.2 , when n = 1 , H = aW , where scalar a is the single element of the 1 × 1 matrix A1 and elements of W ∈ Rk×d come from the concatenation of S1 , . . . , Sd ∈ R1×k . Since learning a and W separately is equivalent to learning their multiplication jointly , scalar a can be dropped , which is learning the single weight matrix in an FC layer . Therefore , a R2D2 layer is degenerated to an FC layer when n = 1 , where there is no parameter saving . 2.4 RELATIONSHIP WITH HAMILTON PRODUCT IN HYPERCOMPLEX SPACE . Next , we explore the relationship between the R2D2 layer and Hamilton product in hypercomplex space . For background , in hypercomplex space , inputs are represented by multiple imaginary components . For example , a Quarternion has 1 real component and 3 imaginary components , while an Octonion has 1 real component and 7 imaginary components . For the sake of illustration , let us take a Quarternion Q as an example , which is a 4-dimensional hypercomplex number Q = Qr +Qxi +Qyj +Qzk , ( 2.5 ) where Qr , Qx , Qy , Qz are real numbers , ijk = i2 = j2 = k2 = −1 , ij = k , jk = i , ki = j , ji = −k , kj = −i , ik = −j . The Hamilton product , which represents the multiplication of two Quaternions S = Sr + Sxi + Syj + Szk and Q equation 2.5 , is defined as Sr −Sx −Sy −SzSx Sr −Sz SySy Sz Sr −Sx Sz −Sy Sx Sr QrQxQy Qz , ( 2.6 ) where the 4 output elements are the real values for the Quaternion unit basis [ 1 , i , j , k ] > . Denoting that s = [ Sr , Sx , Sy , Sz ] > , equation 2.6 can be rewritten as 1 0 0 00 1 0 00 0 1 0 0 0 0 1 s ; 0 −1 0 01 0 0 00 0 0 1 0 0 −1 0 s ; 0 0 −1 00 0 0 −11 0 0 0 0 1 0 0 s ; 0 0 0 −10 0 1 00 −1 0 0 1 0 0 0 s QrQxQy Qz . ( 2.7 ) It can be observed that , when n = k = d = 4 , the R2D2 layer may perform Hamilton Product of Quarternions . Specifically , partitioned parameter blocks A1 , . . . , A4 in equation 2.3 parameterize the 4 permutation blocks ( matrices composed of −1 , 0 , 1 in equation 2.7 ) that reflect the rules of Hamilton product , while S1 in equation 2.4 are s in equation 2.7 , and the layer input x in equation 2.2 is [ Qr , Qx , Qy , Qz ] > in equation 2.7 . Likewise , Hamilton product of Octonions or Sedenions in hypercomplex space may also be performed by the R2D2 layer when n , k , d are equally set to 8 or 16 . 2.5 GENERALIZING HYPERCOMPLEX MULTIPLICATION WITH NEURAL APPROXIMATION . In fact , weight reuse by component-wise partitioning in Quaternion space has demonstrated reasonable success ( Parcollet et al. , 2018 ; Zhu et al. , 2018 ; Parcollet et al. , 2019 ; Tay et al. , 2019 ) . However , one key problem is that Hypercomplex algebra can not be generalized to arbitrary n values simply because algebraic axioms can not hold when n takes values apart from power of 2 . At this point , it is easy to see that 5-dimensional complex numbers can not exist due to a clear inability to multiply two 5- dimensional complex numbers . Within the context of hypercomplex space , specialized multiplication rules , such as Hamilton product , have to be devised and encoded in the network as a fixed inductive bias . In sharp contrast , the R2D2 layer learns such specialized multiplication rules from data , as manifested in partitioned parameter blocks Ai ( i = 1 , . . . , n ) in equation 2.3 . Thus , R2D2 subsumes and is degenerated to hypercomplex multiplication when such partitioned parameter blocks are set to reflect those predefined multiplication rules in hypercomplex space . Moreover , the proposed R2D2 layer can be seen as a trainable and parameterized form of n-dimensional hypercomplex multiplication , where n can be values other than power of 2 . To sum up , the R2D2 layer reduces parameterization cost by reusing partitioned parameter blocks . It also offers a way to bridging multiplication between real space and hypercomplex space .

Summary: This paper proposes a new Reuse and Reduce with Dynamic weight Diffusion (R2D2) layer as an alternative to feed-forward layers in neural networks. The layer is inspired by the Hamilton Product in a hypercomplex space (where numbers have multiple imaginary components). The main idea is to have two smaller parameter blocks that are partitioned, multiplied, and concatenated together in order to form the full weight matrix of the feed-forward layer.