Designing Compact Neural Architectures via Neuron Gating and Mixed Activation

Abhishek Shukla; Ankur Sinha; Faiz Hamid

cs.LGAug 14, 2026

Designing Compact Neural Architectures via Neuron Gating and Mixed Activation

Authors: Abhishek Shukla, Ankur Sinha, Faiz Hamid

Organizations: Department of Management Sciences, IIT Kanpur, India

Abstract

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.

Figures & tables

Figure 1: Progressive construction of MLP architectures via neuron- and activation-level modifications.

Figure 2: Schematic representation of a general MLP architecture with an input layer, $L$ hidden layers, and an output layer.

Notation	Description
$S^{T},\,S^{V}$	Training and validation sets
$\mathbf{x}_{i}\in\mathbb{R}^{n_{0}},\,y_{i}\in\{0,1\}$	Input vector and true output label for the $i^{\text{th}}$ sample
$j=0$	Index of the input layer
$j=1,\dots,l,\;\;l+1,\dots,L$	Indices of the hidden layers
$j=L+1$	Index of the output layer
$n_{j}$	Number of neurons in the $j^{\text{th}}$ layer

Table 1: Notations used for the MLP model.

Figure 3: Neural network with one hidden layer.

Figure 4: Neural network architecture illustrating the neuron gating mechanism employed for MLP architecture search.

Figure 5: Modified neural network architecture for mixed activation-based architecture search.

Figure 6: Neural network incorporating neuron gating mechanism and mixed activation function for MLP architecture search.

g_{A}^{k}=\nabla_{A}\mathcal{L}_{v}\Bigl(A^{k},\;W^{k}-\xi\nabla_{W}\mathcal{L}_{t}(A^{k},W^{k})\Bigr)

Algorithm 1 NAS-NGMA

Figure 7: An overview of NAS-NGMA. (a) The gating states and activation functions of neurons are initially unknown. (b) Continuous relaxation of the search space through neuron gating and a weighted mixture of candidate activation functions for each neuron. (c) Joint optimization of the gating parameters, mixing probabilities, and network weights via a bilevel optimization framework. (d) Derivation of the final neural architecture from the learned gating states and activation-function mixing probabilities.

1–4 Hidden Layers $\times$ {1024, 2048, 3072} Neurons/Layer
Search Method	Val Acc (%)		Time (hrs)		Params (M)
	Mean	Std Dev	Mean	Std Dev	Mean	Std Dev
NAS-NG	97.59	0.04	0.45	0.14	1.90	1.33
NAS-MA	97.27	0.12	18.80	12.04	4.64	4.38
NAS-NGMA	96.20	0.31	18.81	11.89	3.54	3.56

Table 2: Architecture search performance for MLP models on the MNIST dataset.

Figure 8: Comparison of architecture search methods for MLP models on the MNIST dataset.

Search Method	Mean Test Acc.	Max Test Acc.	Mean Params	Param
	(%)	(%)	(M)	Reduction (%)
Baseline	98.56	98.69	8.97	—
NAS-NG	98.53	98.60	1.90	78.8
NAS-MA	98.50	98.65	4.64	48.3
NAS-NGMA	98.55	98.68	3.54	60.5

Table 3: Test performance of discovered MLP models on MNIST dataset.

Figure 9: Comparison of architecture search methods for CNN models on the MNIST dataset.

Search Method	3 Layer Conv [64, 128, 256]			4 Layer Conv [32, 64, 128, 256]
	Val Acc (%)	Time (hrs)	Params (M)	Val Acc (%)	Time (hrs)	Params (M)
NAS-NG	99.31	0.50	0.26	99.24	0.44	0.12
NAS-MA	99.22	5.46	0.45	99.24	4.64	0.24
NAS-NGMA	99.26	5.51	0.25	99.24	4.59	0.14

Table 4: Architecture search performance for CNN models on the MNIST dataset.

Search Method	3 Layer Conv [64, 128, 256]		4 Layer Conv [32, 64, 128, 256]
Search Method	Test Acc (%)	Params $\downarrow$ (%)	Test Acc (%)	Params $\downarrow$ (%)
Baseline	99.60	-	99.62	-
NAS-NG	99.63	73.5	99.50	74.7
NAS-MA	99.49	55.1	99.55	50.5
NAS-NGMA	99.61	75.4	99.56	72.3

Table 5: Test performance of discovered CNN models on MNIST dataset.

Search Method	Val Acc (%)		Time (hrs)		Params (M)
	Mean	Std Dev	Mean	Std Dev	Mean	Std Dev
NAS-NG	77.59	0.31	0.56	0.02	0.49	0.04
NAS-MA	74.16	0.13	10.75	0.39	0.83	0.10
NAS-NGMA	74.88	0.24	10.19	0.08	0.41	0.03
DARTS	81.93	0.31	26.30	0.56	3.40	0.10

Table 6: Architecture search performance for CNN models on the CIFAR-10 dataset.

Figure 10: Comparison of NAS methods for CNN models on the CIFAR-10 dataset.

Search Method	Test Acc (%)		Eval. Time (hrs)		Params $\downarrow$ (%)
	Mean	Std Dev	Mean	Std Dev
Baseline	86.14	0.22	1.48	0.05	–
NAS-NG	85.60	0.21	1.09	0.02	91.4
NAS-MA	84.55	0.31	4.61	0.22	85.4
NAS-NGMA	84.73	0.18	3.08	0.08	92.8
DARTS	82.03	0.50	20.20	0.47	–

Table 7: Test performance of discovered CNN models on CIFAR-10 dataset.

Figure 11: Test accuracy vs model complexity (measured in millions of parameters) for all evaluated configurations on the MNIST dataset.

Method	Layers	Config.	Test Acc.	Params	Search
			(%)	(M)	Time (hrs)
Baseline	4	[3072, 3072, 3072, 3072]	98.69	30.76	—
NAS-NG	4	[3072 4 ] $\rightarrow$ Pruned	98.60	4.29	0.76
NAS-MA	4	[1024 4 ] $\rightarrow$ Pruned	98.65	1.75	16.16
NAS-NGMA	3	[3072 3 ] $\rightarrow$ Pruned	98.68	7.69	35.80

Table 8: MLP best performing models identified by each search method on MNIST dataset.

Method	Layers	Config.	Test Acc.	Params	Search
			(%)	(M)	Time (hrs)
Baseline	4	[32, 64, 128, 256]	99.62	0.49	—
NAS-NG	3	[64, 128, 256] $\rightarrow$ Pruned	99.63	0.26	0.50
NAS-MA	4	[32, 64, 128, 256] $\rightarrow$ Pruned	99.55	0.24	4.64
NAS-NGMA	3	[64, 128, 256] $\rightarrow$ Pruned	99.61	0.25	5.51

Table 9: CNN best performing models identified by each search method on MNIST dataset.

Method	Layers/Cells	Config.	Test Acc.	Params	Search
			(%)	(M)	Time (hrs)
Baseline	3	[128, 256, 512]	86.34	5.69	—
NAS-NG	3	[128, 256, 512] $\rightarrow$ Pruned	85.77	0.51	0.53
NAS-MA	3	[128, 256, 512] $\rightarrow$ Pruned	84.83	0.79	10.64
NAS-NGMA	3	[128, 256, 512] $\rightarrow$ Pruned	84.85	0.38	10.15
DARTS	20	Stacked Normal–Reduction Cells	82.57	3.38	25.88

Table 10: CNN best performing models identified by each search method on CIFAR-10 dataset.

Figure 12: Mean efficiency (test accuracy per million parameters) achieved by each search method for MLP models on the MNIST dataset.

Rank	MNIST			CIFAR-10
	Method	Config.	Efficiency	Method	Config.	Efficiency
1	NAS-NG	[32, 64, 128, 256]	803.97	NAS-NGMA	[128, 256, 512]	225.65
2	NAS-NG	[32, 64, 128, 256]	803.89	NAS-NGMA	[128, 256, 512]	199.23
3	NAS-NGMA	[32, 64, 128, 256]	732.76	NAS-NG	[128, 256, 512]	195.28
4	NAS-NGMA	[32, 64, 128, 256]	732.32	NAS-NGMA	[128, 256, 512]	193.98
5	NAS-MA	[32, 64, 128, 256]	410.58	NAS-NG	[128, 256, 512]	166.88

Table 11: Top five most efficient models (highest accuracy-to-parameter ratios) on MNIST and CIFAR-10 datasets.

Model	Test Acc. (%)	Model	Test Acc. (%)
CNN [ 45 ]	99.45	LSTM [ 45 ]	99.22
LeNet-5 [ 23 ]	99.05	Consensus Clustering [ 46 ]	95
LR-based [ 24 ]	93.83	SVM [ 47 ]	97.83
Spiking-CNN [ 32 ]	96.89	Multi-Column DNN [ 38 ]	99.77
PCA/LDA [ 26 ]	86.60	EfficientDet-D4 [ 35 ]	99.83
CNN-SVM [ 34 ]	99.30	RMDL [ 36 ]	99.82

Table 12: Comparison of MNIST classification accuracies from various models.

Model	Architecture	Baseline		NAS-NG
		Test Acc. (%)	Params (M)	Test Acc. (%)	Params (M)
MLP
	[3072, 3072, 3072, 3072, 3072, 3072]	98.68	49.64	98.65	4.80
	[4096, 4096, 4096, 4096, 4096, 4096]	98.64	87.16	98.67	7.38
	[6144, 6144, 6144, 6144, 6144, 6144]	98.64	193.66	98.73	12.97
CNN Backbone: [128, 256, 512, 512, 1024]

Table 13: Evaluation experiments for NAS-NG with larger baseline architectures.

Figure 13: Applying NG to literature-optimal DARTS cells simultaneously improves classification accuracy and reduces the number of trainable parameters, demonstrating that neuron-level optimization complements topology search.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Activation Function	Definition and Mathematical Formula
None (Inactive Neuron)	$\displaystyle f_{1}(z)=0$ Produces a zero output for any input, effectively deactivating the neuron to enable soft pruning within the activation search framework
Identity	$\displaystyle f_{2}(z)=z$ A linear activation that preserves the pre-activation value, allowing the network to retain linear transformations when nonlinearity is unnecessary
Rectified Linear Unit (ReLU)	$\displaystyle f_{3}(z)=\max(0,z)$ Introduces nonlinearity without saturation for positive inputs, enabling effective gradient propagation in deep networks
Leaky Rectified Linear Unit (Leaky ReLU)	$\displaystyle f_{4}(z)=\max(0.01z,z)$ A ReLU variant that retains a small gradient for negative inputs to mitigate the dead-neuron problem
Hyperbolic Tangent	$\displaystyle f_{5}(z)=\tanh(z)=\frac{e^{z}-e^{-z}}{e^{z}+e^{-z}}$ A smooth, zero-centered nonlinear activation mapping inputs to $(-1,1)$ to enhance optimization stability in shallow and medium-depth networks
Sigmoid	$\displaystyle f_{6}(z)=\sigma(z)=\frac{1}{1+\exp(-z)}$ A bounded, smooth activation mapping inputs to $(0,1)$ , widely used for probability estimation and gating mechanisms

Appendix

Table 14: Candidate activation functions are used for neuron activation-based architecture search.

Figure 14: Test accuracy heatmaps showing performance on MNIST dataset across different network configurations.

Figure 15: Test accuracy vs number of hidden layers, averaged across all hidden dimension configurations.

Figure 16: Test accuracy versus total hidden dimensions, averaged across configurations with the same total capacity.

Figure 17: Distribution of parameter reduction percentages achieved by each search method compared to baseline dense networks of equivalent initial size.

Neural Architecture Search

Neural Architecture Search automates neural network design for target tasks, resource budgets, and hardware constraints.

See all 132 papers

Apr 1, 2026

G-ICSO-NAS: Shifting Gears between Gradient and Swarm for Robust Neural Architecture Search

Neural Architecture Search (NAS) has become a pivotal technique in automated machine learning. Evolutionary Algorithm (EA)-based methods demonstrate superior search quality but suffer from prohibitive computational costs, while gradient-based approaches like DARTS offer high efficiency but are prone to premature convergence and performance collapse. To bridge this gap, we propose G-ICSO-NAS, a hybrid framework implementing a three-stage optimization strategy. The Warm-up Phase pre-trains supernet weights (

w

) via differentiable methods while architecture parameters (

α

) remain frozen. The Exploration Phase adopts a hybrid co-optimization mechanism: an Improved Competitive Swarm Optimizer (ICSO) with diversity-aware fitness navigates the architecture space to update

α

, while gradient descent concurrently updates

w

. The Stability Phase employs fine-grained gradient-based search with early stopping to converge to the optimal architecture. By synergizing ICSO's global navigation capability with differentiable methods' efficiency, G-ICSO-NAS achieves remarkable performance with minimal cost. In the context of the DARTS search space, an accuracy of 97.46% is achieved on CIFAR-10 with a computational budget of just 0.15 GPU-Days. The method also exhibits strong transfer potential, recording accuracies of 83.1% (CIFAR-100) and 75.02% (ImageNet). Furthermore, regarding the NAS-Bench-201 benchmark, G-ICSO-NAS is shown to deliver state-of-the-art results across all evaluated datasets.

Xingbang Du, Enzhi Zhang, Rui Zhong +2

Information Systems Design Laboratory, Hokkaido University, Sapporo, Japan · Information Initiative Center, Hokkaido University, Sapporo, Japan

Neural Architecture Search Particle Swarm Optimization

Ask PDF

Aug 14, 2026

LP-NAS: Linear Programming-based Neural Architecture Search

Neural Architecture Search (NAS) aims to automate neural network architecture design, reducing reliance on human expertise. Among the various NAS methods, differentiable NAS has gained prominence due to its efficiency and accuracy compared to conventional NAS approaches. Since differentiable NAS relaxes the architecture search space into a continuous domain, it is possible to apply principles from continuous optimization to NAS. In this paper, we propose Linear Programming-based NAS (LP-NAS), a mathematical programming-based framework for differentiable NAS that is applicable to a wide range of continuous search spaces. LP-NAS formulates a linear program (LP) using the validation-loss gradient and the training-loss Hessian to compute an architecture update direction that improves generalization while preserving the optimality of the model parameters. By following this LP-derived descent direction, LP-NAS efficiently navigates the architecture search space, leading to faster and more effective architecture optimization. We introduce two computationally efficient variants of LP-NAS, namely S-LP-NAS and R-LP-NAS. Applying LP-NAS to the Differentiable Architecture Search (DARTS) search space results in two algorithmic variants, S-LP-DARTS and R-LP-DARTS. Both variants achieve faster convergence and significantly higher validation performance during the early search iterations than the standard DARTS algorithm. Extensive experiments on CIFAR-10 and CIFAR-100 show that LP-DARTS outperforms standard DARTS in both the architecture search and evaluation phases. Additionally, we compare our approach with several DARTS variants (P-DARTS, PC-DARTS, and STO-DARTS) on the CIFAR-10 dataset and demonstrate its effectiveness. Furthermore, we validate the transferability of the discovered architectures through experiments on the ImageNet dataset.

Abhishek Shukla, Ankur Sinha, Faiz Hamid

Department of Management Sciences, IIT Kanpur, India · Krishnamurthy Tandon School of AI, IIM Ahmedabad, India

Neural Architecture Search

Ask PDF

Jun 28, 2026

Bilevel Optimization for Neural Architecture Search

Bilevel optimization has become an influential and widely adopted framework for addressing hierarchical optimization problems in machine learning, providing an effective approach to modeling the interaction between two levels of optimization, with applications such as hyperparameter tuning, meta-learning, adversarial training, and data poisoning. Neural Architecture Search (NAS), a subfield of hyperparameter optimization, is a prime example of a bilevel optimization problem, with architecture parameters optimized at the outer-level and network weights optimized at the inner level. This paper presents a structured overview of NAS through the lens of bilevel optimization. We categorize existing NAS approaches into two main classes: sampling-based methods, which search optimal architectures using different architecture samplers, and bilevel theory-based methods, which solve the architecture search problem using bilevel optimization principles. We further highlight our current research direction, wherein the bilevel NAS formulation is addressed through an auxiliary mathematical programming framework. This framework enables the systematic integration of second-order information from the model's training loss function and ensures the optimality of the model parameters while modifying architecture parameters. By simultaneously updating the architecture and model parameters along their respective optimal descent directions derived from the auxiliary mathematical program, these methods achieve more principled and theoretically consistent results. The same auxiliary program can also be used for simultaneous hyperparameter and model fine-tuning. A comparative analysis shows that bilevel theory-based approaches generally outperform sampling-based methods, both in accuracy and efficiency.

Abhishek Shukla, Ankur Sinha, Faiz Hamid

Department of Management Sciences, IIT Kanpur, India · Krishnamurthy Tandon School of AI, IIM Ahmedabad, India

Bilevel Optimization Neural Architecture Search Hyperparameter Optimization

Ask PDF

Abstract

Neural Architecture Search

Explore similar work

G-ICSO-NAS: Shifting Gears between Gradient and Swarm for Robust Neural Architecture Search

LP-NAS: Linear Programming-based Neural Architecture Search

Bilevel Optimization for Neural Architecture Search