TaylorShift: Shifting the Complexity of Self-Attention from Squared to Linear (and Back) using Taylor-Softmax
This addresses a key bottleneck for processing long sequences in Transformers, offering a linear-time alternative that maintains performance, though it appears incremental as it builds on existing softmax approximations.
The paper tackles the quadratic complexity problem of self-attention in Transformers by introducing TaylorShift, a reformulation of Taylor softmax that enables full token-to-token interactions in linear time and space, showing memory efficiency for sequences as short as 800 tokens and faster inference for inputs of about 1700 tokens without accuracy loss in classification benchmarks.
The quadratic complexity of the attention mechanism represents one of the biggest hurdles for processing long sequences using Transformers. Current methods, relying on sparse representations or stateful recurrence, sacrifice token-to-token interactions, which ultimately leads to compromises in performance. This paper introduces TaylorShift, a novel reformulation of the Taylor softmax that enables computing full token-to-token interactions in linear time and space. We analytically determine the crossover points where employing TaylorShift becomes more efficient than traditional attention, aligning closely with empirical measurements. Specifically, our findings demonstrate that TaylorShift enhances memory efficiency for sequences as short as 800 tokens and accelerates inference for inputs of approximately 1700 tokens and beyond. For shorter sequences, TaylorShift scales comparably with the vanilla attention. Furthermore, a classification benchmark across five tasks involving long sequences reveals no degradation in accuracy when employing Transformers equipped with TaylorShift. For reproducibility, we provide access to our code under https://github.com/tobna/TaylorShift.