Pretraining Without Attention (2023)

Abstract: Transformers have been essential to pretraining success in NLP. While other architectures have been used, downstream accuracy is either significantly worse, or requires attention layers to match standard benchmarks such as GLUE. This work explores pretraining without attention by using recent advances in sequence routing based on state-space models (SSMs). Our proposed model, Bidirectional Gated SSM (BiGS), combines SSM layers with a multiplicative gating architecture that has been effective in simplified sequence modeling architectures. The model learns static layers that do not consider pair-wise interactions. Even so, BiGS is able to match BERT pretraining accuracy on GLUE and can be extended to long-form pretraining of 4096 tokens without approximation. Analysis shows that while the models have similar average accuracy, the approach has different inductive biases than BERT and scales more efficiently to longer sequences.

## introduction transformers are the de facto model architecture for nlp pretraining @xcite . since bert @xcite , they have proven central to nlp tasks with their ability to learn effectively on large unlabeled datasets. specifically, the use of attention as a central routing component seems to be critical to empirical success on downstream tasks. other architectures have been proposed but require attention layers for high-accuracy @xcite . is the centrality of attention in pretraining due to inductive bias or computational convenience? this question is complicated by the properties of common sequence routing layers: recurrent neural network (rnn) models do not scale as well as attention, whereas convolutional neural networks (cnns) can not easily model long-distance dependencies. state-space models (ssms) for deep learning provide a promising alternative. recent works show that ssms are a competitive architecture for long-range sequence modeling @xcite . ssms achieve strong results on speech generation @xcite and on the long range arena benchmark @xcite outperform standard and long-range transformer architectures @xcite @xcite . in addition to improving accuracy, ssm-based routing does not have quadratic complexity as the length of the sequence grows. concretely, the model provides a way to achieve rnn-like long-range dependencies with cnn-like training speed. this work proposes an architecture for applying ssms using a bidirectional gated ssm (bigs) model for bert-style pretraining. bigs uses ssmrouting at its core as a replacement for attention. however, this change alone significantly degrades the representational capacity of the model. to target this issue, we develop a multiplicative gating architecture @xcite @xcite . in combination, this leads to a simpler routing approach that remains surprisingly effective at modeling necessary interactions. experiments compare ssms to standard nlp pretraining. while we find that ssms by themselves underperform on nlp pretraining tasks, bigs is able to match the performance of a bert model when trained on the same data in a controlled setting. by additionally pretraining on longerlength instances, the model is able to grow without approximation to extend to input sequences of length 4,096. analysis shows the importance of multiplicative gating in fixing specific issues of variable-length textual input. all models from this work will be available open-source (apache 2.0 license) upon release. ## related work prior to bert, promising pretraining approaches for learning contextual representations were learned using rnn-based models @xcite . while important precursors, their accuracy did not scale with data or compute as well as transformers. this gap remains even when back-porting best-practices from transformer pretraining @xcite . recently @xcite explored pretraining with several convolutional (cnn) variants. results show that cnn without attention does not perform well, although they note benefits in routing speed. @xcite propose fnet which replaces the attention layer with a fourier transform. without attention, this achieves 92-97% results on glue @xcite . other works have used cnn-based models with multiplicative gating for nlp tasks such as machine translation @xcite . we believe bigs is the first model to achieve bert-level transfer learning on the glue benchmark without attention. researchers have begun to use state-space models for nlp tasks, and have primarily focused on auto-regressive language modeling. in s4 @xcite and its variants @xcite , researchers experimented with language modeling, achieving promising results, though slightly worse than transformers. gated state space adapts a ssm plus gating approach to language modeling @xcite . concurrent to this work, @xcite propose h3 which closes the gap in auto-regressive language modeling, and with two attention layers outperforms transformers on openwebtext. finally, a related method, mega @xcite combines exponential moving average routing with a simple attention unit to outperform transformer baselines. our approach instead focuses on bidirectional masked language modeling and questions of downstream generalization. ## bigs model we consider two different architectures for ssm pretraining: a stacked architecture (stack) and a multiplicative gated architecture (gated) shown in figure 2 . transformer architecture the stack architecture with self-attention is equivalent to the bert / transformer model. we replace the attention block with two sequential ssm blocks to mimic the na-ture of bi-directional self-attention. ## experimental setup experiments compare the performance of ssmbased models to attention-based models on several standard fine-tuning benchmarks. experiments control for total parameter-size and amount of pretraining in terms of the number of tokens. all models are on the order of magnitude of bert-large at around 350m parameters; all gated ssm models use 23 layers and stack models 24 to match parameter count. in order to run ablation tests, we consider three different pretraining scales: 11b (short), 29b (medium), and 97b (full) tokens. models and architectures are roughly similar in training speed at this length. the 11b (short) training scale is roughly equivalent to the "24h bert" setting typically used in research studies @xcite . full training is closer to the original bert model which was trained on 128b tokens. for all pretraining, we follow the training data and masking strategy of @xcite . since roberta @xcite shows it does not hurt accuracy, we use only masked language modeling and not next-sentence prediction. we preprocess and mask tokens offline for all models for consistency, with maximal sequence length to be 128. we use a grid search on perplexity to select configurations of weight decay and learning rate; other hyperparameters follow @xcite . for ssm, we use a cosine decay learning rate scheduler, which starts at 0, warms up to the peak learning rate, and then decays back @xcite . pretraining is done with length 128 token sequences. in order to adapt to longer sequences we apply continued pretraining. to adapt to 512 tokens for the squad dataset, we follow the protocol of @xcite and train on longer sequences of the same pretraining dataset. to adapt to 4,096 tokens, we follow the longformer (beltagy et al., 2020) protocol and continue training the bigs model on the text of length up to 4,096 tokens long, for 10k more steps using their proposed training corpus of longer documents. for 4,096 tokens, we also use a smaller bigs model (around 130m) so that it is comparable in size longformer-base and bart-base models. we note that longformer (led) and bart are based on superior underlying models that are trained significantly longer. for downstream task fine-tuning, we average the embeddings of non-padding tokens to create the sentence representation and add a classification head. this approach yields superior performance @xcite . see figure 2 for details. we fine-tune rte, mrpc, and sts-b from a mnli checkpoint following the convention by @xcite . we average results of six runs and report accuracy for mnli, qnli, rte, sst-2 and f1 score for qqp, mrpc and matthew's correlation for cola and spearman's correlation for sts-b. all models are comparable to bert-large in size. (bottom) reported comparable results for other non-attention-based pretraining models based on cnns, lstms and fnet @xcite @xcite . bert 1 represents the official bert result @xcite , and bert 2 represents the result using an mnli checkpoint for other nli tasks @xcite . we useto denote those results were not reported by previous research. our ssm implementation is based on the annotated s4 1 @xcite , and our pretraining uses the template from hugging face transformers 2 @xcite . we experimented with variants of ssms and found they performed similarly; experiments use s4d @xcite for simplicity. note that for a fair comparison, we keep the size of the gated architecture comparable to a stacked architecture and our bert implementation. 1 https://srush.github.io/annotated-s4 2 https://github.com/huggingface/transformers 6 results ## conclusion we propose bigs as a model for pretraining without attention. bigs makes use of ssm-based routing and multiplicative gating. results show that ssms alone perform poorly in a stacked architecture, but gating helps them to generalize. as far as we are aware, this architecture is the first to replicate bert results without attention. this work opens up many interesting questions. we experimented with adapting to longer text, but ssm-based models could be pretrained fully on much longer sequences. combining ssms with reductions in feed-forward costs could give further optimizations. finally, we took the steps in exploring the syntactic properties of ssms, but need further probing of how their internal representations lead to these properties.