Preventing Topological Collapse in One-Shot Architecture Search via Forman-Ricci Curvature

One-Shot Neural Architecture Search relies heavily on channel reordering to ensure smaller sub-architectures inherit critical parameters from weight-sharing Supernets. However, traditional magnitude and activation-aware metrics evaluate channels in isolation, ignoring global topology and risking "topological collapse" where critical information bridges are inadvertently pruned. In this work, we propose a novel Activation-Aware Geometric Reordering strategy that bridges the gap between signal strength and topological robustness. We model the neural network as a Traffic Graph and apply Forman-Ricci Curvature (FRC) to distinguish between redundant paths and critical structural bottlenecks geometrically. Furthermore, we integrate these geometric flow features into a comprehensive search space pruning pipeline. By training a machine learning performance predictor, we stratify and filter candidate architectures prior to the final one-shot training phase. Extensive experiments on Vision Transformers (ViT, DeiT) across CIFAR-10, CIFAR-100, and Tiny-ImageNet demonstrate that our generated subnets establish a superior Pareto frontier. Our approach strictly dominates standard baselines, achieving up to 7.51% improvement in accuracy on CIFAR-100 and up to 43.3% parameter reduction on CIFAR-10. Finally, we provide a highly efficient GPU-accelerated FRC implementation that eliminates computational bottlenecks, enabling geometric NAS for large-scale search spaces. 

Committee: Ali Jannesari (major professor), Chris Quinn and Hongyang Gao