Attention Alignment and Flexible Positional Embeddings Improve Transformer Length Extrapolation (2024)

Abstract: An ideal length-extrapolatable Transformer language model can handle sequences longer than the training length without any fine-tuning. Such long-context utilization capability relies heavily on a flexible positional embedding design. Upon investigating the flexibility of existing large pre-trained Transformer language models, we find that the T5 family deserves a closer look, as its positional embeddings capture rich and flexible attention patterns. However, T5 suffers from the dispersed attention issue: the longer the input sequence, the flatter the attention distribution. To alleviate the issue, we propose two attention alignment strategies via temperature scaling. Our findings show improvement on the long-context utilization capability of T5 on language modeling, retrieval, multi-document question answering, and code completion tasks without any fine-tuning. This suggests that a flexible positional embedding design and attention alignment can go a long way toward Transformer length extrapolation. The code is released at: https://github.com/chijames/T5-Attention-Alignment

## introduction pre-training large transformer language models on long sequences is inherently expensive due to self-attention's quadratic complexity w.r.t the input sequence length @xcite . even with the help of memory-efficient attention @xcite , the maximum supported input length of current open-source pretrained transformer language models are capped at 4,096 tokens @xcite , limiting their efficacy in handling long-context tasks. one notable research topic aiming to lift the input length restriction is length extrapolation @xcite . ideally, a lengthextrapolatable transformer language model is trained on short sequences and can perform equally to circumvent the recency bias, we sift through the positional embeddings of existing open-source large pre-trained transformer language models, shown in transformer length extrapolation existing research on transformer length extrapolation is mostly confined to the task of natural language modeling @xcite @xcite . unfortunately, the reported positive results do not carry over to long-context retrieval @xcite . this contrastive observation can be explained by models' short empirical receptive field @xcite . in short, the strong decaying prior of positional embeddings prevents models from accessing distant tokens that may be necessary for retrieval tasks. in this work, we improve the flexible positional embeddings of t5 to get around this limitation. transformer position interpolation instead of performing direct length extrapolation, a different line of research conducts model fine-tuning on long input sequences @xcite , where the main focus is to identify the most efficient fine-tuning scheme that can improve long-context utilization. positive results have been reported on retrieval tasks @xcite . however, we argue that fine-tuning incurs additional costs since it needs 1) gpu resources to perform long sequence finetuning with large models and 2) a pre-defined target sequence length, which still imposes a sequence length upper limit. our proposed methods can circumvent these two limitations. ## proposed methods a natural solution to the dispersed attention issue described in § 3 is to sharpen the self-attention distribution. this can be achieved by reducing the temperature τ during extrapolation. we set the extrapolation temperature τ ex such that the sharpness during training with τ @xmath0. as a measurement of sharpness, we explore the maxi- ## experiments we compare the two alignment strategies against the length-only softmax temperature scaling scheme τ = log lex l tr @xcite and longchat-13b-16k @xcite . note that longchat-13b-16k @xcite , the best baseline, was fine-tuned from llama (touvron et al., 2023) on long sequences of length 16k while our proposed methods do not need any fine-tuning. our experiments are conducted on an a6000 gpu. ## discussion the goal of this section is to explain the observations made in § 5 via the lens of temperature analysis. we visualize proposition 1 and 2 by plotting the temperature curves in figure 3 and 4. we evaluate p tr max and σ tr at the training length and σ ex at every extrapolation length considering only the 0-th layer. you may find the temperature curves for the other tasks in appendix a.4. first, while both proposed strategies lower the temperature when the input sequence length increases, the entropy alignment strategy does so more aggressively, possibly leading to its inferior performance observed in second, deciding the optimal temperature based on sequence lengths, e.g. τ = log lex l tr , is not the most robust method. it gives too high of a temperature in figure 3 compared to algorithm 1. in other words, it does not sharpen the distribution enough, possibly explaining its perplexity spike in ## conclusion in this paper, we show that the t5 model family has great potential when it comes to transformer length extrapolation. we propose the maximum probability and entropy alignment strategies to fix t5's dispersed attention issue without model fine-tuning. we conduct experiments on natural language modeling, retrieval, multi-document question answering, and code completion tasks to demonstrate the effectiveness of our proposed methods. finally, we present a simplified theoretical analysis to elucidate how the temperature is scaled to achieve attention alignment. we hope that our work can inspire future length-extrapolatable transformer designs.