Whitening reveals stable neuronal embedding geometry
Abstract
Deep neural networks provide a powerful framework for learning functional embeddings of neurons from their responses to visual stimuli. These embeddings can reveal structure in neuronal function, but biological interpretations require the geometry of the embedding space to be reproducible across model fits. We show that a common assumption underly- ing distance-based comparisons is violated in Sensorium-like predictive models: the learned functional basis is not orthonormal. Consequently, computing Euclidean distances with a standard formula can distort relationships between neuronal embeddings. We show that post-hoc whitening of the learned functional basis corrects this distortion while leaving model predictions unchanged. Across models trained with different random seeds, whitening improves the consistency of local neighborhoods and global structure, measured by mutual kNN consistency, adjusted rand index, linear centered kernel alignment, and representational similarity analysis. These results demonstrate that the coordinate basis of learned neuronal representations is an important determinant of their geometry and that whitening provides a simple way to obtain more stable embedding spaces.