PaperlayerView exact revision on arXiv

Source-linked AI summary

Designing Compact Neural Architectures via Neuron Gating and Mixed Activation

Abhishek Shukla, Ankur Sinha, Faiz Hamid

arXiv:2608.14443v1cs.LGcs.AI

TL;DR

Finding effective neural architectures remains challenging because expert-designed networks are not necessarily optimal. This paper applies relaxed bilevel optimization to neuron- and activation-level choices, and reports that its methods match or exceed literature baselines across MNIST and CIFAR-10 while using fewer parameters.

  • Problem

    Determining effective architectures remains a fundamental deep-learning challenge, since expert-designed networks are not necessarily optimal as model complexity grows.

    Source: §1 · p. 2

  • Method

    The paper develops three relaxed bilevel formulations that replace discrete neuron- and activation-level choices with continuous variables for differentiable architecture search.

    Source: §1 · p. 2

  • Results

    Across MNIST and CIFAR-10, the methods matched or exceeded literature baselines with fewer parameters; on MNIST, CNN parameter reductions were 50.5%–75.4%.

    Source: §5.9 · p. 23, §6 · p. 24

  • Takeaways & Limitations

    The authors recommend NAS-NG as a practical starting point, reporting up to 86% compression in under one hour with minimal hyperparameter tuning.

    Source: §5.8 · p. 22, §6 · p. 24

    MNIST’s near-ceiling accuracy may obscure differences on harder datasets, and evaluation on larger benchmarks such as ImageNet remains outstanding.

    Source: §6 · p. 24

Abstract

from arXiv · show

Neural Architecture Search (NAS) is naturally formulated as a bilevel optimization problem, where the upper-level optimizes the architecture using validation performance and the lower-level trains network parameters using training loss. However, NAS is computationally expensive due to discrete architectural decisions, exponentially growing search spaces, and the high cost of training candidate architectures. This work develops a general bilevel optimization framework for NAS across diverse architectures, including MLPs, CNNs, RNNs, and Transformers, to identify compact architectures with strong predictive performance. We propose three scalable formulations that replace discrete neuron- and activation-level decisions with continuous relaxations, enabling differentiable optimization over otherwise combinatorial architecture spaces. These formulations give rise to three NAS methods: NAS based on Neuron Gating (NAS-NG), NAS based on Mixed Activation (NAS-MA), and NAS based on Neuron Gating and Mixed Activation (NAS-NGMA). Experiments on MLPs and CNNs using MNIST and CIFAR-10 show that the proposed methods consistently identify compact architectures with competitive or improved predictive performance. On MNIST, NAS-NGMA achieves 98.68% test accuracy with 7.69M MLP parameters, while NAS-NG achieves 99.63% accuracy with only 0.26M CNN parameters. On CIFAR-10, the proposed methods consistently outperform vanilla DARTS. Further experiments demonstrate that NAS-NG can optimize substantially over-parameterized and literature-optimal architectures, improving accuracy while reducing parameters. These results establish relaxed bilevel optimization as a scalable alternative to discrete NAS and provide a general framework for efficient neuron- and activation-level architecture optimization.

1 Introduction

The paper proposes differentiable bilevel NAS methods that relax neuron- and activation-level choices, aiming to find compact, accurate architectures. It reports competitive MNIST results and outperformance of vanilla DARTS on CIFAR-10.

1 Introduction

  • NAS is framed as bilevel optimization because discrete, high-dimensional architecture choices and costly candidate training make search computationally challenging.The upper level selects architectural parameters, while the lower level trains network weights.Source: §1 · p. 2
  • The evaluation covers MLP and CNN architectures on MNIST and CNN architectures on CIFAR-10 to assess the proposed formulations.Source: §1 · p. 3
  • Three relaxed bilevel formulations replace discrete neuron- and activation-level choices with continuous variables, yielding differentiable optimization through NAS-NG, NAS-MA, and NAS-NGMA.The framework is developed using MLPs and extends to CNNs and RNNs through appropriate tensor generalizations.Source: §1 · pp. 2–3
  • 98.68% test accuracy with 7.69M parameters for MLPs (NAS-NGMA), and 99.63% test accuracy with 0.26M parameters for CNNs (NAS-NG); the approaches outperform vanilla DARTS on CIFAR-10.The proposed methods also achieve approximately 40–75% parameter reductions on MNIST, with similar trends on CIFAR-10.Source: §1 · p. 3
  • NAS-NG can optimize substantially over-parameterized architectures and improve literature-optimal architectures, showing that neuron-level optimization complements existing topology search methods.Source: §1 · p. 3

2 Related Works

The related-work section surveys handwritten-digit recognition and differentiable NAS, covering classical and deep learning, ensembles, and extensions of DARTS.

2.1 Neural Architectures and Learning Methods

  • Related work spans gradient-based learning and CNNs, classical ML methods [24–27], CNN architectures and training factors, and ensemble learning [34].The passage notes that CNNs have outperformed MLPs in image tasks through their ability to capture spatial hierarchies.Source: §2.1 · p. 4

2.2 Differentiable Architecture Search

  • DARTS [9] inspired extensions addressing NAS efficiency, stability, robustness, attention, collapse, and search effectiveness.P-DARTS [10] uses search-space approximation and regularization; DARTS+ [41] applies early stopping to mitigate collapse, while Fair DARTS [13] addresses unfair competition among candidate operations. Further approaches include PC-DARTS [11], Att-DARTS, DE-DARTS, EG-DARTS, and Relax-DARTS [15], STO-DARTS, OSTR-DARTS [16], semantic DARTS, LMD-DARTS, HN-DARTS, and DARTS-EAST.Source: §2.2 · p. 4

3 Preliminaries

NAS is framed as bilevel optimization: network weights minimize training loss for a fixed architecture, while architecture selection minimizes validation loss. The preliminaries describe approximate hypergradients for updating architecture parameters and how one-step unrolling determines their approximation.

3 Preliminaries

  • The lower-level problem trains weights W to minimize training loss for a fixed architecture A, while the upper level selects architectures by validation loss.Lt and Lv denote training and validation losses, and A and W denote the architecture and network weight spaces.Source: §3 · p. 5
  • Approximate hypergradients update architecture parameters to improve validation performance while accounting for weights after one training-loss gradient step [21].The hypergradient is the gradient of the upper-level validation-loss objective with respect to architecture A.Source: §3 · p. 5
  • Setting the unrolling learning rate ξ to 0 gives a first-order approximation, whereas ξ ≠ 0 adds a second-order correction.ξ is the small learning rate used for one-step unrolled optimization of the training problem.Source: §3 · p. 5
  • The second-order term captures architecture–weight interactions through the training-loss mixed Hessian and can be approximated efficiently by finite differences; further DARTS details are cited [9].The one-step updated weights are W′ = W − ξ∇W Lt(A, W).Source: §3 · p. 5

4 Proposed Method

The method formulates MLP architecture search as bilevel optimization over continuous architecture variables. It covers neuron selection, activation choice, their joint optimization, and a procedure for deriving a discrete architecture.

4 Proposed Method

  • An MLP architecture is represented by the neuron counts in its input, hidden, and output layers, allowing hidden layers to have varying widths.Source: §4 · p. 6

4.1 Bilevel Formulations for MLP Architecture Search

  • The framework defines three MLP search settings: neuron gating, mixed activation, and their joint optimization within a bilevel formulation.Source: §4.1 · p. 7

4.1.1 Neuron Gating (NG)-based MLP Architecture Search

  • Neuron gating assigns each hidden neuron a real-valued architecture parameter, transformed by a sigmoid into an effective gate for bilevel architecture search.Source: §4.1.1 · pp. 7–8

4.1.2 Mixed Activation (MA)-based MLP Architecture Search

  • Mixed activation assigns each hidden neuron parameters over candidate functions, uses softmax to obtain mixing weights, and defines its activation as their convex combination.Source: §4.1.2 · pp. 9–10

4.1.3 Neuron Gating and Mixed Activation (NGMA)-based MLP Architecture Search

  • NGMA jointly optimizes neuron selection and activation function, using sigmoid-mapped gates and softmax-mapped mixing weights for each hidden neuron.Source: §4.1.3 · p. 10

4.2 Algorithm

  • The NGMA algorithm initializes a sufficiently large network, continuously optimizes weights and architecture variables via bilevel optimization, then prunes neurons by gating and selects each retained neuron's highest-weight activation.The framework extends to CNNs by associating architecture variables with convolutional channels and final linear layers; its relaxation and optimization structure remain unchanged.Source: §4.2 · pp. 11–12

5 Experimental Results and Discussion

The experiments evaluate NAS-NG, NAS-MA, and NAS-NGMA on MLPs and CNNs for MNIST and CNNs for CIFAR-10, comparing searched models with dense baselines and vanilla DARTS. They examine model complexity and accuracy, including performance on literature-optimal DARTS architectures.

5.3.1 Architecture Search

  • NAS-NG completes MLP search in an average of 0.45 CPU hours, about 40 times faster than NAS-MA and NAS-NGMA, whose mean search time is approximately 18.8 hours.The paper attributes the speed difference to gating-based importance metrics versus the extensive forward passes needed to analyze activation patterns; the methods use neuron importance, activation magnitudes, or both to guide pruning.Source: §5.3.1 · p. 14

5.4.2 CIFAR-10

  • On CIFAR-10, the proposed methods outperform vanilla DARTS while discovering CNN architectures with substantially smaller network sizes.The comparison evaluates a three-convolutional-layer configuration with one fully connected layer.Source: §5.4.2 · p. 17, Table 7 · p. 17

5.5 Overall Performance Comparison

  • For MLPs, all three methods’ mean test accuracies are within 0.06 percentage points of the baseline; NAS-NG reduces average size from 8.97M to 1.90M while achieving 98.53% accuracy.NAS-NGMA attains the highest maximum accuracy at 98.68% with 60.5% fewer parameters, while NAS-MA reduces parameters by 48.3%.Source: §5.5 · p. 18

5.5.2 Additional Analysis

  • The paper reports competitive accuracy with reduced complexity across network depths and hidden dimensions, and compares searched architectures with literature models on MNIST.The depth-and-width analysis is reported in Appendix E; the supplied passage describes the literature comparison as competitive while noting simpler model structures.Source: §5.5.2 · p. 19, §5.6.3 · p. 21, Table 12 · p. 21

5.6.1 Best Performing Architectures

  • 98.68% accuracy with 7.69M parameters is the best MLP searched result, just 0.01 percentage points below the 98.69% best baseline while using 75% fewer parameters.The NAS-NGMA result comes from a 3-layer network; NAS-MA reaches 98.65% with 1.75M parameters, while NAS-NG reaches 98.60% with 4.29M parameters.Source: §5.6.1 · p. 19, Table 8 · p. 19
  • NAS-NG’s best three-convolutional-layer CNN on MNIST exceeds the best CNN baseline by 0.01% while using approximately 47% fewer parameters.The paper presents the best-performing MNIST CNN models in Table 9.Source: §5.6.1 · p. 19, Table 9 · p. 20

5.6.2 Efficiency Metrics

  • NAS-NG achieves the highest mean MNIST MLP efficiency at 86.05 accuracy points per million parameters, 2.7 times the baseline’s 31.94 and 1.8 times NAS-MA’s 48.75.NAS-NGMA’s efficiency is 60.66; the paper defines efficiency as test accuracy percentage divided by parameter count in millions.Source: §5.6.2 · p. 20, Figure 12 · p. 20

5.7 Additional Experiments for NAS-NG

  • NAS-NG is tested on substantially larger MLP and CNN baselines, where it reduces parameters while maintaining or sometimes improving test accuracy; a further experiment applies neuron-level optimization to literature-optimal DARTS architectures.For the DARTS experiment, optimal cell structures are held fixed and only the number of stacked cells is varied during initialization; the supplied passage does not report its outcome.Source: §5.7.1 · p. 21, §5.7.2 · p. 21, §5.7 · p. 21

5.7.2 Better Initialization

  • NAS-NG applied to complete 10-cell and 14-cell DARTS macro-architectures simultaneously improved classification accuracy and reduced trainable parameters.This supports neuron-level optimization as complementary to topology search.Source: §5.7.2 · p. 22, Figure 13 · p. 22

5.8 Discussion, Implications, and Recommendations

  • For MLPs, NAS-NG offered up to 86% compression in under an hour, while NAS-MA retained about 60% of parameters and took 40× longer to search.NAS-NGMA matched baseline accuracy with 75% fewer parameters, but inconsistent average gains and longer search times limited its broad applicability; network size from 0.45M–30.76M parameters showed no strong correlation with accuracy.Source: §5.8 · p. 22

5.8.1 Implications for Neural Network Design

  • The 98.40%–98.69% MLP and 99.49%–99.63% CNN accuracy ranges suggest efficiency and ease of use can matter more than small accuracy differences; simpler methods can match complex search.The reported parameter reductions without accuracy loss align with network pruning and the lottery ticket hypothesis [48].Source: §5.8.1 · p. 23

5.8.2 Practical Recommendations

  • For limited compute or tight deadlines, the recommendations favor starting with NAS-NG and considering more cautious pruning for tasks more complex than MNIST.The authors also recommend trying multiple methods when resources permit and starting from medium-sized networks, such as 1024–2048 neurons per layer and 2–3 layers.Source: §5.8.2 · p. 23
  • Validate final models on distinct representative datasets because search-phase accuracy may not predict test generalization, and compare accuracy-to-parameter efficiency rather than accuracy alone.The paper’s example favors a 98.5%-accurate model with 0.5M parameters over a 98.6%-accurate model with 5M parameters for resource-limited devices.Source: §5.8.2 · p. 23

5.9 Summary

  • All three methods maintained competitive accuracy while reducing MLP parameters by 40%–69% and MNIST CNN parameters by 50.5%–75.4% versus their respective baselines.For MLPs, test accuracy ranged from 98.40% to 98.68%; for CNNs, it ranged from 99.49% to 99.63%.Source: §5.8 · p. 22, §5.9 · p. 23
  • NAS-NGMA reached 98.68% MLP accuracy with 75% fewer parameters, while NAS-NG achieved 99.63% CNN accuracy; the most efficient MLP used 0.45M parameters.That MLP achieved 98.49% accuracy and 218.88 accuracy points per million parameters, approximately 68 times the best baseline’s efficiency; CNN NAS-NGMA reached approximately 4.06 times baseline efficiency.Source: §5.9 · pp. 23–24
  • All proposed CIFAR-10 architectures outperformed vanilla DARTS with smaller model complexity, and NAS-NG scaled to larger models and refined literature-optimal DARTS architectures.The latter neuron-level optimization improved accuracy while reducing parameters, complementing topology search.Source: §5.9 · p. 24

6 Conclusions · Appendix

The paper presents three relaxed bilevel NAS methods that find compact architectures with competitive performance, and recommends evaluating them by accuracy, model size, and deployment constraints. It also identifies limitations and directions for extending and improving the methods.

6 Conclusions

  • Across MNIST and CIFAR-10, NAS-NG, NAS-MA, and NAS-NGMA match or exceed literature baselines with substantially fewer parameters.The methods relax neuron- and activation-level choices for differentiable optimization, offering an alternative to discrete combinatorial search.Source: §6 · p. 24
  • NAS-NG is the recommended practical starting point for resource-constrained deployment, with sub-hour searches and aggressive compression, but discovered architectures require held-out validation.Evaluation should prioritize accuracy-to-parameter trade-offs over marginal accuracy gains.Source: §6 · p. 24
  • NAS-NG remains effective from substantially over-parameterized initial networks and can refine DARTS-optimized architectures, complementing rather than replacing topology search..Source: §6 · p. 24
  • Limitations include near-ceiling MNIST MLP accuracy (98.40–98.69%), limited architectural components, fixed hyperparameters, scratch retraining, and no ImageNet-scale evaluation.CIFAR-10 partly addresses dataset simplicity; interactions with dropout, batch normalization, and skip connections, hyperparameter tuning, and weight inheritance or distillation remain unexplored.Source: §6 · p. 24
  • Future work includes harder vision tasks, broader architectures such as RNNs and Transformers, adaptive hyperparameter tuning, and theory explaining when simple gating metrics work.Theoretical understanding could clarify which problem classes best suit these relaxations.Source: §6 · p. 24
  • Data augmentation, stronger regularization, and improved optimization may help close the remaining state-of-the-art accuracy gap without sacrificing efficiency.AutoAugment is given as an example of a data augmentation strategy.Source: §6 · p. 24
  • The study challenges the need for heavy overparameterization, suggesting targeted task-specific adjustments can rival elaborate NAS pipelines.It advocates weighing accuracy jointly with model size and deployment constraints rather than optimizing marginal accuracy gains alone.Source: §6 · p. 25

A Loss Function

This section formulates the network’s forward computation from hidden-layer pre-activations and activations through output prediction, then defines a binary cross-entropy objective for labeled samples.

A Loss Function

  • The forward pass specifies hidden-layer pre-activations and activations, then derives the output pre-activation and predicted output in a single expression.Source: A Loss Function · p. 28
  • Binary cross-entropy defines the per-sample loss from the true binary label and predicted probability, with a further expression substituting the explicit prediction.Source: A Loss Function · p. 28
  • The loss is rewritten using the collection of network weights and biases, W = {W(1), b(1), W(2), b(2)} ∈ R13.Source: A Loss Function · p. 28

B Modified Loss Function for Architecture Search · C Activation Functions

Section B describes sigmoid gating after ReLU, then specifies the output pre-activation and a modified per-sample binary cross-entropy loss; Section C points to Table 14 for the activation functions used.

B Modified Loss Function for Architecture Search

C Activation Functions

D Configuration-Specific Performance Patterns

Across configurations, the methods show similar accuracy patterns, while searched architectures achieve this performance with fewer parameters. NGMA reaches its highest peak in a favorable 3-layer, high-dimensional region.

D Configuration-Specific Performance Patterns

  • Searched architectures achieve similar accuracy patterns to the other methods with fewer parameters.Figure 14 heatmaps report mean test accuracy for each layer-count and total-hidden-dimension combination.Source: D Configuration-Specific Performance… · p. 29
  • All methods have similar accuracy distributions, with the highest accuracies generally in deeper networks with larger hidden dimensions.This pattern suggests the search methods identify efficient points within a broadly similar architecture-performance landscape.Source: D Configuration-Specific Performance… · p. 29
  • Search-method heatmaps show more uniform performance across configurations, suggesting neuron selection can partly compensate for suboptimal initial choices.The baseline shows clearer accuracy drops in shallower or narrower networks.Source: D Configuration-Specific Performance… · pp. 29–30
  • 98.68% is the highest peak accuracy shown by NGMA in the 3-layer, high-dimensional region.This is consistent with NGMA achieving its best individual results when the initial configuration is favorable.Source: D Configuration-Specific Performance… · p. 30

E Additional Analysis

The section analyzes how network depth and hidden dimensions affect classification performance, and how parameter reductions compare with baseline dense architectures.

E Additional Analysis

  • The analysis examines depth and hidden-dimension effects on classification performance alongside parameter reductions relative to dense baselines.These analyses provide insight into the trade-off between model complexity and predictive performance.Source: E Additional Analysis · p. 30

A Impact of Network Depth

Deeper networks generally perform better, though gains diminish beyond 3 layers; search methods remain competitive across depths and vary less than the baseline.

A Impact of Network Depth

  • Deeper networks generally perform better, with 3- and 4-layer networks performing best, but gains diminish beyond 3 layers and search methods remain competitive across depths.Error bars represent standard deviation across hidden dimensions of 1024, 2048, and 3072.Source: A Impact of Network Depth · p. 30
  • Baseline MNIST accuracy rose from 98.51% (1 layer) to 98.63% (4 layers), a 0.12 percentage-point improvement.The passage characterizes the gain as modest given MNIST’s relative simplicity.Source: A Impact of Network Depth · p. 31
  • NAS-NG accuracy stayed within 98.50% to 98.57% across depths, while NAS-NGMA reached its best accuracy, 98.65%, at 3 layers.The passage suggests gating and activation criteria may favor intermediate depths where both metrics provide clear signals.Source: A Impact of Network Depth · p. 31

B Influence of Hidden Dimensions

Accuracy generally improves as total hidden dimensions increase, but gains diminish at larger capacities. The search methods use fewer parameters at each dimension level, and NAS-NG remains competitive at lower dimensions.

B Influence of Hidden Dimensions

  • At each dimension level, search methods use fewer parameters, while NAS-NG achieves competitive accuracy at lower dimensions by identifying informative neurons.The passage contrasts this with retaining neurons in proportion to the initial network size.Source: B Influence of Hidden Dimensions · p. 31
  • Accuracy improves from 1,024 to 4,096 total dimensions, but gains diminish beyond 6,000, and 12,288-dimension configurations do not consistently outperform smaller ones.Error bars show standard deviation across layer configurations with identical total dimensions; some large configurations show slight performance degradation.Source: B Influence of Hidden Dimensions · p. 31, Figure 16 · p. 32

C Parameter Reduction Analysis

The analysis compares parameter reductions from NAS-NG, NAS-MA, and NAS-NGMA against baseline configurations. NAS-NG has the highest and most consistent average reduction; NAS-NGMA is more conservative, while NAS-MA varies more across configurations.

C Parameter Reduction Analysis

  • 68.67% was NAS-NG’s mean parameter reduction, the highest of the methods, with a 43.88%–86.06% range and 15.64% standard deviation.Reductions are measured against baseline dense networks of equivalent initial size.Source: Figure 17 · p. 33, C Parameter Reduction Analysis · p. 33
  • 54.49% was NAS-NGMA’s mean parameter reduction, with a 9.77% standard deviation; agreement between gating and activation metrics makes pruning more conservative.The combined approach is less likely to make extreme pruning decisions that benefit one criterion but perform poorly on the other.Source: C Parameter Reduction Analysis · p. 33
  • 40.98% was NAS-MA’s mean parameter reduction, with a 15.82%–56.76% range and 14.61% standard deviation, indicating configuration-dependent performance.A similar reduction in model complexity was observed for CNN architectures, as reported in Tables 5 and 7.Source: C Parameter Reduction Analysis · p. 33

Related papers

Loading 2608.14443v1…