Learning Multi-Agent Coordination via Sheaf-ADMM
Abstract
We present a differentiable optimization framework for multi-agent coordination. An input is decomposed into overlapping local views, each processed by an agent that solves a convex subproblem parameterized by a neural encoder. Agents coordinate through the Alternating Direction Method of Multipliers (ADMM) with inter-agent constraints specified by a cellular sheaf. The sheaf specifies which aspects of neighboring solutions must agree, allowing for heterogeneous notions of global consensus. Backpropagating through the unrolled optimization jointly trains all components of the multi-agent system. We evaluate on maze pathfinding, image classification, and Sudoku, where agents with individually insufficient local views learn to coordinate to produce correct global outputs. On MNIST, the local-view decomposition yields improved robustness to distribution shifts relative to a standard CNN. On Sudoku, the optimization-derived structure yields markedly higher solve rates than parameter-matched MPNN baselines. Finally, the ADMM structure exposes distinct primal, consensus, and dual state variables, opening the coordination dynamics to direct analysis and intervention—a property unavailable in standard message-passing architectures.
Lay Summary
Standard AI systems use a single large neural network to process inputs as a whole. By contrast, intelligence in nature—ant colonies, neural circuits, immune systems—emerges from many small agents that each see only part of the world and coordinate to solve global problems. We ask whether neural networks can be trained to operate the same way. We build a system where each agent observes only a small local patch of the input (for example, a 3×3 pixel window of an image). Individually, no agent has enough information to solve the task. The agents must communicate with their neighbors and reach a shared interpretation. We use a classical optimization algorithm (ADMM) combined with cellular sheaves—a mathematical structure that specifies which features neighbors must agree on while leaving the rest private. The entire system is learned end-to-end from data. We evaluate on maze pathfinding, MNIST image classification, and Sudoku. The system solves all three tasks despite no single agent ever seeing more than a small fragment of the input. It is more robust than a standard convolutional network when inputs are perturbed, and substantially outperforms similarly-sized message-passing baselines on Sudoku, where solving the puzzle requires reconciling constraints across the entire board.