Aug 22, 2026 - 5 min read

Bridging Categorical Deep Learning and Trace Theory: A Foundation for Distributed AI

Two Sentences Summary

In this post, we share a possible extension to Abbott & Zardini's categorical formalization of deep learning. More specifically, we propose to connect the internal broadcasting formalism with the external concurrency of Diekert & Muscholl's Trace Theory to formalize distributed training and multi-agent systems.

Introduction

From the early Large Scale Distributed Deep Networks, deep learning (DL) is fundamentally parallel. Yet, our mathematical descriptions of it are often sequential or ad-hoc. The recent paper by Abbott & Zardini takes a giant leap forward by using category theory to formalize the internal parallelism of a single model. Specifically, they identify broadcasting as a core mechanism in DL programming frameworks. Indeed, it is parallelized computing that gives DL its ability to cut through massive datasets, leaving classical statistical learning theory struggling to keep pace.

However, as I dug deeper, I realized that the formalism stops at the boundaries of a single model. It describes a single model perfectly, but what happens when we have a society of brains? What about distributed training, federated learning, or multi-agent systems?

This is where Trace Theory—beautifully summarized by Diekert & Muscholl (2011)—might come into play. I believe it provides the exact missing layer to extend Abbott & Zardini's work into the realm of concurrent, distributed, and interactive AI systems.

The View from Abbott & Zardini (The Internal View)

The Abbott & Zardini paper introduces new categorical structures—axis-stride and array-broadcasted—to describe how a computation defined on a single vector extends to every row of a matrix in parallel.

In their framework:

  • Morphisms represent tensor operations,
  • Composition represents sequential layers,
  • Categorical products represent parallel constructs.

This formalism elegantly captures the "physics" of a deep learning model. It tells us how a computation is parallelized across tensor axes. Crucially, it identifies broadcasting as the mathematical prerequisite for parallelized computations and especially well-suited for GPU kernels.

But it is centralized. It seems to assume a single, static computation graph on a single device (or a homogeneous cluster abstracted as one).

The View from Diekert & Muscholl (The External View)

Trace theory (introduced by Mazurkiewicz in 1977), as outlined in Diekert & Muscholl (2011), approaches concurrency from a different angle. It starts with an alphabet of atomic actions and defines an independence relation (I). Two sequences of actions are considered equivalent (i.e., they belong to the same trace) if they can be transformed into each other by swapping independent actions:

ab = ba \quad \text{for } (a, b) \in I

This simple rule generates free partially commutative monoids. Here are the key takeaways from the survey:

  1. Dependence Graphs: Traces are visualized as directed acyclic graphs (DAGs) where edges represent dependencies (the complement of independence). If two operations don't have a path between them, they are concurrent.
  2. (Foata) Normal Form: This provides a canonical representation of a trace in terms of maximal parallel steps (e.g., ({a,d};{b,c};{a,e};{f})).
  3. Asynchronous Automata: Zielonka's theorem proves that any recognizable trace language can be accepted by a distributed automaton with finite-state processes synchronizing over shared variables (CREW, EREW models).
  4. Regular vs. Rational Languages: A fascinating result: while Kleene's theorem holds for words, it fails for traces. Rational expressions like ((ab)^*) (with (a) and (b) independent) represent non-regular languages (equality of counts). Recognizability requires the star operation to be restricted to connected languages. Finally, because it's very close to the language of words, one may expect to re-use/transfer the results of that body of work.

The Bridge: The Internal and External Meet

Here is my core thesis: Abbott & Zardini gives us the "letters" (the tensor kernels), and Trace Theory gives us the "schedules" (the communication and concurrency).

We can visualize how this forms a unified stack for distributed training:

1. The Node (Traced Morphism)

Each GPU or client in a distributed system is represented using the Abbott & Zardini formalism. It defines the local model, the forward pass, and the backward pass. To make it stateful (i.e., an optimizer with momentum), we wrap it in a Trace (categorical trace, not trace theory yet—though they are related). This gives us a morphism F: X \otimes W \rightarrow Y \otimes W , where W represents the persistent weights.

2. The Communication (Trace Monoid)

Now, consider the communication actions:

  • (a): Compute gradient on Node 1
  • (b): Compute gradient on Node 2
  • (c): All-Reduce over the network

In a synchronous setup, (a) and (b) are independent (they can happen in any order). However, (a) and (c) are dependent because the All-Reduce must wait for the local gradient computation to finish.

Using Trace Theory, we can model the global training step as a trace: t = (a \parallel b) \rightarrow c \rightarrow (d \parallel e) . This directly maps to the Foata normal form! The maximal parallel execution of a distributed training step is precisely a trace.

3. Federated Learning as a Trace Automaton

Federated Learning is particularly interesting. Different clients have different data distributions. They compute locally (a_i ) and then aggregate (FedAvg, (agg)).

If the aggregation is synchronous, the system is a finite trace. But if we consider asynchronous updates (where clients participate at different times), we enter the realm of infinite traces or rational trace languages.

Using the CREW (Concurrent-Read Exclusive-Write) model from trace theory:

  • Multiple clients can read the global model simultaneously.
  • The write (update) must be exclusive or follow a conflict-free schedule.

Zielonka's theorem—stating that recognizable trace languages correspond to asynchronous automata—suggests that we can rigorously design a distributed training algorithm where the convergence properties are independent of the specific scheduling of independent operations.

Why This Matters for Multi-Agent Systems

This unification isn't just for training. If we extend this to multi-agent systems:

  • Each agent is a traced morphism (stateful).
  • The interaction between agents (e.g., message passing, negotiation) defines the independence alphabet.
  • An agent's strategic decision (e.g., policy distribution) can be mapped to a game sheaf on top of this trace structure.

The trace monoid captures what can happen in parallel. The sheaf or game-theoretic layer captures why an agent chooses a specific action within that concurrent space.

Practical Implications

  1. Formal Verification: We can use recognizable trace languages to verify properties of distributed training algorithms (e.g., "Does this asynchronous schedule guarantee global gradient descent?").
  2. Automatic Kernel Generation: The Foata normal form of a trace gives you the optimal parallel schedule of communication and computation. A compiler could automatically derive the communication schedule from a high-level description of a distributed model.
  3. Decidability: Trace theory provides clear boundaries. For example, inclusion and equality of rational trace languages are undecidable, but for transitive independence alphabets (free products of free monoids), they are decidable. This tells us exactly which distributed topologies are "safe" to analyze automatically.

What's Next?

I sketched out this idea over the past few days, and it feels like a promising formal foundation. The Abbott & Zardini paper gives us the microscope (the internals of a model). Mazurkiewicz's trace theory—through Diekert & Muscholl's reference—gives us the telescope (the concurrency between models).

The next step for someone (maybe you) would be to:

  1. Define a functor from the Abbott & Zardini tensor categories into a trace monoid.
  2. Implement a small prototype in pyncd that can output the Foata normal form of a distributed model graph.
  3. Explore the connection between asynchronous automata and gradient synchronization protocols.

Formalizing distributed AI is a hard problem, but with tools like category theory and trace theory, I think we may have a solid, right language for it.

If you found this useful, consider supporting my work.