Luna: Linear Unified Nested Attention

Introduction

Transformers (Vaswani et al., 2017) are surprisingly versatile models that preform well on a wide range of language and vision tasks, including machine translation (Vaswani et al., 2017; Ott et al., 2018), language understanding (Devlin et al., 2019), image recognition (Dosovitskiy et al., 2020) and bioinformatics (Madani et al., 2020). Attention (Bahdanau et al., 2015) provides the key mechanism that captures contextual information from the entire sequence by modeling pairwise interactions between the inputs at every timestep. However, a common weakness of Transformers is their quadratic time and memory complexity within the attention mechanism w.r.t the length of the input sequence, which prohibitively restricts their potential application to tasks requiring longer input sequences.

A number of techniques have been recently introduced to improve the time and memory efficiency of Transformer models (‘xformers’) (Tay et al., 2020b, 2021). One popular technique is using sparsity to restrict the attention field range, such as local attention (Parmar et al., 2018), blockwise attention (Qiu et al., 2019), strided attention patterns (Child et al., 2019; Beltagy et al., 2020), compressed attention (Liu et al., 2018), and attention with learnable patterns (Kitaev et al., 2020; Tay et al., 2020a; Roy et al., 2021). Another emerging approach is to improve efficiency by leveraging low-rank approximations of the attention matrix. Linformer (Wang et al., 2020), for example, projects the length dimension of key and value matrices to a fixed-dimensional representation by assuming low-rank structure in the full-rank attention matrix. Recently, some kernel-based methods, such as Linear Transformer (Katharopoulos et al., 2020), Performer (Choromanski et al., 2020) and Random Feature Attention (Peng et al., 2021), attempt to efficiently approximate regular (softmax) full-rank attention through kernelization. Although these models demonstrate better asymptotic complexity for long sequences, their efficiency gains are less prominent for moderate length sequences and their performance remains behind Transformers with regular attention.

In this work, we propose a linear unified nested attention mechanism (Luna), which uses two nested attention functions to approximate the regular softmax attention in Transformer (§2). Specifically, with the first attention function, Luna packs the input sequence into a sequence of fixed length. Then, the packed sequence is unpacked using the second attention function (§3.1). As compared to a more traditional attention mechanism, Luna introduces an additional sequence with a fixed length as input and an additional corresponding output. Importantly, the extra input allows Luna to perform attention operation linearly as efficiently as Linformer (Wang et al., 2020), while also storing adequate contextual information. Unlike Linformer, Luna is capable of modeling variable-length sequences and autoregressive (causal) attention (§3.4). We perform extensive experiments on three sequence modeling tasks, including long-context sequence modeling, neural machine translation, and masked language modeling for large-scale pretraining and downstream task finetuning. Compared to a variety of strong baseline models, Luna achieves competitive or even better performance, while acquiring prominent gains of efficiency in both speed and memory (see Figure 1). More importantly, Luna manages to obtain superior performance with small projection lengths such as 16 (§4).

Background

The traditional attention mechanism is a function:

2 Transformer Layers

The other two key components of Transformer, besides attention, are position-wise feed-forward networks (FFN) and layer normalization (Ba et al., 2016). Technically, the position-wise feed-forward layer operates on each position independently and layer normalization plays a crucial role in controlling the gradient scales (Xiong et al., 2020). Each Transformer layer can be expressed as:

where $X$ and $C$ are the two input sequences and $X^{\prime}$ is the output of the Transformer layer. The Transformer layer in (2) adopts the original post-layer normalization architecture (Vaswani et al., 2017; Devlin et al., 2019) that places layer normalization after residual connection, rather than pre-layer normalization (Vaswani et al., 2018; Wang et al., 2019).

Linear Unified Nested Attention (Luna)

Our goal is to design an efficient attention mechanism to solve the quadratic complexity problem of full attention. We first introduce the proposed linear unified nested attention mechanism, named Luna attention (§3.1), and the architecture of each Luna layer (§3.2). Then, we present the variant of Luna for causal attention, named Luna causal attention (§3.3). Finally, we discuss the differences between Luna and three closely related models: Linformer (Wang et al., 2019), Set Transformer (Lee et al., 2019) (§3.4) and Shared Workspace (Goyal et al., 2021).

To unpack the sequence back to the length of the original query sequence $X$, Luna leverages its second attention, named unpack attention:

where the computation of $Y_{P}$ and $Y_{X}$ is in (3) and (4). By stacking multiple layers of Luna attention, the output $Y_{P}$ from the previous layer, which captures contextual information of $C$, is employed as the input $P$ of the next layer. For the first layer of Luna, we formulate $P$ as learnable positional embeddingsWe also experimented with sinusoidal positional embeddings, and obtained similar results. (Vaswani et al., 2017).

Reducing the Number of Parameters.

Due to the two nested attention operations, there are two sets of parameters ($W_{Q},\,\,W_{K},\,\,W_{V}$) in a single Luna attention function. There are several techniques to reduce the number of parameters, such as parameter sharing (Xia et al., 2019). In this work, we follow Wang et al. (2020) to share $W_{K}$ and $W_{Q}$ in each layer, and conduct experiments to analyze performance decline against Luna with full sets of parameters (§4.2).

2 Luna Layers

The Luna attention is used as a drop-in-replacement for the regular attention. We incorporate the position-wise feed-forward network and layer normalization into Luna layers. Concretely, layer normalization is applied to both $Y_{X}$ and $Y_{P}$, while FFN only to $Y_{X}$:

where $X^{\prime}$ and $P^{\prime}$ are the two outputs of the Luna layer. The graphical specification of one Luna layer is illustrated in Figure 2.

3 Luna Causal Attention

As discussed in Tay et al. (2020b), the ability to support causal autoregressive decoding, i.e. attending solely to the past and current tokens, is required when designing efficient self-attention mechanisms. However, due to the pack attention that packs the long sequence $X$ into a fixed (shorter) length, it is not straight-forward to support causal attention in Luna.

The complexity of the causal attention in Luna is still linear: $O(ln)$. One drawback of Luna causal attention, similar to the causal attention in Random Feature Attention (RFA) (Peng et al., 2021) and Linear Transformer (Katharopoulos et al., 2020), is its sequential computation for each timestep $t$.

In the formulation of causal attention, $P$ is expected to contain no information about $X$. Thus, we need to formulate $P$ based on the usage mode of the causal attention. For the encoder-decoder mode in sequence-to-sequence modeling (e.g. for machine translation), we can use packed output from the Luna encoder as $P$. For the decoder-only mode (e.g. for language modeling), $P$ might be formulated as a learnable parameter of each layer.

4 Discussion

Relation to Set Transformer.

The additional input $P$ in Luna can be regarded as a side memory module that can access the entire sequence to gather contextual information. From this view of point, Luna is also closely related to Set Transformer (Lee et al., 2019), an early model to integrate side memory module in Transformers. Similar to the projection matrix in Linformer, the inducing points in Set Transformer are learnable parameters. Thus, these inducing points might be formulated as the non-contextual version of $P$ in Luna. Moreover, Set Transformer is designed for set-input problems, which are problems wherein the input is a set of features and the model is thereby invariant to permutation or ordering of the input features (Tay et al., 2020b), while Luna attention is used as a drop-in replacement for regular softmax attention.

Experiments

We evaluate the effectiveness and efficiency of Luna on the Long Range Arena (LRA) benchmark recently introduced by Tay et al. (2021), which is designed for the purpose of evaluating efficient Transformer models under the long-context scenario. They collect five tasks in this benchmark which are ListOps (Nangia and Bowman, 2018), byte-level text classification (Text; Maas et al., 2011), byte-level document retrieval (Retrieval; Radev et al., 2013), image classification on sequences of pixels (Image; Krizhevsky et al., 2009) and Pathfinder (Linsley et al., 2018). These tasks consist of input sequences ranging from 1K to 8K tokens and span across a variety of data types and modalities.

To ensure fair comparisons, for all tasks except for the task Retrieval, we closely follow the model configurations in Tay et al. (2021) such as data preprocessing, data split, model architecture, etc. For the task of Retrieval, we find that models are not fully converged when being trained for 5K steps as stated in Tay et al. (2021). Therefore, we train models for 20K steps for this task and obtain much better results. For a direct comparison, besides the average performance of models across all tasks, we also report the average accuracy on tasks excluding Retrieval. We run each experiment for five times with different random seeds and report the average accuracy. The hyper-parameters for each task are shown in Appendix A.1.

The results of various models on the LRA benchmark are presented in Table 1. For our proposed method, we report results from models of three different projected dimensions (16, 128 and 256). First, we note that Luna achieves good results on all tasks consistently compared to the Transformer model and significantly outperforms all the other baseline methods in terms of the average accuracy. By taking a closer look at the accuracy for each individual task, Luna wins over baseline models on three out of five tasks and performs comparably with the best performed model on the other two tasks, i.e. ListOps and byte-level text classification. Notably, Luna improves over the Transformer model on image classification and pathfinder by a large margin. Second, we observe that although Luna achieves the best average performance with a projection dimension of 256, it also performs considerably well with smaller projection dimensions (16 and 128). This demonstrates the effectiveness of Luna even with small projected dimensions.

Memory and Speed Efficiency.

Luna employs two nested linear attention functions to reduce the time and memory complexity compared to the vanilla softmax attention. Here, we examine the speed and memory footprint of various models with varying input lengths (1K, 2K, 3K and 4K). Following Tay et al. (2021), all models are evaluated on the byte-level classification task with the same batch size. The result is shown in Table 2.

Considering the memory efficiency, Luna with a projected dimension of 16 is highly memory-efficient, which is only 10% of the vanilla Transformer at 4K input sequence length. With larger projected dimensions, i.e. 128 and 256, Luna requires more memory but is still competitive compared to other efficient Transformer models. In terms of time efficiency, Luna-16 speeds up over the standard Transformer by 1.2-5.5 times, varying by the sequence length. Compared to other efficient Transformers, Luna-16 performs comparably with the fastest models, i.e. Performer and Linformer. Overall, our models achieve competitive advantage both in time- and memory-efficiency over other models, while attaining the best performance on the LRA benchmark (see Figure 1).

In addition, we plot the trade-off among memory, time and averaged LRA score without task Retrieval in Figure 1. Models such as Linformer and Performer have faster speed and small memory requirement with the sacrifice of performance. However, besides competitive time- and memory-efficiency, Luna models retain superior performance even with a small projected dimension ($l$=16).

Contextual information in P𝑃P of Luna.

Recently, a popular method to model the classification task using Transformer-based models is to prepend a special symbol, [CLS], to every input example. The last hidden state of this symbol is regarded as the aggregate sequence representation. In Luna, we introduce an extra model input $P$ which not only allows us to efficiently compute the attention mechanism but learn contextual information as well. Theoretically, the $P$ from the last layer is capable of learning the representation of the input sequence. To validate this, we extract $P$ at the last layer and employ the mean pooling strategy over positions to obtain the final feature for classification. We test its performance on three long-text modeling tasks in LRA (Tay et al., 2021), i.e., ListOps, Text and Retrieval and report results in Table 3. We find that $P$-based methods obtain better scores across all tasks against the [CLS]-based one, validating the powerful ability of $P$ to encode contextual information of the input sequence.

2 Machine translation

To evaluate Luna on sequence-to-sequence modeling, we conduct experiments on a standard machine translation benchmark, i.e. WMT’14 English-German (EN$\rightarrow$DE) dataset (4.5M sentence pairs). The data split and preprocessing steps follow those of Vaswani et al. (2017), using the scripts from FairSeq (Ott et al., 2019). We share the source and target vocabularies within the language pair, with 37K byte pair encoding (BPE) types (Sennrich et al., 2016). The Luna models closely follow the architecture of Transformer-base: 6 encoder and decoder layers with 8 attention heads and $d_{\textrm{model}}/d_{\textrm{hidden}}=512/2048$. We train the Transformer-base model with two optimization methods: Adam (Kingma and Ba, 2015) and Apollo (Ma, 2020), and find Apollo achieves better performance. Therefore, we use Apollo as the optimizer for all Luna models. For each experiment, we conduct distributed training across eight NVIDIA Tesla V100 GPUs with maximum batch size of 8192 tokens per GPU. Further details are provided in Appendix A.2.

Effect of Encoding Contextual Information into P𝑃P.

3 Masked Language Modeling for Large-Scale Pretraining

One popular application of Transformer is to pretrain a large-scale language model on a large amount of data which can then be fine-tuned on a wide range of downstream tasks, such as BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), etc. Therefore, we pretrain a Luna-based language model with RoBERTa-base model configuration on two versions of data as our pretraining set: 1) BERT version with BookCorpus (Zhu et al., 2015) and English Wikipedia (totally 16GB), 2) RoBERTa version with BookCorpus, English Wikipedia, CC-News (Nagel, 2016), OpenWebText (Gokaslan and Cohen, 2019) and Stories (Trinh and Le, 2018) (totally 160GB). For Luna models, we set $l=128$. On the larger training corpus (160GB), we train models w./w.o parameter sharing, respectively. We compare our models with RoBERTa-base, BERT-base and Linformer which are trained on the same training data. Experimental details are provided in Appendix A.3.

After obtaining the pretrained Luna-based language model, we finetune it on various natural language processing tasks, including sentiment classification (SST-2; Socher et al., 2013), natural language inference (QNLI; Rajpurkar et al., 2016), textual similarity (QQP; Chen et al., 2018, question answering (RACE (Lai et al., 2017) and CommonsenseQA (CSQA; Talmor et al., 2019). For GLUE tasks, following Liu et al. (2019), we consider a limited hyperparameter sweep for each task, with batch sizes $\in\{16,32\}$ and learning rate $\in\{5e^{-6},1e^{-5},2e^{-5}\}$, with a linear warmup for the first 6% of steps followed by a linear decay to . Finetuning is performed for 20 epochs with early stopping based on each task’s evaluation metric on the dev setWe observed that Luna finetuning requires more epochs than vanilla Transformer (20 vs. 10). We also finetuned RoBERTa with 20 epochs but did not obtain better results.. For QA tasks, we concatenate each candidate answer with the corresponding question and passage. We then encode every candidate and pass the [CLS] output at the last layer through a fully-connected layer, which is used to predict the correct answer. We truncate question-answer pairs that are longer than 128 tokens and, if needed, the passage so that the total length is at most 512 tokens. Following Liu et al. (2019), we try a small range of possible values for hyperparameters, i.e., batch size $\in\{16,32\}$, learning rate $\in\{1e^{-5},2e^{-5},3e^{-5}\}$ and dropout $\in\{0.0,0.1,0.2\}$. For other configurations such as warm-up steps, optimizer, we follow thoses in Liu et al. (2019).

The result is reported in Table 6. We observe that on the smaller dataset (16GB) our Luna model has similar or slightly better downstream results compared to other pretrained language models. On QNLI and SST-2, Luna models obtain the best performance among all models, reaffirming the effectiveness of Luna in pre-training. This demonstrates the strong ability of Luna for language representations. On the larger dataset (160GB), however, the performance of Luna is slightly worse than RoBERTa with vanilla Transformer architecture. One possible reason is that the capacity of Luna is not as sufficient as vanilla Transformer, due to the efficient attention mechanism. This is supported by the evidence that Luna with full sets of parameters achieves better performance than that with parameter-sharing, because Luna with full sets of parameters has better capacity.

Related Work

There has been signficiant prior work on improving the efficiency of Transformers, besides the three closely related works discussed in §3.4. The common techniques include, but are not limited to, weight sharing (Dehghani et al., 2018), quantization (Shen et al., 2020; Fan et al., 2020), sparse attention (Parmar et al., 2018; Kitaev et al., 2020), side memory module (Lee et al., 2019; Gupta and Berant, 2020; Goyal et al., 2021), and low-rank or compressed context (Wang et al., 2019; Ainslie et al., 2020). In this section, we briefly review some recently proposed methods. For a detailed overview we refer the readers to Tay et al. (2020b).

The general idea of these methods is that, instead of attending to the whole sequence, each token only access to a fixed, predefined range such as local neighborhoods and strided or “dilated” windows. Popular methods include local attention (Parmar et al., 2018), blockwise attention (Qiu et al., 2019), strided attention patterns (Child et al., 2019; Beltagy et al., 2020), and compressed attention (Liu et al., 2018). To make this range more flexible, Reformer (Kitaev et al., 2020) employs a hash-based similarity measure to efficiently cluster tokens into chunks and Routing Transformer(Roy et al., 2021) employ online k-means clustering on the tokens. The Sinkhorn sorting Network (Tay et al., 2020a) exposes the sparsity in attention weights by learning to sort blocks of the input sequence.

Kernel Methods.

A recently popular method to improve the efficiency of Transformers is to avoid explicitly computing the $m\times n$ attention matrix $A$ in (1) by re-writing it with kernels. Typical models leveraging kernelization are Linear Transformer (Katharopoulos et al., 2020), Performer (Choromanski et al., 2020) and Random Feature Attention (Peng et al., 2021). Since kernels are a form of approximation of the attention matrix, they can be also viewed as a form of low-rank method (Choromanski et al., 2020) that compresses the context to a shorter length, such as Linformer (Wang et al., 2019) and the proposed Luna model.

Recurrence.

The simplest technique to reduce the complexity of Transformer is to chunk input sequences into fixed blocks, with the obvious disadvantage of losing contextual information from past chunks. As discussed in Tay et al. (2020b), these models can be regarded as fixed pattern models. Transformer-XL (Dai et al., 2019) proposed a natural extension to the blockwise method to connect these blocks via a recurrence mechanism. Compressive Transformer (Rae et al., 2020) further extends Transformer-XL by maintaining a fine-grained memory of past chunk activations, which are discarded in Transformer-XL. Technically, Luna can be adapted to a recurrence method, by simply using $P$ as an inherent memory module to maintain the recurrence across segments.

Conclusion

We have introduced Luna, a simple, efficient and effective linear attention mechanism used as a drop-in substitute for regular softmax attention. By introducing an extra input with the fixed length, Luna is capable of capturing adequate contextual information while performing attention operations linearly. On three sequence modeling tasks, i.e., long-context sequence modeling, neural machine translation, and large-scale pretraining and finetuning, Luna achieves comparable or even better performance than a variety of strong baselines, while acquiring prominent gains of efficiency in both speed and memory. In future work, we are interested in combining Luna with recurrence methods where $P$ can be used as a running memory across segments of inputs. Another interesting direction would be to apply Luna to other tasks with long input sequences, such as document-level summarization and translation.

Acknowledgments and Disclosure of Funding

This material is based on research sponsored by Air Force Research Laboratory (AFRL) under agreement number FA8750-19-1-1000. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation therein. Xiang Kong was supported by U.S. DARPA AIDA Program No. FA8750-18-2-0014. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of Air Force Laboratory, DARPA or the U.S. Government.

References

Appendix: Luna: Linear Unified Nested Attention

Appendix A Experimental Details

For all tasks except Retrieval, we closely follow the model configurations in Tay et al. (2021) such as data preprocessing, data split, model architecture, batch size etc. To guarantee convergence, we train models for the Retrieval task with 20k steps instead of the 5k steps prescribed inTay et al. (2021). The hyperparameters of models in these tasks are listed in Table 7. We mainly tune three hyperparameters: learning rate, dropout and attention dropout. For the other main hyperparametrs such as batch size, number of layers and number of warmup steps, we follow the guidance of Tay et al. (2021).

A.2 Neural Machine Translation

Our experiments on WMT 2014 English-German are based on the Transformer-base model (Vaswani et al., 2017), with implementation from the FairSeq package (Ott et al., 2019). This dataset contains 4.5M parallel sentence pairs for training. We following the standard setting (Vaswani et al., 2017), using Newstest2013 as the validation set and Newstest2014 as the test set. The dataset is pre-processed following (Ma, 2020), using the scripts from FairSeq packagehttps://github.com/pytorch/fairseq. Specifically, we use word embedding with 512 dimension and 6-layer encoder/decoder with 8 multi-head attention and 2048 feed-forward dimensions. We apply 0.1 label smoothing (Szegedy et al., 2016), and perform totally $500,000$ updates to train each model. For Adam, we use start learning rate $0.0005$, set $\beta=(0.9,0.98)$, and apply the decoupled weight decay technique (AdamW) (Loshchilov and Hutter, 2019). For all the models trained with Apollo, we set the learning rate is $0.1$, $\beta=0.9$ and $\epsilon=1e^{-4}$. For learning rate scheduling, we applied linear warm up the learning rate for both Adam, and Apollo — $4000$ updates for Adam and $1000$ updates and Apollo. After learning rate warming up, we applied the inverse square root decay (Vaswani et al., 2017) to Adam. For Apollo, following Ma (2020), we decayed the learning rate at the $300,000$ and $450,000$ updates by decay rate 0.1. Gradient clips with 1.0 are applied to all the optimization methods, and the dropout ratio are set to $0.1$. Weight decay rates are $1e^{-4}$ for Adam methods and $1e^{-8}$ for Apollo. The decoding beam size is set to 5, and the checkpoints of the last 10 epochs are averaged before evaluation. For each experiment, we conducted distributed training across eight NVIDIA Tesla V100 GPUs with maximum batch size as 8192 tokens per GPU (totally $8192\times 8$ tokens per batch).

A.3 Masked Language Modeling for Large-Scale Pretraining and Finetuing

We pre-trained all the models on 64 Tesla V100 GPUs with the standard masked-language-modeling (MLM) objective and two pre-training corpus: (i)BERT version with BookCorpus (Zhu et al., 2015) and English Wikipedia (totally 16GB); (ii) RoBERTa version with BookCorpus, English Wikipedia, CC-News (Nagel, 2016), OpenWebText (Gokaslan and Cohen, 2019) and Stories (Trinh and Le, 2018) (totally 160GB). We use the standard Adam optimizer with a linear decay learning rate scheduler. Table 8 describes the hyperparameters for pre-training of Luna-128 model. For finetuning stage, we closely follow the training configuration used in released Roberta finetuning script for different tasks and main hyperparameters are listed in Table 9.