GLU Variants Improve Transformer

Introduction

The Transformer [Vaswani et al., 2017] sequence-to-sequence model alternates between multi-head attention, and what it calls "position-wise feed-forward networks" (FFN). The FFN takes a vector $x$ (the hidden representation at a particular position in the sequence) and passes it through two learned linear transformations, (represented by the matrices $W_{1}$ and $W_{2}$ and bias vectors $b_{1}$ and $b_{2}$). A rectified-linear (ReLU) [Glorot et al., 2011] activation function applied between the two linear transformations.

max0𝑥subscript𝑊1subscript𝑏1subscript𝑊2subscript𝑏2\textrm{FFN}(x,W_{1},W_{2},b_{1},b_{2})=\textrm{max}(0,xW_{1}+b_{1})W_{2}+b_{2} (1) Following the T5 codebase [Raffel et al., 2019] 111Also in the interest of ML fairness., we use a version with no bias:

Subsequent work has proposed replacing the ReLU with other nonlinear activation functions such as Gaussian Error Linear Units, $\textrm{GELU}(x)=x\Phi(x)$ [Hendrycks and Gimpel, 2016], and $\textrm{Swish}_{\beta}(x)=x\sigma(\beta x)$ [Ramachandran et al., 2017].

Gated Linear Units (GLU) and Variants

[Dauphin et al., 2016] introduced Gated Linear Units (GLU), a neural network layer defined as the component-wise product of two linear transformations of the input, one of which is sigmoid-activated. They also suggest omitting the activation, which they call a "bilinear" layer and attribute to [Mnih and Hinton, 2007].

𝑥𝑊𝑏𝑥𝑉𝑐Bilinear𝑥𝑊𝑉𝑏𝑐tensor-product𝑥𝑊𝑏𝑥𝑉𝑐\begin{split}\textrm{GLU}(x,W,V,b,c)&=\sigma(xW+b)\otimes(xV+c)\\ \textrm{Bilinear}(x,W,V,b,c)&=(xW+b)\otimes(xV+c)\end{split} (4) We can also define GLU variants using other activation functions:

𝑥𝑊𝑏𝑥𝑉𝑐GEGLU𝑥𝑊𝑉𝑏𝑐tensor-productGELU𝑥𝑊𝑏𝑥𝑉𝑐SwiGLU𝑥𝑊𝑉𝑏𝑐𝛽tensor-productsubscriptSwish𝛽𝑥𝑊𝑏𝑥𝑉𝑐\begin{split}\textrm{ReGLU}(x,W,V,b,c)&=\textrm{max}(0,xW+b)\otimes(xV+c)\\ \textrm{GEGLU}(x,W,V,b,c)&=\textrm{GELU}(xW+b)\otimes(xV+c)\\ \textrm{SwiGLU}(x,W,V,b,c,\beta)&=\textrm{Swish}_{\beta}(xW+b)\otimes(xV+c)\\ \end{split} (5) In this paper, we propose additional variations on the Transformer FFN layer which use GLU or one of its variants in place of the first linear transformation and the activation function. Again, we omit the bias terms.

All of these layers have three weight matrices, as opposed to two for the original FFN. To keep the number of parameters and the amount of computation constant, we reduce the number of hidden units $d_{ff}$ (the second dimension of $W$ and $V$ and the first dimension of $W_{2}$) by a factor of $\frac{2}{3}$ when comparing these layers to the original two-matrix version.

Experiments on Text-to-Text Transfer Transformer (T5)

We test the FFN variants we have described on the transfer-learning setup from [Raffel et al., 2019]. An encoder-decoder transformer model [Vaswani et al., 2017] is trained on a denoising objective of predicting missing text segments, and subsequently fine-tuned on various language understanding tasks.

We use the same code base, model architecture, and training task as the base model from [Raffel et al., 2019]. The encoder and decoder each consist of 12 layers, with $d_{model}=768$. For the attention layers, $h=12$ and $d_{k}=d_{v}=64$. The FFN layers have hidden size $d_{ff}=3072$. As we describe above, for the GLU-variant-based FFN layers, which have thee weight matrices instead of two, we reduce the hidden layer to $d_{ff}=2048$, so as to maintain the same parameter and operation counts as the base model.

2 Pre-Training and Perplexity Results

Identically to [Raffel et al., 2019], we pre-train for 524,288 steps on the span-filling objective on the C4 dataset. Each training batch consists of 128 examples, each of which has an input of 512 tokens and an output of 114 tokens, the output containing multiple spans of tokens which were deleted from the input222Each training step took approximately 0.15 seconds on a 32-core TPUv2 cluster.. Similarly to [Raffel et al., 2019], we use the Adafactor optimizer [Shazeer and Stern, 2018] and an inverse-square-root learning-rate schedule. We also decay the learning rate linearly for the final 10 percent of the training steps. Our main departure from [Raffel et al., 2019] is that we use no dropout during pre-training. We find this to produce superior results. We compute the log-perplexity on the training objective on a heldout shard of C4, which we believe to be a good indicator of model quality. For each model architecture, we also trained four models for a shorter period (65,536 steps) to measure inter-run variability. The results are listed in table 1. The GEGLU and SwiGLU variants produce the best perplexities.

3 Fine-Tuning

We then fine-tune each fully-trained model once on an examples-proportional mixture of the Stanford Question-Answering Dataset (SQuAD) [Rajpurkar et al., 2016] and all the language understanding tasks in the GLUE [Wang et al., 2018] and SuperGlue [Wang et al., 2019] benchmarks.333This departs from [Raffel et al., 2019], who fine-tuned separately on the different tasks. We chose one fine-tuning run for simplicity. Fine-tuning consists of 131072 steps with a learning rate of $10^{-3}$. As in training, the input sequences for each step have a combined length of approximately 65,536 tokens. Following [Raffel et al., 2019], we use a dropout rate of $0.1$ on the layer outputs, feed-forward hidden-layers and attention weights. The embedding matrices are fixed during fine-tuning.

Tables 2, 3 and 4 show results on the development sets. For each task, we report the best score of any of the checkpoints recorded during fine-tuning. While the results are noisy, the new GLU-variants perform best on most of the tasks. For comparison, at the bottom of each of the tables we list the reuslts from [Raffel et al., 2019]. The model is identical to our $\textrm{FFN}_{\textrm{ReLU}}$ model. Their results are notably worse, which we believe was caused by their use of dropout during pre-training. Also listed are the inter-run standard deviations measured by [Raffel et al., 2019].

Conclusions

We have extended the GLU family of layers and proposed their use in Transformer. In a transfer-learning setup, the new variants seem to produce better perplexities for the de-noising objective used in pre-training, as well as better results on many downstream language-understanding tasks. These architectures are simple to implement, and have no apparent computational drawbacks. We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence.

References