arXiv:2505.09326math.CTcs.LG2025-05被引 1

用范畴论设计高效基因调控网络注意力机制,性能超PyTorch三倍。

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

  • 基于范畴论的神经电路图推导出新注意力算法
  • 在A100上比PyTorch快3.6倍,接近FlashAttention性能
  • 适合追求高效率架构自动化的AI系统开发者

如何让人工智能模型自我优化?这对实现可泛化的智能系统至关重要,能高效利用最新硬件处理新问题。当前自动化编译方法效果差,高效算法需多年人工开发。本文基于范畴论的神经电路图,证明了深度学习算法的一般定理,指导设计专用于基因调控网络的新注意力算法,并生成高效内核。所提球面注意力(spherical attention)通过用L²范数替代SoftMax,克服了标准注意力的特殊函数单元瓶颈,同时保持流式计算特性。其衍生的 extit{FlashSign}内核在A100上表现接近最先进微调版FlashAttention,较PyTorch提升3.6倍。该研究验证了神经电路图作为自动化高效架构开发的高层框架的适用性。

原文摘要 · Abstract (English)

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.

注意力机制范畴论基因调控高性能计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。