[Paper Review] KNAS: Green Neural Architecture Search

Jingjing Xu, Liang Zhao+4 more|arXiv (Cornell University)|Nov 26, 2021

Machine Learning in Materials Science4 citations

TL;DR

KNAS proposes a green neural architecture search method that evaluates neural network architectures without training by leveraging the gradient kernel hypothesis—specifically, using the mean of the Gram matrix of gradients (MGM) as a proxy for downstream performance. This approach achieves competitive accuracy with orders of magnitude faster search, outperforming RoBERTA-large on text classification tasks while drastically reducing computational and environmental costs.

ABSTRACT

Many existing neural architecture search (NAS) solutions rely on downstream training for architecture evaluation, which takes enormous computations. Considering that these computations bring a large carbon footprint, this paper aims to explore a green (namely environmental-friendly) NAS solution that evaluates architectures without training. Intuitively, gradients, induced by the architecture itself, directly decide the convergence and generalization results. It motivates us to propose the gradient kernel hypothesis: Gradients can be used as a coarse-grained proxy of downstream training to evaluate random-initialized networks. To support the hypothesis, we conduct a theoretical analysis and find a practical gradient kernel that has good correlations with training loss and validation performance. According to this hypothesis, we propose a new kernel based architecture search approach KNAS. Experiments show that KNAS achieves competitive results with orders of magnitude faster than "train-then-test" paradigms on image classification tasks. Furthermore, the extremely low search cost enables its wide applications. The searched network also outperforms strong baseline RoBERTA-large on two text classification tasks. Codes are available at \url{https://github.com/Jingjing-NLP/KNAS} .

Motivation & Objective

• To address the high computational and environmental cost of existing neural architecture search (NAS) methods that rely on full training for architecture evaluation.
• To investigate whether gradients from randomly initialized networks can serve as a reliable proxy for downstream training performance.
• To develop a green NAS solution that eliminates training during search while maintaining competitive model accuracy.
• To validate the generalization capability of the proposed method across diverse tasks, including image and text classification.

Proposed method

• Propose the gradient kernel hypothesis: gradients from randomly initialized networks can act as a coarse-grained proxy for downstream training performance.
• Identify the Gram matrix of gradients (GM) as a key feature, with its Frobenius norm bounding the convergence rate of optimization.
• Introduce the mean of the Gram matrix (MGM) as a practical, scalable metric for architecture evaluation.
• Use MGM to rank architectures without training, selecting only the top-k candidates for full training.
• Apply the method to NAS-Bench-201 for image classification and to text classification tasks with RoBERTA-based search spaces.
• Conduct theoretical analysis showing that higher MGM correlates with faster convergence and better generalization.

Experimental results

Research questions

• RQ1 Can gradients from randomly initialized networks reliably predict downstream training performance?
• RQ2 Does the Gram matrix of gradients (GM) exhibit strong correlation with training loss and validation accuracy?
• RQ3 Can a training-free architecture evaluation method based on MGM achieve competitive performance compared to full-training NAS?
• RQ4 Does the proposed KNAS method generalize across different tasks, such as image and text classification?
• RQ5 What is the impact of the proposed method on search efficiency and carbon footprint reduction?

Key findings

• MGM shows a Spearman correlation of 0.53 (p ≪ 0.01) with negative training loss and 0.56 (p ≪ 0.01) with validation accuracy on CIFAR-100.
• KNAS achieves competitive accuracy on NAS-Bench-201 with search costs reduced by orders of magnitude compared to full-training NAS.
• The searched architecture outperforms RoBERTA-large by 1.24% on MRPC and 0.24% on RTE in text classification tasks.
• KNAS achieves a search time of only 0.4K seconds on MRPC and 2K seconds on RTE, demonstrating extreme efficiency.
• The method generalizes well across different initialization schemes, normalization types, and network architectures (e.g., CNNs vs. self-attention networks).
• Theoretical analysis confirms that higher F-norm of the GM matrix corresponds to faster convergence, validating the gradient kernel hypothesis.