Proxy-Mamba: Training-Free Architecture Search for Mamba via Gradient-Weight Correlation
摘要
State Space Models (SSMs) such as Mamba have recently emerged as efficient alternatives to Transformers for sequence modeling, especially in long-context and low-resource scenarios. However, the design of Mamba-based architectures still relies on manual heuristics and exhaustive tuning, limiting their scalability and performance. In this work, we propose Proxy-Mamba, a training-free Neural Architecture Search (NAS) framework tailored for Mamba models. At the core of Proxy-Mamba is a novel proxy metric, SigScore, designed to estimate the performance of Mamba-based architectures without any training. By computing a Sigmoid-normalized aggregation of the magnitudes of parameter weights and their corresponding gradients within the SSM module, our proxy exhibits strong correlation with actual model performance, enabling fast and cost-effective architecture search. Experiments on CIFAR10, CIFAR100, and ImageNet16-120 demonstrate that Proxy-Mamba outperforms existing training-free NAS baselines in terms of correlation with test accuracy and the final model’s classification performance. Our findings open a promising path toward scalable and efficient Mamba architecture design.