Nonlinear classification of neural manifolds with contextual information

TL;DR

The paper tackles nonlinear classification of neural manifolds by introducing context-dependent gating to enable nonlinear readouts, addressing the limitations of linear manifold capacity in early layers. It develops a tractable theoretical framework that models neural representations as manifolds embedded in high-dimensional space and gates decisions via 2^K context hyperplanes, yielding an exact capacity formula $\frac{1}{\alpha^*(K,\Phi,\gamma)}=\mathbb{E}_{y,\xi,R}[\max_{c}\min_{H_c} (1/P)\sum_{\mu} ||H^\mu_c-\xi^\mu_c||^2]$, capturing the interplay between manifold geometry, context correlations $\Phi$, and margin $\gamma$. The authors validate the theory on synthetic data (random points, spherical manifolds, Gaussian mixtures) and apply it to ResNet-50 features under supervised and SimCLR objectives, revealing progressive, context-driven untangling across layers and improvements when using informative context directions (e.g., principal components). This framework offers a scalable, data-driven way to quantify context-dependent computation across scales, datasets, and models, with potential applications to biological neural data and beyond.

Abstract

Understanding how neural systems efficiently process information through distributed representations is a fundamental challenge at the interface of neuroscience and machine learning. Recent approaches analyze the statistical and geometrical attributes of neural representations as population-level mechanistic descriptors of task implementation. In particular, manifold capacity has emerged as a promising framework linking population geometry to the separability of neural manifolds. However, this metric has been limited to linear readouts. To address this limitation, we introduce a theoretical framework that leverages latent directions in input space, which can be related to contextual information. We derive an exact formula for the context-dependent manifold capacity that depends on manifold geometry and context correlations, and validate it on synthetic and real data. Our framework's increased expressivity captures representation reformatting in deep networks at early stages of the layer hierarchy, previously inaccessible to analysis. As context-dependent nonlinearity is ubiquitous in neural systems, our data-driven and theoretically grounded approach promises to elucidate context-dependent computation across scales, datasets, and models.

Figures (8)

• Figure 1: Neural manifolds are composed by the collection of the neural responses elicited by the same concept, either "cat" (red) or "mouse" (blue) in this illustration. These two concepts are expressed by different representations according to the different contexts: auditory or visual signals (stimulus modality as different contexts), realistic or cartoon-like images (visual styles as different contexts).
• Figure 2: Three hyperplanes shatter the input space into different contexts, marked by different colors. Manifold shapes are ellipsoids, while labels are encoded by the black/white coloring. The vectors $\bm{r}_1$ and $\bm{r}_2$ define uncorrelated contexts ($\Phi_{12}\sim 0$), while $\bm{r}_1$ and $\bm{r}_3$ are highly correlated ($\Phi_{13}\sim 1$). (a) Correlations between manifold directions. (b) Correlation between manifold centers, where $\bf O$ denotes the origin. (c-d) Manifolds can be "cut" by context hyperplanes and lie into multiple contexts.
• Figure 3: Capacity as a function of the context correlation parameter $\phi$ in the case of uncorrelated random points. Curves for different number of contexts are depicted with different colors. Full lines mark our theoretical predictions from Eq. \ref{['eq:formula_randomPoints']}, symbols mark numerical simulations at $N=5000$.
• Figure 4: Capacity as a function of the manifold correlation, that for visibility purposes we take uniform $\Sigma^{\mu i}_{\nu j}=\sigma$ for all pairs of manifolds $(\mu,\nu)$ and directions $(i,j)$ for synthetic spherical manifolds. Subplot in different rows represent different values of context correlation $\phi\in [0,0.8]$, while different columns represent different latent dimension $D\in [5,15]$, embedded in ambient dimension $N=4000$. We consider $P=50$ spherical manifolds, and for each we draw $M=50$ points. Each panel depicts the capacity for $2^K=1,2,4,8,16,32$ contexts, represented by different colors. Full lines mark theoretical predictions while dots mark simulations.
• Figure 5: Capacity from Eq. \ref{['eq:formula_manifolds']} for spherical synthetic manifolds of latent dimension $D=10$, embedded in ambient dimension $N=8000$, different combinations of uniform manifold and context correlations $(\sigma,\Phi_{12})$, and four contexts.