Understanding and controlling the geometry of memory organization in RNNs

TL;DR

This paper investigates the geometric restructuring in RNNs during abrupt learning phases in memory tasks, introducing a regularization method that accelerates training and promotes attractor formation, with implications for neuroscience and AI.

Contribution

It uncovers geometric restructuring phenomena in RNNs during abrupt learning and proposes a temporal consistency regularization to enhance training efficiency and attractor formation.

Findings

Geometric restructuring occurs before abrupt learning drops.

Regularization accelerates training and attractor formation.

Predictions for neuroscience and artificial network learning mechanisms.

Abstract

Training recurrent neural networks (RNNs) is a high-dimensional process that requires updating numerous parameters. Therefore, it is often difficult to pinpoint the underlying learning mechanisms. To address this challenge, we propose to gain mechanistic insights into the phenomenon of \emph{abrupt learning} by studying RNNs trained to perform diverse short-term memory tasks. In these tasks, RNN training begins with an initial search phase. Following a long period of plateau in accuracy, the values of the loss function suddenly drop, indicating abrupt learning. Analyzing the neural computation performed by these RNNs reveals geometric restructuring (GR) in their phase spaces prior to the drop. To promote these GR events, we introduce a temporal consistency regularization that accelerates (bioplausible) training, facilitates attractor formation, and enables efficient learning in strongly connected networks. Our findings offer testable predictions for neuroscientists and emphasize the need for goal-agnostic secondary mechanisms to facilitate learning in biological and artificial networks.