Paper deep dive
Learnable Low-Rank Polynomial Sketch for Effective Linear Attention
Zongyue Qin, Haoran Deng, Jason Cong, Yizhou Sun
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 10/14/2026, 9:41:27 AM
Summary
The paper introduces Learnable Low-Rank Polynomial Sketch (LLoPS), a framework for approximating polynomial kernels in linear attention mechanisms. It addresses the high computational cost and approximation errors of existing methods by proposing a learnable, low-rank projection that provably minimizes approximation error compared to data-independent sketches like PolySketchFormer and BASED. LLoPS uses CP decomposition to efficiently compute these projections, achieving superior performance on benchmarks.
Entities (8)
Relation Signals (7)
Softmax Attention → hasproblem → Quadratic Complexity
confidence 95% · Softmax attention in Transformers suffers from quadratic complexity in sequence length...
LLoPS → isa → Linear Attention Framework
confidence 95% · We propose Learnable Low-Rank Polynomial Sketch (LLoPS), a principled and flexible framework for approximating polynomial kernels with linear attention.
Linear Attention → solves → Quadratic Complexity
confidence 95% · Linear attention alleviates this issue by replacing the exponential kernel with alternative functions that enable linear-time computation.
Polynomial Kernel → approximates → Softmax Attention
confidence 90% · Recent studies have shown that polynomial kernels are particularly effective, as they exhibit sparse and spiky behavior similar to softmax attention...
LLoPS → uses → CP Decomposition
confidence 90% · We further propose a structured decomposition of the optimal linear mapping using low-rank tensor rank decompositions (also known as CP decompositions)...
LLoPS → outperforms → PolySketchFormer
confidence 85% · Our method learns a low-rank polynomial sketch that provably achieves a strictly smaller approximation error than existing approaches.
LLoPS → outperforms → BASED
confidence 85% · Experiment results show that LLoPS achieves the highest performance across extensive benchmarks, comparing with various linear attention baselines.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.14219v1
- Canonical: https://arxiv.org/abs/2610.14219v1
Trouble viewing inline? Open PDF directly →
Full Text
62,401 characters extracted from source content.
Expand or collapse full text
Learnable Low-Rank Polynomial Sketch for Effective Linear Attention Zongyue Qin Haoran Deng Jason Cong Yizhou Sun Affiliation: Department of Computer Science Affiliation: University of California Los Angeles Affiliation: CA, USA Email: qinzongyue,denghaoran,cong,yzsun@cs.ucla.edu 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 Learnable Low-Rank Polynomial Sketch (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. 1 Introduction Transformers have become the dominant architecture for sequence modeling across natural language processing, vision, and multimodal tasks. Despite their empirical success, the quadratic time and memory complexity of the attention mechanism with respect to sequence length severely limits scalability to long-context settings. This bottleneck is particularly problematic in emerging applications such as long-document understanding, code modeling, and retrieval-augmented generation, where context lengths can easily reach tens or hundreds of thousands of tokens. Linear attention methods (Xiong et al., 2021; Team et al., 2025; Schlag et al., 2021) address this challenge by retaining the core attention formulation while reducing its complexity from quadratic to linear in sequence length. By rewriting softmax attention as a kernelized dot-product and exploiting associativity, linear attention enables streaming computation with constant memory. Recent large-scale deployments (Team et al., 2025) demonstrate that linear attention can attain a strong empirical performance while supporting extremely long contexts, thus making it a promising direction for next-generation foundation models. Existing work on linear attention can be broadly categorized into two complementary directions. The first focuses on improving the state update rule, designing better ways to accumulate and normalize the running key-value statistics over time (Team et al., 2025; Yang et al., 2024). The second direction, which this paper focuses on, aims to improve the feature map of query and key vectors so that the induced kernel more faithfully approximates softmax attention (Arora et al., 2024; Zhang et al., 2024). Prior studies on feature map design have observed that high-order polynomial kernels can approximate the softmax kernel surprisingly well (Kacham et al., 2023; Arora et al., 2024; Zhang et al., 2024). Intuitively, both exponential and polynomial kernels can exhibit sparse, spiky behavior that emphasizes large dot-products while suppressing irrelevant interactions. However, computing exact polynomial kernels is infeasible for high degree kernel and high-dimensional vectors, as the computation complexity grows exponentially with the polynomial degree. To address this challenge, PolySketchFormer (Kacham et al., 2023) exploits a data-independent sketching technique to approximate exact polynomial kernels. More recently, BASED (Arora et al., 2024) realizes state-of-the-art performance in language modeling by combining linear attention and sliding window attention. In their linear attention formulation, queries and keys are first compressed into a low-dimensional space and then the exact polynomial kernel is computed in that space. In this paper, we show that from the perspective of polynomial kernel approximation, the previous methods that approximate polynomial kernel are suboptimal. Our key insight is that, when the output feature dimension is fixed, there exists a theoretically optimal sketch that minimizes the approximation error to the target polynomial kernel. We prove that a linear projection for the p-th order outer products of input vectors can achieve this optimal approximation, whereas existing approaches such as BASED (Arora et al., 2024) and PolySketchFormer (Kacham et al., 2023) can be viewed as constrained or data-independent approximations to this optimal mapping. Building on this observation, we further propose a structured decomposition of the optimal linear mapping using low-rank tensor rank decompositions (also known as CP decompositions) and column-wise Kronecker products, which significantly reduces computational costs. We provide theoretical analysis showing that this structured approximation is strictly tighter than prior feature-map constructions under the same dimensional budget. Our contributions are summarized as follows: (1) We provide a theoretical characterization of the optimal sketching strategy for approximating polynomial kernels in the context of linear attention. (2) Motivated by this analysis, we propose a novel Learnable Low-Rank Polynomial Sketch (LLoPS), a linear attention kernel that accurately approximates polynomial attention while remaining computationally efficient. (3) We further derive approximation error analysis for LLoPS and show that several existing linear attention kernels can be interpreted as special cases within our unified framework. (4) We demonstrate through experiments that LLoPS consistently improves the effectiveness of linear attention transformers across a broad range of benchmarks. 2 Preliminaries 2.1 Softmax-based Full Attention Given a sequence of queries ∈ℝn×d Q ^n× d, keys ∈ℝn×d K ^n× d, and values ∈ℝn×dv V ^n× d_v, standard Transformers compute attention as Attn(,,)=softmax(⊤d),Attn( Q,K,V)=softmax\! ( QK d ) V, where the softmax is applied row-wise. Because this formulation has (n2)O(n^2) time and memory complexity due to the explicit computation of the attention matrix ⊤ QK , it is prohibitively expensive for long input sequences. 2.2 Linear Attention Previous studies have observed that softmax attention can be interpreted as kernelized dot-product attention with an exponential kernel. By approximating the exponential function with a feature map ϕ:ℝd→ℝmφ:R^d ^m and exploiting associativity, attention can be rewritten as Attn(,,)=(ϕ()ϕ()⊤)=ϕ()(ϕ()⊤),Attn( Q,K,V)= (φ( Q)φ( K) ) V=φ( Q) (φ( K) V ), where ϕ(),ϕ()∈ℝn×mφ( Q),φ( K) ^n× m are mapped features with dimension m and ϕ()ϕ()⊤φ( Q)φ( K) is the inner product in the new feature space that approximates softmax(⊤d)softmax\! ( QK d ).11 1 We use ϕ()φ( Q) to denote applying ϕ(⋅)φ(·) to each row vector in Q. Similar for ϕ()φ( K) and later usages. Crucially, the overall complexity is reduced from quadratic to linear in the sequence length n. This enables streaming computation and a constant memory state, which are key advantages of linear attention. Recent work reveals that polynomial kernels provide a strong approximation to the exponential kernel underlying softmax attention. In particular, a degree-p polynomial kernel, denoted as κ(⋅,⋅)κ(·,·), takes the form κ(,)=(⊤)p,κ( q,k)=( q k)^p, where ,∈ℝd q,k ^d are query and key vectors. However, computing the exact polynomial feature map required by linear attention involves expanding all degree-p monomials, i.e., ψp()=vec(⊗p), _p( x)=vec\! ( x p ), where ⊗p p is p-th order outer product, whose dimensionality scales as (dp)O(d^p). This makes exact polynomial kernels impractical for real-world models. BASED addresses this issue by first compressing the query and key vectors into a low-dimensional space using a linear projection, ~=q,~=k,q,k∈ℝr×d,r≪d, q= W_q q, k= W_k k, W_q, W_k ^r× d,\;r d, (1) and then computing exact polynomial features in the compressed space. In practice, BASED utilizes a second-order polynomial kernel κ(,)=ψ2()⊤ψ2(),κ( q,k)= _2( q) _2( k), which corresponds to the quadratic term in the Taylor expansion of the exponential kernel.22 2 BASED proposes to use the Taylor expansion 1+x+x2/21+x+x^2/2. For simplicity, we omit the constant and linear terms, as they can be handled by concatenation and do not affect the core analysis. By combining this polynomial linear attention with sliding-window attention to capture local interactions, BASED achieves state-of-the-art language modeling performance while maintaining favorable computational efficiency. 2.3 Related Work We broadly categorize prior studies on linear attention into two directions: (i) the design of feature maps for query and key vectors, and (i) the design of state update rules. Feature Map Design. This line of work focuses on constructing efficient and expressive feature maps ϕ(⋅)φ(·) such that the softmax kernel exp(⊤) ( q k) can be approximated by an inner product ϕ()⊤ϕ()φ( q) φ( k) (Katharopoulos et al., 2020; Choromanski et al., 2020; Qin et al., 2022; Keles et al., 2023; Meng et al., 2025). Among the many design choices, several recent studies have identified polynomial kernels as particularly effective approximations, as they share important inductive biases with softmax attention, such as sparsity and sharp selectivity (Han et al., 2023; Zhang et al., 2024; Kacham et al., 2023; Arora et al., 2024). As discussed by (Arora et al., 2024), many existing feature maps are substantially less effective than polynomial-kernel-based approaches when evaluated in realistic settings. Moreover, for current methods that explicitly approximate polynomial kernels (Han et al., 2023; Kacham et al., 2023; Arora et al., 2024), their approximation accuracy is still limited. Our work finds that these approximations can be further improved through a more principled and flexible feature-map construction. Overall, there remains significant room for boosting the effectiveness of linear attention feature maps. Additionally, despite the extensive exploration of feature map design, many prior works do not assess performance on large-scale, real-world language modeling tasks. In some cases, linear attention variants are reported to outperform softmax attention, which contrasts with empirical observations in large language models where linear attention usually has sub-optimal effectiveness. State Update Rules. Another line of work focuses on designing improved state update rules for linear attention that aim to stabilize training and better capture long-range dependencies (Schlag et al., 2021; Yang et al., 2024; Team et al., 2025). These methods often resemble state space models (Gu and Dao, 2024; Peng et al., 2023) in that they maintain and update a recurrent state, rather than strictly adhering to the standard attention formulation. Importantly, these approaches are largely orthogonal to feature map design: their update rules can be combined with different feature maps by replacing q and k with ϕ()φ( q) and ϕ()φ( k). Since our work focuses exclusively on improving feature maps for polynomial kernel approximation, we do not directly compare with methods that primarily modify the state update mechanism. Polynomial Sketches for Neural Networks. Polynomial sketches have also been incorporated into neural network modules in prior work. For example, PolySketchFormer (Kacham et al., 2023) uses polynomial sketches to approximate attention, and EUGens (Kim et al., 2026) applies polynomial sketches within feedforward layers. While these works demonstrate the practical usefulness of polynomial sketching, their theoretical analyses are based on data-independent sketches, where the sketching map is fixed independently of the data distribution. This setting differs fundamentally from ours: LLoPS learns a data-dependent sketch tailored to the target representations, which can yield substantially better approximation under a fixed sketch dimension. To the best of our knowledge, our work is the first to provide theoretical analysis and approximation guarantees for such data-dependent polynomial sketches. 3 Methodology Polynomial kernel is considered to be a good candidate to replace softmax kernel (Kacham et al., 2023; Zhang et al., 2024; Arora et al., 2024). A key reason for the effectiveness of softmax attention is its sparse and spiky attention distribution: for a given query, only a small subset of keys receives large attention weights, while the majority of keys are assigned values close to zero. This behavior arises from the exponential function, which sharply amplifies differences in inner products, effectively separating keys that are well aligned with the query from those that are not. Figure 1: Illustration of LLoPS with an example where p=2,d=2,m=3p=2,d=2,m=3. (a) Given the optimal linear projection matrix W, each column of W can be written as a 2 by 2 matrix, and be approximated by its rank-1 CP decomposition; (b) The rank-1 CP decomposition can be stacked into two weight matrices , _1, _2; (c) In this way, the linear projection ⊤vec(⊤) W vec( x ) can be approximated as (⊤)∗(⊤)( _1 x)*( _2 x) where ∗* is an element-wise product. Polynomial kernels exhibit a similar spikiness property, as higher-order powers increasingly emphasize large inner products while suppressing smaller ones. Besides, polynomial functions are often used to approximate exponential function (e.g., Taylor expansion) (Zhang et al., 2024; Arora et al., 2024). Thus, a polynomial kernel is a natural and effective alternative to the exponential kernel in the context of linear attention. However, as discussed in Sec 2.2, an exact computation of polynomial kernels becomes impractical for high degrees or high-dimensional vectors due to the exponential growth of dimensionality. Therefore, our work focuses on designing feature maps for queries and keys separately, denoted by ϕQ(),ϕK():ℝn×d→ℝn×m _Q( Q), _K( K):R^n× d ^n× m, that accurately approximate polynomial kernels while remaining computationally efficient. Specifically, we aim to minimize the following approximation error: ‖ϕQ()ϕK()⊤−(⊤)∗p‖F2, \| _Q( Q) _K( K) -( Q K )^*p \|_F^2, (2) where ∗p*p is the element-wise p-th power and ∥⋅∥f\|·\|_f is Frobenius norm. To achieve this goal, we introduce Learnable Low-Rank Polynomial Sketch (LLoPS). We first demonstrate that, when the output dimension m is fixed, complicated feature maps are unnecessary to achieve the best compression. Instead, linear projections of ⊗p q p and ⊗p k p are sufficient to attain the theoretically optimal approximation to the target polynomial kernel. Building on this, we further decompose the optimal linear projection and approximate it using a low-rank CP (CANDECOMP/PARAFAC) decomposition (Papalexakis, 2016), which is also known as tensor rank decomposition, that enables a practical computation with linear attention while preserving the approximation quality. Finally, we provide a detailed error analysis of the proposed feature maps and elucidate their connections to existing linear attention kernels. 3.1 Optimal Sketch of Polynomial Kernel We first consider the optimal compression of polynomial kernel when the output dimension m is fixed. For clarity, let ψp:ℝd→ℝdp _p:R^d ^d^p, where ψp()=vec(⊗p) _p( x)=vec( x p); let =(⊤)∗p=ψp()ψp()⊤ A=( Q K )^*p= _p( Q) _p( K) 33 3 Function ψp _p is applied to each row of the matrices.; and let ⊤ U V be the SVD decomposition of A where =diag(σ1,⋯,σn) =diag( _1,·s, _n). Since ϕQ() _Q( Q) and ϕK() _K( K) are both an n×mn× m matrix (n≫mn m), the rank of ϕQ()ϕK()⊤ _Q( Q) _K( K) is at most m. By the property of SVD decomposition, the minimum possible loss of Eq. 2 is ∑i=m+1nσi2 _i=m+1^n _i^2 where σi _i is the i-th singular value. It is achieved when ϕQ()ϕK()⊤=⊤ _Q( Q) _K( K) = U _mV (3) where m=diag(σ1,⋯,σm,0,⋯,0) _m=diag( _1,·s, _m,0,·s,0). The following theorem shows how to achieve the minimum approximation error. Theorem 3.1 (Optimal Polynomial Sketch). Let ψp():ℝd→ℝdp=vec(⊗p) _p( x):R^d ^d^p=vec( x p). There exists Q∗,K∗∈ℝdp×m W^*_Q, W^*_K ^d^p× m such that ϕQ∗()ϕK∗()⊤=⊤ _Q^*( Q) _K^*( K) = U _mV where ϕQ∗()=ψp()Q∗ϕK∗()=ψp()K∗ _Q^*( Q)= _p( Q) W^*_Q _K^*( K)= _p( K) W^*_K (4) The proof is illustrated in Appendix A.1. Theorem 3.1 suggests that we could compress ψp(),ψp() _p( Q), _p( K) into a m-dimensional space via linear mapping, while achieving the minimal approximation error stated in Eq 3. 3.2 Learnable Low-Rank Polynomial Sketch Although Theorem 3.1 illustrates how to achieve the optimal approximation, since Q∗ W_Q^* and K∗ W_K^* have dpd^p rows, the linear projection is still too computationally expensive. To address this problem, we propose to decompose and approximate the linear projection matrices utilizing their low-rank CP decompositions (also known as tensor rank decomposition (Papalexakis, 2016)). Figure 1 depicts an example of our decomposed sketch for a toy example where p=2,d=2,m=3p=2,d=2,m=3. Let W denote Q∗ W^*_Q or K∗ W^*_K, and let i∈ℝdp W_i ^d^p (1≤i≤m1≤ i≤ m) denote the i-th column of W. We can write i W_i as a (p)(p)-way tensor i∈ℝd×⋯×d W_i ^d×·s× d. According to CP decomposition, i W_i can be represented, with a sufficiently large RiR_i, as a linear combination of RiR_i rank-1 tensors: i=∑r=1Riλr(i)1,r(i)⊗2,r(i)⊗⋯⊗p,r(i) W_i= _r=1^R_i _r^(i) w_1,r^(i) w_2,r^(i) ·s w_p,r^(i) (5) where λr(i)∈ℝ _r^(i) , j,r(i)∈ℝd w_j,r^(i) ^d (1≤j≤p1≤ j≤ p), and ⊗ denotes outer product. Inspired by the CP decomposition in Eq. 5, the key idea of LLoPS is as follows. Each row of W is expressed as a sum of rank-one outer products of learnable vectors, as illustrated in Figure 1(a). These vectors are trained implicitly to approximate the low-rank CP decomposition of each row vector i W_i. We then organize the learned factors into a collection of projection matrices, as shown in Figure 1(b), which enables an efficient computation of the feature map, avoiding explicitly constructing the high-dimensional polynomial features ψp() _p( Q) and ψp() _p( K) through a sequence of lightweight matrix multiplications and element-wise operations instead, as depicted in Figure 1(c). Specifically, iW_i can be approximated with its rank-t CP decomposition, i.e., i≈∑r=1tλr(i)1,r(i)⊗2,r(i)⊗⋯⊗p,r(i)≜^i W_i≈ _r=1^t _r^(i) w_1,r^(i) w_2,r^(i) ·s w_p,r^(i) W_i (6) By reshaping the tensor back to vector, we have an approximate column vector ^i=vec(^i) W_i=vec( W_i) and thus the approximate linear projector ^=(1^,⋯,m^)∈ℝdp×m W=( W_1,·s, W_m) ^d^p× m. Next, we demonstrate how to efficiently compute ψp() _p( Q) W and ψp() _p( K) W without explicitly computing ψp() _p( Q), ψp() _p( K), and W, which are with extremely high dimensions. Notice that the inner product between the outer product of vectors 1,r,…,p,r\ w_1,r,…, w_p,r\ and vec(⊗p)vec( x p) equals the product of ⟨j,r,⟩:1≤j≤p\ w_j,r, x :1≤ j≤ p\, i.e., ⟨vec(1,r⊗⋯⊗p,r),vec(⊗p)⟩=∏j=1p⟨j,r,⟩ ( w_1,r ·s w_p,r),vec( x p) = _j=1^p w_j,r, x (7) So we have ⟨^i,vec(⊗p)⟩=∑r=1tλr(i)∏j=1p⟨j,r(i),⟩ W_i,vec( x p) = _r=1^t _r^(i) _j=1^p w_j,r^(i), x . Thus ψp() _p( Q) W (8) = = (ψp()^1,⋯,ψp()^m) ( _p( Q) W_1,·s, _p( Q) W_m) = = (∑r=1tλr(1)∏j=1pj,r(1),⋯,∑r=1tλr(m)∏j=1pj,r(m)) ( _r=1^t _r^(1) _j=1^p Q w_j,r^(1),·s, _r=1^t _r^(m) _j=1^p Q w_j,r^(m) ) = = ∑r=1t(λr(1)∏j=1pj,r(1),⋯,λr(m)∏j=1pj,r(m)) _r=1^t ( _r^(1) _j=1^p Q w_j,r^(1),·s, _r^(m) _j=1^p Q w_j,r^(m) ) With Eq 8, we have the following theorem to approximate the optimal projection matrix W with a set of matrices . Theorem 3.2. Let j,r∈ℝd×m _j,r ^d× m be the stack of j,r(i) w_j,r^(i) for i=1,…,mi=1,…,m, i.e., j,r=(λr(1)j,r(1),⋯,λr(m)j,r(m)) _j,r=( _r^(1) w_j,r^(1),·s, _r^(m) w_j,r^(m)) (9) where 1≤j≤p1≤ j≤ p, 1≤r≤t1≤ r≤ t. We have ψp()^=∑r=1t∏j=1p,, _p( Q) W= _r=1^t _j=1^p Q _j,r, (10) ^=∑r=1t(,⊙col,⊙col⋯⊙col,) W= _r=1^t ( _1,r _col _2,r _col·s _col _p,r ) where ⊙col _col is the column-wise Kronecker product. Similarly, ψp()^=∑r=1t∏j=1p, _p( K) W= _r=1^t _j=1^p K _j,r. More importantly, this formulation avoids explicitly constructing the exponentially large polynomial feature map ψp() _p( Q): we first compute the projected matrices j,r Q _j,r for all j=1,…,pj=1,…,p and r=1,…,tr=1,…,t using standard matrix multiplications. Then, for each rank component r, we combine the p projected terms via element-wise multiplication, corresponding to the product over j. Finally, we sum the resulting features across all t components. Overall, the total computation cost is reduced from O(nmdp)O(nmd^p) to O(ndmtp)O(ndmtp). In practice, we find that the setting t=1t=1 already attains an accurate approximation of the polynomial kernel. Setting t>1t>1 will lead to too many learnable parameters. Therefore, we fix t=1t=1 and our proposed feature maps is formulated as below: Definition 3.3 (Learnable Low-Rank Polynomial Sketch). ϕQ()=∏j=1p(Q),ϕK()=∏j=1p(K) _Q( Q)= _j=1^p Q _j^(Q), _K( K)= _j=1^p K _j^(K) (11) where (Q) _j^(Q) and (Q) _j^(Q) (1≤j≤p1≤ j≤ p) are learnable. The computation complexity is O(ndmp)O(ndmp), and the number of trainable parameters are O(dmp)O(dmp). During training, are trained implicitly to approximate the low-rank CP decomposition of the optimal projecion matrix Q∗ W_Q^* and K∗ W_K^*. The following theorem analyze the error of our proposed learnable polynomial sketch. Theorem 3.4 (LLoPS Error). Let ∗∈ℝdp×m W^* ^d^p× m be the optimal matrix suggested in Theorem 3.1. Then min,…,p‖∗−1⊙…⊙p‖F2=‖∗‖F2−∑i=1mτi2 _ ,…, _p\| W^*- _1 … _p\|_F^2=\| W^*\|_F^2- _i=1^m _i^2 (12) where τi=max‖1‖=…=‖p‖=1⟨∗[:,i],1⊗…⊗p⟩ _i= _\| u_1\|=…=\| u_p\|=1 W^*[:,i], u_1 … u_p (13) is the tensor spectral norm (i.e., the largest tensor singular value) of iW_i, and i∈ℝd×…×dW_i ^d×…× d is the order p tensor by reshaping the iith column of ∗∈ℝdp W^* ^d^p. The proof is given in Appendix A.2. Some previous studies found that making the attention score of linear attention to be non-negative is helpful for training stability (Kacham et al., 2023; Han et al., 2023). If non-negativity is required, we could adopt the approach used in (Kacham et al., 2023) to apply another square to the feature map, i.e., ϕQ()=(∏j=1p(Q))∗2 _Q( Q)= ( _j=1^p Q _j^(Q) )^*2, and similarly for ϕK() _K( K). 3.3 Relationship with Existing Feature Maps In this subsection, we discuss the relationship between our proposed feature map and several existing approaches for approximating polynomial kernels in linear attention. Element-wise Polynomial Feature Maps Some earlier linear attention methods (Han et al., 2023) adopt element-wise polynomial feature maps of the form ϕ()=()∗pφ( x)=( x)^*p where ∗p*p denotes the element-wise p-th power. This construction can be viewed as a special case of our formulation in which all projection matrices are shared, i.e., 1=⋯=p _1=·s= _p. Under this constraint, our kernel reduces exactly to an element-wise polynomial feature map. BASED This method (Arora et al., 2024) approximates the second-order polynomial kernel by first projecting queries and keys into a low-dimensional space via a linear mapping, and then applying the quadratic feature map ψ2(⋅) _2(·). Specifically, for ∈ℝm′×d W ^m × d and ∈ℝd x ^d, BASED computes ψ2()=vec(⊗)=(⊗kron)vec(⊤), _2( W x)=vec( W x W x)=( W _kron W)\,vec( x x ), (14) where ⊗kron _kron is the full Kronecker product. This construction can also be recovered as a special case of our kernel. In particular, let 1⊤=⊙col,2⊤=⊙col, _1 = W _col 1, _2 = 1 _col W, where ∈ℝm′×d 1 ^m × d is an all-ones matrix. Under this choice, we have 1⊙col2=(⊗kron)⊤. _1 _col _2=( W _kron W) . Therefore, with these specific parameter settings, our feature map exactly recovers the BASED kernel. PolySketchFormer This approach (Kacham et al., 2023) approximates polynomial kernels using randomized sketching techniques. Its feature map can be expressed as multiplication by a randomized sketching matrix ∈ℝdp×m R ^d^p× m satisfying the (ϵ,p)(ε,p)-AMM property (Woodruff and others, 2014). Following prior constructions (Ahle et al., 2020), R can be formulated as a column-wise Kronecker product of p sparse randomized matrices, permitting efficient approximation of high-order polynomial kernels. From our perspective, PolySketchFormer corresponds to a special case of our framework in which the projection matrices ii=1p\ _i\_i=1^p are fixed, randomized, and sparse, rather than learned from data. Taken together, these observations reveal that several existing polynomial-kernel-based linear attention methods can be interpreted as special or constrained instances of our proposed feature map. By allowing more general and learnable projections, our formulation strictly generalizes these prior approaches. As a consequence, for a fixed output dimension, our method can achieve an approximation error that is no worse than these existing methods, and in practice often substantially better, as shown in the next section. This is equivalent to enforcing a constraint where half of the learnable matrices share parameters with the other half. 4 Experiments In this section, we report the main experiment results comparing our method with state-of-the-art linear attention methods. Additional experiments and ablation studies are reported in Appendix B. 4.1 Zero-Shot NLP Benchmarks Following prior work on alternative sequence modeling architectures (Gu and Dao, 2024; Arora et al., 2024), we evaluate our method on a suite of zero-shot NLP benchmarks. We compare against the standard softmax Transformer, Mamba, and BASED, using the publicly released 360M and 1B parameter checkpoints from Arora et al. (2024). Given that these architectures require extensive pre-training from scratch, additional linear attention baselines were excluded from this evaluation. Since the BASED attention kernel can be viewed as a special case of our proposed kernel, we initialize our model from the BASED checkpoint and continue training it to obtain our final model to reduce training time (BASED+LLoPS). For LLoPS-360M, the polynomial degree p=4p=4; for LLoPS-1B p=2p=2. Importantly, the original baselines in Arora et al. (2024) were pretrained on the Pile dataset (Gao et al., 2020), which is no longer publicly available. Therefore, we perform continued training utilizing the SlimPajama-6B dataset (Soboleva et al., 2023)44 4 https://huggingface.co/datasets/DKYoon/SlimPajama-6B, which is 50 times smaller in size compared to Pile. For a fair comparison under this updated data setting, the BASED models are also post-trained with SlimPajama-6B for the same number of iterations. This allows us to isolate the benefits of our proposed kernel from differences arising purely from additional training data or compute. The evaluated benchmarks are implemented in lm_eval (Gao et al., 2023), including: WinoGrande (Sakaguchi et al., 2021), PIQA (Bisk et al., 2020), LAMBADA (Paperno et al., 2016), HellaSwag (Zellers et al., 2019), ARC-challenge (Clark et al., 2018), SQUAD (Rajpurkar et al., 2016), SWDE (Lockard et al., 2019; Arora et al., 2023), FDA (Arora et al., 2023). The tasks include information extraction and commonsense reasoning QA. The results are reported in Table 1. Notably, for 360M models, we observe that after continued training on SlimPajama-6B, both BASED and our method exhibit a substantial accuracy drop on the FDA benchmark. We hypothesize that this degradation arises from a mismatch between the SlimPajama-6B and the FDA task distribution, which may induce catastrophic forgetting of task-relevant knowledge. Since this phenomenon occurs for both BASED and our method, it is more likely due to the choice of training data rather than a limitation of our proposed kernel. Despite this issue, our method produces a higher average accuracy than BASED. These results further support the effectiveness of our approach in improving the performance of linear attention models across diverse zero-shot NLP benchmarks. Table 1: Zero-shot accuracy (%) across diverse NLP benchmarks for pretrained models at two scales. Bold numbers indicate the best performance among linear-complexity methods within each model size. BASED+CT is BASED under the same continual training setting as LLoPS. 360M Models Method Wino PIQA LAM Hella ARC-E ARC-C FDA SWDE SQuAD Avg Softmax Attn 51.78 63.71 29.32 33.60 42.85 24.32 57.89 56.08 27.88 43.05 Mamba 50.36 63.76 27.89 33.88 42.42 24.66 5.90 17.37 24.83 32.34 BASED 51.30 63.55 26.64 33.23 41.88 24.23 17.42 27.18 27.41 34.76 BASED+CT 50.91 65.40 27.73 36.85 43.10 25.51 8.35 30.87 26.27 35.00 BASED+LLoPS 53.04 66.70 29.23 37.07 43.94 25.51 7.53 31.41 31.94 36.26 1B Models Method Wino PIQA LAM Hella ARC-E ARC-C FDA SWDE SQuAD Avg Softmax Attn 52.49 67.30 38.07 39.12 51.05 22.10 68.78 68.68 35.79 49.26 Mamba 51.30 66.81 38.87 39.34 46.80 26.11 11.07 28.17 29.46 37.55 BASED 51.62 66.70 37.86 39.10 45.54 25.60 24.05 42.84 34.58 40.88 BASED+CT 51.22 68.44 36.43 41.59 47.01 27.30 24.41 46.80 36.23 42.16 BASED+LLoPS 51.62 68.72 38.11 42.69 47.31 27.13 26.32 47.97 36.36 42.91 4.2 Long Range Arena Benchmark Following prior work (Zhang et al., 2024; Meng et al., 2025; Han et al., 2023), we evaluate LLoPS on the Long Range Arena (LRA) benchmark (Tay et al., 2020). Instead of the original train-from-scratch protocol, we adopt the pretraining-and-finetuning protocol of Amos et al. (2023), which better reflects the behavior of pretrained transformers. We apply the same protocol to all reproduced baselines for a fair comparison, while reporting the best numbers from prior papers and our training results. All methods share the same architecture except for the attention layer, and we compare against softmax attention, Nyström (Xiong et al., 2021), Skyformer (Chen et al., 2021), Hedgehog (Zhang et al., 2024), PolaFormer (Meng et al., 2025), Flatten Transformer (Han et al., 2023), PolySketchFormer (Kacham et al., 2023), and BASED (Arora et al., 2024). We set the polynomial degree of LLoPS to p=4p=4. Table 2 summarizes the results. Under this protocol, softmax attention consistently outperforms linear attention variants, aligning with empirical observations from large-scale language modeling. Among linear attention methods, LLoPS achieves top-3 performance on every task and obtains the highest average accuracy, improving over the strongest linear attention baseline by 2.28%2.28\%. These results show that LLoPS substantially improves linear attention effectiveness while preserving its efficiency advantages. Table 2: Performance comparison across Long Range Arena (LRA) benchmarks. Bold numbers indicate the best performance in linear attention methods, underlined numbers indicate the second-best, and italicized numbers indicate the third-best. Method Image Text Retrieval ListOps Pathfinder Avg Softmax Attn 64.36 82.96 83.23 48.49 70.45 69.90 Nyström 58.66 66.19 83.90 45.46 69.71 64.78 Skyformer 58.54 65.25 79.55 46.77 72.73 64.57 Hedgehog 52.68 66.59 82.24 41.03 74.16 63.34 PolaFormer 47.74 73.06 80.50 38.86 70.53 62.14 Flatten 51.48 66.13 78.72 38.51 70.27 61.02 PolySketchFormer 51.40 64.60 81.10 43.00 66.20 61.26 BASED 51.13 65.73 80.76 43.55 72.98 62.83 LLoPS 57.54 73.94 82.17 48.29 73.84 67.16 4.3 Polynomial Kernel Approximation Error Furthermore, we directly assess the error of our proposed sketch in approximating polynomial kernels. To that end, we obtain a set of query and key vectors with the following steps: (1) we train a transformer with the exact polynomial kernel (Attn(,,)=(⊤)pAttn( q,k,v)=( q k)^p v, p=2,3p=2,3) on the synthetic associative recall (AR) task introduced in Arora et al. (2024); (2) we select the first (layer 0), middle (layer 8), and last attention layers (layer 16) of the trained model and sample the corresponding query and key vectors with synthetic AR data. After obtaining the query and key vectors, for each layer, we apply different sketching methods to approximate the polynomial kernel on these sampled vectors. We compare our sketch against the data-independent sketching approach used in PolySketchFormer (Kacham et al., 2023) (PSF), as well as neural network-based sketches (N). For the neural sketch baseline, the feature maps ϕq _q and ϕk _k are parameterized as two-layer MLPs and then trained to approximate the exact polynomial kernel. To quantify approximation quality, we report the relative mean absolute error (RMAE), defined as RMAE=,|(⊤)p−ϕq()⊤ϕk()(⊤)p|.RMAE=E_ q, k | ( q k)^p- _q( q) _k( k)( q k)^p |. (15) The results are summarized in Table 3. We observe that our sketch consistently achieves substantially smaller approximation error than PSF. This aligns with the intuition that learnable methods are usually better than data-independent methods. Furthermore, our learnable sketch also outperforms the neural network-based feature maps. We attribute this to two main reasons. First, Theorem 3.1 suggests that a linear projection is sufficient to achieve optimal compression for polynomial kernels, making more complex nonlinear parameterizations unnecessary. Second, our structured and decomposed linear formulation is easier to optimize than MLP-based sketches, which may suffer from additional training instability. Overall, these results demonstrate that our proposed sketch provides an effective and accurate approximation of polynomial kernels in practical transformer query and key distributions. Table 3: Relative mean absolute error of different sketches approximating polynomial kernels with degree p=2,3p=2,3. p Method Layer 0 Layer 8 Layer 15 Avg 22 PSF 0.186 0.247 0.445 0.293 N 0.035 0.070 0.058 0.054 LLoPS 0.008 0.032 0.025 0.022 33 PSF 0.740 0.738 0.745 0.741 N 0.059 0.114 0.143 0.105 LLoPS 0.018 0.040 0.058 0.038 4.4 Efficiency Comparison We further benchmark the runtime of our proposed linear attention kernel against softmax attention and the BASED kernel. We measure forward-pass latency over 1,000 repetitions with sequence lengths from 1,000 to 10,000 tokens, using the same PyTorch implementation environment. The input dimension is 128. For LLoPS, we set the output feature dimension to 256 and p=2p=2; for BASED, we compress queries and keys to 16 dimensions, yielding a comparable final feature dimension of 162+16+1=27316^2+16+1=273. Figure 2 reports the total forward-pass runtime. As expected, both linear attention methods are much faster than softmax attention at long sequence lengths. Moreover, LLoPS is consistently faster than BASED, since it mainly uses matrix multiplications and element-wise products, which are highly optimized on modern accelerators, whereas BASED involves less efficient outer-product-style operations. These results show that LLoPS improves approximation quality while retaining strong practical efficiency. Figure 2: Time cost of different attention kernels at different input length with p=2p=2. 5 Conclusion We studied polynomial-kernel-based linear attention through the lens of sketching. We provided a theoretical characterization of the optimal sketch for approximating polynomial kernels in the linear attention setting, revealing that structured and learnable projections can significantly improve approximation quality over previous data-independent feature constructions (Kacham et al., 2023). Building on this insight, we proposed LLoPS, a Learnable Decomposed Polynomial Sketch that offers an accurate and efficient approximation to polynomial attention while retaining the favorable computational properties of linear attention. Empirically, LLoPS consistently improves the effectiveness of linear attention Transformers across a diverse set of benchmarks. Limitation. Since the original pre-training dataset used by BASED is no longer publicly available, we are unable to evaluate LLoPS under the same from-scratch training setting. Nevertheless, the results in Table 1 show that, when BASED and LLoPS are continually trained from the same initialization for the same number of iterations, LLoPS consistently achieves higher accuracy. We believe this controlled comparison is sufficient to demonstrate the effectiveness of LLoPS. References Ahle et al. (2020) T. D. Ahle, M. Kapralov, J. B. Knudsen, R. Pagh, A. Velingker, D. P. Woodruff, and A. Zandieh Oblivious sketching of high-degree polynomial kernels. In Proceedings of the Fourteenth Annual ACM-SIAM Symposium on Discrete Algorithms, p. 141–160. Cited by: §3.3. Amos et al. (2023) I. Amos, J. Berant, and A. Gupta Never train from scratch: fair comparison of long-sequence models requires data-driven priors. arXiv preprint arXiv:2310.02980. Cited by: §4.2. Arora et al. (2024) S. Arora, S. Eyuboglu, M. Zhang, A. Timalsina, S. Alberti, D. Zinsley, J. Zou, A. Rudra, and C. Ré Simple linear attention language models balance the recall-throughput tradeoff. arXiv preprint arXiv:2402.18668. Cited by: §B.1, §B.2, §B.3, §1, §1, §1, §2.3, §3.3, §3, §3, §4.1, §4.1, §4.2, §4.3. Arora et al. (2023) S. Arora, B. Yang, S. Eyuboglu, A. Narayan, A. Hojel, I. Trummer, and C. Ré Language models enable simple systems for generating structured views of heterogeneous data lakes. arXiv preprint arXiv:2304.09433. Cited by: §4.1. Bisk et al. (2020) Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. Piqa: reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, Vol. 34, p. 7432–7439. Cited by: §4.1. Chen et al. (2021) Y. Chen, Q. Zeng, H. Ji, and Y. Yang Skyformer: remodel self-attention with gaussian kernel and nystr\ " om method. Advances in Neural Information Processing Systems 34, p. 2122–2135. Cited by: §4.2. Choromanski et al. (2020) K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, L. Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794. Cited by: §2.3. Clark et al. (2018) P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457. Cited by: §4.1. Gao et al. (2020) L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al. The pile: an 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027. Cited by: §4.1. Gao et al. (2023) L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac’h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou A framework for few-shot language model evaluation. Zenodo. External Links: Document, Link Cited by: §4.1. Gu and Dao (2024) A. Gu and T. Dao Mamba: linear-time sequence modeling with selective state spaces. In First conference on language modeling, Cited by: §2.3, §4.1. Han et al. (2023) D. Han, X. Pan, Y. Han, S. Song, and G. Huang Flatten transformer: vision transformer using focused linear attention. In Proceedings of the IEEE/CVF international conference on computer vision, p. 5961–5971. Cited by: §2.3, §3.2, §3.3, §4.2. Kacham et al. (2023) P. Kacham, V. Mirrokni, and P. Zhong PolySketchFormer: fast transformers via sketching polynomial kernels. arXiv preprint arXiv:2310.01655. Cited by: §1, §1, §1, §2.3, §2.3, §3.2, §3.3, §3, §4.2, §4.3, §5. Katharopoulos et al. (2020) A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret Transformers are rnns: fast autoregressive transformers with linear attention. In International conference on machine learning, p. 5156–5165. Cited by: §2.3. Keles et al. (2023) F. D. Keles, P. M. Wijewardena, and C. Hegde On the computational complexity of self-attention. In International conference on algorithmic learning theory, p. 597–619. Cited by: §2.3. Kim et al. (2026) S. M. Kim, B. Kim, A. Sehanobish, S. B. R. Chowdhury, R. Kidambi, D. Shim, A. Dubey, S. Chaturvedi, M. Oh, and K. Choromanski EUGens: efficient, unified, and general dense layers. arXiv preprint arXiv:2601.22563. Cited by: §2.3. Lockard et al. (2019) C. Lockard, P. Shiralkar, and X. L. Dong Openceres: when open information extraction meets the semi-structured web. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), p. 3047–3056. Cited by: §4.1. Meng et al. (2025) W. Meng, Y. Luo, X. Li, D. Jiang, and Z. Zhang PolaFormer: polarity-aware linear attention for vision transformers. arXiv preprint arXiv:2501.15061. Cited by: §2.3, §4.2. Papalexakis (2016) E. E. Papalexakis Automatic unsupervised tensor mining with quality assessment. In Proceedings of the 2016 SIAM International Conference on Data Mining, p. 711–719. Cited by: §3.2, §3. Paperno et al. (2016) D. Paperno, G. Kruszewski, A. Lazaridou, N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández The lambada dataset: word prediction requiring a broad discourse context. In Proceedings of the 54th annual meeting of the association for computational linguistics (volume 1: Long papers), p. 1525–1534. Cited by: §4.1. Peng et al. (2023) B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, M. Grella, et al. Rwkv: reinventing rnns for the transformer era. arXiv preprint arXiv:2305.13048. Cited by: §2.3. Qin et al. (2022) Z. Qin, W. Sun, H. Deng, D. Li, Y. Wei, B. Lv, J. Yan, L. Kong, and Y. Zhong Cosformer: rethinking softmax in attention. arXiv preprint arXiv:2202.08791. Cited by: §2.3. Rajpurkar et al. (2016) P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250. Cited by: §4.1. Sakaguchi et al. (2021) K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi Winogrande: an adversarial winograd schema challenge at scale. Communications of the ACM 64 (9), p. 99–106. Cited by: §4.1. Schlag et al. (2021) I. Schlag, K. Irie, and J. Schmidhuber Linear transformers are secretly fast weight programmers. In International conference on machine learning, p. 9355–9366. Cited by: §1, §2.3. Soboleva et al. (2023) D. Soboleva, F. Al-Khateeb, R. Myers, J. R. Steeves, J. Hestness, and N. Dey SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. External Links: Link Cited by: §4.1. Tay et al. (2020) Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler Long range arena: a benchmark for efficient transformers. arXiv preprint arXiv:2011.04006. Cited by: §4.2. Team et al. (2025) K. Team, Y. Zhang, Z. Lin, X. Yao, J. Hu, F. Meng, C. Liu, X. Men, S. Yang, Z. Li, et al. Kimi linear: an expressive, efficient attention architecture. arXiv preprint arXiv:2510.26692. Cited by: §1, §1, §2.3. Woodruff et al. (2014) D. P. Woodruff et al. Sketching as a tool for numerical linear algebra. Foundations and Trends® in Theoretical Computer Science 10 (1–2), p. 1–157. Cited by: §3.3. Xiong et al. (2021) Y. Xiong, Z. Zeng, R. Chakraborty, M. Tan, G. Fung, Y. Li, and V. Singh Nyströmformer: a nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI conference on artificial intelligence, Vol. 35, p. 14138–14148. Cited by: §1, §4.2. Yang et al. (2024) S. Yang, J. Kautz, and A. Hatamizadeh Gated delta networks: improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464. Cited by: §1, §2.3. Zellers et al. (2019) R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi Hellaswag: can a machine really finish your sentence?. arXiv preprint arXiv:1905.07830. Cited by: §4.1. Zhang et al. (2024) M. Zhang, K. Bhatia, H. Kumbong, and C. Ré The hedgehog & the porcupine: expressive linear attentions with softmax mimicry. arXiv preprint arXiv:2402.04347. Cited by: §1, §2.3, §3, §3, §4.2. Appendix A Proofs A.1 Proof of Theorem 3.1 Proof. For simplicity, let ~=ψp() Q= _p( Q) and ~=ψp() K= _p( K). Notice that ~~⊤=(⊤)∘p Q K =( QK ) p (16) Let ϕQ()=~1 _Q( Q)= Q W_1, ϕK()=~2 _K( K)= K W_2, where 1,2∈ℝdp×m W_1, W_2 ^d^p× m. We need to prove that there exist 1,2 W_1, W_2 such that ~12⊤~⊤=⊤ Q W_1 W_2 K = U_m _mV_m (17) where m=⊤ S_m= U_m _mV_m is the truncated SVD decomposition of =~~⊤ S= Q K . Notice that the column and row spaces of m S_m are subspaces of the column and row spaces of S. Since =~~⊤ S= Q K , the columns of S are linear combinations of the columns of ~ Q. Therefore, the range (column space) of S is a subspace of the range of ~ Q: range()⊆range(~)range( S) ( Q) (18) Similarly, the rows of S are linear combinations of the rows of ~⊤ K : range(⊤)⊆range(~)range( S ) ( K) (19) As a result, range()⊆range(~)range( U_m) ( Q) and range()⊆range(~)range( V_m) ( K) Therefore, there must exist 1,2 W_1, W_2 such that ~1=1/2 Q W_1= U_m _m^1/2 and ~2=1/2 K W_2= V_m _m^1/2. So ~1(~2)⊤=m Q W_1( K W_2) = S_m ∎ A.2 Proof of Theorem 3.4 Proof. min1,…,p∥∗−1⊙…⊙p∥F2=∑i=1mmin1,i,…,p,i∥1,i⊙…⊙p,i−∗[:,i]∥F2 _ _1,…, _p\| W^*- _1 … _p\|_F^2= _i=1^m _ _1,i,…, _p,i\| _1,i … _p,i- W^*[:,i]\|_F^2 (20) where j,i _j,i is the i-th row of j _j, 1≤j≤p1≤ j≤ p, 1≤i≤m1≤ i≤ m. By the definition of tensor spectral norm and the best Rank-1 CP approximation, for each i, min1,i,…,p,i∥1,i⊙…⊙p,i−∗[:,i]∥F2=∥∗[:,i]∥22−τi2 _ w_1,i,…, w_p,i\| w_1,i … w_p,i- W^*[:,i]\|_F^2=\| W^*[:,i]\|_2^2- _i^2 (21) ∎ Appendix B Additional Experiment Results B.1 Synthetic Associative Recall Task We compare our method with Based on the synthetic associative reacall task introduced in Arora et al. [2024]. For both methods, the hidden dimension of the model is set to d∈64,128d∈\64,128\. The query and key dimensions of BASED is 16. The sketch dimension of LLoPS m=256m=256 and p=4p=4. Table 4 shows the result. We can see that LLoPS achieves higher accuracy than BASED, further proving its effectiveness. Table 4: Comparion of Based and LLoPS on synthetic associative recall task. For LLoPS, m=256m=256, p=4p=4. For both methods d=128d=128. d 128 192 BASED 91.9 98.1 LLoPS 93.2 99.0 B.2 Polynomial Degree Ablation Here, we study how the polynomial degree p affects the effectiveness of LLoPS. We consider the synthetic associative recall task Arora et al. [2024] and train three 2-layer transformers from scratch using polynomial kernels with degrees p∈2,4,6p∈\2,4,6\. All models use a hidden dimension of 128, and the output sketch dimension is fixed to 256 across settings. The results are summarized in Table 5. We observe that increasing the degree from p=2p=2 to p=4p=4 leads to improved accuracy. This suggests that moderate higher-order interactions can better distinguish key vectors that align well with the queries from those that do not, thereby enhancing retrieval performance. However, when further increasing the degree to p=6p=6, accuracy declines. We hypothesize that this drop arises because the sketch dimension m remains fixed: as p grows, the polynomial kernel becomes more difficult to approximate accurately, and the resulting increase in approximation error can negatively impact model effectiveness. Another possible reason is that, since the number of parameters increase with p, a higher value of p make the model easier to overfit the training data, which would also cause the test accuracy to drop. Overall, these results highlight an important trade-off: choosing p requires balancing representational power and approximation accuracy, and higher degree is not always better in practice. Table 5: Accuracy of synthetic association recall task under different polynomial degree p with d=128d=128. p 2 4 6 Acc 88.7 93.2 86.6 (a) m vs Acc (b) m vs efficency Figure 3: Results of how accuracy and forward time changes with sketch dimension m∈64,128,192,256m∈\64,128,192,256\ and d=128d=128, p=4p=4. B.3 Sketch Dimension Ablation Here we demonstrate the effect of sketch dimension m to effectiveness and effciency. We fix the model dimension d=128d=128 and vary m∈64,128,192,256m∈\64,128,192,256\. We measure the accuracy for synthetic association recall task [Arora et al., 2024] and the forward computation time. The results show that when m increases, the effective improves but the efficiency drops, which fits the intuition. Nevertheless, we note that when the sketch dimension drops from 256 to 192, the efficiency improves while the accuracy does not drop. This suggests that our method has great potential in approximating polynomial kernels with limited budget of state space size. NeurIPS Paper Checklist The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit. Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist: • You should answer [Yes] , [No] , or [N/A] . • [N/A] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available. • Please provide a short (1–2 sentence) justification right after your answer (even for [N/A] ). The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will also be asked to include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper. The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While [Yes] is generally preferable to [No] , it is perfectly acceptable to answer [No] provided a proper justification is given (e.g., error bars are not reported because it would be too computationally expensive” or “we were unable to find the license for the dataset we used”). In general, answering [No] or [N/A] is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found. IMPORTANT, please: • Delete this instruction block, but keep the section heading “NeurIPS Paper Checklist", • Keep the checklist subsection headings, questions/answers and guidelines below. • Do not modify the questions and only use the provided macros for your answers. 1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? Answer: [Yes] Justification: The main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope. 2. Limitations Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The limitation is discussed in conclusion. 3. Theory assumptions and proofs Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: The proofs are provided in the main paper and appendix. 4. Experimental result reproducibility Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: All information is disclosed. 5. Open access to data and code Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [No] Justification: The code will be released upon acceptance. 6. Experimental setting/details Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer) necessary to understand the results? Answer: [Yes] Justification: All information is disclosed 7. Experiment statistical significance Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: We include error bars in Figure 2. For other experiments, due to the high cost of pre-training, we cannot use different random seeds to train it multiple times. 8. Experiments compute resources Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: Information is disclosed Guidelines: 9. Code of ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.c/public/EthicsGuidelines? Answer: [Yes] Justification: This paper conform with the NeurIPS code of ethics. 10. Broader impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: The impact of the paper is discussed in the main paper. 11. Safeguards Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pre-trained language models, image generators, or scraped datasets)? Answer: [N/A] Justification: The paper has no such risk. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: They are credited. 13. New assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [N/A] Justification: the paper does not release new assets. 14. Crowdsourcing and research with human subjects Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [N/A] Justification: the paper does not involve crowdsourcing nor research with human subjects. 15. Institutional review board (IRB) approvals or equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [N/A] Justification: the paper does not involve crowdsourcing nor research with human subjects. 16. Declaration of LLM usage Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigor, or originality of the research, declaration is not required. Answer: [N/A] Justification: the core method development in this research does not involve LLMs as any important, original, or non-standard components.