R2D2: R EUSE & R EDUCE VIA D YNAMIC W EIGHT D IFFUSION FOR T RAINING E FFICIENT NLP M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
We propose R2D2 layers, a new neural block for training efﬁcient 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 ﬂexibility 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.
I NTRODUCTION
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 cannot 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 ﬁnal 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 4 D 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 ﬂagship 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 ﬁnd that R2D 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).
O UR P ROPOSED M ETHOD
This section introduces our proposed R2D2 components.
2. V ANILLA FC L AYERS
Before we delve into our proposed method, the standard FC transform layer is deﬁned as:
[Equation]
The FC layer (equation 2.1) is fundamental to many modern and traditional neural network architectures.
Owing to its ubiquity and widespread usage, an efﬁcient 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.
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. R EUSE AND R EDUCE WITH D YNAMIC W EIGHT D IFFUSION (R2D2)
Preserving the same notation from the FC layer equation 2.1, a R2D2 layer transforms an input x ∈ R d into an output y ∈ R k with H ∈ R k × d and b ∈ R k :
[Equation]
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:
[Equation]
This section describes our proposed Dynamic Weight Diffusion mechanism.
The central idea is to operate on partitioned 1 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 ∈ R n × n × n controls the weight diffusion process and learns the soft interaction between S partitions.
Here, n is a user deﬁned 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 S j ∈ R n × k n .
A i ∈ R n × 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 ( ; ):
[Equation]
where each segment s ( A i ) is also formed by column-wise concatenation:
[Equation]
In equation 2.4, function ψ : R p × q → R pq , where ψ ( X ) ﬂattens the matrix X ∈ R p × 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, ψ ( A i S j ) ∈
R k , s ( A i ) ∈ R k × dn , thus H ∈ R k × d .
A
S
A
S
A
S
A
S
A
S
A
S
A
S
A
S
Columns of H :
A A
S S
S S
Parameters for H :
Parameter size of H :
Size of H :
As illustrated in Figure 1, the reuse of partitioned parameter blocks S 1 , . .
, S dn in the n segments s ( A i ) ( i = 1 , . . . , n ) in equation 2.3 and the reuse of the partitioned parameter block
A i in each individual s ( A i ) segment in equation 2.4 may reduce the degree of freedom for H .
Remark 2.2 It is partitioned parameter blocks
A i ( i = 1 , . . . , n ) and S j ( 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 (cid:39) 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 approximatelyn of that of an FC layer under mild assumptions.
2. R ELATIONSHIP WITH M ATRIX M ULTIPLICATION IN R EAL S PACE
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 = a W , where scalar a is the single element of the 1 × 1 matrix A 1 and elements of W ∈ R k × d come from the concatenation of S 1 , . .
, S d ∈ R 1 × 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. R ELATIONSHIP WITH H AMILTON P RODUCT IN H YPERCOMPLEX S PACE
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
[Equation]
where Q r , Q x , Q y , Q z are real numbers, ijk = i
2 =
j 2 = k 2 = − 1 , ij =
k , jk =
i , ki
= j , ji = − k , kj = − i , ik = − j .
The Hamilton product, which represents the multiplication of two
Quaternions S =
S r + S x
i + S y j + S z k and Q equation 2.5, is deﬁned as  
S r − S x − S y − S z S x S r − S z S y S y
S z
S r − S x
S z − S
y
S x
S r 
   Q r Q x Q y Q z   , (2.6) where the 4 output elements are the real values for the Quaternion unit basis [1 , i , j , k ] (cid:62) .
Denoting that s =
[ S r , S x , S y , S z ] (cid:62) , equation 2.6 can be rewritten as     1 0 00 1 00
10 0 0  s ;   0 −            − 1  s ;   0 0 − 0 −1  0  
 s ;    0 −    0 − 1     
 s 
  
Q r Q x
Q y Q z   .
It can be observed that, when n = k = d = 4 , the R2D2 layer may perform Hamilton Product of
Quarternions.
Speciﬁcally, partitioned parameter blocks A 1 , . . .
, A 4 in equation 2.3 parameterize the 4 permutation blocks (matrices composed of − 1 , 0 , 1 in equation 2.7) that reﬂect the rules of
Hamilton product, while S 1 in equation 2.4 are s in equation 2.7, and the layer input x in equation 2. is [ Q r , Q x , Q y ,
Q z ] (cid:62) 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. G ENERALIZING H YPERCOMPLEX M ULTIPLICATION WITH N EURAL A PPROXIMATION
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 cannot be generalized to arbitrary n values simply because algebraic axioms cannot hold when n takes values apart from power of 2.
At this point, it is easy to see that 5-dimensional complex numbers cannot exist due to a clear inability to multiply two 5dimensional 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 ﬁxed inductive bias.
In sharp contrast, the R2D2 layer learns such specialized multiplication rules from data, as manifested in partitioned parameter blocks A i ( i = 1 , . . . , n ) in equation 2.3.
Thus, R2D2 subsumes and is degenerated to hypercomplex multiplication when such partitioned parameter blocks are set to reﬂect those predeﬁned 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.
E FFICIENT N EURAL NLP M ODELS WITH R2D2 L AYERS
To demonstrate the applicability of the R2D2 layers, we develop two popular neural network models,
LSTM and Transformer, which are able to beneﬁt from R2D2 layers.
3. R2D2-LSTM S
Recurrent neural networks such as LSTMs are gated recurrent networks where gating functions are parameterized by linear transformations.
We introduce the our efﬁcient LSTM (R2D2-LSTM), which replace such linear transformations in the LSTM with the R2D2 layers:
[Equation]
where σ s is the sigmoid activation function and σ t is the tanh activation function, φ : R 1 × d
→ R 4 × dis a four way split on the last dimension, c t , h t are the cell state and hidden state of the R2D2 -LSTM unit at time step t .
3. R2D2 T RANSFORMERS
Transformers (Vaswani et al., 2017), the state-of-the-art model for sequence transduction task, is a stacked neural network architecture that aggressively exploits linear transformations.
Each selfattention layer comprises of three Q , K , V transformations, along with N h heads per self-attention layer.
Each layer also comprises of a two layered FC layer with nonlinearities.
A large majority of the
Transformer parameters stem from the FC layers.
In R2D2 Transformers, we replace all FC layers with R2D2 layers.
The single-headed self-attention module is rewritten as:
[Equation]
where d k is the key dimension, Φ : R 1 × d → R 3 ×
d 3 is a three way split on the last dimension, X is the input sequence, and A is the self-attentive representation.
For multi-headed, using R2D2 layers also enables weight sharing not only among Q , K , V but also among heads.
attention, the transformation of multiple heads are also projected with R2D2 layers:
[Equation]
where N h is the number of heads and (; ) is the column-wise concatenation.
Finally, the position-wise
FC layer is now deﬁned as
[Equation]
which transforms X with 2 R2D2 layers.
E XPERIMENTS
This section reports the experimental results of R2D2-LSTMs and R2D2 Transformers that are equipped with R2D2 layers.
Overall, we conduct 4 main experiments in which all experimental settings and hyperparameters (e.g., parameter initialization) remain consistent with the vanilla version.
We evaluate R2D2 LSTMs on natural language inference (NLI) and compare them with standard LSTMs.
Next, we evaluate R2D2 Transformers and compare them with standard Transformers on neural machine translation (NMT), text style transfer, and subject verb agreement (SVA) tasks.
4. N ATURAL L ANGUAGE I NFERENCE
The task of natural language inference is to determine if two sequences entail or contradict with each other (MacCartney, 2009).
NLI is a fundamental task pertaining to language understanding.
To this end, they serve as a suitable benchmark for evaluating recurrent models.
[Table]
Datasets and Setup We run experiments on 5 NLI datasets, i.e. , (1) SNLI (Bowman et al., 2015), (2) MultiNLI
(Williams et al., 2017), (3) Dialogue NLI (Welleck et al., 2018), (4) QNLI (Quora) (Wang et al., 2017), and (5) Scitail (Science Entailment) (Khot et al., 2018).
We implement 300dimensional unidirectional encoders with shared parameters for both premise/hypothesis.
We take the concatenation of max and mean pooled representations as input to a two-layered 300-dimensional
MLP for prediction.
Our model is trained with the Adam with a learning rate of 0 .
0004 with a batch size of 256 .
Word embeddings are initialized with GloVe (Pennington et al., 2014) and are ﬁxed.
No cross sentence attention (Parikh et al., 2016) is used, mainly to observe the effectiveness of standalone encoders.
For R2D2-LSTM, we use n = { 2 , 5 , 10 } and report the results accordingly.
In this task, since word embeddings are 300-dimensional, we select multiples of 5 instead of 4 for simplicity and ease of divisibility.
Experimental Results Table 1 reports the results on all the 5 natural language inference datasets.
All in all, results are extremely encouraging, showing that the R2D2 layer can not only reduce the parameterization cost but also improve performance (4 out of 5 datasets show reasonable improvement).
On QNLI (the only exception), the performance drop is marginal ( < 1% ), which is decent considering the parameter savings.
It is also noteworthy that on SNLI, DNLI, and SciTail, all the
R2D2 -LSTM variants outperform the base LSTM model.
Overall, we think that the parameter reuse properties, in addition to learning to share such reused parameter blocks amongst recurrent gating functions, may contribute to a regularizing effect.
4. N EURAL M ACHINE T RANSLATION
Machine translation (MT) is concerned with translating between source-target language pairs.
To this end, sequence transduction models are central to this problem domain.
In this experiment, the key goal is to compare R2D2 Transformers against the standard Transformer model.
[Figure] [Table]
[Table]
and Decoding Time (seconds to decode test set) with Beam size of 4 and length penalty of 0 .
6 on
IWSLT 2014 German-English.
[Table]
Datasets and Setup We run experiments on 8 NMT datasets.
The datasets are (1) IWSLT’
English-Vietnamese (En-Vi) † , (2)
IWSLT’17 English-Indonesian (En-Id) † , (3) IWSLT’14 GermanEnglish (En-De) † , (4) IWSLT’14 Romanian-English
(Ro-En) † , (5) WMT’18 English-Estonian (EnEt) ∗ ,
(6) Setimes English-Macedonian
(En-Mk) ∗ , (7) WMT’18 English-Romanian (En-Ro) ∗ and (8)
WMT’16 English-German
(En-De).
Datasets with † are run with 50K steps while datasets with ∗ are trained for 100K steps also on a single GPU.
For all tasks, we specify that Transformers have 4 layers, 8 heads, and a hidden size 512.
We use beam size of 5 and α = 0 .
6 (length penalty) for evaluating all the models.
For all R2D models, we benchmark several settings for the hyperparameter n = { 2 , 4 , 8 , 16 } .
Likewise, we also re-scale the parameters by increasing the dimensions to match the original number of parameters.
This is denoted by hyperparameter U .
In the experiments, the upsampling factor is ﬁxed at 2.
We benchmark against two baselines that are also concerned with reducing parameterization of the
FC layer.
The ﬁrst, is the Factorized Transformer which approximates each FC layer with low-rank approximation.
In this case, the number of latent factors is set to d 4 to enable direct comparison to the
Quaternion Transformer.
The second is the Hamilton product based Transformer proposed in (Tay et al., 2019).
Experimental Results Table 3 reports our results on machine translation.
Across 5 out of 7 NMT benchmarks, R2D2 Transformer outperforms the standard Transformer model.
Overall, the empirical results are promising.
On one hand, we observe that increasing n all the way to 16 does not cause signiﬁcant degradation in performance on datasets such as En-Vi.
On the other hand, for most datasets, even with signiﬁcant parameter reduction, we ﬁnd that the decrease in the BLEU score is overall manageable ( ≈ 1–3 BLEU points).
However, we also note a rare occurrence where n =results in signiﬁcant decrease in the BLEU score, such as on the En-Id dataset.
On several datasets, the R2D2 Transformer model improves the performance.
For example, datasets such as En-Vi and En-Et enjoy a performance boost of about 1 BLEU point with n = 2 .
Generally, there is only marginal performance degradation at n = 4 or n = 8 for most datasets.
Finally, by upsampling with hyperparameter U , we are able to improve the performance of three datasets: En-Vi,
En-Id, and En-Mk.
Additionally, we note that Factorized Transformer performs substantially worse than the vanilla
Transformer or our R2D2-Transformer.
On the other hand, our proposed R2D2 also makes reasonable gains over the Quaternion Transformer (Tay et al., 2019), signifying that ﬂexible approximation of
Hamilton products are effective.
Table 4 reports the decoding and training time for each Transformer variant.
We observe that R2D with n = 8 has the fastest decoding speed amongst all variants.
All in all, the training speed is also approximately similar.
As such, this ascertains that our proposed dynamic weight diffusion method is not computationally expensive, despite its complexity.
4. T EXT S TYLE T RANSFER
We experiment with sequence transduction for text style transfer.
The goal of this task is to convert text of a certain style to another style.
Datasets and Setup We use the Modern → Shakespeare corpus 2 in the experiments.
The key goal here is to convert modern writing into Shakespeare writing.
This dataset comprises of 18395 parallel sentences for training, 1218 parallel sentences for evaluation (dev set), and 1462 parallel sentences for testing.
We still specify that Transformers have 4 layers, 8 heads, and a hidden size 512.
Similar to NMT, we experiment with n = { 2 , 4 , 8 , 16 } .
We train all the models for 10K steps.
[Table]
Experimental Results Table 5 reports the results on text style transfer.
We observe that the best performance is achieved with R2D2 Transformer ( n = 4 ).
Notably, all except the n = 16 variant outperforms the standard Transformer model.
This ascertains the effectiveness of the proposed R2D layer.
This not only enables parameter savings but improves the performance of Transformer.
4. S UBJECT V ERB A GREEMENT
We conduct additional experiments on subject-verb agreement task (Linzen et al., 2016).
The task is a binary classiﬁcation problem.
It predicts if the sentence, e.g. , ‘The keys to the cabinet .’ follows by a plural or a singular.
Dataset and Setup This dataset can be found online (Linzen et al., 2016).
We evaluate the R2D
Transformer encoder against the standard Transformer encoder for the subject verb agreement task.
In contrast to the previous experimental settings, we use a smaller Transformer architecture.
Speciﬁcally,
Transformers here have 2 layers, 4 head, and a hidden size 128.
Since the hidden size is smaller than those in the previous experimental settings, we experiment with n = { 2 , 4 , 8 } .
[Table]
Experimental Results Table 6 reports the results on the SVA task.
Results are promising, demonstrating that all variants with R2D2 layers outperform the standard Transformer.
The best performance peaks at n = 8 , despite parameter saving to up to 1 / 8 .
R2D2 also outperforms
Quaternion Transformers.
R ELATED W ORK
FC layers are ubiquitous in deep learning research (Goodfellow et al., 2016).
Many, if not all neural building blocks incorporate this functionality, i.e. , transforming input vectors to an output space using a parameterized linear transformation matrix.
Recurrent neural networks (Hochreiter & Schmidhuber, 1997), multilayer perceptrons, attention (Bahdanau et al., 2014), and Transformer models (Vaswani et al., 2017) are all heavily grounded in these transformation operations.
Notably, recent state-of-theart models in language domains aggressively exploit transformation layers, achieving very exceptional results (Devlin et al., 2018;
Radford et al., 2019).
Our work can be interpreted as a form of soft parameter sharing, albeit dynamic and learned.
A recent work (Savarese & Maire, 2019) proposed soft-weight sharing across stacked convolution layers, effectively simulating recurrence, i.e. , RNNs can be considered to be weight sharing across time.
Quaternion networks (Zhu et al., 2018; Parcollet et al., 2018; 2019) are also known to inhibit weight sharing qualities and have demonstrated reasonable success despite having fewer parameters.
Low-rank approximations (Markovsky, 2012) and factorization methods are known techniques to reduce parameters and improve generalization.
Majority of techniques aimed at reducing parameters of neural networks belong to this class (Sainath et al., 2013;
Chen et al., 2018; Si et al., 2017).
Within the context of language, GroupReduce (Chen et al., 2018) applies low-rank approximations to compress neural language models.
(Si et al., 2017) proposed memory efﬁcient kernels using low-rank approximation.
Block diagonal approaches (Xie et al., 2017), structured sparsity (Gale et al., 2019) are also highly notable alternatives for improving the FC layer.
Methods such as knowledge distillation (Hinton et al., 2015; Freitag et al., 2017) are also capable of reducing model complexity and can be considered an orthogonal direction from ours.
In a similar vein, sharing parameters by the means of a meta network (HyperNetworks Ha et al.
(2016) also perform a form of weight sharing.
C ONCLUSION
We propose R2D2 layers, which can reduce the parameterization cost of their FC counterparts without signiﬁcantly compromising performance.
The R2D2 layer also offers a way to bridging multiplication between real space and hypercomplex space by subsuming and generalizing both matrix multiplication in real space and Hamilton product in hypercomplex space.
It is highly modular and applicable to many dominant models such as LSTMs and Transformers.
We evaluate such models equipped by R2D2 layers on comprehensive tasks, demonstrating substantial parameter savings with minimal degradation or improvement in performance.