Hierarchical Shift Mixing -- Beyond Dense Attention in Transformers
Robert Forchheimer
TL;DR
The paper tackles the quadratic complexity of dense attention in Transformers by proposing Hierarchical Shift Mixing (HSM), a framework that distributes pairwise token interactions across layers to achieve linear-time complexity while preserving causality. It demonstrates a spectrum of HSM variants, from simple scalar weighting to nonlinear gated and fusion mechanisms, and shows that many variants reach performance close to GPT. Additionally, it introduces hybrid HSM-GPT architectures that replace selected GPT layers with HSM layers, achieving equal or better predictive performance with faster training. The findings suggest that efficient token mixing, coupled with careful parameter allocation (e.g., reallocating capacity to FFNs), can yield practical, scalable improvements for Transformer models, though scaling to larger datasets remains an open question.
Abstract
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.