Emergent Generalization by Representation Learning in Artificial Neural Networks

Technical Summary: Emergent Generalization by Representation Learning in Artificial Neural Networks

Problem Statement
The paper addresses a longstanding debate in neuroscience regarding the functional relevance of "neural manifolds"—low-dimensional structures underlying high-dimensional neural activity. While dimensionality reduction techniques have successfully identified these manifolds across various brain regions and species, it remains contested whether these structures are merely descriptive epiphenomena of neuron-level activity or if they confer a causal, functional advantage for learning and generalization. Specifically, the authors investigate whether the ability of a learning system to extract low-dimensional representations from high-dimensional inputs provides a functional benefit for generalizing to out-of-distribution (OOD) data, and whether such representations exhibit "causal emergence."

Methodology
The authors employ a hybrid approach combining artificial neural network (ANN) simulations with the analysis of biological neural data.

1. Synthetic Task and Architecture:

   • Task: A time-series prediction task where the network must predict the next time step of a high-dimensional input $X_t$. This input is generated by projecting a low-dimensional latent dynamical system $x_t$ (e.g., Lorenz, Rössler, Sprott attractors) into a high-dimensional space via a random orthogonal projection $R$, with added noise.
   • Architecture: A modified Reservoir Computing (RC) architecture. It consists of a fixed recurrent reservoir layer, a trainable bottleneck layer (the focus of representation learning), and a linear ridge regression readout layer. The bottleneck layer is trained via backpropagation to learn a low-dimensional predictive latent representation, while the readout layer is trained to minimize Mean Squared Error (MSE).
   • Generalization Tests: The model is trained on a subset of dynamical systems and tested on: (a) random orthogonal rotations of the training attractors, and (b) held-out dynamical systems that are OOD relative to the training data.
   • Complexity Control: Task complexity is modulated by subsampling the time series to retain $N_{\tau}$ timesteps before the autocorrelation drops below $1/e$.

2. Information-Theoretic Measures:

   • Causal Emergence ($\Psi$): The authors utilize the $\Psi$ measure (Rosas et al., 2020) to quantify causal emergence. It is defined as the difference between the macro predictability (mutual information between the coarse-grained state $V_t$ and its future $V_{t+1}$) and the sum of micro predictabilities (sum of mutual information between individual reservoir neurons $Y^i_t$ and their futures). A positive $\Psi$ indicates the macro state is more predictive than the sum of its parts.
   • Information Bottleneck (IB): The IB loss ($I(Y_t; V_t) - \beta I(V_t; X_{t+1})$) is used to track the trade-off between compressing input information and preserving predictive information.

3. Biological Data Analysis:

   • Dataset: Neural spiking activity recorded from the dorsal CA1 and medial Prefrontal Cortex (PFC) of mice performing a W-maze spatial alternation task (8 mice, 8+ training sessions).
   • Analysis: Preprocessed firing rates were fed into a simple linear bottleneck encoder with a ridge decoder to predict the next timestep of behavioral variables (position and velocity). The $\Psi$ measure was estimated for the bottleneck states across training sessions to track its evolution relative to behavioral performance and prediction error.

Key Results

1. Necessity of Bottlenecks for Generalization:

   • Networks without a bottleneck layer could learn the training data but failed to generalize to random rotations of the training attractors or to held-out OOD dynamical systems.
   • The presence of a bottleneck layer enabled the network to learn a rotation-invariant representation of the latent states, allowing for successful generalization to both rotated training data and entirely new dynamical systems.
   • Generalization performance was robust across different dynamical systems and depended primarily on task complexity (controlled by $N_{\tau}$) rather than the specific attractor being learned.

2. Non-Monotonic Trajectory of Causal Emergence:

   • While the prediction loss (MSE) decreased monotonically during training, the measure of causal emergence ($\Psi$) exhibited a non-monotonic trajectory.
   • Phase 1 (Memorization): $\Psi$ initially decreased during early training.
   • Phase 2 (Transition): $\Psi$ reached a minimum.
   • Phase 3 (Generalization): $\Psi$ subsequently rose to a maximum during the generalization phase.
   • This trajectory mirrors the "grokking" phenomenon observed in other contexts, where a transition from memorization to generalization is marked by a structural change in the representation. The magnitude of this emergent structure reliably predicted generalization performance.

3. Correlation with Generalization:

   • The degree of causal emergence was positively correlated with the network's ability to generalize to OOD data.
   • Decomposition of the $\Psi$ measure revealed that generalization is associated with high macro predictability (stability of the representation) and low micro predictability (highly distributed representations where no single component predicts the macro).
   • The transition to high emergence preceded the minimization of the Information Bottleneck loss, suggesting that the emergence of structure is a precursor to optimal compression and generalization.

4. Biological Validation:

   • Analysis of mouse CA1 and medial PFC data revealed a similar non-monotonic behavior in the causal emergence measure during spatial learning.
   • The emergence measure decreased initially and then increased in later training sessions. Crucially, the minimum of the emergence measure significantly preceded the minimum of the prediction error (MSE) and the peak in behavioral performance, mirroring the dynamics observed in the artificial network.

Significance and Claims
The paper argues that the ability to learn compact, distributed, and emergent representations confers a functional advantage for generalization in learning systems. By demonstrating that causal emergence is a necessary condition for OOD generalization in a controlled synthetic setting, and observing analogous non-monotonic emergence dynamics in biological neural networks during learning, the authors support a causal role for learned representations in cognition.

The findings suggest that neural manifolds are not merely descriptive summaries but may represent causally emergent structures that enhance predictive power. The study posits that the learning of these emergent manifolds is a key feature of learning in biological neural networks, providing a functional explanation for why brains might evolve to organize high-dimensional activity into low-dimensional, emergent structures. The authors remain modest regarding mechanistic explanations, noting that while they identify the functional advantage, the specific plasticity mechanisms enabling this in biological brains remain an open question for future work.