Fast and Simplex: 2-Simplicial Attention in Triton

2025-07-03

AI Summary

Large language models (LLMs) suffer from low token efficiency due to their reliance on massive internet-scale corpora. To address this, we propose the 2-simplicial Transformer—a novel architecture that generalizes dot-product attention to a trilinear form for the first time, thereby provably altering the scaling exponent of knowledge acquisition and reasoning tasks. We further design a high-performance, Triton-based GPU kernel that optimizes memory access patterns and computational throughput, significantly improving token utilization. Empirical evaluation across mathematical reasoning, code generation, logical deduction, and multi-step inference shows consistent superiority over standard Transformers at equivalent parameter counts, demonstrating enhanced modeling capacity and generalization under fixed token budgets. Our core contributions are: (i) a theoretically grounded attention paradigm with improved asymptotic scalability, and (ii) its efficient, system-level implementation enabling practical deployment.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Recent work has shown that training loss scales as a power law with both model size and the number of tokens, and that achieving compute-optimal models requires scaling model size and token count together. However, these scaling laws assume an infinite supply of data and apply primarily in compute-bound settings. As modern large language models increasingly rely on massive internet-scale datasets, the assumption that they are compute-bound is becoming less valid. This shift highlights the need for architectures that prioritize token efficiency. In this work, we investigate the use of the 2-simplicial Transformer, an architecture that generalizes standard dot-product attention to trilinear functions through an efficient Triton kernel implementation. We demonstrate that the 2-simplicial Transformer achieves better token efficiency than standard Transformers: for a fixed token budget, similarly sized models outperform their dot-product counterparts on tasks involving mathematics, coding, reasoning, and logic. We quantify these gains by demonstrating that $2$-simplicial attention changes the exponent in the scaling laws for knowledge and reasoning tasks compared to dot product attention.

Problem

• Improving token efficiency in large language models
• Generalizing dot-product attention to trilinear functions
• Enhancing performance in math, coding, and reasoning tasks

Innovation

• 2-simplicial Transformer for token efficiency
• Trilinear attention via Triton kernel
• Improved scaling laws for reasoning tasks