A Wasserstein Geometric Framework for Hebbian Plasticity

TL;DR

This paper introduces a geometric framework for Hebbian plasticity using Wasserstein geometry, modeling memory states as evolving probability measures with a variational structure.

Contribution

It formalizes Hebbian learning as a Wasserstein gradient flow, connecting internal memory dynamics with observable neural quantities through geometric projections.

Findings

Memory states evolve along Wasserstein geodesics in a latent space.

Classical neural network schemes are shown as flat projections of the curved dynamics.

The framework accommodates complex distributional representations and spectral extensions.

Abstract

We introduce the Tan-HWG framework (Hebbian-Wasserstein-Geometry), a geometric theory of Hebbian plasticity in which memory states are modeled as probability measures evolving through Wasserstein minimizing movements. Hebbian learning rules are formalized as Hebbian energies satisfying a sequential stability condition, ensuring well-posed fiberwise JKO updates, optimal-transport realizations, and an energy descent inequality. This variational structure induces a fundamental separation between internal and observable dynamics. Internal memory states evolve along Wasserstein geodesics in a latent curved space, while observable quantities, such as effective synaptic weights, arise through geometric projection maps into external spaces. Simplicial projections recover classical affine schemes (including exponential moving averages and mirror descent), while revealing synaptic competition and pruning as geometric consequences of mass redistribution. Hilbertian projections provide a geometric account of phase alignment and multi-scale coherence. Classical neural networks appear as flat projections of this curved dynamics, while the framework naturally accommodates richer distributional representations, including structural weights and embedding memories, and their spectral extensions in complex internal spaces. Under mild Lipschitz regularity assumptions, including a quasi-stationary "sleep-mode" regime, we establish the existence of continuous-time limit curves. This yields a variational formulation of memory consolidation as a perturbed Wasserstein gradient flow. The framework thus provides a unified geometric foundation for synaptic plasticity, representation dynamics, and context-dependent computation.