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

G-ICSO-NAS: Shifting Gears between Gradient and Swarm for Robust Neural Architecture Search
Apr 1, 2026Xingbang Du, Enzhi Zhang, Rui Zhong +2Neural Architecture Search Particle Swarm Optimization
Information Systems Design Laboratory, Hokkaido University, Sapporo, Japan · Information Initiative Center, Hokkaido University, Sapporo, Japan
LP-NAS: Linear Programming-based Neural Architecture Search
Aug 14, 2026Abhishek Shukla, Ankur Sinha, Faiz HamidNeural Architecture Search
Department of Management Sciences, IIT Kanpur, India · Krishnamurthy Tandon School of AI, IIM Ahmedabad, India
Bilevel Optimization for Neural Architecture Search
Jun 28, 2026Abhishek Shukla, Ankur Sinha, Faiz HamidBilevel Optimization Neural Architecture Search
Department of Management Sciences, IIT Kanpur, India · Krishnamurthy Tandon School of AI, IIM Ahmedabad, India

Apr 1, 2026J/K move · Enter open · S save

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

Xingbang Du, Enzhi Zhang, Rui Zhong, Yang Cao+1

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

Neural Architecture Search Particle Swarm Optimization

Open paper AskPDF arXiv

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.

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