Gated Linear Attention Transformers with Hardware-Efficient Training

Understand

Transformers with linear attention allow for efficient parallel training but can simultaneously be formulated as an RNN with 2D (matrix-valued) hidden states, thus enjoying linear-time inference complexity.

• However, linear attention generally underperforms ordinary softmax attention.
• Moreover, current implementations of linear attention lack I/O-awareness and are thus slower than highly optimized implementations of softmax attention.
• This work describes a hardware-efficient algorithm for linear attention that trades off memory movement against parallelizability.