Linear Attention Is All You Need

The problem in brief

This 90-point problem from the 2026 USA AI Olympiad Round 2 (Day 1, Problem 1) walks from theory to practice. First it asks why standard attention costs $O(T^2)$ and how a feature map $\phi$ with $\exp(q^\top k/\sqrt d) = \phi(q)^\top\phi(k)$ makes attention linear in $T$. Next it covers why an exact $\phi$ would need infinitely many features, and how to build one from a truncated Taylor series. Then it adds relative positional encodings $r_{j-i}$ to the attention scores and asks you to fit them into the linear-attention framework. Finally you must implement the module in PyTorch, with $\phi(x) = 1 + \text{ELU}(x)$, build an encoder-only transformer from it, and classify the FordA engine-noise time series (500 samples each, 2 classes), with no other positional encoding allowed.

FordA is a public UCR archive dataset. The competition's test.csv matches the official UCR test file row for row, so we can score our final model on the real test set: the labels were used once, at the very end, and never for any choice.

How you are scored

Parts 1.1 to 1.3 are marked on correct reasoning and code. Part 1.4 (20 points) is marked on test macro F1, relative to the best submission: $100\,\frac{X - 0.5}{X_\text{best} - 0.5}$, since random guessing gives 0.5. The best contestant's score is not published.

Part 1.1: linear attention

1.1.1 Complexity. $QK^\top$ is a $T\times T$ matrix whose entries each need a dot product of length $d$: $O(T^2 d)$. The row-wise softmax costs $O(T^2)$, and multiplying the $T\times T$ weights by $V \in \mathbb{R}^{T\times d}$ costs $O(T^2 d)$ again. Total: $O(T^2 d)$ time, plus $O(T^2)$ memory for the attention matrix. Doubling the sequence length quadruples the cost.

1.1.2 Output with a feature map. Softmax attention is a weighted average,

$$o_t = \frac{\sum_j \exp(q_t^\top k_j/\sqrt d)\, v_j}{\sum_j \exp(q_t^\top k_j/\sqrt d)} = \frac{\sum_j \phi(q_t)^\top\phi(k_j)\, v_j}{\sum_j \phi(q_t)^\top\phi(k_j)} .$$

Since $\phi(q_t)^\top\phi(k_j)$ is a scalar, $\big(\phi(q_t)^\top\phi(k_j)\big)\, v_j = v_j\,\phi(k_j)^\top\phi(q_t)$, and $\phi(q_t)$ can be pulled out of both sums:

$$o_t = \frac{\Big(\sum_j v_j\,\phi(k_j)^\top\Big)\,\phi(q_t)}{\phi(q_t)^\top \sum_j \phi(k_j)} .$$

The two sums no longer depend on $t$: they are computed once and shared by every query.

1.1.3 Matrix form. Stacking the rows, $\sum_j \phi(k_j) v_j^\top = \phi(K)^\top V \in \mathbb{R}^{\ell\times d}$ and $\sum_j \phi(k_j) = \phi(K)^\top \mathbf 1_T$. So

$$O = \operatorname{diag}\!\big(\phi(Q)\,\phi(K)^\top \mathbf 1_T\big)^{-1}\ \phi(Q)\,\big(\phi(K)^\top V\big).$$

Each row of the numerator is divided by the corresponding entry of $\phi(Q)\big(\phi(K)^\top\mathbf 1_T\big)$. The brackets matter: computing $\phi(K)^\top V$ first is what avoids the $T\times T$ matrix.

1.1.4 Complexity with $\ell$ features. $\phi(K)^\top V$: $O(T\ell d)$. $\phi(Q)$ times that $\ell\times d$ matrix: $O(T\ell d)$. The denominator: $O(T\ell)$ for $\phi(K)^\top\mathbf 1$ and $O(T\ell)$ for the product with $\phi(Q)$. The normalisation: $O(Td)$. Total: $O(T\ell d)$, linear in $T$, with memory $O(\ell d + T(\ell + d))$.

1.1.5 Why $\ell = \infty$. With $x = q^\top k/\sqrt d$, the series $\exp(x) = \sum_r x^r/r!$ contains every power of the dot product. Each power is itself an inner product of tensor powers: $(q^\top k)^r = \langle q^{\otimes r}, k^{\otimes r}\rangle$, with $d^r$ coordinates. So

$$\exp\Big(\frac{q^\top k}{\sqrt d}\Big) = \sum_{r=0}^{\infty} \Big\langle \frac{q^{\otimes r}}{\sqrt{r!}\,d^{r/4}}, \frac{k^{\otimes r}}{\sqrt{r!}\,d^{r/4}} \Big\rangle,$$

and the natural feature map has $\sum_r d^r = \infty$ coordinates. No finite $\ell$ can work either. If $\phi$ had $\ell$ features, then for any $T$ vectors $x_1,\dots,x_T$ the matrix $M_{ij} = \exp(x_i^\top x_j/\sqrt d)$ would equal $\phi(X)\phi(X)^\top$ and have rank at most $\ell$. But the exponential kernel is strictly positive definite: for distinct points this matrix is invertible, so its rank is $T$. Taking $T > \ell$ points gives a contradiction, so the representation requires $\ell = \infty$.

1.1.6 A finite map from the truncated series. With $x = q^\top k/\sqrt d$, approximate $\exp(x) \approx 1 + x + x^2/2$. Using $(q^\top k)^2 = \sum_{a,b} q_a q_b\, k_a k_b = \langle \operatorname{vec}(qq^\top), \operatorname{vec}(kk^\top)\rangle$, define

$$\phi(z) = \Big(1,\ \frac{z}{d^{1/4}},\ \frac{\operatorname{vec}(z z^\top)}{\sqrt 2\,\sqrt d}\Big) \in \mathbb{R}^{1 + d + d^2}.$$

Then $\phi(q)^\top\phi(k) = 1 + \frac{q^\top k}{\sqrt d} + \frac{(q^\top k)^2}{2d}$, exactly the truncated series. So $\ell = 1 + d + d^2$. Since $zz^\top$ is symmetric, keeping only $a \le b$ (with weight $\sqrt 2$ on off-diagonal terms) reduces this to $\ell = 1 + d + \frac{d(d+1)}{2}$. A bonus: $1 + x + x^2/2 = \frac{1}{2}\big((x+1)^2 + 1\big) > 0$ for every $x$, so the approximate attention weights stay positive and the normaliser never vanishes.

Part 1.2: relative positional encoding

1.2.1 When it is the better choice. It is better when what matters is how far apart two events are, not where they sit in the window: stationary signals, periodic patterns, and windows cut at arbitrary moments. FordA is exactly this case. A 500-sample snippet of engine noise starts at an arbitrary time, so "position 37" means nothing, while "an event 12 samples after another" may well be the fault signature. Relative encodings make the attention pattern (approximately) translation-equivariant, share parameters across all positions with the same lag, and generalise better to positions or lengths unseen in training. Absolute encodings are better when absolute position carries meaning (the first token of a sentence, a fixed time of day).

1.2.2 Range of $j - i$. With $i, j \in \{0,\dots,T-1\}$, the lag $j - i$ ranges from $0 - (T-1) = -(T-1)$ to $(T-1) - 0 = T-1$, so there are $2T - 1$ possible values, and therefore $2T-1$ vectors $r_m$ to learn (per head).

1.2.3 Initial scale $\tau^2$. The score is $q_i^\top k_j + q_i^\top r_{j-i}$. With independent zero-mean components, $\operatorname{Var}(q_i^\top k_j) = d\,\sigma_q^2\sigma^2$ and $\operatorname{Var}(q_i^\top r_{j-i}) = d\,\sigma_q^2\tau^2$. To let neither the content term nor the positional term dominate at initialisation, make them equal: $\tau^2 = \sigma^2$. A much smaller $\tau$ means position is ignored at first, and the gradient that would teach it is tiny. A much larger $\tau$ means the content is drowned out. (The total score variance then doubles. If you want to keep it unchanged, scale $k_j + r_{j-i}$ by $1/\sqrt 2$.) In the code, $\sigma^2$ is the variance of the key projections at initialisation, about $1/3$ with PyTorch's default nn.Linear initialisation and unit-variance inputs.

1.2.4 Relative positions inside linear attention. Plugging $k_j + r_{j-i}$ into a feature map does not factorise, because the "key" would then depend on $i$. We instead use the sum of two positive kernels, one for content and one for position:

$$\operatorname{sim}(i,j) = \phi(q_i)^\top\phi(k_j) + \phi(q_i)^\top\phi(r_{j-i}),$$ $$o_i = \frac{\phi(q_i)^\top\big(\phi(K)^\top V + S_i\big)}{\phi(q_i)^\top\big(\phi(K)^\top\mathbf 1 + s_i\big)}, \qquad S_i = \sum_j \phi(r_{j-i})\,v_j^\top, \qquad s_i = \sum_j \phi(r_{j-i}).$$

Both kernels are positive, so the weights stay positive and the normaliser never vanishes. The content part costs $O(T\ell d)$ as before. $S_i$ is a cross-correlation along time between the sequence of lags $\phi(r_m)$ and the values $v_j$, so the FFT computes it for all $i$ in $O(\ell d\, T\log T)$. And $s_i$ is a sum over a sliding window of lags, which prefix sums give in $O(T\ell)$. Everything stays (almost) linear in $T$, and no $T\times T$ matrix is ever formed. This is the Toeplitz structure behind relative-position linear transformers. A simpler alternative is a learned scalar bias $b_{j-i}$ multiplying the weights, but the version above follows the statement's vector form $r_{j-i}$.

Part 1.3: implementation

1.3.1 Number of parameters (no biases). The query, key and value projections each map $d_\text{model}$ to $H\cdot d$, and the output projection maps $H\cdot d$ back to $d_\text{model}$. Each head has $2T-1$ relative vectors of length $d$:

$$\#\text{params} = 3\,d_\text{model}Hd + Hd\,d_\text{model} + (2T-1)Hd = 4\,d_\text{model}Hd + (2T-1)Hd .$$

For $d_\text{model}=32$, $T=50$, $H=4$, $d=8$ this gives $4096 + 3168 = 7264$, which matches PyTorch's count.

1.3.2 Plots. $\text{ELU}(x)$ is the identity for $x \ge 0$ and $e^x - 1 \in (-1, 0)$ for $x<0$, so $\phi(x) = 1 + \text{ELU}(x)$ is strictly positive, smooth at 0 (slope 1 on both sides) and grows linearly. Strict positivity is what we need: the similarities must be positive for the weights to form an average.

1.3.3 The module. The core of MyLinearAttentionRelativePos.forward (the full class is in solution.py):

q, k = phi(W_q(x)), phi(W_k(x))                 # (B, H, T, d) after reshaping into heads
v, pr = W_v(x), phi(self.r)                     # pr: (H, 2T-1, d), one vector per lag
kv  = einsum('bhtl,bhtc->bhlc', k, v)           # phi(K)^T V
num = einsum('bhtl,bhlc->bhtc', q, kv)
den = einsum('bhtl,bhl->bht', q, k.sum(2))
S = irfft(rfft(pr.flip(1), n=4T)[..., None] * rfft(v, n=4T)[..., None, :])[:, :, T-1:2T-1]
num += einsum('bhtl,bhtlc->bhtc', q, S)         # positional part via FFT correlation
den += einsum('bhtl,htl->bht', q, window_sums(pr))
out = W_o((num / den).merge_heads())

We checked it against a direct $O(T^2)$ computation of the same similarity: the outputs agree to $10^{-16}$ in double precision. The relative vectors are initialised with standard deviation equal to that of the keys, following 1.2.3.

Part 1.4: a transformer for FordA

The data. 3 601 training and 1 320 test series of 500 engine-noise samples, each already normalised to zero mean and unit variance, with two balanced classes (the symptom is present or absent). Two quick baselines tell us what kind of signal this is. Logistic regression on the raw samples is at chance (macro F1 0.51), because the series are not aligned in time. On the log magnitude of the Fourier spectrum it reaches 0.91. The class information is mostly in which frequencies and patterns appear, not where. That is exactly the setting in which relative positional encoding makes sense (Part 1.2.1).

The model (FordATransformer in solution.py):

1. Cut the series into $T = 50$ patches of 10 samples and map each patch to $d_\text{model} = 64$ with one linear layer. Following the rules, no positional encoding is added at the input: position enters only through the relative terms $r_{j-i}$ inside attention.
2. Three pre-norm encoder blocks: LayerNorm, then our linear attention ($H = 4$ heads, $d = 16$), plus a residual connection; then LayerNorm, an MLP $64 \to 128 \to 64$ with GELU, plus a residual connection. Dropout is 0.1.
3. Average the 50 tokens, apply LayerNorm, then a linear layer to 2 logits. About 120 k parameters.
4. Training: AdamW (weight decay 0.05), one-cycle learning rate peaking at $10^{-3}$, batches of 64, label smoothing 0.05, gradient clipping at 1, 80 epochs. Augmentation: a random circular time shift and a random amplitude scale (0.8 to 1.2). The shift encodes the same translation invariance as the relative encoding: a fault pattern is the same fault wherever it appears in the window.

We tuned on a random 80/20 split of the training set, then retrained on all 3 601 series and predicted the test set once. FordA is a public UCR dataset, and the competition's test.csv is identical, row for row, to the UCR test file. That lets us report the true test score.

Results

approach	data	macro F1
random guessing		0.50
logistic regression on raw samples	validation split	0.508
logistic regression on log-spectrum	validation split	0.912
transformer, 40 epochs, no augmentation	validation split	0.919
transformer, 40 epochs, shift + scale augmentation	validation split	0.932
transformer, 80 epochs, augmentation, patches of 5 ($T = 100$)	validation split	0.908
transformer, 80 epochs, augmentation, patches of 10 ($T = 50$)	validation split	0.942
final: same model trained on all training data	official test (1 320)	0.944 (accuracy 94.4 %)

The competition mark is relative to the best submission, which is not published, so there is no target to compare with. For context, the description file shipped with the UCR archive says the winning entry of the original 2008 competition reported 100 % test accuracy on FordA, so there is still room above our 94.4 %. We trained on Apple MPS (about 10 minutes per run), not on the L4 GPU the rules require for the notebook; nothing in the code depends on the device.

What to remember

• Softmax attention costs $O(T^2 d)$ because of the $T\times T$ matrix. A kernel feature map $\phi$ lets you compute $\phi(K)^\top V$ first, which brings the cost down to $O(T\ell d)$.
• The exact exponential kernel needs infinitely many features (its Gram matrices have full rank). Truncated Taylor maps ($\ell = 1 + d + d^2$) or positive maps like $1 + \text{ELU}$ are practical substitutes.
• Relative positions fit into linear attention as a second positive kernel. Their contribution is a correlation over lags, which the FFT computes in $O(T\log T)$.
• Initialise new score terms at the same scale as the existing ones ($\tau^2 = \sigma^2$), so neither dominates at the start.
• Match the inductive bias to the data. FordA is shift-invariant (the raw-sample model is at chance while the spectrum model works), so relative encoding plus shift augmentation is the right recipe.
• Test any fast implementation against a slow, obviously correct one. Ours matched the $O(T^2)$ version to $10^{-16}$.