Bridging Abstract Mathematics and Systems Engineering: Category Theory for Tiny ML in Rust
The intersection of high-level mathematical abstraction and low-level systems programming is rarely a fertile ground for practical software engineering. Usually, category theory is relegated to the realm of purely functional languages like Haskell, while machine learning (ML) is viewed as a series of numerical linear algebra operations. However, a new project, Category Theory for Tiny ML in Rust, seeks to bridge this gap by treating ML not just as computation, but as a structured pipeline of objects, transformations, and constraints.
Developed by Hamze Ghalebi and Farzad Jafarranmani, this working draft proposes a framework where the mathematical rigor of category theory serves as an engineering tool for building reliable, auditable, and maintainable ML systems in Rust. By mapping mathematical concepts directly to Rust's type system, the authors aim to make the "abstract nonsense" of category theory executable and concrete.
The Conceptual Mapping: Math to Rust
At the heart of the project is the translation of categorical concepts into the idioms of the Rust programming language. The goal is to move away from treating ML as a "black box" of tensors and instead view it as a composition of morphisms.
Domain Objects as Rust Types
In category theory, a category consists of objects and morphisms. In this framework, domain objects are mapped directly to Rust types. This ensures that the data flowing through an ML pipeline is strictly typed, reducing runtime errors and making the architecture of the system explicit in the code.
Morphisms as Typed Transformations
Morphisms (the arrows between objects) are implemented as typed transformations. In the context of ML, a morphism could be a layer in a neural network or a preprocessing step. By defining these as morphisms, the system emphasizes the composition of these functions, ensuring that the output of one transformation is mathematically and programmatically compatible with the input of the next.
Training as an Endomorphism
One of the more provocative ideas in the book is the treatment of training as an endomorphism. An endomorphism is a morphism that maps an object back to itself. In this framework, training is viewed as a repeated transformation of the model's state, where the state is the object being transformed to improve its performance over time.
From Theory to Production
The project is driven by two distinct perspectives: the mathematical foundation provided by Farzad Jafarranmani (specializing in proof theory and denotational semantics) and the production engineering perspective of Hamze Ghalebi (focused on GenAI and auditable AI systems).
This duality is intended to solve a common problem in AI development: the gap between a prototype and a production-ready system. By using Rust, the authors leverage memory safety and performance, while category theory provides a blueprint for creating systems that can be evaluated, monitored, and kept under human accountability.
Critical Perspectives and Community Debate
As with any project blending high-level math with systems code, the community response has been mixed, raising important questions about the utility of such abstractions.
The "Decorative Abstraction" Critique
Some critics argue that category theory may be "stapled on" to the implementation. One commenter noted that regular typed programming already uses types for domain objects and functions for transformations, suggesting that the categorical terminology might be descriptive rather than prescriptive:
"I don't understand, this looks to me like regular Rust, or regular programming for that matter... Category Theory terminology can be used to describe the structure of a regular typed program."
The Missing Link: HKTs and Theorems
From a more technical standpoint, some developers pointed out the limitations of Rust's type system. Specifically, the lack of Higher-Kinded Types (HKTs) in Rust makes it difficult to implement certain categorical patterns (like Functors or Monads) as elegantly as one could in Haskell.
Furthermore, some researchers argued that for category theory to be truly useful in ML, it should allow for the derivation of theorems about the implementation. They suggested that the framework should move toward Markov categories or use adjunctions to characterize ML algorithms, which would provide a stronger mathematical guarantee of the system's behavior.
Conclusion
Category Theory for Tiny ML in Rust represents an ambitious attempt to formalize the "plumbing" of machine learning. While skeptics question whether the added mathematical layer provides tangible benefits over standard typed programming, the project's focus on auditability and structure offers a compelling vision for the future of reliable AI. By turning mathematical structures into executable Rust code, the authors are challenging developers to think of ML pipelines not as a series of matrix multiplications, but as a rigorous composition of typed transformations.