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

NAS remains difficult because architectural choices are discrete and costly to search. This paper relaxes neuron and activation decisions into differentiable bilevel formulations, finding compact MLP and CNN architectures that match or exceed baselines with fewer parameters.

  • Problem

    Architectural design remains challenging because expert-driven choices do not necessarily yield optimal neural network architectures.

  • Method

    The paper introduces three relaxed bilevel NAS formulations that replace discrete neuron- and activation-level choices with continuous variables for differentiable optimization.

  • Results

    Across MNIST and CIFAR-10, the methods identified compact architectures with competitive or improved accuracy, reducing parameters by 40%–69% for MLPs and 50.5%–75.4% for CNNs on MNIST.

  • Takeaways & Limitations

    Continuous relaxations provide an effective and scalable alternative to discrete combinatorial NAS for neuron- and activation-level architecture optimization.

  • Takeaways & Limitations

    MNIST’s simplicity may obscure differences on harder datasets, and the study does not evaluate larger-scale benchmarks such as ImageNet.

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 addresses the difficulty of designing effective neural architectures by replacing discrete neuron- and activation-level NAS decisions with continuous variables in relaxed bilevel formulations. It proposes NAS-NG, NAS-MA, and NAS-NGMA, which identify compact architectures with competitive or improved performance across MNIST and CIFAR-10.

  • Proposed framework: Relaxed bilevel optimization replaces discrete neuron- and activation-level design choices with continuous variables, enabling differentiable and scalable NAS.The formulation separates architecture evaluation from weight optimization while making the combinatorial search tractable.
  • Proposed methods: The proposed NAS strategies are NAS-NG, NAS-MA, and NAS-NGMA.These methods target neuron gating, mixed activation, or both.
  • Comparisons and extensions: The proposed approaches outperform vanilla DARTS on CIFAR-10.NAS-NG also remains effective from substantially over-parameterized initial architectures and can improve literature-optimal architectures through architecture fine-tuning.
  • Experimental results: 40–75% reduction in model parameters is achieved on MNIST while maintaining competitive or improved predictive performance, with similar trends on CIFAR-10.The study evaluates compactness and predictive performance across the proposed formulations.
  • Experimental results: 98.68% test accuracy with 7.69M parameters is achieved by NAS-NGMA for MLPs on MNIST, while NAS-NG achieves 99.63% with 0.26M CNN parameters.The experiments cover MLP and CNN architectures on MNIST and CNN architectures on CIFAR-10.

2 Related Works

Related work spans handwritten-digit recognition methods and differentiable neural architecture search, from classical and deep learning approaches to ensemble techniques and DARTS-based methods. DARTS introduced efficient continuous-space NAS, followed by extensions targeting computational cost, search stability, and robustness.

  • The related-work discussion covers handwritten-digit recognition and differentiable NAS, including classical machine learning, deep architectures, ensemble techniques, and DARTS-based methods.
  • Earlier research established gradient-based learning and CNNs for document recognition, while subsequent work explored enhanced logistic regression, privacy-preserving learning, and dimensionality reduction.
  • DARTS introduced an efficient, competitive NAS approach using a continuous search space, compared with reinforcement-learning and evolutionary-computation methods.
  • P-DARTS extended DARTS by addressing computational challenges and improving search stability through search-space approximation and regularization.

3 Preliminaries

Bilevel NAS trains network weights by minimizing training loss while optimizing architectures for validation performance. Approximate hypergradients enable architecture updates through one-step unrolled optimization, with first- and second-order variants.

  • The lower level minimizes training loss over weights W for fixed architecture A, while the upper level selects architectures using validation loss to promote generalization.
  • The hypergradient is the gradient of the validation objective with respect to architecture A and updates architecture parameters using weights obtained from training optimization.
  • Setting ξ = 0 produces a first-order hypergradient approximation, whereas ξ ≠ 0 introduces a second-order correction through one-step unrolled optimization.
  • The unrolled weights W′ = W −ξ∇W Lt(A, W) capture architecture–weight interaction through the training-loss mixed Hessian, which can be approximated by finite differences.

4 Proposed Method

Section 4 formulates MLP architecture search as bilevel optimization under neuron gating, mixed activation, and their joint use. Continuous sigmoid and softmax parameterizations enable differentiable neuron selection and activation-function optimization, followed by discrete configuration derivation.

  • MLP Architecture Representation: The MLP architecture is represented by input, heterogeneous hidden, and output-layer neuron counts, allowing hidden layers with varying widths.The architecture is specified as (Input = [n0], Hidden = [n1, . . . , nL], Output = [nL+1]).
  • Bilevel Search Formulations: The framework develops three bilevel MLP search settings: neuron gating, mixed activation, and joint neuron gating with mixed activation.Neuron gating controls individual hidden-neuron participation, while mixed activation learns activation functions from a predefined candidate set.
  • Neuron Gating: Neuron gating uses unconstrained architecture parameters mapped through sigmoids to control each hidden neuron’s effective contribution.The parameters are indexed by neuron and hidden layer, and the resulting gates determine neuron participation in the network.
  • Mixed Activation: Mixed activation assigns each hidden neuron a convex combination of candidate activation functions, with mixing weights obtained through a softmax mapping.For the illustrated example, M = 7 candidates are None, Identity, ReLU, Leaky ReLU, tanh, σ, and SiLU.
  • Joint Neuron Gating and Mixed Activation: The joint NGMA formulation associates each hidden neuron with neuron-gating and activation parameters, optimizing selection and activation jointly.The gating uses a sigmoid mapping, and activation mixing uses a softmax mapping.
  • NGMA Search Procedure: Algorithm 1 initializes architecture parameters and network weights, alternates hypergradient-based architecture and weight updates, and derives the final configuration by thresholding gates and selecting the largest activation component.The procedure uses the approximate hypergradient from equation (2) before updating architecture parameters and network weights.

5 Experimental Results and Discussion

Experiments on MNIST and CIFAR-10 show that NAS-NG, NAS-MA, and NAS-NGMA produce compact MLP and CNN architectures with accuracy matching or approaching baselines. NAS-NG is substantially faster and more parameter-efficient, while the proposed methods outperform vanilla DARTS on CIFAR-10.

  • Search efficiency: 0.45 CPU hours is NAS-NG’s average search time, approximately 40 times faster than NAS-MA and NAS-NGMA.The speed advantage results from simpler gating-based importance metrics rather than comprehensive activation-pattern analysis.
  • CNN architecture evaluation: NAS-NG outperforms vanilla DARTS on the three-convolutional-layer configuration while discovering substantially smaller networks.The proposed methods also achieve performance higher than or close to baseline with substantial parameter reductions in two-convolutional-layer configurations.
  • Overall evaluation: The proposed methods match or closely approach baseline accuracy while substantially reducing parameters across MLPs and CNNs on MNIST and CIFAR-10.On CIFAR-10, they outperform vanilla DARTS.
  • MLP architecture evaluation: 98.68% accuracy with 7.69M parameters is achieved by the best searched MLP, versus 98.69% with 30.76M parameters for the best baseline.This represents a 75% parameter reduction with nearly identical accuracy.
  • MLP architecture evaluation: 98.65% accuracy with 1.75M parameters is achieved by NAS-MA, a 94.3% reduction versus the best baseline.NAS-NGMA achieves the highest search-method accuracy, 0.01 percentage points below the best baseline, using 75% fewer parameters.
  • CNN architecture evaluation: NAS-NG achieves the best MNIST CNN result, outperforming the best CNN baseline by 0.01% while using approximately 47% less model size.CNN search also requires less time than MLP search for activation-driven methods.
  • Efficiency analysis: 86.05 accuracy points per million parameters is NAS-NG’s mean efficiency, compared with 31.94 for the baseline and 48.75 for MA-based NAS.NAS-NGMA reaches 60.66, between NAS-NG’s aggressive compression and MA’s more conservative approach.

6 Conclusions · Appendix

The paper presents relaxed bilevel NAS formulations that efficiently optimize neuron- and activation-level choices, achieving competitive performance with fewer parameters. It also identifies limitations and future directions, while advocating evaluation that balances accuracy, model size, and deployment constraints.

  • 6 Conclusions: Three relaxed formulations—NAS-NG, NAS-MA, and NAS-NGMA—enable differentiable neuron- and activation-level architecture optimization for MLPs and CNNs.The formulations provide a scalable alternative to discrete architecture search.
  • 6 Conclusions: Across MNIST and CIFAR-10, the proposed methods match or exceed literature baselines while using substantially fewer parameters.These results support continuous relaxations as an effective and scalable alternative to discrete NAS.
  • 6 Conclusions: 98.40–98.69% is the near-ceiling MNIST MLP accuracy range, which may obscure architecture differences on the simple dataset.CIFAR-10 experiments partially address this limitation, while ImageNet remains an important next benchmark.
  • 6 Conclusions: The study varies MLP depth and width, leaving additional architectural variation as a limitation.The supplied passage identifies this scope limitation but truncates the specific omitted variation.
  • 6 Conclusions: Future work should address harder vision tasks, broader architectural families including RNNs and Transformers, adaptive hyperparameter tuning, and deeper theory for gating metrics.These directions aim to improve robustness and clarify when relaxed formulations are most suitable.
  • 6 Conclusions: The findings challenge the necessity of heavy overparameterization and suggest targeted, task-specific adjustments can rival more elaborate NAS pipelines.The conclusion relates this observation to lottery-ticket-style findings.
  • 6 Conclusions: The paper advocates evaluating accuracy jointly with model size and deployment constraints rather than optimizing designs solely for predictive performance.This framework prioritizes efficient architectures that deliver practical value.

A Loss Function

The section formulates a binary-classification network’s forward pass from hidden-layer pre-activations through the predicted output. It then defines binary cross-entropy loss and rewrites it using the network’s explicit prediction and parameter collection.

  • The forward pass defines the hidden-layer pre-activation vector and corresponding hidden-layer activations.
  • The output neuron’s pre-activation and predicted output are specified, then combined into a single explicit prediction expression.
  • Binary cross-entropy loss is defined for true labels yi ∈ {0, 1} and predicted probabilities ŷi ∈ (0, 1).
  • The loss is further expressed by substituting the network’s explicit predicted output and collecting all weights and biases as W = {W(1), b(1), W(2), b(2)} ∈ R13.

B Modified Loss Function for Architecture Search · C Activation Functions

The method applies sigmoid-based gating after ReLU and defines a modified binary cross-entropy loss for binary labels. The activation functions used in the work are summarized in Table 14.

  • B Modified Loss Function for Architecture Search: Sigmoid-based gating is applied after the ReLU activation in the hidden layers.
  • B Modified Loss Function for Architecture Search: The hidden-layer activations are expressed using the gated activation formulation.
  • B Modified Loss Function for Architecture Search: The output pre-activation is specified by a corresponding expression.
  • B Modified Loss Function for Architecture Search: For each sample with ground-truth label y_i ∈ {0, 1}, a modified binary cross-entropy loss is defined.
  • B Modified Loss Function for Architecture Search: Additional displayed expressions provide the component terms used in the modified loss formulation.
  • C Activation Functions: The work summarizes its activation functions in Table 14.

D Configuration-Specific Performance Patterns

Across layer-count and hidden-dimension configurations, all methods show similar accuracy landscapes, while searched architectures achieve comparable performance with fewer parameters. Search methods also produce more uniform performance and can compensate for suboptimal initial configurations, with NGMA reaching 98.68% under favorable settings.

  • Configuration-specific accuracy: Searched architectures match similar accuracy patterns across layer counts and hidden dimensions while using fewer parameters.Each heatmap cell reports mean test accuracy for one configuration.
  • Configuration-specific accuracy: Highest accuracies generally occur in deeper networks with larger hidden dimensions, especially in the bottom-right heatmap regions.The search methods preserve rather than fundamentally alter the architecture-performance landscape.
  • Configuration robustness: Search-method heatmaps are more uniform than the baseline, which shows larger accuracy drops in shallower or narrower networks.This suggests architecture search partially compensates for suboptimal initial choices by retaining the most valuable neurons.
  • NGMA performance: 98.68% is the highest peak accuracy, achieved by NGMA in the 3-layer, high-dimensional region.This peak is consistent with NGMA’s strongest individual results when the initial configuration is favorable.

E Additional Analysis

Additional empirical analysis examines how network depth and hidden dimensions affect classification performance, while also assessing parameter reductions relative to dense baselines. These analyses characterize the trade-off between model complexity and predictive performance.

  • Architectural Configuration: The analysis evaluates proposed NAS methods across different architectural configurations.It focuses on how architectural choices shape method behavior.
  • Architectural Configuration: Network depth and hidden dimensions are examined for their influence on classification performance.The analysis studies these architectural factors directly.
  • Model Complexity: Parameter reduction is assessed relative to baseline dense architectures.This comparison provides evidence about changes in model complexity.

A Impact of Network Depth

Increasing depth generally improves accuracy, with the strongest results at 3–4 layers and diminishing benefits beyond 3 layers. Search methods remain competitive across depths, with NAS-NG showing stable performance and NAS-NGMA favoring an intermediate depth.

  • A Impact of Network Depth: 3-layer and 4-layer networks perform best, while deeper networks generally improve performance but yield diminishing benefits beyond 3 layers.Error bars summarize standard deviation across hidden dimensions 1024, 2048, and 3072.
  • A Impact of Network Depth: 98.63% accuracy at 4 layers versus 98.51% at 1 layer gives the baseline a 0.12 percentage-point improvement.The gain demonstrates a modest depth benefit on MNIST.
  • A Impact of Network Depth: 98.50% to 98.57% is NAS-NG’s performance range across depths, indicating remarkably consistent accuracy as architecture configurations change.The pruning process adapts across different depths.
  • A Impact of Network Depth: 98.65% is NAS-NGMA’s best performance at 3 layers, rather than 4, suggesting that gating and activation criteria favor intermediate depth.Its depth preference differs from the general best performance at 3–4 layers.

B Influence of Hidden Dimensions

Increasing total hidden dimensions generally improves accuracy, but gains diminish beyond 6,000 dimensions. NAS-NG maintains competitive accuracy at lower capacities while using fewer parameters at each dimension level.

  • B Influence of Hidden Dimensions: Figure 16 averages test accuracy across configurations with identical total hidden-dimension capacity and displays standard deviations across their layer configurations.The analysis evaluates total capacity as the sum of hidden dimensions across layers.
  • B Influence of Hidden Dimensions: Beyond 6,000 total hidden dimensions, accuracy gains diminish, with some configurations showing slight performance degradation.Increasing capacity from 1,024 to 4,096 total dimensions produces noticeable improvements across all methods.
  • B Influence of Hidden Dimensions: NAS-NG achieves competitive accuracy at lower total dimensions while reducing overall parameter counts at each dimension level.This suggests NAS-NG identifies informative neurons instead of retaining neurons in proportion to the initial network size.

C Parameter Reduction Analysis

The parameter-reduction analysis shows that NAS-NG provides the strongest and most consistent compression, while NAS-NGMA is more conservative and NAS-MA is smaller and more variable. Comparable complexity reductions are also observed for CNN architectures.

  • Parameter Reduction Analysis: 54.49% average parameter reduction: NAS-NGMA is more conservative, with 9.77% standard deviation because gating and activation metrics must agree.This joint criterion reduces the likelihood of extreme pruning decisions optimized for only one criterion.
  • Parameter Reduction Analysis: 40.98% average parameter reduction: MA-based pruning is smaller and more variable, spanning 15.82%–56.76% with 14.61% standard deviation; similar complexity reductions occur for CNNs.The variability indicates configuration-dependent performance across architectures.
Loading 2608.14443v1…