Paper: Rapid Learning without Catastrophic Forgetting in the Morris Water Maze (ICLR-2024)

Reviewer comment (Novelty):
It is not clear that the proposed task is of practical importance.

Author response:
Our proposed water maze task is designed for testing neuroscience tasks using machine learning methods to reduce the gap between biological and artificial neural networks. We believe that this can give important insights into overcoming catastrophic forgetting, something biological agents do not suffer from. Our results indicate that, in contrast to standard continual learning methods which significantly struggle with catastrophic forgetting, our neuro-inspired method demonstrates substantially improved performance.

Reviewer comment (Experiments):
Comparison is extremely limited and considers old methods such as EWC.

Author response:
Although EWC was published in 2016, it is still widely used due to its simplicity and fewer constraints (e.g., no need for replay buffer, multiple backbone) (Wołczyk et al. 2021, Wołczyk et al.  2022). We have also compared our method with other recent methods: DER, DER++, Experience Replay, A-GEM and included our results in the appendix Table 5.  

For convenience, here are the average success rates on all 5 environments after training on all 5, with the exception of our method, which we only trained our action-policy in the first environment and zero-shot test across the rest of the environments. 

| Training Scheme   | Average |
|-------------------|---------|
| Fine-Tune         | 19.5    |
| EWC (Kirkpatrick et al., 2017) | 23.2    |
| Replay Buffer     | 4.0     |
| DER (Buzzega et al., 2020) | 2.05    |
| DER++ (Buzzega et al., 2020) | 3.39    |
| A-GEM (Chaudhry et al., 2019) | 7.31    |
| ER (Rolnick et al., 2019) | 16.92   |
| Ours (MSN)        | 99.2    |



[1] Wołczyk, Maciej, et al. "Continual world: A robotic benchmark for continual reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 28496-28510.

[2] Wolczyk, Maciej, et al. "Disentangling transfer in continual reinforcement learning." Advances in Neural Information Processing Systems 35 (2022): 6304-6317.

Reviewer comment (Reproducibility):
The code is not provided which makes judgment about reproducibility challenging.

Author response:
We attached the repository in the revision. https://anonymous.4open.science/r/water_maze-ED17/

Reviewer comment (Novelty):
The major novelty that I see in this paper is building connections between biological systems and ML, however, it is not clear how relevant the proposed task is and how effective it will be in the context of CL.

Author response:
Our proposed water maze task is designed for testing neuroscience tasks using machine learning methods to reduce the gap between biological and artificial neural networks. We believe that this can give important insights into overcoming catastrophic forgetting, something biological agents do not suffer from. Our results indicate that, in contrast to standard continual learning methods which significantly struggle with catastrophic forgetting, our neuro-inspired method demonstrates substantially improved performance.

Reviewer comment (Experiments):
Given that CL is an extremely well-established field, it is unclear why the paper has not performed an evaluation according to precedent by including SOTA methods.

Author response:
Although EWC was published in 2016, it is still widely used due to its simplicity and fewer constraints (e.g., no need for replay buffer, multiple backbone) (Wołczyk et al. 2021, Wołczyk et al.  2022). We have also compared our method with other recent methods: DER, DER++, Experience Replay, A-GEM and included our results in the appendix Table 5.  

For convenience, here are the average success rates on all 5 environments after training on all 5, with the exception of our method, which we only trained our action-policy in the first environment and zero-shot test across the rest of the environments. 

| Training Scheme   | Average |
|-------------------|---------|
| Fine-Tune         | 19.5    |
| EWC (Kirkpatrick et al., 2017) | 23.2    |
| Replay Buffer     | 4.0     |
| DER (Buzzega et al., 2020) | 2.05    |
| DER++ (Buzzega et al., 2020) | 3.39    |
| A-GEM (Chaudhry et al., 2019) | 7.31    |
| ER (Rolnick et al., 2019) | 16.92   |
| Ours (MSN)        | 99.2    |



[1] Wołczyk, Maciej, et al. "Continual world: A robotic benchmark for continual reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 28496-28510.

[2] Wolczyk, Maciej, et al. "Disentangling transfer in continual reinforcement learning." Advances in Neural Information Processing Systems 35 (2022): 6304-6317.

Reviewer comment (Novelty):
I don't find this work compelling and mostly find it to be a proof-of-concept level work which is not clear whether it will be of practical relevance.

Author response:
Our proposed water maze task is designed for testing neuroscience tasks using machine learning methods to reduce the gap between biological and artificial neural networks. We believe that this can give important insights into overcoming catastrophic forgetting, something biological agents do not suffer from. Our results indicate that, in contrast to standard continual learning methods which significantly struggle with catastrophic forgetting, our neuro-inspired method demonstrates substantially improved performance.

Reviewer comment (Writing):
Enhancing the paper's readability, especially for readers less familiar with neuroscience, would improve its significance. This could involve providing clearer (maybe more intuitive) explanations of certain concepts, such as the entorhinal cortex, the neocortical-entorhinal-hippocampal circuit, the memory scaffold, Hebbian learning, and grid cell patterns.

Author response:
Memory scaffold is a memory structure with predefined fixed points that allow arbitrary sensory inputs to be remembered without catastrophic forgetting [1, 2]. 

Grid cells are specialized cells in the entorhinal cortex that encode spatial information in a hexagonal pattern [3]. This cell is known for its crucial role in spatial navigation [4, 5]. 

A neocortical-entorhinal-hippocampal circuit is a key neural circuit for spatial memory and navigation that combines three different components. The neocortex corresponds to sensory cells in MESH, the entorhinal cortex corresponds to grid cells and the hippocampal corresponds to place cells. Please refer to Sharma et al. (2022) [6] for details.

[1] Mulders, D., et al.. A structured scaffold underlies activity in the hippocampus. bioRxiv, 2021.

[2] Yim M. Y., et al., Place-cell capacity and volatility with grid-like inputs. eLife, 2021.

[3] Moser EI, Moser MB. A metric for space. Hippocampus. 2008

[4] Bush D, Barry C, Manson D, Burgess N. Using Grid Cells for Navigation. Neuron. 2015

[5] Moser EI, Kropff E, Moser MB. Place cells, grid cells, and the brain's spatial representation system. Annu Rev Neurosci. 2008

[6] Sharma S., et al. Content addressable memory without catastrophic forgetting
by heteroassociation with a fixed scaffold. In ICML, 2022

Reviewer comment (Writing):
The description of the Morris Water Maze environment lacks clarity. What are the observations and sensory-cells in Figure 2. Is the observation meant to represent the agent's view as high dimensional vector? What are place and grid cells, and what are their dimensionalities? Does the memory store grid cells? What are the differences between various environments, apart from variations in the goal positions?

Author response:
The observations are the markings on the wall in the agent’s field of view, and they are passed into the sensory cells. The memory only stores the grid cell representation for the goal location once the agent reaches out there. The agent’s location is not stored.
In each environment, the surrounding wall pattern is different as shown in Figure 1 (b). It leads to catastrophic forgetting of previous environments.

Reviewer comment (Reproducibility):
The description of the proposed method requires further clarity. An explanation of the design architecture for the policy is needed. Additionally, specify which parts of the MESH incorporate attention mechanisms. How is the entire network trained? Is it trained using Reinforcement Learning? Define the objective function for training. Does the statement "The policy requires no further training" imply that the displacement network also doesn't require training?

Author response:
The entire network is trained through supervised learning. The observations and place cells are learned through hebbian-like association without backpropagation. The random projections from place cells to grid cells give us the grid cells. Please refer Sharma et al. (2022) for more details on associating observations and grid cells. The fixed convolutional neural network that takes the current grid code and the goal grid code and produces the invariant representation is not trained. This representation is then passed into the attention mechanism, and this is associated with the correct action through supervision. Essentially we have two networks, one for getting the grid code and another to get the correct action. The observation to grid code is learned in each environment, however the action network is only trained in the first environment, and zero-shot in the rest.

Reviewer comment (Experiments):
The experimental evaluation lacks comprehensiveness. It would be valuable to compare the proposed method with more state-of-the-art continual learning approaches, such as ER [1] and A-GEN [2]. Moreover, in the paper [3], a strategy involving the replay of similar experiences is used for continual learning. It would be insightful to discuss the relationship between the proposed method and the paper [3] on memory replay.

Author response:
Our method does not rely on replaying previous trajectories as the methods mentioned. Using the high capacity of MESH, we are able to achieve high accuracy without having replay. In-addition, our method does not require any further training for the network mapping the grid-cell displacement to the action after training in the first environment, which is both more biological and more efficient. Other approaches such as the ones mentioned need training in every environment encountered. 

We have added 4 more modern methods for comparison in the appendix Table 5. We included DER, DER++, A-GEM, and Experience replay. For convenience, here are the average success rate on all 5 environments after training on all 5, with the exception of our method, which we only trained our action-policy in the first environment and zero-shot test across the rest of the environments. 

| Training Scheme   | Average |
|-------------------|---------|
| Fine-Tune         | 19.5    |
| EWC (Kirkpatrick et al., 2017) | 23.2    |
| Replay Buffer     | 4.0     |
| DER (Buzzega et al., 2020) | 2.05    |
| DER++ (Buzzega et al., 2020) | 3.39    |
| A-GEM (Chaudhry et al., 2019) | 7.31    |
| ER (Rolnick et al., 2019) | 16.92   |
| Ours (MSN)        | 99.2    |

Reviewer comment (Writing):
It is beneficial to provide clearer (maybe more intuitive) explanations of some concepts, such as neocortical-entorhinal-hippocampal circuit, memory scaffold, and grid cell patterns.

Author response:
Memory scaffold is a memory structure with predefined fixed points that allow arbitrary sensory inputs to be remembered without catastrophic forgetting [1, 2]. 

Grid cells are specialized cells in the entorhinal cortex that encode spatial information in a hexagonal pattern [3]. This cell is known for its crucial role in spatial navigation [4, 5]. 

A neocortical-entorhinal-hippocampal circuit is a key neural circuit for spatial memory and navigation that combines three different components. The neocortex corresponds to sensory cells in MESH, the entorhinal cortex corresponds to grid cells and the hippocampal corresponds to place cells. Please refer to Sharma et al. (2022) [6] for details.

[1] Mulders, D., et al.. A structured scaffold underlies activity in the hippocampus. bioRxiv, 2021.

[2] Yim M. Y., et al., Place-cell capacity and volatility with grid-like inputs. eLife, 2021.

[3] Moser EI, Moser MB. A metric for space. Hippocampus. 2008

[4] Bush D, Barry C, Manson D, Burgess N. Using Grid Cells for Navigation. Neuron. 2015

[5] Moser EI, Kropff E, Moser MB. Place cells, grid cells, and the brain's spatial representation system. Annu Rev Neurosci. 2008

[6] Sharma S., et al. Content addressable memory without catastrophic forgetting
by heteroassociation with a fixed scaffold. In ICML, 2022

Reviewer comment (Presentation):
What are the observations and sensory cells in Figure 2? Does the memory store grid cells? What are the differences between various environments, apart from variations in the goal positions?

Author response:
The observations are the markings on the wall in the agent’s field of view, and they are passed into the sensory cells. The memory only stores the grid cell representation for the goal location once the agent reaches out there. The agent’s location is not stored.
In each environment, the surrounding wall pattern is different as shown in Figure 1 (b). It leads to catastrophic forgetting of previous environments.

Reviewer comment (Reproducibility):
How is the entire network trained? Is it trained through Reinforcement Learning? What is the training objective? Additionally, specify which part of the network utilizes attention mechanisms.

Author response:
The entire network is trained through supervised learning. The observations and place cells are learned through hebbian-like association without backpropagation. The random projections from place cells to grid cells give us the grid cells. Please refer Sharma et al. (2022) for more details on associating observations and grid cells. The fixed convolutional neural network that takes the current grid code and the goal grid code and produces the invariant representation is not trained. This representation is then passed into the attention mechanism, and this is associated with the correct action through supervision. Essentially we have two networks, one for getting the grid code and another to get the correct action. The observation to grid code is learned in each environment, however the action network is only trained in the first environment, and zero-shot in the rest.

Reviewer comment (Novelty):
Consider highlighting the unique aspects of the proposed method compared to replay-based approaches, such as ER [1], A-GEM [2], and Feedforward [3].

Author response:
Our method does not rely on replaying previous trajectories as the methods mentioned. Using the high capacity of MESH, we are able to achieve high accuracy without having replay. In-addition, our method does not require any further training for the network mapping the grid-cell displacement to the action after training in the first environment, which is both more biological and more efficient. Other approaches such as the ones mentioned need training in every environment encountered. 

We have added 4 more modern methods for comparison in the appendix Table 5. We included DER, DER++, A-GEM, and Experience replay. For convenience, here are the average success rate on all 5 environments after training on all 5, with the exception of our method, which we only trained our action-policy in the first environment and zero-shot test across the rest of the environments. 

| Training Scheme   | Average |
|-------------------|---------|
| Fine-Tune         | 19.5    |
| EWC (Kirkpatrick et al., 2017) | 23.2    |
| Replay Buffer     | 4.0     |
| DER (Buzzega et al., 2020) | 2.05    |
| DER++ (Buzzega et al., 2020) | 3.39    |
| A-GEM (Chaudhry et al., 2019) | 7.31    |
| ER (Rolnick et al., 2019) | 16.92   |
| Ours (MSN)        | 99.2    |

Reviewer comment (Experiments):
The empirical evaluation is the main weakness of the paper, as the authors compare their method to only two continual learning baselines, both of which are quite old.

Author response:
While we don't disagree that the two baseline methods are a bit dated, EWC is still one of the main continual learning methods evaluated in the current literature [1, 2]. The replay buffer size we find is relatively proportionate to the environment, but as we show, increasing the replay buffer size doesn't have a significant difference. 

We have added 4 more modern methods for comparison in the appendix Table 5. We included DER, DER++, A-GEM, and Experience replay. For convenience here are the average success rate on all 5 environments after training on all 5, with the exception of our method, which we only trained our action-policy in the first environment and zero-shot test across the rest of the environments. 

| Training Scheme   | Average |
|-------------------|---------|
| Fine-Tune         | 19.5    |
| EWC (Kirkpatrick et al., 2017) | 23.2    |
| Replay Buffer     | 4.0     |
| DER (Buzzega et al., 2020) | 2.05    |
| DER++ (Buzzega et al., 2020) | 3.39    |
| A-GEM (Chaudhry et al., 2019) | 7.31    |
| ER (Rolnick et al., 2019) | 16.92   |
| Ours (MSN)        | 99.2    |

Reviewer comment (Experiments):
The replay buffer sizes used in the experiments are rather small (200-1800).

Author response:
The replay buffer size is usually set to 2,000 in iCIFAR10, iCIFAR100, and ImageNet100 for 5 or 10 tasks [1, 2, 3]. This is almost similar to the maximum size of the replay buffer we showed in Figure 6 as the reviewer noticed. Appendix A.2 shows that the larger memory size does not improve the performance.

[1] Rebuffi, Sylvestre-Alvise, et al. "icarl: Incremental classifier and representation learning." CVPR. 2017.

[2] Yan, Shipeng, Jiangwei Xie, and Xuming He. "Der: Dynamically expandable representation for class incremental learning." CVPR. 2021.

[3] Zhou, Da-Wei, et al. "A model or 603 exemplars: Towards memory-efficient class-incremental learning." ICLR 2023.

Reviewer comment (Novelty):
The proposed method seems to be custom-designed for the navigation task, so while it can serve as a model of how rodents learn to navigate in new environments, it is not a practical continual learning method that could be applied to an arbitrary task.

Author response:
Our method is specifically designed for spatial navigation tasks, rather than being a versatile continual learning tool applicable to any arbitrary task. The aim is to bridge the gap between natural and artificial neural networks, focusing on the specialized area of spatial navigation.

Reviewer comment (Experiments):
To give existing baselines a fairer chance, the reviewer recommends having two separate networks: one mapping observations and goal position into some latent representation, and another mapping these to actions. In the first environment, train both networks. For each new environment, re-train only the first network.

Author response:
In our method, we actually only train the association between the latent representation and action in one environment, this network is zero-shot tested in all other environments, where, exactly as you mention, the "first network" (finding the latent representation of the displacement between the goal and current position) is the only one re-trained in multiple environments.

Reviewer comment (Reproducibility):
It makes sense that the policy network is invariant across tasks, since given true localization and goal location it only needs to learn a good search algorithm. But the place cell and grid cell may still suffer from catastrophic forgetting. Does the model get another set of newly initialized place cells and grid cells when switching to a different environment? Otherwise, how does it prevent forgetting? I would appreciate further clarification on this part.

Author response:
No, place cells and grid cells are not initialized for each environment. When initializing the agent, we initialize one set of place cells and grid cells that the agent uses for all environments. The high capacity of the MESH network allows for multiple environments, and when the agent is placed in a new environment, only the associations between the sensory cells and places cells are learned through Hebbian-like associations. The projections from the grid cell and place cell are predefined random projections that are initialized once during the initialization of the agent. Critically, this setup *prevents place cells and grid cells from catastrophic forgetting*. Even as more and more sensory-grid cell associations are learned, the network retains separate grid and place representations for distinct sensory inputs: it avoids interference.

Reviewer comment (Experiments):
I understand that the proposed method is tailored to a path finding task, however, to make it more generalizable, it would be better to test on other types of maze tasks (perhaps with more complex visual features and map topology). Moreover, I don’t see why the task needs to be a water maze (with no walls) vs. a real maze.

Author response:
We hope to test more types of maze tasks in the future, but we chose Morris WaterMaze as an example for 2 primary reasons. First, this is a task that standard deep learning and continual learning methods struggle with as we show experimentally. Second, this task is a popular neuroscience task that remains to be modeled. [1, 2] On this neuroscience task, we find that a structured neural model inspired by biological circuits can outperform standard deep learning. This task is also difficult in the sense that the agent must localize based solely on the markings of the walls: the agent receives no additional location-related signals. Note that the observation is not a bird-eye view of the 2D map.

Reviewer comment (Reproducibility):
I would appreciate more clarity on model training and loss functions. An algorithm block can also strengthen the presentation clarity of the paper.

Author response:
We have added a new algorithm block in the appendix. We do not use any gradient descent in our method, rather we are using heteroassociation between sensory input and grid cell representation. For the baselines, we use cross entropy loss with the ground truth action.