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

TL;DR

This paper introduces a category theory-based framework using neural circuit diagrams to develop a novel spherical attention algorithm for gene regulatory networks, achieving high efficiency and performance comparable to state-of-the-art methods.

Contribution

It presents a new theoretical approach using neural circuit diagrams to systematically design efficient deep learning algorithms, specifically a spherical attention method for gene regulatory networks.

Findings

The spherical attention algorithm overcomes SoftMax bottleneck.

The FlashSign kernel matches state-of-the-art performance on A100.

The approach enables automated development of efficient AI architectures.

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.