Add a training-free, closed-form MNIST robustness baseline

[kayuksel]

Hi — I really admired the Sheaf-ADMM paper; using a cellular sheaf to make heterogeneous agreement precise is elegant. Reading the solver, I noticed the per-agent x-update (Q + rho*I) x = rho*(z - y) - q is a closed-form ridge, and got curious how much of the reported MNIST robustness a single such ridge already delivers — with no consensus iteration and no training. This PR adds that as a baseline. Happy to adjust the framing or scope.

TL;DR

An evolutionary search — no gradient descent, no backprop — discovered a weightless collective of
12 random-feature "genes" that, fused by a single closed-form ridge solve, matches a strong CNN on
clean MNIST and dramatically exceeds trained models under distribution shift. It fits in seconds on a
CPU. We offer it as a baseline that maps the boundary between what needs learned coordination and what
emerges from structure — which we think sharpens exactly where Sheaf-ADMM's iterative machinery earns
its keep.

Results on the paper's own splits (full 50k/10k, via ImageDataset)

split	evolved closed-form (ours)	Sheaf-ADMM (trained)	CNN
clean	0.990	0.985	0.993
pad 4 / 8 / 16	0.990 / 0.990 / 0.990	0.984 / 0.975 / 0.863	0.708 / 0.441 / 0.114
noise σ .1 / .2	0.989 / 0.986	0.740 / 0.319	0.540 / 0.099
noise σ .3 / .4 / .5	0.981 / 0.960 / 0.893	(not reported)	(not reported)

No training. The ridge sees only clean train; every test* split is a corruption of the
disjoint held-out test images — no image is shared between fit and eval, so there is no leakage.

Relationship to Sheaf-ADMM (why this is the natural baseline)

The two methods are the same family — decompose the input into many local views, solve each in
closed form, and fuse — and reading your solver made that concrete:

• Your per-agent x-solver is a closed-form ridge. DenseQuadraticXSolver solves
  (Q + rho*I) x = rho*(z - y) - q via jnp.linalg.solve — a Tikhonov / ridge step. The encoder emits
  the local quadratic (Q, q); the z-step is the sheaf-constrained consensus; the whole thing is
  iterated and trained end-to-end.
• EvoForest is the single-shot, weightless version of the same primitive. It's a pool of expert
  feature-views (the "genes"), each a different random-conv projection of the image, fused by exactly
  that closed-form ridge — computed once, with no consensus iteration and no learned encoder. Evolution
  plays the role of expert selection (competing feature families, survivors kept); ridge is the
  provably-optimal linear fusion.

So this baseline is a clean ablation of what the trained ADMM iteration buys: the same ridge primitive,
with vs. without the learned spatial-consensus loop.

More broadly it reflects a search-first stance — the learned object is a discovered computation (an
evolved feature program), not a trained weight matrix, and a parameter-light closed-form head is all that's
needed once the right computation exists. Where a mixture-of-experts routes to fixed experts, and your
agents share a fixed per-patch architecture refined by training, here the experts themselves are
searched and invented.

Two axes of decomposition — and we quietly use both

Robustness is redundancy, but there are two orthogonal kinds:

• Spatial (yours): one agent per patch — corrupt or drop part of the image and the clean patches keep
  voting → padding / occlusion robustness.
• Feature-space (ours): many feature-views of the same input — if one representation is misled, others
  compensate → a different robustness (noise, via the low-pass views).

The shipped model does both: the 12 genes are the feature-space axis, and each gene's adaptive-max
pool over a 2×2 grid is a (coarse) spatial axis — which is exactly why one closed-form solve inherits
your padding robustness (size-invariant spatial pooling) and strong noise robustness (blurred
feature-views). The natural extension — a full patch × feature-view expert grid — is strictly richer
than spatial-only agents and still closed-form; we'd be keen to explore whether it closes the gap on the
hard shifts (rotation / heavy occlusion) where your iterated agents still lead.

Honest by construction

On every reported condition the result is invariant to the random draw — re-seeding the kernels
(--seed-offset) leaves clean / pad / noise≤0.2 unchanged (verified in __main__); we embed the specific
evolved kernels so the extreme-σ tail is exact too. The robustness is a property of the architecture the
search discovered, not a lucky initialization.

What's in the PR

Purely additive, under scripts/ — no changes to existing code, configs, or training paths:

• robust_mnist_jax.py — the model (JAX, documented). The exact evolved kernels are embedded at the
  bottom of the file (float16, zlib+base64; ~72 KB of text, no binary, no RNG regeneration).
  RobustMNIST().fit(X, y).predict(X).
• eval_closed_form_baseline.py — evaluates on every robustness split through the repo's ImageDataset,
  printing accuracy next to the paper's Table-4 numbers.

Run

python -m sheaf_admm.data.build_mnist --output-dir datasets/mnist --with-robustness-splits
python scripts/eval_closed_form_baseline.py --dataset-dir datasets/mnist

Provenance

The 12-gene architecture was discovered by EvoForest [arXiv:2604.19761], an evolutionary feature-map
search whose fitness is held-out robustness — selecting from a vocabulary of conv scales, dilations,
pooling geometries, and Gaussian-blur bandwidths. The exact evolved kernels are embedded (float16,
zlib+base64) in champion_weights.py; re-deriving them from the genome seeds (--seed-offset, numpy)
reproduces every reported-condition number, confirming the robustness is a property of the discovered
architecture rather than the specific draw.