Table of Contents
Fetching ...

L-SWAG: Layer-Sample Wise Activation with Gradients information for Zero-Shot NAS on Vision Transformers

Sofia Casarin, Sergio Escalera, Oswald Lanz

Abstract

Training-free Neural Architecture Search (NAS) efficiently identifies high-performing neural networks using zero-cost (ZC) proxies. Unlike multi-shot and one-shot NAS approaches, ZC-NAS is both (i) time-efficient, eliminating the need for model training, and (ii) interpretable, with proxy designs often theoretically grounded. Despite rapid developments in the field, current SOTA ZC proxies are typically constrained to well-established convolutional search spaces. With the rise of Large Language Models shaping the future of deep learning, this work extends ZC proxy applicability to Vision Transformers (ViTs). We present a new benchmark using the Autoformer search space evaluated on 6 distinct tasks and propose Layer-Sample Wise Activation with Gradients information (L-SWAG), a novel, generalizable metric that characterizes both convolutional and transformer architectures across 14 tasks. Additionally, previous works highlighted how different proxies contain complementary information, motivating the need for a ML model to identify useful combinations. To further enhance ZC-NAS, we therefore introduce LIBRA-NAS (Low Information gain and Bias Re-Alignment), a method that strategically combines proxies to best represent a specific benchmark. Integrated into the NAS search, LIBRA-NAS outperforms evolution and gradient-based NAS techniques by identifying an architecture with a 17.0% test error on ImageNet1k in just 0.1 GPU days.

L-SWAG: Layer-Sample Wise Activation with Gradients information for Zero-Shot NAS on Vision Transformers

Abstract

Training-free Neural Architecture Search (NAS) efficiently identifies high-performing neural networks using zero-cost (ZC) proxies. Unlike multi-shot and one-shot NAS approaches, ZC-NAS is both (i) time-efficient, eliminating the need for model training, and (ii) interpretable, with proxy designs often theoretically grounded. Despite rapid developments in the field, current SOTA ZC proxies are typically constrained to well-established convolutional search spaces. With the rise of Large Language Models shaping the future of deep learning, this work extends ZC proxy applicability to Vision Transformers (ViTs). We present a new benchmark using the Autoformer search space evaluated on 6 distinct tasks and propose Layer-Sample Wise Activation with Gradients information (L-SWAG), a novel, generalizable metric that characterizes both convolutional and transformer architectures across 14 tasks. Additionally, previous works highlighted how different proxies contain complementary information, motivating the need for a ML model to identify useful combinations. To further enhance ZC-NAS, we therefore introduce LIBRA-NAS (Low Information gain and Bias Re-Alignment), a method that strategically combines proxies to best represent a specific benchmark. Integrated into the NAS search, LIBRA-NAS outperforms evolution and gradient-based NAS techniques by identifying an architecture with a 17.0% test error on ImageNet1k in just 0.1 GPU days.
Paper Structure (26 sections, 2 theorems, 22 equations, 11 figures, 8 tables, 1 algorithm)

This paper contains 26 sections, 2 theorems, 22 equations, 11 figures, 8 tables, 1 algorithm.

Key Result

Theorem 1

Given the linear regressor $f({\bm{a}}, {\bm{x}})$ with trainable parameters ${\bm{a}} = (a_j)_{j = 1}^M$, let $g({\bm{x}}_i) = (g_j({\bm{x}}_i))_{j = 1}^d$ be the gradient of ${\bm{a}}$ w.r.t. to ${\bm{x}}_i$, and $\hat{{\bm{a}}} = {\bm{a}} -\eta\sum_i g_j({\bm{x}}_i)$ the updated parameters with l

Figures (11)

  • Figure 1: Our approach applies to different task types of architectures. L-SWAG takes as input a batch of images and a DNN, extracts the gradient statistics, and counts the # of linear regions in a layer-wise fashion. The relevant layers are identified una-tantum, before running the metric and are specific for each benchmark. L-SWAG outputs a rank of the architectures. LIBRA takes as input the pre-computed ZC-proxy metrics for a given benchmark. It has three steps: (i) selects the best performing one according to their correlation $\rho$. (ii) Computes the information we gain over the validation accuracy $y$ given $z_{best}$ and each other $z_i$, and selects the $z$ leading to the lowest validation accuracy. (iii) Select $z_3$ with the closest bias to $y$. LIBRA outputs the 3 identified metrics.
  • Figure 2: Empirical motivation for our layer selection strategy.
  • Figure 3: Spearman rank correlation coefficient between ZC proxy values and validation accuracies. Results were obtained from 5 multiple runs. Rows and columns are ordered based on the mean scores.
  • Figure 4: Average Spearman $\rho$ coefficient of ZC proxies across different search spaces.
  • Figure 5: Toy example for the positive correlation of $\mu$ and the loss $\mathcal{L}$ for 1000 linear networks trained for one epoch on M = 1000 samples with different $\eta$.
  • ...and 6 more figures

Theorems & Definitions (6)

  • Theorem 1
  • proof
  • Definition 1
  • Definition 2
  • Theorem 1
  • proof