Parameter efficient dendritic-tree neurons outperform perceptrons

Multiple works have proposed ANNs based on neuron models that simulate dendritic input. These alternatives include dendrite morphological neural networks (DMNNs) [9], dendritic neural networks (DENNs) [14], the single dendritic neuron model (DNM) [12], and most recently the model by Jones and Kording based on a balanced tree structure [4]. Our model expands on the work of Jones and Kording by: (1) exploring the effect of generalizing the dendritic tree structure to allow tunable branching, dropout, and activations without k-tree redundancy; (2) using layers of dendritic neurons for non-binary classification tasks and evaluate overfitting; (3) attaching the layer of neurons to a CNN to evaluate performance as a perceptron replacement.

Other studies have connected dendritic neurons in MLP-like architectures [9, 14], including a hybrid DNM-CNN adaptation of the DNM model [13]. Each of these multi-neuron architectures uses a fundamentally different neuron model. DNM neurons connect each dendritic branch to each input and rely on logical operations, while DMNNs utilize a different underlying mathematical structure. By contrast, our model enforces sparse, localized connections between dendrites and inputs in a tree structure, which more closely models the spatially-limited connections between biological neurons. The DENN model enforces dendrite-input sparsity, but uses one-layer dendrite trees [14]. Our dendritic trees are deeper, to more closely replicate the anatomy and complexity of biological neural networks.

Using PyTorch (1.10.0+cu111) [7], we implemented a dendritic tree neuron based on the Jones-Kording single neuron model [4] (Figure 2 in the Appendix). Our implementation generalizes the original architecture by allowing users to specify the branching factor and number of neurons in the DendriticLayer. Each DendriticLayer has a constant branching factor, but multiple instances of the module can be stacked together to achieve different branching per layer. To vectorize the DendriticLayer, we treat the inputs to all the neurons as a single input tensor, and we compute the next layer of all dendrite trees simultaneously. The tree structure is enforced by constructing a mask for the weight matrix at each layer (Figure 1). Table 1 describes the full architecture of the DendriticLayer module.

We used Jones and Kording’s modified density gain Kaiming He initialization scheme to account for tree sparsity and stabilize training [4, 3]. Each weight in the weight matrices is initially sampled: $\mathbf{W}^{(i)}_{jk}\overset{iid}{\sim}\mathcal{N}(0,\frac{2}{I\cdot b^{-i+1}\cdot\text{density}})$, where $\text{density}=\frac{I\cdot b^{-i+1}}{I\cdot b^{-i+1}\cdot I\cdot b^{-i}}$, and $I\cdot b^{-i+1},I\cdot b^{-i}$ are the number of input and output units in the $i^{th}$ layer of the dendritic tree.

Using the DendriticLayer as our base, we built several multi-layer neuron (MLN) architectures on top of it. (1) An MLNBinaryClassifier, which is a single neuron that predicts a binary output activated using a logistic sigmoid function. (2) An MLNClassifier, which repeats the same set of inputs to a specified number of neurons (one for each class) and returns the output of each neuron activated through a softmax function. (3) A ConvMLN, which runs the inputs through a CNN (described in Section 9.3 in the Appendix), then flattens the output of the convolutional layers and feeds it to a single layer of neurons for classification. All models are equipped with a tunable input dropout layer to be robust against overfitting [11].

We tested our single neuron models on binary classification tasks, using a subset of MNIST that consisted only of images labeled as ”4” or ”9” (referred to as 4-9 MNIST) [2]. We used the full MNIST dataset to test the classification capabilities of our multi-neuron architectures [2].

For controls, we constructed multi-layer perceptrons (MLPs) that performed the same tasks as each of our dendritic models. The architecture of the MLPs is described in Table 2. The number of hidden units ($h$) was modified in order to match the number of parameters in the dendritic models. The number of parameters in each model is listed in Tables 9, 10, and 11 in the Appendix. Each model was initialized using the Kaiming He method, but without density gain since MLPs are not sparse [3].

The standard MNIST dataset images are $28\times 28$. Nearest neighbour upsampling was used to scale inputs to $32\times 32$, to better fit the dendritic branching factors. For non-CNN architectures, input flattening was applied to create a $1024$-dimension tensor.

To aggregate results from multiple trials, the epoch with the lowest validation loss was taken from each trial as the best performance to evaluate at. Post-experiment statistical analysis was conducted using R [8].

All models were trained on the Google Colab environment with CUDA for 100 epochs, using a batch size of 128. Every model was re-initialized and trained 10 times. The learning rates used for dendritic MLN and MLP models were 0.05 and 0.001, respectively. All models used either binary or categorical cross entropy loss. All models were trained using the Adam Optimizer [5].

We modified the Jones-Kording single neuron model to investigate the effects of branching factors ($b$) and dropout probabilities ($p$) (Table 3). The control MLPs with similar numbers of parameters are in Table 4. Dropout models were tested with $p$ = 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, but only the best-performing set of models are reported in each experiment.

To test non-binary classification, we integrated multiple neurons in a layer (MLNClassifier). Each neuron connected to the same set of inputs and was expected to predict a single MNIST digit. The experimental setup identical to the binary case, as described in Sections 3.4, 4.1. The average performance of each multi-neuron model and each control MLP is reported in Tables 5 and 6.

To explore the potential as a tunable, modular MLP replacement in existing architectures, we connected 10 neurons to a simple CNN architecture for MNIST classification (ConvMLN) as a replacement to the usual perceptrons. This experimental setup is described in Sections 3.4, 4.1. The performance of each CNN-MLN and control CNN-MLP is listed in Tables 7and 8.

Both increasing branching factor and introducing dropout significantly improved performance. The best performing dendritic models were two layers deep, with $b=\sqrt{i}$, where $i$ is number of inputs to each neuron and $p\geq 0.1$. We hypothesize deeper trees led to vanishing gradients; shallow dendritic trees trained and performed better despite having fewer parameters. In all experiments, dendritic neuron models with $b=2$ perform worse than their control MLPs, but neurons with $b>2$ surpass their control MLPs in terms of validation performance and robustness to overfitting, both with and without dropout. This shows dendritic trees can improve performance while maintaining parameter efficiency relative to MLP counterparts.

PyTorch implementation and experiment data linked here: github.com/zw123han/DendriticNeuralNetwork

An identical and simple CNN was attached to all classifiers. The CNN consisted of three sets of convolutions. Each convolutional layer had kernel size of 5 and a padding of 2. The number of filters in the first layer was 4, the number of filters in the second layer was 8, the number of filters in the third layer was 16. Each convolutional layer was followed by a MaxPool layer (kernel size = 2, stride = 2), a BatchNorm layer, and a ReLU activation layer.

The initial input to the CNN was a $B\times 1\times 32\times 32$ tensor, where $B$ was the batch size, $1$ is the number of input channels (grayscale), and $32\times 32$ are the dimensions of the upsampled MNIST digits. Each set of convolutions and pooling cut the size of the image in half ($32\rightarrow 16\rightarrow 8\rightarrow 4$). The final output of the CNN was 16 4$\times$4 maps, which were flattened into a single 256-unit input before it was given to the dendritic classifier.