Adaptive Scale-Tuned Transformer: Extrapolating the Context Window of LLMs Through Adaptive Attention Scale
摘要
Positional encoding in the Transformer architecture provides crucial positional information that helps the model understand word order and semantic relationships. Conventional positional encodings struggle with sequences that exceed training limits due to fixed-size context windows. Although position extrapolation methods theoretically overcome the limitations of processing longer text sequences, they often need to improve in long-context tasks. Positional interpolation methods, while demonstrating great performance in long-context applications, show notable limitations in maintaining a balanced attention distribution. In order to reduce the limitations of the positional interpolation methods, we propose the Adaptive Scale-Tuned Transformer (ASTT). The ASTT employs an adaptive temperature factor matrix to adjust the scale of the attention. The matrix consists of trainable parameters tailored for each attention head in different transformer blocks. It can be fine-tuned during training to allow the model to adapt its attention mechanism to the varying lengths and complexities of the input sequences, thereby optimizing performance in long-context. Experimental results show that the ASTT performs well in a specific benchmark for long-context processing. In addition, we demonstrate that ASTT can extrapolate beyond the context window of a fine-tuning dataset.