ICML2026

Axiomatic Atlas: A Prescriptive Framework for Neural Architecture Design

Minghao Guo, Wojciech Matusik

Abstract

Neural architecture design lacks first principles: innovations are discovered empirically and justified post-hoc, with no systematic way to diagnose why an architecture fails or derive what repair will succeed. We introduce the Axiomatic Atlas, encoding requirements as composable axioms over graph connectivity, operator contracts, numerical stability, and information preservation. Given an operator library and wiring conventions, the Atlas constructs certificates lower-bounding output variation via min-cut analysis and diagnoses failures by locating axiom violations. Crucially, the framework is prescriptive: each violation implies a targeted repair, reducing architecture design to constraint satisfaction. We prove variation bounds under exact and finite-precision arithmetic, enabling modular verification across transformers, MoEs, SSMs, and GNNs. Four Atlas-derived interventions validate the approach: +46 percentage points on GNN bottlenecks, 3×3\times robustness to MoE quantization, 83% gap closure with adaptive expert budgets, and 0%\to100% retrieval via orthogonal keys---each against matched negative controls.