Prefix-Sum PSD Attention Kernel
Implementation & benchmark of arXiv:2608.15125 — Entrywise Loewner Preservers on Min and Max Matrix Cones
Source paper: Entrywise Loewner Preservers on Min and Max Matrix Cones arXiv:2608.15125 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace or augment a dense attention similarity matrix with a Min-cone matrix generated by a monotone scalar sequence. The resulting matrix is positive semidefinite by construction, has only O(n) learned scalar parameters, and can be multiplied by values in O(n d) time using cumulative sums rather than forming an n-by-n matrix.
Formulas
Mathematical statement
For x=(x_1,...,x_n), the Min matrix is A_min(x)_{ij}=x_{min(i,j)}. Set x_0=0 and d_k=x_k-x_{k-1}. If d_k>=0, define the prefix indicator u_k in R^n by (u_k)_i=1 when i>=k and 0 otherwise. Then A_min(x)=sum_{k=1}^n d_k u_k u_k^T, so A_min(x) is positive semidefinite. The same construction can be applied to suffix indicators for A_max(x)_{ij}=x_{max(i,j)}. The paper also states that an entrywise function f preserves PSD on all PSD Min or Max matrices exactly when f is nonnegative and nondecreasing on [0,infinity). The transferable asset is an explicit PSD factorization plus a scalar parameterization that enforces all required inequalities.
Implementation notes
Integrate this as a structured attention branch or structured covariance kernel. Store unconstrained parameters a_1,...,a_n and define d_k=softplus(a_k), then compute x_k=sum_{r=1}^k d_r. The kernel is K_ij=x_min(i,j), or equivalently K=U diag(d) U^T. Do not materialize K for the value path: compute U^T V using suffix cumulative sums, multiply each row by d_k, and compute U times the result using prefix cumulative sums. This costs O(n d_model) time and O(n d_model) working memory. To use it in attention, either add beta K to dense QK^T/sqrt(d_model), or use K V as a parallel linear-attention branch with a learned scalar beta. If an entrywise transform is needed, use f(t)=softplus(c t+b) with c>=0, or f(t)=t^alpha with alpha>=1, and verify nonnegativity numerically. The paper provides the PSD guarantee from d_k>=0 and the scalar PSD-preserver condition; normalization, causal masking, beta selection, and quality improvements are empirical. First test on a 2-layer Transformer with 4 heads on WikiText-2 or synthetic long-range sequence prediction. Compare dense attention, standard linear attention, and one Min-kernel head at matched hidden size. Measure validation perplexity versus sequence length, wall-clock time, peak memory, and the smallest eigenvalue on short sequences. Success means stable nonnegative spectra and comparable perplexity with subquadratic value computation.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.