Neural population geometry and optimal coding of tasks with shared latent structure (2402.16770v2)

Abstract: Humans and animals can recognize latent structures in their environment and apply this information to efficiently navigate the world. However, it remains unclear what aspects of neural activity contribute to these computational capabilities. Here, we develop an analytical theory linking the geometry of a neural population's activity to the generalization performance of a linear readout on a set of tasks that depend on a common latent structure. We show that four geometric measures of the activity determine performance across tasks. Using this theory, we find that experimentally observed disentangled representations naturally emerge as an optimal solution to the multi-task learning problem. When data is scarce, these optimal neural codes compress less informative latent variables, and when data is abundant, they expand these variables in the state space. We validate our theory using macaque ventral stream recordings. Our results therefore tie population geometry to multi-task learning.

• The paper develops an analytical framework linking neural population geometry with optimal generalization performance in tasks sharing latent structure.
• It quantifies neural activity using key geometric measures—such as neural-latent correlation and spectral flattening—to predict performance across artificial and biological systems.
• It validates the theoretical predictions on both multi-layer perceptrons and macaque neural data, demonstrating improved disentangled representations and signal-noise separation.

Neural Population Geometry and Optimal Coding of Tasks with Shared Latent Structure

Introduction

This paper addresses the analytical theory connecting neural population geometry to generalization performance in multi-task learning scenarios. The authors posit that geometric properties of neural activity patterns, specifically four key measures, determine the ability to generalize across tasks that share a latent structure. Through this framework, they identify the mechanisms by which disentangled neural representations emerge optimally, facilitating the navigation and application of latent environmental variables to various tasks.

Theory of Multi-Task Learning

The paper examines the role of neural population geometry in supporting tasks reliant on a shared latent structure. Specifically, it develops a model where binary classification tasks are generated by stimuli formed from latent variables, subjected to linear separation in the latent space (Figure 1).

Figure 1: Schematic of the task and model setup using images from the d-sprites dataset as an example.

The generalization error of a linear readout tasked with classifying these stimuli is linked to the statistical properties of the neural activity. The analytical theory employs Gaussian model simplifications to approximate generalization error, validating these predictions against non-linear neural activations in artificial and biological data (Figures 2, 3).

Figure 2: Schematic of the geometric terms. Geometric patterns elicited by stimuli are shown with varying levels of correlation and factorization.

The model incorporates geometric aspects such as neural-latent correlation, signal-signal factorization, signal-noise factorization, and neural dimension, demonstrating that these measures can predict generalization error with high fidelity.

Optimal Representation of Latent Variables

An intriguing result from the theory is the emergence of disentangled representations as optimal codes for multi-task learning. The paper elucidates that optimal neural representations possess orthogonal subspaces where each latent direction corresponds to orthogonal neuronal directions (Figure 3).

Figure 3: Optimal representational geometry as a function of training samples and latent structure.

These optimal representations adaptively expand or compress less informative latent variables depending on the availability of training samples, affirming that higher-dimensional neural activity correlates with improved generalization as data increases. The eigenstructure of the neuron-neuron covariance thus evolves with the sample size, demonstrating significant spectral flattening in optimal conditions (Figure 4).

Figure 4: Theory predicts empirical generalization error in Gaussian model with power law covariance spectra.

Geometry of Multi-Task Learning in Artificial Networks

The study extends its analysis to non-linear MLPs, demonstrating the theory’s robustness even under non-Gaussian conditions. Validation against trained and random MLPs reveals consistent agreement between theoretical predictions and empirical generalizations (Figures 5, 6).

Figure 5: Theory predicts generalization error in random and trained MLPs.

Trained networks exhibited improved geometric organization through sequential layers, optimizing signal-noise and signal-signal factorization while raising dimensionality, particularly through relu layers, aligning with optimal spectral strategies (Figure 6).

Figure 6: Evolution of generalization error through training stages in MLPs.

Predicting Readout Performance in Biological Systems

Further empirical validation is conducted on macaque V4 and IT neural data, forming task labels from latent variables tied to visual stimulus categories. The application of the theoretical framework accurately predicts generalization error in biological neural responses (Figure 7).

Figure 7: Theory predicts multi-task error in macaque V4 and IT neural data.

Comparative analysis across raw pixels, V4, and IT reveals superior generalization performance in brain regions compared to raw data, with IT optimally factorizing signal directions relative to V4, consistent with improved signal-noise separation strategies.

Conclusion

This research delineates an analytical pathway from neural population geometry to multi-task learning efficacy. Through exploring the geometry of neural activity, the study articulates a coherent mechanism underlying both artificial and biological systems’ ability to generalize across shared latent tasks. The results not only provide insight into disentangled representation emergence as optimal learning codes but also hypothesize future developments in AI inspired by biological efficiency, offering predictive tools for decoding neural dynamics and facilitating targeted cognitive computational models.