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, robustness to MoE quantization, 83% gap closure with adaptive expert budgets, and 0%100% retrieval via orthogonal keys---each against matched negative controls.