Linearized 2-Simplicial Attention

The paper presents a linearized form of 2-simplicial attention, called LinSimp, which rewrites the trilinear score as an inner product between a composite query and a key, enabling linear cost in sequence length while maintaining global reach. The authors implement this with custom Triton kernels and combine it with Kimi Delta Attention (KDA) to build a model with no softmax attention at all. Under matched compute, this model achieves the highest mean downstream accuracy among compared architectures, and at 16k context it improves mean accuracy over a KDA hybrid while lowering LAMBADA perplexity from 715.6 to 602.6.

The key contributions are: (1) deriving a one-mode kernelization of 2-simplicial attention, (2) providing a linear-time implementation with a positive orthogonal random-feature bank and a short explicit anchor window, (3) providing initial custom Triton forward and backward kernels including numerical tests against an fp32 reference, and (4) comparing the model across standard attention, an exact windowed 2-simplicial baseline, and a KDA hybrid under both iso-Token and iso-FLOP settings.

The method works by rewriting the trilinear score s ijc = tau i, j, c as s ijc = ((tau i) c) j. For fixed query i and anchor c, the sum over j becomes an ordinary exponential dot-product kernel, which is approximated with positive random features. The entire prefix is stored in a fixed-size linear-attention state (matrices M i and vectors a i), while the anchor axis stays explicit inside a 64-token window. The output is computed as o i = N i / (Z i + epsilon) where N i = sum c in C i phi(z ic) M i u c and Z i = sum c in C i phi(z ic) a i, with C i being the recent anchor window.

The architecture uses 24 pre-norm residual blocks with hidden size 1024, RMSNorm, SwiGLU, tied embeddings, and a Llama 2 tokenizer. The main no-softmax model uses 18 KDA layers and 6 LinSimp layers. LinSimp layers have six heads with D = Dv = 64, five input projections (q, k, r, v, u), feature rank m = 128, and anchor width w = 64. Parameter counts are matched within ±0.5% of the 330.35M standard-attention model by adjusting SwiGLU widths.

Experimental results show: In web pretraining at 2k context (350M-token iso-FLOP budget), the KDA hybrid has the best validation loss (3.4563) and perplexities, while KDA+LinSimp is close but does not win. Six LinSimp layers improve validation loss over 24 standard-attention layers (3.6528 vs 3.6618). In the iso-token setting, replacing three of six attention slots with LinSimp raises mean accuracy by 0.0080. In math pretraining at 2k context (3B tokens), KDA+LinSimp has the highest mean accuracy (0.3900) and highest OpenBookQA accuracy (0.188). At 16k context under iso-FLOP, KDA+LinSimp improves mean accuracy from 0.3809 to 0.3888 and lowers LAMBADA perplexity from 715.6 to 602.6, winning five of seven downstream tasks.

Kernel performance measurements on an RTX 6000 Ada show one LinSimp layer (m=128, w=64) takes 37 ms for forward and backward, about 1.03 times the softmax-attention layer time. The full no-softmax model processes 28.4k tokens/s at 2k context and 20.2k tokens/s at 16k, while the 16k KDA-softmax hybrid reaches 24.5k tokens/s. The authors note the backward pass has substantial room for optimization.

The paper includes detailed derivations in appendices covering the algebra, random-feature bank construction with orthogonal Gaussian marginals, log-space evaluation for stabilization, kernel algorithms, FLOPs counting, and partial ablations. The authors acknowledge limitations: experiments are preliminary with a single seed, no confidence intervals, models are near 330M parameters, scaling behavior is unknown, the backward pass is not highly optimized, and iso-FLOP budgets use analytic operation counts rather than energy or latency measurements.

Improvements to AI systems:

1. Replace softmax attention entirely with linear-time, kernelized 2-simplicial attention (LinSimp) in long-context models. This removes the quadratic bottleneck of standard attention while preserving global token interactions, enabling processing of sequences far beyond 16k tokens without the perplexity degradation seen in KDA hybrids (LAMBADA PPL drops from 715.6 to 602.6).

2. Use the positive orthogonal random-feature bank with a short explicit anchor window (w=64) as a drop-in attention layer. This provides a fixed-size linear state (matrices M i, vectors a i) that can be updated in O(1) per token, allowing streaming inference over arbitrarily long inputs with bounded memory—critical for real-time dialogue, code completion, or document analysis.

3. Adopt the composite query–key rewriting s ijc = ((tau i) c) j to fuse multi-head relational reasoning into a single inner product. This enables models to capture higher-order interactions (e.g., entity A relates to entity B in context C) at linear cost, improving performance on tasks requiring relational or compositional understanding (e.g., OpenBookQA accuracy rose to 0.188, highest among compared models).

4. Implement custom Triton forward/backward kernels for LinSimp with log-space evaluation for numerical stability. This allows training and inference on standard GPUs with only 1.03× the time of softmax attention per layer, making the improvement practical without specialized hardware. The backward pass can be further optimized for speedups.

5. Combine LinSimp with KDA layers in a hybrid architecture (e.g., 18 KDA + 6 LinSimp) to balance short-range precision and long-range efficiency. This yields the highest mean downstream accuracy (0.3888 at 16k context) and wins 5/7 tasks, outperforming both pure softmax and KDA-only hybrids under matched compute.

6. Use the iso-FLOP and iso-token evaluation methodology to guide architecture search. The paper shows that replacing 3 of 6 attention slots with LinSimp improves mean accuracy by 0.0080 under iso-token, suggesting a principled way to allocate compute between linear and quadratic attention layers for given budgets.

What the improved AI system can do:

• Process 100k+ token contexts (e.g., entire books, long codebases, multi-hour audio transcripts) with linear memory and time, while maintaining or improving downstream task accuracy compared to softmax-based models.

• Perform real-time streaming inference on unbounded sequences (e.g., live transcription, continuous sensor data) with a fixed-size state, without forgetting distant context.

• Answer complex relational questions that require combining multiple pieces of information across a long document (e.g., How does the author's argument in Chapter 3 relate to the evidence in Chapter 10?).

• Achieve lower perplexity on language modeling benchmarks (e.g., LAMBADA) than hybrid softmax-linear models, reducing hallucination and improving coherence in generated long-form text.

• Run efficiently on consumer GPUs (e.g., RTX 6000 Ada) with only 3% overhead per layer compared to standard attention, making long-context AI accessible without massive compute clusters.

We present a linearized form of 2-simplicial attention by rewriting the trilinear score as an inner product between a composite query and a key, so that the sum over one token axis takes the same form as ordinary softmax attention. We then approximate this sum with positive random features and store the entire past in a fixed-size state, while the second axis stays explicit over a short window of recent tokens. This enables us to achieve linear cost in sequence length combined with a global reach that windowed 2-simplicial attention lacks. We implement it with custom Triton kernels and combine it with Kimi Delta Attention to build a model with no softmax attention at all. Under matched compute, this model achieves the highest mean downstream accuracy among the compared architectures, and at 16k context it improves mean accuracy over a KDA hybrid while lowering LAMBADA perplexity from 715.6 to 602.6.