Learning Successor Features with Distributed Hebbian Temporal Memory

Thank you for your detailed review and valuable comments! We updated the manuscript taking into account your suggestions. The most important changes are highlighted and tagged for your convenience.

W1 Misleading introduction

Thank you for pointing this out! We revised the introduction, adding more emphasis on the type of the memory we present in this work.

W2 DHTM’s description

We substantially revised section 3.2 to make it easier to follow. We also moved most of the derivations to Appendix B, since they are too heavy for the first read.

W3 Implementation details

Due to paper size limits, we had to put most of the implementation details into Appendix E.

W4 EC vs DHTM

The EC version of the algorithm lacks some important features such as the ability to process distributed encoding, factoring hidden state space as well as probabilistic estimates, which might be crucial in stochastic environments. It’s also possible to modify the EC version, to add these features, however, it seems like the resulting algorithm would be very similar to full DHTM. At the current stage of DHTM’s development, there is, probably, no decisive experiment that would show DHTM’s advantage over EC. However, newly added experiments in Appendix H.2 might give an idea that, at this point, DHTM already has better scaling capacity and faster convergence.

W5 Related work

We added more on related work in the introduction.

Q1 Temporal Memory

We agree that the term “Temporal Memory” might be confusing in the abstract, so we removed it.

Q2 Complete data sequences

There is a common challenge in continual learning that algorithms should be able to incorporate new sequential data into the system by small fragments, and without access to previously seen data (Khetarpal et al. 2022). Training only on small amounts of new data (at the extreme, just one sequence element) usually results in overfitting for deep learning methods and loss of previously acquired knowledge, known as catastrophic forgetting. Therefore, contemporary deep learning methods rely on huge experience buffers that are sampled to calculate parameter updates. That’s what we understand by “complete data sequences”.

Q3 Factor graph formalism

We added more information about factor graphs in the introduction and background sections.

Q4 Fast Hebbian learning

Here “faster” means that algorithms relying on Hebbian learning are usually more sample efficient, and require less training steps to converge in comparison to stochastic gradient descent (see for example Journé et al. 2023).

Q5 “Eq. (3) is unclear”

We added corresponding clarification.

Q6 Meaning of $f_l$

It simply means that $f_l$ is an exponential average of the segment's activation frequency, given that its root cell is also active.

Q7, Q8 VAE encoder with other models

We updated AnimalAI experiments by adding other models with VAE encoder for comparison. Since EC and CSCG have only one categorical variable as an input, each unique VAE representation should be mapped to the corresponding state of one categorical variable. For Categorical VAE with $5$ categorical variables with $50$ states each, it means that the resulting input variable should have $50^5$ states to have the same expressive power as distributed encoding. However, in practice, the total number of unique representations is much less for a particular environment and task. For example, in AnimalAI, it’s only about $1200$ VAE patterns. Still, CSCG with such a large state space would be infeasible to train. But since EC doesn’t require allocating all states in advance, it’s possible to train EC baseline trained on VAE representations, which, however, can’t use possible factorization in VAE encoding space in contrast to DHTM.

Khetarpal, Khimya, Matthew Riemer, Irina Rish, and Doina Precup. 2022. ‘Towards Continual Reinforcement Learning: A Review and Perspectives’. Journal of Artificial Intelligence Research 75 (December):1401–76. https://doi.org/10.1613/jair.1.13673.

Journé, Adrien, Hector Garcia Rodriguez, Qinghai Guo, and Timoleon Moraitis. 2023. ‘Hebbian Deep Learning Without Feedback’. arXiv. https://doi.org/10.48550/arXiv.2209.11883.

Thank you for your thoughtful questions and commentaries! We uploaded a revised version of our paper that, hopefully, addresses your concerns and takes into account your suggestions.

W1 Model description

We substantially revised DHTM’s description, simplifying it and moving the most difficult part into Appendix B.

W2 Background

We added more information about Factor graphs into the background section about HMM.

Q1 Scalability

We made additional experiments to investigate DHTM’s scaling properties and compare it to EC agent in Appendix H.2.

Q2 Aliased observations

DHTM was specifically designed to disentangle aliased observations, since it's a common problem in POMDP environments. DHTM is a high-order memory that is able to assign different hidden states to identical observations. We show it in our POMDP Gridworld navigation experiments, where many positions have identical colours. Only memory that is able to handle aliased observations will successfully solve the task.

Q3 Generalization

We revised our introduction section to stress more clearly that, in contrast to CSCG, DHTM is closer to episodic-like memory, and it is not able to generalize experience forming schemes. By comparison to CSCG, we aimed to show that, in certain tasks, it’s more advantageous to use fast episodic types of memories than slowly learning semantic memories like CSCG.

Thank you for your questions and valuable comments! We did our best to address your concerns and highlighted important changes in the new revision for your convenience.

Q1 Generalization

We think that hierarchical approach is the most promising for online generalization in DHTM. We see current DHTM implementation as a first level in the hierarchy, which aims to remember as much experience as possible. We are in the process of developing the second level, which shares some properties with DHTM, but its main purpose is to connect hidden states from the first level into clusters according to their n-step predictions. Furthermore, we aim to make the second level non-destructive, so that in case of generalization error we don’t lose potentially valuable experience on the first level.

Q2 DHTM vs EC

The EC version of the algorithm lacks some important features such as the ability to process distributed encoding, factoring hidden state space as well as probabilistic estimates, which might be crucial in stochastic environments. It’s also possible to modify the EC version, to add these features, however, it seems like the resulting algorithm would be very similar to full DHTM. At the current stage of DHTM’s development, there is, probably, no decisive experiment that would show DHTM’s advantage over EC. However, newly added experiments in Appendix H.2 might give an idea that, at this point, DHTM already has better scaling capacity and faster convergence.

Q3 VAE encoder with other models

We updated AnimalAI experiments by adding other models with VAE encoder for comparison. Since EC and CSCG have only one categorical variable as an input, each unique VAE representation should be mapped to the corresponding state of one categorical variable. For Categorical VAE with $5$ categorical variables with $50$ states each, it means that the resulting input variable should have $50^5$ states to have the same expressive power as distributed encoding. However, in practice, the total number of unique representations is much less for a particular environment and task. For example, in AnimalAI, it’s only about $1200$ VAE patterns. Still, CSCG with such a large state space would be infeasible to train. But since EC doesn’t require allocating all states in advance, it’s possible to train EC baseline trained on VAE representations, which, however, can’t use possible factorization in VAE encoding space in contrast to DHTM.

Q4 Adding new baselines

We updated the experiments in the new revision of the manuscript by adding also RWKV for comparison. We found that, in our setup, there is not much performance difference between RWKV and LSTM, except that RWKV of the same hidden state size as LSTM is a bit slower to train. Though their losses are decreasing with training, the performance stays approximately the same in AnimalAI. Apparently, LSTM and RWKV are not apt for probabilistic multistep predictions and, perhaps, require more architectural adjustments that are beyond the scope of this work.

Q5 Related papers

Thank you for pointing this out! We added a corresponding paragraph to the introduction section.

Q6 Oracle and EC

We updated Gridworld experiments by adding EC with Oracle. Interestingly, EC with complete memory clean-up converges a bit slower. Apparently, it loses some experience that it could reuse otherwise.

We are pleased to inform you that we have updated the manuscript by adding another experiment addressing your question about the difference between EC and DHTM (see Appendix H.3 "Noise Tolerance"). If you have any further questions, please let us know.

Thank you for such a detailed review! We updated the manuscript and hope that we managed to improve its quality by incorporating your suggestions and addressing your questions. We added highlights of important changes to the revised text, and you can use text search by your id to find them.

W1 RESET_MEMORY

Though, indeed, the name of the procedure might seem counterintuitive, we elaborate in section 3.4 that RESET_MEMORY procedure concerns only memory state or its current activity, but not its parameters (synaptic weights), which are usually considered the main carriers of the experience in the models we tested.

W2 Biological interpretation

About spikes: Thank you for noting that! Indeed, we have only one timescale, so we agree that calling cells activity as spikes might be confusing here, therefore we removed it.

Connection to HTM: Certainly, you are right, we’ve changed the reference to a more suitable paper. We also added clarification on DHTM’s connection to HTM in Appendix C.

Columnar structure: This form of the emission factor means that one feature state (input cell) influences several hidden states (cells in column) and only them. That is, several hidden states share input. We use the same emission factor that CSCG uses, which also refers to columnar structure of the neocortex, and it’s in concordance with HTM’s implementation of columnar structure.

W3 Presentation style

Thank you for your suggestions! We substantially revised Background and Neural Implementation sections in order to simplify them. We agree that the details of the derivation might be confusing and are difficult to comprehend from the first read, so we decided to move most of the details to Appendix B.

W4 Scope of the paper

“The experiments are barely non-stationary” - there are different types of non-stationarity, here we refer to piecewise non-stationarity of the environment’s transition function according to taxonomy used in Khetarpal et al. 2022. Though the task might appear trivial, it’s still challenging for most of the contemporary deep learning methods, which we aimed to show in our research.

Our method indeed has some scaling challenges, which, however, are common for episodic-like memories. We added additional experiments in Appendix H.2 and discussed DHTM’s scaling properties there. We also provided theoretical analysis of DHTM’s capacity In Appendix G.

About segment efficiency: We meant that in contrast to conventional HMM, we use sparse representation of transition matrix in the form of segments. In that case, prediction step time complexity grows linearly with the amount of segments and with hidden state space size. In contrast, if a full transition matrix is used, the time complexity of the prediction step grows quadratically with the state space size. For example, CSCG with the same state space size as DHTM in our experiments would be infeasible to train.

W5 Figures

Thank you for your feedback! We agree that Figure 1 could be improved and we will take it into account in future revisions of the manuscript. Good point, about Figure 6, we added annotations.

W6 Others

Thank you for your questions, we did our best to add more clarifications where it’s possible. However, there are some misunderstandings, which we should address directly:

“neuron is also a factor” - it’s not true, a factor is represented by a collection of segments, so we need to grow only segments to store new experiences. We hope that in the revised text it’s more clear. Populations of neurons are set up at the onset of learning and correspond to RV’s state space sizes.

“291,292, please explain this step” - we refer to the fact that if one of the terms in the sum under logarithm $\log(x_1+x_2+x_3+...)$ is much larger than others, then we can approximate this sum by taking a maximum instead. Indeed, let’s assume that $x_m = \mathrm{max}(x_1, x_2, x_3, …)$, then \log(x_1+x_2+x_3+...) =$$ \log x_m(1+x_1/x_m+x_2/x_m+x_3/x_m + ...) = $$\log x_m + \log (1+x_1/x_m+x_2/x_m+x_3/x_m + ...)$$ \approx \log \mathrm{max}(x_1, x_2, x_3, …), when $x_m >> x_1, x_2, x_3, ...$. Interestingly, that this operation is similar to logical OR which is used to aggregate segments' activity in the HTM neuron and also has been found in some biological neurons as noted in Stuart & Spruston, 2015.

Khetarpal, Khimya, Matthew Riemer, Irina Rish, and Doina Precup. 2022. ‘Towards Continual Reinforcement Learning: A Review and Perspectives’. Journal of Artificial Intelligence Research 75 (December):1401–76. https://doi.org/10.1613/jair.1.13673.

Stuart, Greg J., and Nelson Spruston. 2015. ‘Dendritic Integration: 60 Years of Progress’. Nature Neuroscience 18 (12): 1713–21. https://doi.org/10.1038/nn.4157.