Learning Multi-Agent Coordination via Sheaf-ADMM

Learning Multi-Agent Coordination via Sheaf-ADMM

Learning Multi-Agent Coordination via Sheaf-ADMM

We introduce Sheaf-ADMM, a different way to build a neural network based on the notion of multi-agent consensus. The framework is built on the intersection of sheaf theory and ADMM for distributed consensus.

Resources

Paper

Code

Authors

Published

July 2026

Limited-view agents negotiating a global answer.

Introduction

AI systems are increasingly composed of many interacting agents rather than a single monolithic model. In current practice, multi-agent systems are typically centralized, such as with an orchestrator delegating and assigning subtasks. In many systems of interest, however, no such central coordinator exists. The nodes of a sensor network, the ants of a colony, or the neurons of a nervous system each observe a small part of the environment and communicate only with their neighbors, and coherent global behavior arises from these local interactions alone .

We wish to study the mechanisms of collective coordination. Our approach is to focus on the problem of distributed consensus — how multiple agents with individual views of data agree on a global state — and to look for inspiration in existing fields that have studied distributed consensus from different angles.

In distributed optimization, the alternating direction method of multipliers (ADMM) splits a global problem into per-agent subproblems; each agent solves its subproblem, reconciles its solution with neighbors, and repeats until the system reaches global consensus. The algorithm is rooted in the theory of convex optimization, but each step admits an elegant interpretation in terms of multi-agent coordination.

A complementary question is what neighboring agents must agree on. Full consensus is often too restrictive; an alternative is to ask agents to agree on only linear projections of their state. Incidentally, this coincides precisely with an object from applied algebraic topology: a network sheaf , which offers topological tools for studying distributed systems, such as harmonic states, topological obstructions to coordination, and more .

Both ADMM and sheaves offer complementary framings for local-to-global coordination. We develop Sheaf-ADMM, which utilizes both in a learnable system. ADMM supplies coordination and negotiation dynamics, and the sheaf structure supplies the notions of inter-agent agreement. In Sheaf-ADMM, coordination evolves by running an ADMM solver across the latent space of hundreds of communicating agents. No agent sees enough of the input to solve the task on its own. The global solution emerges only from the agents' local negotiation.

We establish the method in simple settings: image classification, multi-agent Sudoku, and maze pathfinding. By focusing on simple tasks, we are able to isolate these coordination mechanisms clearly, and make them amenable to investigation.

Paper: arxiv.org/abs/2605.31005

Code: github.com/SakanaAI/sheaf-admm

Sheaves

We first introduce the sheaf component of the framework. Typical message-passing neural networks (MPNNs) use arbitrary learnable nonlinear maps (e.g. MLPs) to pass messages between agents (i.e. between nodes of a communication graph) . Alternatively, a network sheaf gives a simple, linear, and—importantly—highly interpretable implementation of message passing .

The sheaf consensus condition. Two neighboring agents each hold a private state in $\mathbb{R}^2$. Restriction maps $F_{ij}, F_{ji}$ project both into a shared 1-D public communication channel. Consensus is reached when the projections agree, $F_{ij}x_i = F_{ji}x_j$.

In a sheaf, each agent $i$ holds a private state vector $x_i\in\mathbb{R}^d$ (a decision, or a latent representation). Two neighboring agents try to reach consensus — not by agreeing on their entire state vector, but only on a learned linear projection, $F_{ij} x_i = F_{ji} x_j$. Gradient descent on $\|F_{ij} x_i - F_{ji} x_j\|^2$ yields a message passing update:

$$x_i \leftarrow x_i - \eta\, F_{ij}^\top (F_{ij} x_i - F_{ji} x_j)$$

for each agent. Importantly, this update requires only knowledge of local state and pairwise prediction error ($F_{ij}x_i-F_{ji}x_j$) of adjacent agents. When cast across hundreds of agents with local connectivity, this amounts to a decentralized algorithm — sheaf diffusion — that iteratively updates each agent's private state vector to reach global consensus (defined as pairwise $F_{ij}x_i=F_{ji}x_j$ for all communicating agents).

Sheaf-ADMM

The Sheaf-ADMM architecture. A shared encoder maps each local view to a small convex problem; the ADMM layer alternates local optimization (x), consensus via sheaf diffusion (z), and dual accumulation (u) for $K$ rounds; a shared decoder turns the final states into local predictions, fused into the global answer. Every step is differentiable.

Sheaf diffusion gives a method for reaching inter-agent consensus—but what solution should each agent propose to begin with?

In Sheaf-ADMM, each local decision vector comes from the agent solving a local subproblem. Each agent does a small amount of work given what portion of the input they can see, then compares that to the results given by their neighbors, iterating until convergence. This back and forth — propose and compare — is precisely the kind of process codified in ADMM.

Concretely, each agent's local objective is a small convex program, specified by $f(x_i; \theta_i)$, and a shared encoder reads agent $i$'s local view of the input to produce its parameters $\theta_i$. Treating that per-agent $\arg\min$ as a layer of the network is the idea behind differentiable optimization layers like OptNet .

We iterate for $K$ ADMM rounds . At the first round $k=0$, initialize a consensus state $z_i=0$ for each agent and local error $u_i=0$:

One ADMM round — all agents $i$ in parallel, repeat for $k = 1,\dots,K$

x-updatelocal solve $\displaystyle \textcolor{#2C6FB3}{x_i} \leftarrow \arg\min_{x_i}\ \underbrace{f(x_i;\theta_i)}_{\text{local objective}} + \underbrace{\tfrac{\rho}{2}\big\|\textcolor{#2C6FB3}{x_i} - \textcolor{#2E9E6F}{z_i} + \textcolor{#E07A1F}{u_i}\big\|^2}_{\text{stay near consensus}}$

z-updateconsensus $\displaystyle \textcolor{#2E9E6F}{z_i} \leftarrow \textcolor{#2E9E6F}{z_i} - \eta \sum_{j \sim i} F_{ij}^\top \underbrace{\big(F_{ij}\,\textcolor{#2E9E6F}{z_i} - F_{ji}\,\textcolor{#2E9E6F}{z_j}\big)}_{\text{prediction error}}$ sheaf diffusion$T$ steps, from $z_i = x_i + u_i$

u-updatedual update $\displaystyle \textcolor{#E07A1F}{u_i} \leftarrow \textcolor{#E07A1F}{u_i} + \underbrace{\textcolor{#2C6FB3}{x_i} - \textcolor{#2E9E6F}{z_i}}_{\text{disagreement}}$

Each agent proposes their state based on their local information, diffusion pulls the proposals into agreement, and $u$ records the leftover difference. Because $u$ accumulates over rounds, an agent that stays far from consensus gets pushed harder each round towards consensus until the entire system is in a globally consistent state.

A local view

One of the benefits of this extra machinery is that it specifies a highly localized notion of an agent. Each agent uses only a local decision $x_i$, a local target $z_i$ updated by its neighbors' prediction errors, and a local memory $u_i$.

Experiments

Multi-agent Sudoku

Sudoku: full-puzzle solve rate and per-cell accuracy (%), 3-seed mean.

Recurrent reasoning models achieve high solve rates on Sudoku , using self-attention to couple every cell to every other across a fully connected token graph. We instead use Sudoku as a multi-agent benchmark and deliberately make it harder: each agent sees only a local region of the board and communication only occurs between agents who share cells. A problem like Sudoku splits into local subproblems (each agent's nine cells must all be distinct) coupled by pairwise consensus (two agents that share a cell must agree on its digit), meaning Sudoku can be solved by divide and concur approaches . We found that MPNN baselines struggled (strongest baseline; 34.7%) in the multi-agent Sudoku setting compared to Sheaf-ADMM (92.6%) using the Ritvik19 Sudoku dataset . This is likely due to Sheaf-ADMM using local subproblems in place of learned MLPs, more naturally capturing the “divide and concur” nature of Sudoku.

“Multi-agent” Sudoku. Each row, column and 3×3 box corresponds to one agent, with agent communication defined by whether two agents share cells.

Maze pathfinding with small communication channels

Maze exact-solve rate (%), 3-seed mean.

Maze pathfinding can also serve as a local-to-global testbed. While standard GNNs and MPNNs readily solve maze pathfinding, in some cases replicating exact Bellman-Ford updates , here we show that Sheaf-ADMM can do so while communicating over a much smaller channel. Since each Sheaf-ADMM agent carries distinct state variables $x,z,u$, this local memory eases the burden on communication: in our custom maze data, Sheaf-ADMM solves 2× OOD mazes at a fraction of the public communication bandwidth used by MPNN baselines, which must carry all internal state in a single vector $h$. This suggests that the ADMM splitting allows for more bandwidth-efficient communication at the cost of additional local memory.

Interpretability

Finally, the Sheaf-ADMM framework is especially transparent and open to intervention and interpretability. Each agent keeps three inspectable quantities, $x,z,u$, and all negotiation dynamics occur in these latent states. And all communication is linear, even opening the global network to sheaf-theoretic analysis (harmonic states, topological obstructions) . While this isn't the focus of the present work, the system's linearity makes it especially amenable to such analysis, and we see it as a promising direction for future work. The figure below shows maze pathfinding agents on small mazes, each with a 2-dimensional latent state and 1-dimensional communication channels. The local view specifies the parameters of each agent's local objective, and a shared decoder maps the latent state to a binary class (path or not-path) per pixel. We show the decision boundary in these latent states for path/not-path of the agent's center pixel, along with the trajectories of the $x$ (local greedy) and $z$ (consensus) variables. Over iterations (ADMM rounds), the $x$ and $z$ variables coincide at solutions corresponding to a global maze solution.

$x$/$z$ dynamics on a maze. A 3×3 block of agents, each seeing only a 3×3 patch. Right panels show agents' $x$-state (blue), $z$-state (red) over ADMM negotiation rounds. The decision boundary for path/not-path is shown in the latent state.

Future directions

The Sheaf-ADMM framework opens several directions of future research by drawing on ongoing developments in both the ADMM and applied topology literatures. This is especially important since distributed consensus can face several difficulties and limitations at scale; both the ADMM and sheaf-theoretic literature develop tools for addressing them, and offer directions for future extensions. To name a few: asynchronous updates , Byzantine-tolerant formulations in the presence of adversarial agents , adaptive penalties for improved convergence , and nonconvex extensions . The sheaf formalism offers tools for analyzing local-to-global structure, and exposes global features that are otherwise opaque in generic learned nonlinear systems, e.g. measuring the net effect of stubborn agents , notions of topological obstructions to coordination , and more. While we focus in this work on simple tasks (that could easily have been solved with a “monolithic” approach), we anticipate that some of the lessons learned in these settings may inform research on inherently distributed problems, where no single vantage point exists and coordination must be local.

Citation

@inproceedings{seely2026sheafadmm, title =

[truncated for AI cost control]