Learnable Low-Rank Polynomial Sketch for Effective Linear Attention

AI summary

Overview

Research area: Efficient sequence modeling — specifically feature-map design for linear attention in Transformers, viewed through the lens of polynomial kernel approximation.

Technical level: Advanced. The paper leans on tensor rank (CP) decompositions, column-wise Kronecker products, SVD-based optimality arguments, and tensor spectral norms. The empirical sections are readable on their own, but the core methodology requires comfort with kernel methods and tensor algebra.

Scope in one sentence: The paper derives the theoretically optimal way to compress a degree-p polynomial kernel into a fixed-size feature map, turns that result into a practical learnable low-rank parameterization called LLoPS, and shows it beats prior linear attention kernels on zero-shot NLP, Long Range Arena, kernel-approximation error, and runtime.

What This Paper Is About

Softmax attention costs quadratic time and memory in sequence length, and linear attention fixes this by replacing the exponential kernel with a feature map φ so that attention can be computed as φ(Q)(φ(K)ᵀV). Polynomial kernels, (qᵀk)^p, are known to mimic softmax's sparse, "spiky" behavior, but their exact feature map vec(x^⊗p) has dimension d^p and is unusable. Existing fixes — BASED's project-then-expand scheme and PolySketchFormer's data-independent random sketch — are, the authors argue, suboptimal compressions of the same target kernel. The goal is to characterize the best possible compression at a fixed output dimension m and then build a computationally feasible kernel that realizes it.

Key Contributions

1. A theory of optimal polynomial sketching. The authors prove (Theorem 3.1) that when the output dimension m is fixed, the best possible feature map is a plain linear projection of the exact polynomial features, φ*_Q(Q) = ψ_p(Q)W*_Q and φ*_K(K) = ψ_p(K)W*_K, with W*_Q, W*K ∈ ℝ^{d^p × m}. The minimum achievable error is exactly the tail of the SVD of A = (QKᵀ)^{*p}, namely ∑{i=m+1}^{n} σ_i².

2. LLoPS: a learnable low-rank realization. Because the optimal projector still has d^p rows, each of its m columns is reshaped into an order-p tensor and approximated by a rank-t CP decomposition. Reshaping back yields the feature map φ_Q(Q) = ∏_{j=1}^{p} QΘ_j^{(Q)} (and analogously for K), which costs O(ndmp) time with O(dmp) parameters instead of O(nmd^p).

3. Error analysis and a unifying view of prior kernels. Theorem 3.4 gives an exact expression for the residual of the low-rank approximation in terms of tensor spectral norms τ_i. The paper then shows that element-wise polynomial feature maps (Θ₁ = … = Θ_p), BASED (a specific Kronecker-product parameter choice), and PolySketchFormer (fixed, randomized, sparse Θ) are all special or constrained cases of LLoPS, so LLoPS is strictly more general at equal dimension.

4. Empirical validation across four evaluation regimes. Zero-shot NLP at 360M and 1B parameters, Long Range Arena, direct kernel-approximation error on real query/key vectors, and forward-pass latency — LLoPS leads among linear-complexity methods in all of them.

Main Findings

• Zero-shot NLP, 360M scale: BASED+LLoPS reaches 36.26 average accuracy versus 35.00 for BASED+CT (BASED under the same continued-training budget), 34.76 for the original BASED, 32.34 for Mamba, and 43.05 for softmax attention. LLoPS is best on WinoGrande (53.04), PIQA (66.70), LAMBADA (29.23), HellaSwag (37.07), ARC-E (43.94), SWDE (31.41), and SQuAD (31.94); it ties BASED+CT on ARC-C (25.51) and trails on FDA (7.53 vs 8.35).

• Zero-shot NLP, 1B scale: BASED+LLoPS averages 42.91 versus 42.16 for BASED+CT, 40.88 for BASED, 37.55 for Mamba, and 49.26 for softmax attention. It leads on HellaSwag (42.69), ARC-E (47.31), FDA (26.32), SWDE (47.97), and SQuAD (36.36), and ties BASED on WinoGrande (51.62).

• A documented data artifact, not a kernel failure: At 360M, both BASED and LLoPS lose substantial FDA accuracy after continued training on SlimPajama-6B. The authors attribute this to distribution mismatch and possible catastrophic forgetting, noting the dataset is 50 times smaller than the Pile used for the original checkpoints.

• Long Range Arena: LLoPS attains the best average among linear attention methods at 67.16, a 2.28% improvement over the strongest linear baseline, and places in the top three on every one of the five tasks (Image 57.54, Text 73.94, Retrieval 82.17, ListOps 48.29, Pathfinder 73.84). Softmax attention still leads overall at 69.90, consistent with large-scale language modeling observations.

• Approximation error: Measured as relative mean absolute error on query/key vectors sampled from a transformer trained on synthetic associative recall. For p = 2, LLoPS averages 0.022 versus 0.054 for a two-layer MLP neural sketch and 0.293 for the PolySketchFormer-style data-independent sketch. For p = 3, LLoPS averages 0.038 versus 0.105 (neural) and 0.741 (PSF). LLoPS wins in every reported layer (0, 8, and the final layer, labeled Layer 15 in the table).

• Efficiency: Over 1,000 timed forward passes with input dimension 128 and sequence lengths from 1,000 to 10,000 tokens, both linear methods are far faster than softmax attention at long lengths, and LLoPS is consistently faster than BASED, which the authors credit to LLoPS relying mainly on matrix multiplications.

• Why complex nonlinear sketches don't help: The authors argue the MLP sketch underperforms because Theorem 3.1 shows a linear projection suffices for optimal compression, and because the structured decomposition is easier to optimize than an MLP.

Methodology in Plain English

The argument proceeds in three moves.

First, frame the goal precisely: pick feature maps φ_Q and φ_K of output size m that minimize the Frobenius error between φ_Q(Q)φ_K(K)ᵀ and the exact polynomial kernel matrix (QKᵀ)^{*p}. Since any product of two n×m matrices has rank at most m, the Eckart–Young theorem says the best you can possibly do is keep the top m singular components of the target matrix. That bound is achievable by linearly projecting the exact polynomial features — so elaborate nonlinear feature maps are unnecessary, and the only real question is how to make that linear projection affordable.

Second, attack the cost. The optimal projector has d^p rows, which is the whole problem. Each of its m columns, reshaped into a p-way tensor, is approximated by a rank-t CP decomposition: a short sum of rank-one outer products of learnable vectors. The crucial algebraic fact is that the inner product between such an outer product and vec(x^⊗p) equals the product of the p scalar projections ⟨w_{j,r}, x⟩. So instead of ever materializing d^p-dimensional features, you compute a handful of d×m matrix products (QΘ_j), multiply the results element-wise across j, and sum over the t components. Cost drops from O(nmd^p) to O(ndmtp); in practice the authors fix t = 1, giving O(ndmp) and a strong approximation.

Third, connect the pieces. Because BASED corresponds to a constrained Kronecker-structured choice of Θ and PolySketchFormer to fixed random sparse Θ, LLoPS contains both as special cases and therefore cannot do worse at equal dimension. A theorem quantifies the remaining gap via tensor spectral norms. If non-negative attention scores are desired for stable training, the authors note the feature map can be squared, following PolySketchFormer.

Experiments use a deliberate shortcut: since BASED is a special case of LLoPS, models are initialized from released BASED checkpoints and continually trained on SlimPajama-6B, with BASED given an identical continued-training budget (BASED+CT) so the comparison isolates the kernel rather than extra compute. Polynomial degree is p = 4 for the 360M model and for all LRA runs, and p = 2 for the 1B model.

Why This Matters

Impact on research. The paper shifts the linear attention feature-map conversation from ad hoc construction to a stated optimality target. If a single linear projection of polynomial features is optimal, then the design space collapses to "how well can you approximate that projector," which is a cleaner and more falsifiable research question. The unification of element-wise polynomial maps, BASED, and PolySketchFormer under one parameterization also gives the field a common vocabulary for comparing kernels rather than architectures.

Real-world applications:

• Long-document understanding, where contexts reach tens or hundreds of thousands of tokens and quadratic attention is infeasible.
• Code modeling, where repository-scale context windows are valuable.
• Retrieval-augmented generation, where many retrieved passages must be attended over at once.
• Streaming or on-device inference, where linear attention's constant-memory recurrent state matters.

Industry relevance. The method is a drop-in replacement for the attention layer inside an otherwise unchanged architecture, adds only O(dmp) parameters, and is reported as faster than BASED in the same PyTorch environment — an attractive profile for teams already running hybrid linear-attention models at scale and looking for accuracy gains without retraining from scratch.

Future Directions

• Scaling the evidence. The largest reported models are 360M and 1B parameters, and the 1B result comes from continued training on SlimPajama-6B rather than from-scratch pretraining on a large corpus. Whether the advantage persists at frontier scale, and under from-scratch training, is untested here.

• The data-availability caveat. The original baselines were pretrained on the Pile, which the authors note is no longer publicly available, forcing a 50-times-smaller substitute and producing the FDA regression at 360M. A clean from-scratch comparison on a stable, large corpus would settle how much of the reported gap is kernel and how much is training recipe.

• Beyond t = 1 and fixed degree. The authors set t = 1 because larger values multiply parameters, and they hand-pick p (2 or 4) per setting. Whether rank or degree should be adaptive, scheduled, or learned per layer is left open.

• Composition with state update rules. The authors explicitly flag state-update research (as in Mamba-like or DeltaNet-style recurrences) as orthogonal and uncombined. Merging LLoPS's feature map with improved state update rules is the obvious next experiment, as is closing the remaining gap to softmax attention on LRA (67.16 vs 69.90).

Target Audience

Researchers and graduate students working on efficient attention and long-context architectures who want both a theoretical handle on why polynomial kernels work and a concrete, implementable kernel. It is also useful for practitioners building or fine-tuning linear-attention models who need to decide whether a kernel swap is worth the effort, and for theorists interested in data-dependent sketching guarantees, which the authors claim are established here for the first time for polynomial sketches. Readers without tensor-algebra background will get the most from the introduction, the special-case analysis in Section 3.3, and the experiments; the full derivations are in the appendices, whose contents are referenced but not included in the material available for this summary.

Authors’ abstract

Softmax attention in Transformers suffers from quadratic complexity in sequence length, making it impractical for long-context applications. Linear attention alleviates this issue by replacing the exponential kernel with alternative functions that enable linear-time computation. Among existing linear attention approaches, recent studies have shown that polynomial kernels are particularly effective, as they exhibit sparse and spiky behavior similar to softmax attention, which emphasizes large dot products while suppressing irrelevant interactions. However, the exact computation of high-degree polynomial kernels is infeasible for high-dimensional representations. As a result, prior work relies on approximate polynomial kernels. This introduces a non-negligible approximation error. In this paper, we show that, from the perspective of polynomial kernel approximation, existing linear attention methods are still suboptimal. We propose \textbf{L}earnable \textbf{Lo}w-Rank \textbf{P}olynomial \textbf{S}ketch (LLoPS), a principled and flexible framework for approximating polynomial kernels with linear attention. Our method learns a low-rank polynomial sketch that provably achieves a strictly smaller approximation error than existing approaches. Experiment results show that LLoPS achieves the highest performance across extensive benchmarks, comparing with various linear attention baselines.