Rank-Collapse Quadratic Token Router
Implementation & benchmark of arXiv:2608.07828 — The Rank-Collapse Principle for Quadratic Optimization
Source paper: The Rank-Collapse Principle for Quadratic Optimization arXiv:2608.07828 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace independent token scores with a query-conditioned positive-semidefinite low-rank quadratic score over a fixed-size selected subset. Repeatedly convert the quadratic objective into a linear exposure vector and apply a cheap top-k oracle, allowing the selector to model joint token interactions without constructing an n-by-n attention matrix. The margin between the current low-dimensional shadow and alternatives provides a practical confidence or early-stopping signal.
Formulas
Mathematical statement
The paper factors a positive-semidefinite matrix as Q=BB^{\mathsf{T}}, where Q is an n-by-n matrix and B is an n-by-r matrix with small rank r. For a feasible vector x in a set X, define its rank-space shadow y=B^{\mathsf{T}}x in R^r; then the quadratic objective is q_Q(x)=x^{\mathsf{T}}Qx=||B^{\mathsf{T}}x||_2^2=||y||_2^2. If x^\star is optimal and y^\star=B^{\mathsf{T}}x^\star, the paper states that y^\star uniquely maximizes the self-generated linear functional y mapsto <y^\star,y> and obeys the margin inequality <y^\star,y^\star-y> >= 1/2 ||y^\star-y||_2^2. In the neural adaptation, x is a binary token-selection vector with sum_i x_i=k, B=B(q,K,V) is a query-conditioned rank-r feature matrix, y=B^{\mathsf{T}}x is the selected-set shadow, and the exposure weights are w=By=BB^{\mathsf{T}}x=Qx. The top-k operation is the exact linear oracle for the cardinality-constrained feasible set; low rank makes forming w cost O(nr) rather than O(n^2).
Implementation notes
Integrate this into an attention token-pruning or retrieval block immediately before expensive value aggregation. Given a query q and n candidate key/value tokens, compute a small rank-r feature matrix B without materializing an n-by-n similarity matrix. One parameterization is B=phi(K)U(q), where phi(K) is an n-by-d token feature matrix and U(q) is a d-by-r query-conditioned projection; another is B=(KW_b)diag(g(q)), with r columns. Maintain a binary mask x containing exactly k selected tokens. Initialize x_0 from ordinary qK^T scores. For each of T=2--5 iterations, compute y_t=B^T x_t, compute w_t=B y_t, set x_candidate=TopK(w_t), and accept the candidate only if ||B^T x_candidate||_2^2 exceeds ||y_t||_2^2. Keep the best mask across one or more cheap starts. The paper-derived computations are q_Q(x)=||B^T x||_2^2 and the self-generated exposure w=Qx; B itself is learned and estimated from the network’s projections. Compute the empirical margin m_t=<y_t,y_t-B^T x_candidate> and the normalized shadow gap ||y_t-B^T x_candidate||_2/(||y_t||_2+epsilon). Stop when the gap is below a threshold or when the objective fails to improve. First test on CIFAR-100 or ImageNet-1k with a ViT-Tiny at 25%, 50%, and 75% token retention. Compare against standard attention pruning and linear top-k routing at equal FLOPs and wall-clock budget. Measure accuracy, selector overhead, latency, retained-token count, and quadratic objective. Success means 1--2 percentage points higher accuracy at fixed retention, or equal accuracy with fewer tokens, while selector overhead stays below 10% of original attention cost.
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.