# Accelerating Machine Learning Systems via Category Theory: Applications to Spherical Attention for Gene Regulatory Networks

arXiv: https://arxiv.org/abs/2505.09326 · HTML: https://mlsanity.com/paper/2505.09326/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 23 references that are in the map.
- First posted on arXiv 14 May 2025.

## Builds on (most cited first, 12 of 12)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Backprop as Functor: A compositional perspective on supervised learning](https://mlsanity.com/paper/1711.10455/) (2017)
- [Categorical Foundations of Gradient-Based Learning](https://mlsanity.com/paper/2103.01931/) (2021)
- [Position: Categorical Deep Learning is an Algebraic Theory of All Architectures](https://mlsanity.com/paper/2402.15332/) (2024)
- [Neural Circuit Diagrams: Robust Diagrams for the Communication, Implementation, and Analysis of Deep Learning Architectures](https://mlsanity.com/paper/2402.05424/) (2024)
- [FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness](https://mlsanity.com/paper/2412.03317/) (2024)
- [Categories of Differentiable Polynomial Circuits for Machine Learning](https://mlsanity.com/paper/2203.06430/) (2022)

## Built on by (most cited first, 1 of 1)

- [Weaves, Wires, and Morphisms: Formalizing and Implementing the Algebra of Deep Learning](https://mlsanity.com/paper/2604.07242/) (2026)

## Abstract (arXiv)

How do we enable artificial intelligence models to improve themselves? This is central to exponentially improving generalized artificial intelligence models, which can improve their own architecture to handle new problem domains in an efficient manner that leverages the latest hardware. However, current automated compilation methods are poor, and efficient algorithms require years of human development. In this paper, we use neural circuit diagrams, based in category theory, to prove a general theorem related to deep learning algorithms, guide the development of a novel attention algorithm catered to the domain of gene regulatory networks, and produce a corresponding efficient kernel. The algorithm we propose, spherical attention, shows that neural circuit diagrams enable a principled and systematic method for reasoning about deep learning architectures and providing high-performance code. By replacing SoftMax with an L2 norm as suggested by diagrams, it overcomes the special function unit bottleneck of standard attention while retaining the streaming property essential to high-performance. Our diagrammatically derived textitFlashSign kernel achieves comparable performance to the state-of-the-art, fine-tuned FlashAttention algorithm on an A100, and 3.6× the performance of PyTorch. Overall, this investigation shows neural circuit diagrams' suitability as a high-level framework for the automated development of efficient, novel artificial intelligence architectures.