On the Computational Hardness of Transformers

TL;DR

This paper proves that computing multiple attention heads in transformers cannot be significantly optimized beyond straightforward methods, establishing fundamental computational lower bounds for different embedding regimes.

Contribution

It provides the first non-trivial lower bounds for multi-head multi-layer transformers, showing optimality of current algorithms under standard complexity assumptions.

Findings

In small embedding regime, computing all attention heads takes near-optimal $LHN^{2+o(1)}$ time.

In large embedding regime, computing all attention heads requires $LHN^{ ext{omega}+o(1)}$ operations, matching known upper bounds.

The paper introduces a novel application of the Baur-Strassen theorem to establish lower bounds in the large embedding regime.

Abstract

The transformer has revolutionized modern AI across language, vision, and beyond. It consists of $L$ layers, each running $H$ attention heads in parallel and feeding the combined output to the subsequent layer. In attention, the input consists of $N$ tokens, each a vector of dimension $m$. The attention mechanism involves multiplying three $N \times m$ matrices, applying softmax to an intermediate product. Several recent works have advanced our understanding of the complexity of attention. Known algorithms for transformers compute each attention head independently. This raises a fundamental question that has recurred throughout TCS under the guise of ``direct sum'' problems: can multiple instances of the same problem be solved more efficiently than solving each instance separately? Many answers to this question, both positive and negative, have arisen in fields spanning communication complexity and algorithm design. Thus, we ask whether transformers can be computed more efficiently than $LH$ independent evaluations of attention. In this paper, we resolve this question in the negative, and give the first non-trivial computational lower bounds for multi-head multi-layer transformers. In the small embedding regime ($m = N^{o(1)}$), computing $LH$ attention heads separately takes $LHN^{2 + o(1)}$ time. We establish that this is essentially optimal under SETH. In the large embedding regime ($m = N$), one can compute $LH$ attention heads separately using $LHN^{\omega + o(1)}$ arithmetic operations (plus exponents), where $\omega$ is the matrix multiplication exponent. We establish that this is optimal, by showing that $LHN^{\omega - o(1)}$ arithmetic operations are necessary when $\omega > 2$. Our lower bound in the large embedding regime relies on a novel application of the Baur-Strassen theorem, a powerful algorithmic tool underpinning the famous backpropagation algorithm.