Paper: The Expressive Leaky Memory Neuron: an Efficient and Expressive Phenomenological Neuron Model Can Solve Long-Horizon Tasks. (ICLR-2024)

Reviewer comment (Writing):
A paragraph defining Branch-ELM would be necessary. Can you elaborate more on the intuition behind this variant? How does it work? When is it more suitable compared to the vanilla ELM? Why is it important to over-sample the input in this case?

Author response:
We now have a dedicated description of the Branch-ELM at the beginning of the paper, together with the definition of ELM (in “The integration mechanism dynamics” paragraph), and generally expand on its discussion.
The intuition: the Branch-ELM is inspired by the biological neuron, whose dendritic tree is composed of smaller hierarchically organized units called dendritic branches. While the dendritic tree overall might compute a highly nonlinear transformation, most inputs are summed linearly, and organization into branches guides where linear and where nonlinear summation is appropriate. We mimic this idea by slicing the input into a number (d_tree) of equally sized “branches” (for our best model, each synapse takes part in four slices, accounting for the mismatch between the ground truth and the branch assignment). Inputs are summed linearly within the slice. However, we let the downstream MLP take care of any further processing nonlinearities. If the slice size is large and the number of branches few, this approach reduces dramatically the number of parameters by reducing the size of MLP.

- However, using this stark simplification, we lose general expressivity, and performance is expected to degrade unless the inputs are assigned to the synapse belonging to the correct ground truth branch identity and branch computations can be well approximated by a rectified sum.

- As we only know that neighboring inputs mostly belong to the same branch identity, we sampled larger overlapping windows than the actual suspected branch sizes and let the model discard the incorrect assignments by modifying the synaptic weights.

- In the appendix Figure S5 we show that randomly assigning synapses to branches or fixing the synaptic weights severely hurts performance, and having the incorrect branch activation like “sigmoid” or “relu” also impairs performance significantly.

- In general, we expect the Branch-ELM to work great on datasets with large input dimensionality with low information density, such as SHD or SHD-Adding (coding for one or two digits per sample). We randomly oversample the input for the Branch-ELM on that dataset.

Reviewer comment (Reproducibility):
In Appendix B, the general training setup is detailed (batch size of 8, etc.). However, later on, different hyperparameters are used for the datasets. It is not clear where this general training setup was used.

Author response:
Thank you for catching this inconsistency. There are parameters that can change across tasks (including the optimizer, batch size and initial learning rate). We report them now individually in the corresponding tasks.

Reviewer comment (Evaluation):
This paper is aimed to balance the trade-off between fidelity, efficiency and biological realism. It did a fair job while still failed to capture some important aspects. For example, using MLP sacrifices the interpretability and biological realism to compensate accuracy. On the other hand, the model still sacrifices the accuracy and has a big performance gap when compared to SOTA models (S4 and Mega).

Author response:
1) *Clarification of MLP design choice of ELM neuron:*

We will include in the manuscript that one of the key modeling goals of the ELM design was to quantify the degree of nonlinearity necessary to model a neuron’s dynamics; hence the choice of MLP as it offers a straightforward way to do this using $d_{mlp}$ and $l_{mlp}$. In contrast, in more biologically realistic detailed biophysical models, there is no straightforward way to quantify the amount of nonlinearity in such models, and therefore offers less interpretability along this dimension compared to the ELM neuron.

Reviewer comment (Experiments):
Scalability of the method: as scaling law plays a big role for improving transformer's predictivity, one concern is that transformer might perform better with increasing number of parameters. However, it might not be the same for ELM. As firstly shown in Fig 3, the accuracy quickly saturates with simply increasing $d_m$ and $d_{mlp}$, and not able to get further improved to minimize the performance gap between ELM and Mega.

Author response:
6) *Addressing the scalability and limitations of the ELM neuron:*

We acknowledge the scaling limitation of individual ELM neurons, which we set out to explore using the LRA benchmark, however, we would like to point out that the training was kept fairly simple, and we did not use e.g. warmup, weight decay, layer-norm, dropout, etc (which the comparison models did). Additionally, we suspect that assembling such complex neurons in more sophisticated multi-layered networks might provide the desired scalability; unfortunately, this is beyond the scope of this work. Nevertheless, we provide an ablation of the ELM neuron and  S5 model on Pathfinder in terms of parameters in Tables S5 and S6 showing that an ELM with 3.5K parameters is outperforming a single layered S5 with ~170K parameters (62.95% ACC vs 58.19% ACC). Furthermore, motivated by the observation that SOTA models on LRA use various regularization techniques, we train ELM neuron with various dropout probabilities on Pathfinder in Table S7, and find that models trained with 0.3 dropout probability for 1000 Epochs can reach ~82\% accuracy, and improvement of 10% accuracy, evidencing that the performance gap might be reduced through additional proper regularization techniques.

Reviewer comment (Experiments):
Only one ML task is evaluated, more evaluations and benchmarks needed to demonstrate contributions in addressing long-range sequence modeling.

Author response:
4) *Number of Long-range ML tasks:*

There might be a misunderstanding concerning the number of tasks. The LRA benchmark contains 6 rather different tasks. Additionally, the Spoken Heidelberg Digits Task (SHD), and particularly SHD-addition, has a context length of 2000 for a bin size of 1, and therefore is twice as long as, e.g., the Image or Pathfinder task from LRA. Therefore, we provide results on a total of 8 long-range tasks.

Reviewer comment (Experiments):
As shown in Table S1, large number of hyper-parameters still needed in advance or be tuned based on prior knowledge from neuroscience literatures.

Author response:
3) *Clarification of tuning ELM parameters:*

We will highlight that the ELM neuron primarily only requires tuning along the memory $d_m$ dimension (and $\tau_s$ for very sparse tasks). We add the reference to the text for the hyper-parameters tuning table in the appendix that highlights how other hyperparameters can be derived from $d_m$, the input dimensionality $d_s$, and the task length $T$, and what default parameters should be used otherwise. A further testament to this is the LRA task, where mostly the same hyper-parameters across all tasks. Other training parameters such as batch size, learning rate, etc. still need tuning of course. Therefore, we believe the ELM neuron does not require more tuning than e.g. the LSTM.

Reviewer comment (Evaluation):
The efficiency side might be over-claimed, as Table 1 shows ELM still requires 100k-200k parameters?

Author response:
5) *Addressing when the ELM neuron can be expected to be efficient:*

We agree that this needs further clarification. The ELM and in particular Branch-ELM neuron can mainly be expected to be much more efficient if the input dimensionality is very large, such as in the NeuronIO task ($d_s=1278$), or somewhat in the SHD and SHD-Adding task ($d_s=700$), and less so e.g. on LRA ($d_s = 8$ to $d_s = 169)$; also note the increased relative performance for language tasks with higher input dimensionality. We now elaborate this condition for efficiency more clearly in the manuscript.

Reviewer comment (Experiments):
Is the model able to be trained with other biological plausible learning rules instead of BPTT? How they might end up with different parameters?

Author response:
7. *Training with biological learning rules:* 

We agree that it would be interesting to study the biologically plausible training, and it might happen that the training will converge to a different solution. However, for this manuscript, we aimed to establish how the model can perform with standard machine learning optimizers (that typically define the upper limit on the performance, seldomly achieved by the biologically plausible rules). Finding how to train this model with, for example, Eq.Prop or even more complicated bio-plausible plasticity rules would be an exciting direction. We added it to future work suggestions.

Reviewer comment (Writing):
There are major clarity issues in the paper. Many aspects of the model and notation are unexplained (e.g. $\lambda$, $1-\kappa_m$ in Fig. 1(c)). The explanation of Branch-ELM comes much later, even though it's referred to multiple times before that, which makes it very hard to read. The role of $w_s$ is also not clear at all (the given explanation on Pg. 3 doesn't help).

Author response:
1)  *Better introduction and explanation of the Branch-ELM neuron model:*

We improve the presentation of our paper by introducing the Branch-ELM together with the ELM neuron in section 2, by additionally expanding on the parameters involved in scaling the MLP output, adding references to the supplementary tables and by clarifying the Branch-ELM processing behavior in detail.

2) *Clarification of the parameters scaling of the MLP output:*

The constant $\lambda$ (float constant) and $1-\kappa_m$ (float vector, dependent on $tau_m$) are factors that scale the output of the MLP. Together, they ensure that the memory values $m_t$ are constrained to stay within a range of [-$\lambda$, $\lambda$] if $tanh$ range is [-1,1], which prevents exploding gradients.
We now added a short explanation to the formulation of the model (in “The integration mechanism dynamics” section)

Reviewer comment (Experiments):
The behaviour of Branch-ELM is unclear -- if the input is shuffled, does it affect performance? Since it depends on the local window to group inputs?

Author response:
3) *Clarification of the Branch-ELM processing of its inputs on NeuronIO:* 

Indeed, this is correct; if the ordering of input channels is shuffled, performance degrades significantly for the Branch-ELM. Neighboring input channels are typically located next to each other in the biophysical model dataset we are using, hence we oversample the input in a moving window fashion to exploit this insight (the resulting overlap between branches is to account for uncertainty in exact ground truth branch bounds). We tested the shuffling effect in Supplementary Figure S5 (purple trace), where random synapse-to-branch assignment makes the model incapable of reaching the spike-prediction threshold. For other datasets, such as SHD, however, we randomly assign synapses to branches. We will highlight this in the main text and refer to additional branch-related ablations in the appendix.

Reviewer comment (Experiments):
It's a bit odd that ELM doesn't perform well for short sequences (large bin size) as seen in Fig. 5 whereas LSTM does. The performance of ELM for short sequences could be explored more, since it sounds like that might be a major shortcoming of the model.

Author response:
4) *Clarification of ELM neuron and LSTM performance for larger bin sizes on SHD:*

The ELM and Branch-ELM neurons consistently outperform the LSTM for all binning sizes. Nevertheless, performance degrades for larger bin sizes for all models, as the input binning is a lossy pre-processing step, where temporal information is eliminated. Performance improves with the bin size for LSTM because, for the small time bins, LSTM could not take advantage of the whole temporal structure but gets into training problems due to the many time bins it has to keep in memory. The information is lost as the bin size increases, but the small number of bins enables LSTM to learn the task. We show that the temporal structure of the task is not a problem for ELM neurons, so they are only facing the drawbacks of the coarser data discretization. Overall, the observed performance dynamics is not a shortcoming of, e.g. the ELM or LSTM models but a consequence of the data preprocessing. We now discuss it in the main text, and include a statement about the hyper-tuning on the dataset in the appendix.

Reviewer comment (Experiments):
It is not clear how the number of parameters for the various cases were chosen.

Author response:
6) *Comparison to matched parameter performance on LRA between:* 

To provide some evaluation results with matched parameters on LRA, we provide an additional ablation on Pathfinder of the number of memory units for the ELM neuron and number of layers for the S5 model respectively; now displayed in Tables S5 and S6 The results show that an ELM with 3.5K parameters is outperforming a single layered S5 with ~170K parameters (62.95% ACC vs 58.19% ACC).