Gumbel-Softmax based Neural Architecture Search for Hierarchical Brain Networks Decomposition

Abstract

Understanding the brain's functional architecture has been an important topic in the neuroimaging field. A variety of brain network modeling methods have been proposed. Recently, deep neural network-based methods have shown a great advantage in modeling the hierarchical and complex functional brain networks (FBNs). However, most of these deep neural networks were handcrafted, making it time-consuming to find the relatively optimal architecture. To address this problem, we propose a novel unsupervised differentiable neural architecture search (NAS) algorithm, named Gumbel-Softmax based Neural Architecture Search (GS-NAS), to automate the architecture design of deep belief network (DBN) for hierarchical FBN decomposition. Specifically, we introduce the Gumbel-Softmax scheme to reframe the discrete architecture sampling procedure during NAS to be continuous. Guided by the reconstruction error minimization procedure, the architecture search can be driven by the intrinsic functional architecture of the brain, thereby revealing the possible hierarchical functional brain organization via DBN structure. The proposed GS-NAS algorithm can simultaneously optimize the number of hidden units for each layer and the network depth. Extensive experiment results on both task and resting-state functional magnetic resonance imaging data have demonstrated the effectiveness and efficiency of the proposed GS-NAS model. The identified hierarchically organized FBNs provide novel insight into understanding human brain function.

Keywords: Deep neural network; Hierarchical brain networks decoding; Neural architecture search; Task fMRI; Unsupervised.