Estimating the Intrinsic Dimension of Neural Manifolds from Variational Autoencoders

Neural population activity is often conceptualized as lying on low-dimensional manifolds. Yet estimating the intrinsic dimension of these manifolds is difficult in noisy, curved neural data, and existing estimators are often not reliable enough to serve as a scientific yardstick. We therefore investigate a method based on fitting a variational autoencoder (VAE) to the neural data and estimating dimension from the effective rank of the decoder Jacobian. We explain the effectiveness of this approach by linking the VAE to local probabilistic principal component analysis, showing that the variational objective penalizes decoders that use unnecessary latent dimensions. In neuroscience-inspired simulations with highly noisy and curved data, our method accurately recovers intrinsic dimension where alternative methods fail. Applied to population recordings from the mouse head direction system, it identifies two intrinsic dimensions. Further analysis suggests that one of these dimensions may reflect gain modulation, even in the absence of explicit experimental manipulation. Applied to population recordings from a grid cell module in mouse medial entorhinal cortex, our method identifies 4 dimensions, related to x and y position, head direction, speed, and gain. Since VAE models are common in neuroscience, our method can be applied to previously fit models. We demonstrate this on a VAE fit to juvenile male zebra finch songs over weeks of vocal practice, identifying an increase in dimensionality over time.