Sakana AI's Sheaf-ADMM Solves Multi-Agent Sudoku at 92.6% Where Rivals Hit 34%

Most neural networks are monoliths: one model, one forward pass, one answer. But a growing number of real-world problems -- sensor grids, robot swarms, decentralized inference -- are fundamentally distributed. No single node sees the whole picture. Sheaf-ADMM, a new framework from Sakana AI accepted at ICML 2026, takes this constraint seriously and builds it into the architecture itself.

The problem with centralized thinking

AI systems are increasingly composed of many interacting agents rather than a single monolithic model. Yet in practice, most multi-agent systems still rely on a central orchestrator to delegate and aggregate. 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.

The question Sheaf-ADMM asks is deceptively simple: how do you build a neural network where coordination is the computation, not an afterthought? The answer pulls from two mature fields that have studied this problem from very different angles.

Two old ideas, one new architecture

The first ingredient is ADMM (Alternating Direction Method of Multipliers), a classical algorithm from distributed optimization. ADMM solves convex optimization problems by breaking them into smaller pieces, each of which are then easier to handle. It was introduced in the 1970s but has seen a recent resurgence in the context of dramatic increases in computing power and the development of widely available distributed computing platforms. In ADMM, each agent solves its own local subproblem, then reconciles with neighbors, and repeats until the whole system converges.

The second ingredient is a network sheaf -- a concept from applied algebraic topology. Sheaf theory provides a systematic framework to track data locally on a topological space and determine how to "glue" or assemble this local data into global structures. In the context of multi-agent systems, sheaf theory provides a framework for consensus on complex data structures. Crucially, a sheaf lets you define what neighboring agents must agree on -- not their entire state, just a learned linear projection of it. This is a much weaker and more realistic form of consensus.

Sheaf-ADMM utilizes both in a learnable system: ADMM supplies coordination and negotiation dynamics, and the sheaf structure supplies the notions of inter-agent agreement.

How it actually works

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.

The pipeline works as follows. A shared encoder maps each agent's local view of the input into the parameters of a small convex optimization problem. Then the ADMM loop runs for K iterations, cycling through three updates per round:

1. x-update (local solve): Each agent minimizes its own local objective while staying close to the current consensus estimate. This is the "local guess" step.
2. z-update (sheaf diffusion): Agents communicate with direct neighbors to smooth out conflicts. They don't need to agree on everything -- only on the learned linear projections defined by the sheaf's "restriction maps." This is the consensus step.
3. u-update (dual accumulation): If an agent and its neighbor still disagree after the z-update, that disagreement is recorded in a memory variable u. This memory grows over rounds, exerting increasing pressure on stubborn agents to compromise.

Because the ADMM loop is unrolled for a fixed number of iterations, the entire pipeline is differentiable and trained end-to-end. The sheaf restriction maps -- the matrices that define what neighbors must agree on -- are learned jointly with the encoder and decoder.

Benchmarks that stress-test coordination

The team deliberately chose tasks where no single agent can succeed alone, to isolate coordination as the variable being tested. Three benchmarks were used:

• Multi-agent Sudoku: Each agent sees only a single row, column, or 3x3 box. Sheaf-ADMM achieved a 92.6% full-puzzle solve rate, while the strongest MPNN baseline (with 4x the parameters) scored only 34.7%. Standard message-passing networks, which rely on learned nonlinear MLPs to pass messages, struggled to capture the structured "divide and concur" nature of the problem.
• Image classification under domain shift: MNIST digits were tested on canvas sizes the model had never seen during training. A standard CNN dropped to 11% accuracy on these out-of-distribution canvas sizes, while Sheaf-ADMM retained 86%.
• Maze pathfinding: Sheaf-ADMM matches a message-passing baseline's accuracy while agents communicate over a 5-dimensional channel -- 8x smaller than the 42-dimensional channel required by the baseline. Because each Sheaf-ADMM agent carries separate x, z, and u state variables locally, less information needs to travel over the wire.

The real breakthrough: you can watch it think

The headline numbers are impressive, but the deeper contribution is interpretability. Traditional message-passing networks hide their reasoning inside opaque, high-dimensional hidden states. 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. All communication is linear, even opening the global network to sheaf-theoretic analysis.

This linearity is not just aesthetically pleasing. It means you can apply tools from algebraic topology -- like measuring "harmonic states" (stable configurations the system naturally settles into) or detecting "topological obstructions" (structural reasons why certain agent networks can never reach consensus) -- directly to the trained model. No black-box probing required.

When to use it (and when not to)

Sheaf-ADMM is a natural fit for problems that are inherently distributed -- where the data itself is fragmented across agents with no single vantage point:

• Sensor networks where each node sees a local patch of the environment
• Multi-robot coordination without a central planner
• Federated inference where data cannot be centralized
• Any structured prediction problem that decomposes into local subproblems with overlap constraints (like Sudoku, scheduling, or constraint satisfaction)

The framework is less suited to problems where a monolithic model already works well. The authors focus on simple tasks that could easily have been solved with a "monolithic" approach, explicitly framing this as a proof-of-concept to isolate and study coordination mechanisms. Scaling to large, noisy, asynchronous agent networks -- and handling non-convex local objectives -- remains open work. The paper itself flags asynchronous updates, Byzantine-tolerant formulations (robustness against adversarial agents), and nonconvex extensions as future directions.

What's available

The official JAX/Flax implementation is available on GitHub at github.com/SakanaAI/sheaf-admm. The repo includes training scripts for all three benchmarks (maze, MNIST, Sudoku), config files for reproducing paper results, and a visualization script for inspecting the x/z/u dynamics of trained agents. Installation uses uv:

uv sync
# For CUDA:
uv pip install -U "jax[cuda12]"
# Train Sheaf-ADMM on Sudoku:
uv run python scripts/train.py +experiment=sudoku_sheaf

The full paper is on arXiv. What makes this work worth watching is not just the benchmark results -- it's the bet that the right inductive bias for distributed AI isn't a bigger hidden state, but a principled theory of what agents should agree on and how they should get there.