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

TL;DR

The paper tackles the bottleneck of designing efficient, self-improving AI systems by introducing neural circuit diagrams grounded in category theory as a framework to derive fusion and streaming properties and to guide hardware-aware algorithm design. It proposes spherical attention, replacing the SoftMax with an $L^{2}$-norm normalization to enable streaming and signed weight signals, backed by fusion-theorem proofs. The approach is applied to gene regulatory networks by modeling gene interactions with signed $QK^{T}$ relations and multiplicities, and is implemented as a low-level CUDA kernel with memory tiling that achieves around 200 TFLOP/s on an A100, delivering up to $3.6\times$ speedup over PyTorch and competitive performance versus FlashAttention-2. The work demonstrates a path toward automated, diagram-driven discovery and optimization of deep learning architectures, potentially enabling self-improving AI systems with hardware-aware compilation pipelines.

Abstract

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 $L^2$ 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 \textit{FlashSign} kernel achieves comparable performance to the state-of-the-art, fine-tuned FlashAttention algorithm on an A100, and $3.6\times$ 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.

Figures (9)

• Figure 1: On the left is a traditional diagram from Attention is All You Needvaswani_attention_2017. On the right, we see a neural circuit diagram, which shows the details of axes and data sizes as well as broadcasting.
• Figure 2: Generic normalized contraction, which we represent with a triple-lined triangle, can be expanded into a loop where the $n$-axis is partitioned into chunks of size $s$, and only $z$ and $o$ are maintained between chunks.
• Figure 3: An attention algorithm with a generic normalized contraction in place of SoftMax can have its $x$-axis streamed at a lower-level, resulting in the $s_x$ relabeling.
• Figure 4: The full model, where spherical attention is combined with iterative feed-forward layers, can be diagrammatically expressed.
• Figure 5: Low-level diagram of spherical attention, adapting Figure \ref{['fig:attention_expansion']} to derive diagrammatic pseudocode.

Theorems & Definitions (6)

• lemma thmcounterlemma
• theorem thmcountertheorem
• definition thmcounterdefinition: Streamable Function
• definition thmcounterdefinition: Normalized Contraction
• lemma thmcounterlemma
• theorem thmcountertheorem