Linearized 2-Simplicial Attention

The evolution of attention mechanisms in deep learning has led to transformers becoming the dominant architecture in NLP. Softmax-based attention provides rich, content-dependent access to tokens but suffers from quadratic complexity, limiting scalability. Recent efforts like linear attention, recurrent models, and state-space approaches aim to reduce this complexity by compressing past information into fixed-size states or sparse representations. High-order attention models, such as 2- simplicial transformers, can capture complex relationships among token triples, but their cubic computational cost restricts practical use. Prior work has attempted windowed or hardware-accelerated solutions, yet a fully global, efficient, high-order attention remains elusive. This paper builds on these foundations, proposing a scalable, high-order attention mechanism that balances expressiveness and efficiency.

The core challenge is to enable higher-order attention, specifically 2-simplicial interactions, with linear or near-linear complexity suitable for long sequences. Traditional formulations involve triple interactions that scale as O(n^3), making them infeasible for large-scale applications. Existing solutions either restrict the attention window or rely on hardware-specific optimizations, which limit the model's ability to capture global dependencies. The key problem is to reformulate the high-order scoring function into a form that allows approximation via random features, enabling scalable, causally consistent modeling of complex relationships across long sequences without sacrificing efficiency or accuracy.

The paper introduces a novel reformulation of 2-simplicial attention by expressing the trilinear score as an inner product between a composite query and a key, enabling kernel approximation via positive random features. This approach reduces the complexity from cubic to linear with respect to sequence length. It combines a fixed-size global state that encodes the entire past with a short, explicit window over recent tokens, allowing the model to maintain long-range dependencies efficiently. Custom Triton GPU kernels further optimize the computation, and the integration with Kimi Delta Attention results in a fully non-softmax, high-capacity architecture that captures complex interactions while remaining computationally feasible.

• �� Reformulate the trilinear score as an inner product between a composite query zic and a key ˆkj. • Approximate the exponential kernel exp(z⊤ic ˆkj) using positive random features ϕ(x). • Store the entire prefix in a fixed-size state Mi and ai, updating incrementally. • Use a short window of recent anchors Ci to incorporate local context. • Implement custom Triton GPU kernels for efficient forward and backward passes, handling the exponential feature computations and prefix updates. • Apply a gating mechanism to modulate the contribution of the trilinear head, enhancing expressiveness. • Ensure causality by restricting updates to past tokens only, maintaining linear complexity. • Train under an iso-FLOP budget to compare computational efficiency across architectures.

Experiments involve large-scale pretraining on datasets like FineWeb-Edu and FineMath-4+, with sequence lengths of 2k and 16k tokens. Baselines include standard softmax attention, windowed 2-simplicial attention, and KDA hybrid models. Metrics include validation loss, accuracy on benchmarks like ARC, BoolQ, and WinoGrande, as well as perplexity on WikiText and LAMBADA. The models are parameter-matched (~330M) with fixed random feature dimension m and window size w. GPU kernel performance is measured, showing inference speeds close to softmax attention. The models are evaluated under both token and FLOP-based compute budgets, demonstrating scalability and efficiency.

The proposed model surpasses baseline architectures in accuracy and perplexity, achieving a 0.0079 accuracy increase at 16k context over KDA hybrid, and reducing LAMBADA perplexity from 715.6 to 602.6. GPU kernel tests show inference speeds of 20-28k tokens/sec at 16k context, close to softmax attention. Ablation studies confirm the importance of the global state and local window mechanisms. The model maintains linear complexity, enabling efficient long-sequence processing without significant performance degradation, marking a breakthrough in high-order, scalable attention.

This architecture is suitable for long-document understanding, large-scale language modeling, and knowledge graph reasoning. Its linear complexity allows deployment in resource-constrained environments, enabling applications like real-time long-form content generation, comprehensive question answering, and multimodal data integration. The model's ability to efficiently capture complex relationships makes it promising for future AI systems requiring deep contextual understanding.

Despite promising results, the model's current GPU kernel optimization is preliminary, with room for speed improvements. Fixed global states may limit adaptability in highly dynamic contexts, and the approach's scalability to billion-parameter models needs further validation. Additionally, the reliance on random feature approximation introduces approximation errors, which could affect performance in certain tasks. Future work should address dynamic memory management, kernel optimization, and robustness in diverse scenarios.