# Architecture-Aware Minimization (A2M): How to Find Flat Minima in Neural Architecture Search

arXiv: https://arxiv.org/abs/2503.10404 · HTML: https://mlsanity.com/paper/2503.10404/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 45 references that are in the map.
- First posted on arXiv 13 Mar 2025.

## Builds on (most cited first, 30 of 38)

- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Once-for-All: Train One Network and Specialize it for Efficient Deployment](https://mlsanity.com/paper/1908.09791/) (2019)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [NAS-Bench-201: Extending the Scope of Reproducible Neural Architecture Search](https://mlsanity.com/paper/2001.00326/) (2020)
- [A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets](https://mlsanity.com/paper/1707.08819/) (2017)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks](https://mlsanity.com/paper/2102.11600/) (2021)
- [PC-DARTS: Partial Channel Connections for Memory-Efficient Architecture Search](https://mlsanity.com/paper/1907.05737/) (2019)
- [Understanding and Robustifying Differentiable Architecture Search](https://mlsanity.com/paper/1909.09656/) (2019)
- [Towards Understanding Sharpness-Aware Minimization](https://mlsanity.com/paper/2206.06232/) (2022)
- [Fair DARTS: Eliminating Unfair Advantages in Differentiable Architecture Search](https://mlsanity.com/paper/1911.12126/) (2019)
- [DARTS+: Improved Differentiable Architecture Search with Early Stopping](https://mlsanity.com/paper/1909.06035/) (2019)
- [Stabilizing Differentiable Architecture Search via Perturbation-based Regularization](https://mlsanity.com/paper/2002.05283/) (2020)
- [Rethinking Architecture Selection in Differentiable NAS](https://mlsanity.com/paper/2108.04392/) (2021)
- [DARTS-: Robustly Stepping out of Performance Collapse Without Indicators](https://mlsanity.com/paper/2009.01027/) (2020)
- [β-DARTS: Beta-Decay Regularization for Differentiable Architecture Search](https://mlsanity.com/paper/2203.01665/) (2022)
- [NAS-Bench-NLP: Neural Architecture Search Benchmark for Natural Language Processing](https://mlsanity.com/paper/2006.07116/) (2020)
- [Understanding Architectures Learnt by Cell-based Neural Architecture Search](https://mlsanity.com/paper/1909.09569/) (2019)
- [Shaping the learning landscape in neural networks around wide flat minima](https://mlsanity.com/paper/1905.07833/) (2019)
- [An Adaptive Policy to Employ Sharpness-Aware Minimization](https://mlsanity.com/paper/2304.14647/) (2023)
- [Deep Networks on Toroids: Removing Symmetries Reveals the Structure of Flat Regions in the Landscape Geometry](https://mlsanity.com/paper/2202.03038/) (2022)
- [Entropic gradient descent algorithms and wide flat minima](https://mlsanity.com/paper/2006.07897/) (2020)
- [AdaSAM: Boosting Sharpness-Aware Minimization with Adaptive Learning Rate and Momentum for Training Deep Neural Networks](https://mlsanity.com/paper/2303.00565/) (2023)
- [Exploring the Loss Landscape in Neural Architecture Search](https://mlsanity.com/paper/2005.02960/) (2020)
- [Lambda-DARTS: Mitigating Performance Collapse by Harmonizing Operation Selection among Cells](https://mlsanity.com/paper/2210.07998/) (2022)
- [GA-SAM: Gradient-Strength based Adaptive Sharpness-Aware Minimization for Improved Generalization](https://mlsanity.com/paper/2210.06895/) (2022)
- [Deep learning via message passing algorithms based on belief propagation](https://mlsanity.com/paper/2110.14583/) (2021)
- [CR-SAM: Curvature Regularized Sharpness-Aware Minimization](https://mlsanity.com/paper/2312.13555/) (2023)

## Built on by (most cited first, 2 of 2)

- [SEAL: Searching Expandable Architectures for Incremental Learning](https://mlsanity.com/paper/2505.10457/) (2025)
- [HERCULES: Hardware-Efficient, Robust, Continual Learning Neural Architecture Search](https://mlsanity.com/paper/2605.04103/) (2026)

## Abstract (arXiv)

Neural Architecture Search (NAS) has become an essential tool for designing effective and efficient neural networks. In this paper, we investigate the geometric properties of neural architecture spaces commonly used in differentiable NAS methods, specifically NAS-Bench-201 and DARTS. By defining flatness metrics such as neighborhoods and loss barriers along paths in architecture space, we reveal locality and flatness characteristics analogous to the well-known properties of neural network loss landscapes in weight space. In particular, we find that highly accurate architectures cluster together in flat regions, while suboptimal architectures remain isolated, unveiling the detailed geometrical structure of the architecture search landscape. Building on these insights, we propose Architecture-Aware Minimization (A2M), a novel analytically derived algorithmic framework that explicitly biases, for the first time, the gradient of differentiable NAS methods towards flat minima in architecture space. A2M consistently improves generalization over state-of-the-art DARTS-based algorithms on benchmark datasets including CIFAR-10, CIFAR-100, and ImageNet16-120, across both NAS-Bench-201 and DARTS search spaces. Notably, A2M is able to increase the test accuracy, on average across different differentiable NAS methods, by +3.60% on CIFAR-10, +4.60% on CIFAR-100, and +3.64% on ImageNet16-120, demonstrating its superior effectiveness in practice. A2M can be easily integrated into existing differentiable NAS frameworks, offering a versatile tool for future research and applications in automated machine learning. We open-source our code at https://github.com/AI-Tech-Research-Lab/AsquaredM.