Hierarchical Shift Mixing -- Beyond Dense Attention in Transformers
This addresses a scalability bottleneck in large language models for AI researchers and practitioners, offering a potential incremental improvement over existing methods.
The paper tackles the quadratic-time computational complexity of softmax-based attention in Transformers by introducing Hierarchical Shift Mixing (HSM), a framework that distributes token interactions across layers for linear-time complexity, achieving performance close to or better than baselines while reducing computational costs.
Since the introduction of the Transformer architecture for large language models, the softmax-based attention layer has faced increasing scrutinity due to its quadratic-time computational complexity. Attempts have been made to replace it with less complex methods, at the cost of reduced performance in most cases. We introduce Hierarchical Shift Mixing (HSM), a general framework for token mixing that distributes pairwise token interactions across Transformer layers rather than computing them densely within each layer. HSM enables linear-time complexity while remaining agnostic to the specific mixing function. We show that even simple HSM variants achieve performance close to softmax attention, and that hybrid architectures combining HSM with softmax attention can outperform a GPT-style Transformer baseline while reducing computational cost during both training and inference.