Formalized Deep Learning Architectures for Automated Low-Level Kernel Optimization

Abstract: Vincent Abbott is a PhD student at the Massachusetts Institute of Technology's (MIT) Zardini Lab who has developed a formal framework for describing the relationship between the mathematical function implemented by a deep learning model, its resource usage, and low-level implementation. These methods are based on category theoretic diagrams [1]. The Zardini Lab has developed these diagrams into a tool for rapidly deriving low-level algorithms, as presented in their recent work FlashAttention on a Napkin [2]. These methods have been put into practice, deriving a FlashAttention-like algorithm for an attention variant from first principles [3].

Recently, he has been working on encoding the underlying mathematics into an automated tool for diagram generation and algorithm optimization. In this talk, Vincent Abbott will cover formal diagrams for deep learning models, show how they can be used to derive low-level algorithms such as FlashAttention and corresponding performance models, and preview work related to automated tools for diagramming and analyzing algorithms.

[1] https://openreview.net/forum?id=RyZB4...
[2] https://openreview.net/forum?id=pF2uk...
[3] https://dl.acm.org/doi/10.1007/978-3-...