---
title: Wasserstein Neurons in Neural Networks
url: https://www.emergentmind.com/topics/wasserstein-neurons
type: topic
---

# Wasserstein Neurons in Neural Networks

Searching arXiv for relevant papers on "Wasserstein neurons" and adjacent usages of the term.
Searching arXiv for: Wasserstein neurons
The expression **Wasserstein neurons** has been used in multiple, non-equivalent ways across recent literature. In one line of work, it denotes a sparse class of **entangled neurons** in transformer MLPs whose pre-activation distributions deviate strongly from Gaussianity, quantified by a one-dimensional Wasserstein distance to a standard Gaussian; these neurons are reported to be disproportionately important for syntax-sensitive language-model behavior [2509.24198]. In a second usage, the term refers to a neuron as a **measure-valued dynamical system** whose internal memory state evolves in Wasserstein space through variational proximal updates and is projected to observable synaptic weights [2604.16052]. Related but broader Wasserstein-based neural constructions include **mean-field neural networks** for learning mappings on the Wasserstein space of probability measures [2210.15179] and **node-as-measure/function** interpretations used in Wasserstein barycenter-based model fusion [2210.06671]. The shared theme is that a neuron, layer, or network component is analyzed or constructed using optimal-transport geometry, but the object being called “Wasserstein” differs substantially across these settings.

## 1. Terminological scope and principal usages

In the transformer-mechanistic usage, a neuron is taken to be a row \(w \in \mathbb R^d\) of \(W_{\text{up}}\) or \(W_{\text{gate}}\) in an MLP block, with pre-activations \(y_i = w^\top x_i\) over a corpus of token activations \(x_1,\dots,x_N\). The empirical distribution of \(\{y_i\}\) is normalized to zero mean and unit variance, and the neuron’s **Wasserstein Distance (WD)** is defined as the \(1\)D Wasserstein-\(1\) distance from that normalized empirical law to \(\mathcal N(0,1)\); neurons with large WD are called Wasserstein neurons [2509.24198]. In that paper, WD serves as a practical proxy for **entanglement**, and an alternate measure, **Mapping Difficulty (MD)**, is also defined from ratios of output-space to input-space pairwise distances.

In the Hebbian-plasticity usage, a Wasserstein neuron is specified by three ingredients: an internal memory \(\mu \in \mathcal P_2(Y)\), a variational update rule in Wasserstein space based on a JKO step, and an observable projection from the latent measure to weights or features [2604.16052]. The term therefore describes not a static feature detector but a neuron whose state is itself a probability measure evolving in a curved metric space.

Two related constructions do not center the exact phrase in the same way but are directly relevant. In **mean-field neural networks**, the network learns mappings \(\mu \mapsto V(\cdot,\mu)\) from \(\mathcal P_2(\mathbb R^d)\) to \(L^2(\mu;\mathbb R^p)\), using either bin-density embeddings or cylindrical moment embeddings [2210.15179]. In **Wasserstein barycenter-based model fusion**, each node is interpreted as a pair of a measure on the previous layer and a weight-function, placing neurons in a transport-based comparison space and enabling layer-wise barycenter fusion [2210.06671].

This multiplicity of meanings implies that “Wasserstein neuron” is not a single standardized object. A plausible implication is that the phrase should be interpreted contextually: in mechanistic interpretability it identifies a sparse empirical neuron subclass, whereas in geometric plasticity it names a model class with measure-valued latent dynamics.

## 2. Wasserstein neurons as entangled transformer neurons

The transformer definition is given for GPT-2-style and GLU-style MLPs. In GPT-2 style,
\[
h = W_{\text{down}}(\mathrm{GELU}(W_{\text{up}} h_0)),
\]
and in GLU style,
\[
h = W_{\text{down}}[\mathrm{SiLU}(W_{\text{gate}} h_0) \odot (W_{\text{up}} h_0)].
\]
Each row of \(W_{\text{up}}\) or \(W_{\text{gate}}\) is treated as a neuron. Given activations \(x_i\), the pre-activations are \(y_i = w^\top x_i\), and the neuron is scored by
\[
WD(w) = W_1(\text{Normalized}(\{y_i\}), \mathcal N(0,1)).
\]
Neurons with large WD are called Wasserstein neurons [2509.24198].

The paper also defines **Mapping Difficulty**:
\[
MD(w) = \frac1M \sum_{(i,j)} \frac{\|y_i-y_j\|}{\|x_i-x_j\|},
\]
where input distances are normalized by their maximum and output distances by their median. WD and MD are reported to correlate strongly, leading the authors to use WD for neuron selection in practice [2509.24198].

A central empirical claim is that, in GELU/SiLU-based models such as **Pythia 1.4B**, the non-Gaussian tail of high-WD neurons concentrates almost entirely in the **negative pre-activation region**, whereas in ReLU-based models such as **OPT 1.3B**, the outlying mass lies in the positive region [2509.24198]. The paper further reports that, for the top \(k\%\) MD neurons in each layer, highly differentiated input pairs can be classified by the signs of \((y_i,y_j)\) into PP, PN, or NN; in **Pythia layer 2**, more than \(50\%\) of the top-100 differentiated pairs are NN, and this NN fraction rises sharply within the first \(25\)K training steps [2509.24198].

The interpretation advanced there is functional rather than merely statistical: pairs of similar inputs, especially involving **syntactic tokens such as determiners and prepositions**, are driven to highly distinct negative values by this sparse subpopulation. This suggests a sign-specific computational specialization in which negative pre-activation magnitudes encode differentiations that are behaviorally consequential for grammar-sensitive processing.

## 3. Identification methodology and sign-specific interventions

The identification pipeline for entangled Wasserstein neurons proceeds by feeding a held-out corpus, specifically **WikiText-2**, through the model, collecting pre-activations for every neuron in \(W_{\text{up}}\) or \(W_{\text{gate}}\), normalizing them to zero mean and unit variance, computing WD for each neuron, and ranking neurons by WD. The top \(p\%\), often \(p=1\%\), are designated Wasserstein neurons [2509.24198]. MD can be computed optionally by sampling input pairs and averaging normalized output-to-input distance ratios.

The paper then studies a **sign-specific intervention** on the pre-activation vector \(a \in \mathbb R^n\). If \(S\) indexes the top \(p\%\) WD neurons in a layer, the altered pre-activations are
\[
a'_k =
\begin{cases}
\max(a_k,0), & k \in S,\\
a_k, & k \notin S.
\end{cases}
\]
Thus only the **negative** pre-activations of the selected Wasserstein neurons are zeroed [2509.24198].

Two control conditions are used. The **random** control selects \(|S|\) neurons uniformly at random in each layer. The **perplexity-matched** control selects the \(m\%\) lowest-WD neurons so that zeroing their negatives yields the same final perplexity as the Wasserstein-group intervention [2509.24198]. These controls are designed to distinguish a generic perturbation effect from a sign-specific effect concentrated in the high-WD subnetwork.

The reported outcomes are strong. In **Llama 3.1 8B** and **Mistral 7B**, clamping only \(2\%\) of negative pre-activations of Wasserstein neurons doubles perplexity; **Qwen 8B** requires \(5\%\) to double. Random controls remain near baseline [2509.24198]. To match the perplexity increase caused by the WD group, the least-entangled-neuron control requires ablating approximately \(50\%\) of the least-entangled neurons in Mistral, \(35\%\) in Llama, and \(20\%\) in Qwen [2509.24198]. On **BLiMP** and **TSE**, only the \(1\%\) WD-group intervention is reported to collapse grammatical accuracy, with scores dropping by up to \(30\)–\(40\%\) in Llama and Mistral, while random and perplexity-matched controls remain near unperturbed levels [2509.24198].

At the token-class level, the added negative-log-likelihood from the \(1\%\) WD ablation, after subtracting the perplexity-matched control, is reported to concentrate on **determiners, auxiliaries, particles, and punctuation**, with negligible change on **nouns, verbs, adjectives, and adverbs** [2509.24198]. This localization motivates the characterization of these neurons as supporting **syntactic scaffolding** rather than generic semantic content.

## 4. Layerwise behavior, training dynamics, and mechanistic interpretation

Layerwise intervention analysis in **Llama 3.1 8B** divides the model into eight groups of four layers. Ablating \(1\%\) WD neurons in only the first group already raises error sharply on **ellipsis**, **subject–verb agreement**, and **negative polarity item licensing**; cumulatively clamping groups \(1,\dots,k\) yields monotonic growth in BLiMP/TSE errors [2509.24198]. The reported interpretation is that small local degradations in early layers accumulate across depth into broad grammatical failure.

Training-dynamics analysis across **Pythia (70M–12B)** shows that larger models exhibit neurons with higher maximum WD, that Wasserstein neurons defined as the top \(1\%\) WD at the final checkpoint acquire large WD within the first \(\sim 25\)K training steps (\(\sim 50\)B tokens) and then plateau, and that their weight directions specialize rapidly early and then stabilize, measured via cosine dissimilarity between successive \(10\)K-step checkpoints normalized to the layer average [2509.24198]. The mean WD of this cohort is reported to correlate strongly with TSE accuracy, with \(r \approx 0.9\), and sign-specific ablation at each checkpoint yields BLiMP/TSE error increases that grow in lockstep with WD, unlike random or perplexity controls [2509.24198].

The synthesis offered in that work is that Wasserstein neurons form a **very sparse, highly entangled subnetwork** whose pre-activation distributions depart most strongly from Gaussianity, and that in GELU/SiLU models this departure resides almost exclusively in the **negative tail** [2509.24198]. Because the intervention zeroes only the negative region and leaves positive values untouched, the findings are presented as evidence that negative pre-activations are not simply an optimization side effect of smooth activations but a functional resource used for syntax.

A common misconception would be to treat these neurons as defined by optimal-transport operations internal to the forward pass. In this usage, the “Wasserstein” descriptor comes from the **diagnostic statistic** \(W_1(\cdot,\mathcal N(0,1))\) applied to empirical pre-activation distributions, not from replacing a conventional neuron with an OT layer. The neuron remains a standard linear unit inside the MLP; what changes is the criterion by which a sparse subset is identified and the interpretation attached to that subset.

## 5. Wasserstein neurons as measure-valued Hebbian dynamical systems

A distinct definition appears in the **Tan–HWG** framework, where a neuron’s memory is a probability measure
\[
\mu \in \mathcal P_2(Y), \qquad \int_Y d_Y(y,y_0)^2\,\mu(dy) < \infty,
\]
on a complete, geodesic Polish metric space \((Y,d_Y)\). Given a target or global signal \(\zeta \in \mathcal P_2(Y)\), the Hebbian energy may be taken as
\[
E(\mu,\zeta) = \frac{\alpha}{2} W_2^2(\mu,\zeta), \qquad \alpha>0,
\]
or more generally by integrating geodesically convex fiber energies \(\mathcal E_x(\mu_x,S_x)\) [2604.16052].

The update rule is a fiberwise JKO step:
\[
\mu_{k+1} = \arg\min_{\nu \in \mathcal P_2(Y)}
\left\{
E(\nu,\zeta) + \frac{1}{2\tau} W_2^2(\nu,\mu_k)
\right\}.
\]
Under AGS-type assumptions—\(E(\cdot,\zeta)\) proper, lower semicontinuous, and geodesically \(\lambda\)-convex—each step admits a minimizer, the discrete sequence is well-defined, and the **Energy Descent Inequality**
\[
E(\mu_{k+1},\zeta) + \frac{1}{2\tau}W_2^2(\mu_{k+1},\mu_k)
\le E(\mu_k,\zeta)
\]
holds [2604.16052].

Observable synaptic quantities arise by projection. The framework specifies a **simplicial projection**, where leaf masses \(\hat p_i = \mu(\{y_i\})\) define a point in \(\Delta^{M-1}\), and a **Hilbertian projection**
\[
G(\mu) = \int_Y V(y)\,\mu(dy) \in \mathbb R^d
\]
or \(\mathbb C^d\), which recovers real or complex weight vectors depending on the choice of \(V\) and \(Y\) [2604.16052]. In the purely quadratic case, the JKO minimizer admits the closed-form interpolation
\[
\mu_{k+1} = \bigl((1-t_\tau)\,\mathrm{Id} + t_\tau\,T_k\bigr)_\# \mu_k,
\qquad
t_\tau = \frac{\alpha\tau}{1+\alpha\tau},
\]
where \(T_k\) is the optimal transport map from \(\mu_k\) to \(\zeta_k\). Projecting to the simplex yields
\[
\hat p_{k+1} = (1-t_\tau)\,\hat p_k + t_\tau\,\hat q_k,
\]
which is described as an exact exponential moving average, and if \(\hat q_k\) is generated by a mirror-descent or softmax update, the projected observable dynamics recover the usual mirror-descent-in-the-simplex rule exactly [2604.16052].

The framework also attributes **synaptic competition and pruning** to simplex constraints, provides a geometric account of **phase alignment** when \(Y=Y_M \times \mathbb C\), and argues for **multi-scale coherence** in which embeddings update on a faster timescale and simplex weights consolidate more slowly under quasi-stationary activations, such as during sleep [2604.16052]. In the continuous-time limit, under mild Lipschitz assumptions, the discrete scheme converges along subsequences to a perturbed Wasserstein gradient flow and, on the simplex, to the ODE
\[
\dot{\hat p}(t) = -\alpha(\hat p(t)-\hat q(t)).
\]
In this usage, a Wasserstein neuron is therefore a variational dynamical system with a latent measure-valued state and a projected external weight update [2604.16052].

## 6. Related Wasserstein-based neuron and network constructions

The most directly related precursor is the **mean-field neural network** for learning mappings from the Wasserstein space of measures to functions [2210.15179]. The target is a map
\[
\mu \in \mathcal P_2(\mathbb R^d) \mapsto V(\cdot,\mu) \in L^2(\mu;\mathbb R^p).
\]
Two architectures are proposed.

The **bin-density** architecture approximates a continuous density \(\rho^\mu\) on a rectangular box \(\Omega \subset \mathbb R^d\) by a piecewise-constant density over a uniform grid of \(K\) bins of volume \(h=|\Omega|/K\), with coefficients
\[
p_k^\mu = \frac{\rho^\mu(x_k)}{\sum_{j=1}^K \rho^\mu(x_j)\,h},
\qquad
\hat \mu := (p_1^\mu,\dots,p_K^\mu)\in \Delta_K.
\]
When \(\rho^\mu\) is unknown, \(\hat\mu\) is estimated from samples by
\[
p_k^\mu \simeq \frac{\#\{X_i \in \mathrm{Bin}(k)\}}{N\cdot h}.
\]
The pair \((x,\hat\mu)\) is then fed into a standard feed-forward network with output \(\Phi(x,\hat\mu)\) [2210.15179].

The **cylindrical** architecture instead computes empirical moments of a learned test-function network \(\phi_\theta:\mathbb R^d \to \mathbb R^k\),
\[
m := \langle \phi_\theta,\mu\rangle = \int \phi_\theta(x')\,\mu(dx') \in \mathbb R^k,
\]
approximated in practice by Monte Carlo, and then outputs \(\Psi_\theta(x,m)\) [2210.15179]. Theoretical support is provided by two universal-approximation theorems: Theorem 2.1 for the bin-density construction under continuity in \(x\) and \(\mu\) on compactly supported measures with continuous density, and Theorem 2.2 for the cylindrical construction in \(L^2(\nu)\) for a training law \(\nu\) on measures [2210.15179].

Training minimizes an empirical \(L^2\)-loss over sampled training measures and sampled points from each measure, optimized by stochastic gradient descent with Adam; in the cylindrical case, back-propagation passes through both the inner network and the Monte Carlo averaging step [2210.15179]. In \(d=1\), \(p=1\) experiments with \(M=1000\) training measures and up to \(N=250\,000\) samples per measure, the paper reports that bin networks work well for smooth \(V\) but require large \(N\) for density estimation, cylindrical networks reach the same training error with smaller \(N\), and both methods struggle on non-smooth \(V\) for distributions with near-Dirac peaks [2210.15179]. For a toy semi-linear PDE on Wasserstein space, embedding these networks into four time-dependent algorithms led to the result that the **local–BSDE scheme with cylindrical networks** was most accurate, with **MSE \(\approx 5 \cdot 10^{-5}\) at \(t=0\)** [2210.15179].

A separate but conceptually allied construction appears in **Wasserstein barycenter-based model fusion**. There, each layer \(N_\ell\) is equipped with the uniform discrete measure
\[
\gamma_\ell = \frac1{|N_\ell|}\sum_{v\in N_\ell}\delta_v,
\]
and each node \(v\in N_\ell\) is interpreted as a pair
\[
v=(\gamma_{\ell-1},w)\in \mathcal M_1^+(N_{\ell-1})\times \mathbb R^{|N_{\ell-1}|},
\]
so that a neuron becomes a point in a transport–\(L^2\) space \(TL^2(N_{\ell-1})\) [2210.06671]. Layer-wise Wasserstein barycenter or Gromov–Wasserstein barycenter problems are then solved by alternating between entropic OT coupling updates, via Sinkhorn, and closed-form weight updates. The framework applies to fully connected networks, CNNs, ResNets, RNNs, and LSTMs, and is used both for model fusion and for analyzing **linear mode connectivity** after permutation alignment [2210.06671].

These adjacent works do not identify a sparse neuron subclass in the mechanistic sense. Instead, they use Wasserstein geometry to represent measures, compare neurons across models, or learn operators defined on measure spaces. The commonality is methodological—optimal transport is the organizing geometry—while the ontological status of the “neuron” differs in each case.

## 7. Conceptual synthesis and points of distinction

The literature grouped under the label **Wasserstein neurons** separates into at least three conceptual layers. First, there is a **diagnostic** layer, where standard neurons are ranked by how non-Gaussian their normalized pre-activations are under a Wasserstein metric to \(\mathcal N(0,1)\), yielding a sparse entangled subnetwork in transformers [2509.24198]. Second, there is a **dynamical-geometric** layer, where a neuron is fundamentally a probability measure updated by JKO steps in \(\mathcal P_2(Y)\) and projected to observable weights [2604.16052]. Third, there are **operator-learning and model-space** constructions, where Wasserstein space is the domain of the learned function or the geometry used to fuse neurons or layers [2210.15179; 2210.06671].

A frequent source of confusion is to conflate these strands because they share the same transport vocabulary. In the transformer study, the neuron is ordinary and the Wasserstein structure is imposed on an empirical distribution of scalar pre-activations. In the Tan–HWG framework, the neuron’s internal state is itself a measure in Wasserstein space. In mean-field neural networks, the object being learned is a function of a measure, not a new neuron primitive. In model fusion, a neuron is interpreted as a measure-plus-function object to permit barycentric alignment across trained networks.

Another possible misunderstanding is to assume that all “Wasserstein neurons” aim at better predictive accuracy or model compression. The available sources support more specific claims. The transformer work ties high-WD neurons to **syntax**, **negative pre-activations**, and **grammatical behavior** under sign-specific ablation [2509.24198]. The Tan–HWG work aims at a **geometric theory of Hebbian plasticity**, recovering EMA, mirror descent, consensus, and phase locking as projections of an underlying curved flow [2604.16052]. The mean-field work targets **learning mappings on Wasserstein space** with universal approximation guarantees and applications to time-dependent mean-field problems [2210.15179]. The barycenter-fusion work targets **layer-wise model fusion** and evidence for permutation-corrected **linear mode connectivity** [2210.06671].

Taken together, these works indicate that “Wasserstein neuron” is best understood not as a single canonical architecture but as a family of transport-geometric viewpoints on neural units and neural computations. This suggests that future usage of the term will remain field-dependent unless a stricter standardization emerges.

Source: https://www.emergentmind.com/topics/wasserstein-neurons