Wasserstein Neurons in Neural Networks

• Wasserstein neurons are neural units analyzed using optimal transport metrics to quantify deviations from Gaussian pre-activations in transformer and Hebbian frameworks.
• They reveal a sparse subnetwork whose high Wasserstein distance values correlate with syntax-sensitive processing and sign-specific computational roles.
• Alternate perspectives interpret these neurons as measure-valued dynamical systems or tools for model fusion, integrating optimal transport geometry with neural computation.

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 (Kong et al., 29 Sep 2025). 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 (Tan, 17 Apr 2026). Related but broader Wasserstein-based neural constructions include mean-field neural networks for learning mappings on the Wasserstein space of probability measures (Pham et al., 2022) and node-as-measure/function interpretations used in Wasserstein barycenter-based model fusion (Akash et al., 2022). 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 (Kong et al., 29 Sep 2025). 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 (Tan, 17 Apr 2026). 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 $W_{\text{up}}$0 from $W_{\text{up}}$1 to $W_{\text{up}}$2, using either bin-density embeddings or cylindrical moment embeddings (Pham et al., 2022). 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 (Akash et al., 2022).

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,

$W_{\text{up}}$3

and in GLU style,

$W_{\text{up}}$4

Each row of $W_{\text{up}}$5 or $W_{\text{up}}$6 is treated as a neuron. Given activations $W_{\text{up}}$7, the pre-activations are $W_{\text{up}}$8, and the neuron is scored by

$W_{\text{up}}$9

Neurons with large WD are called Wasserstein neurons (Kong et al., 29 Sep 2025).

The paper also defines Mapping Difficulty: $W_{\text{gate}}$0 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 (Kong et al., 29 Sep 2025).

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 (Kong et al., 29 Sep 2025). The paper further reports that, for the top $W_{\text{gate}}$1 MD neurons in each layer, highly differentiated input pairs can be classified by the signs of $W_{\text{gate}}$2 into PP, PN, or NN; in Pythia layer 2, more than $W_{\text{gate}}$3 of the top-100 differentiated pairs are NN, and this NN fraction rises sharply within the first $W_{\text{gate}}$4K training steps (Kong et al., 29 Sep 2025).

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{gate}}$5 or $W_{\text{gate}}$6, normalizing them to zero mean and unit variance, computing WD for each neuron, and ranking neurons by WD. The top $W_{\text{gate}}$7, often $W_{\text{gate}}$8, are designated Wasserstein neurons (Kong et al., 29 Sep 2025). 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 $W_{\text{gate}}$9. If $y_i = w^\top x_i$0 indexes the top $y_i = w^\top x_i$1 WD neurons in a layer, the altered pre-activations are

$y_i = w^\top x_i$2

Thus only the negative pre-activations of the selected Wasserstein neurons are zeroed (Kong et al., 29 Sep 2025).

Two control conditions are used. The random control selects $y_i = w^\top x_i$3 neurons uniformly at random in each layer. The perplexity-matched control selects the $y_i = w^\top x_i$4 lowest-WD neurons so that zeroing their negatives yields the same final perplexity as the Wasserstein-group intervention (Kong et al., 29 Sep 2025). 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 $y_i = w^\top x_i$5 of negative pre-activations of Wasserstein neurons doubles perplexity; Qwen 8B requires $y_i = w^\top x_i$6 to double. Random controls remain near baseline (Kong et al., 29 Sep 2025). To match the perplexity increase caused by the WD group, the least-entangled-neuron control requires ablating approximately $y_i = w^\top x_i$7 of the least-entangled neurons in Mistral, $y_i = w^\top x_i$8 in Llama, and $y_i = w^\top x_i$9 in Qwen (Kong et al., 29 Sep 2025). On BLiMP and TSE, only the $x_1,\dots,x_N$0 WD-group intervention is reported to collapse grammatical accuracy, with scores dropping by up to $x_1,\dots,x_N$1–$x_1,\dots,x_N$2 in Llama and Mistral, while random and perplexity-matched controls remain near unperturbed levels (Kong et al., 29 Sep 2025).

At the token-class level, the added negative-log-likelihood from the $x_1,\dots,x_N$3 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 (Kong et al., 29 Sep 2025). 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 $x_1,\dots,x_N$4 WD neurons in only the first group already raises error sharply on ellipsis, subject–verb agreement, and negative polarity item licensing; cumulatively clamping groups $x_1,\dots,x_N$5 yields monotonic growth in BLiMP/TSE errors (Kong et al., 29 Sep 2025). 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 $x_1,\dots,x_N$6 WD at the final checkpoint acquire large WD within the first $x_1,\dots,x_N$7K training steps ($x_1,\dots,x_N$8B tokens) and then plateau, and that their weight directions specialize rapidly early and then stabilize, measured via cosine dissimilarity between successive $x_1,\dots,x_N$9K-step checkpoints normalized to the layer average (Kong et al., 29 Sep 2025). The mean WD of this cohort is reported to correlate strongly with TSE accuracy, with $\{y_i\}$0, and sign-specific ablation at each checkpoint yields BLiMP/TSE error increases that grow in lockstep with WD, unlike random or perplexity controls (Kong et al., 29 Sep 2025).

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 (Kong et al., 29 Sep 2025). 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 $\{y_i\}$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

$\{y_i\}$2

on a complete, geodesic Polish metric space $\{y_i\}$3. Given a target or global signal $\{y_i\}$4, the Hebbian energy may be taken as

$\{y_i\}$5

or more generally by integrating geodesically convex fiber energies $\{y_i\}$6 (Tan, 17 Apr 2026).

The update rule is a fiberwise JKO step: $\{y_i\}$7 Under AGS-type assumptions—$\{y_i\}$8 proper, lower semicontinuous, and geodesically $\{y_i\}$9-convex—each step admits a minimizer, the discrete sequence is well-defined, and the Energy Descent Inequality

$1$0

holds (Tan, 17 Apr 2026).

Observable synaptic quantities arise by projection. The framework specifies a simplicial projection, where leaf masses $1$1 define a point in $1$2, and a Hilbertian projection

$1$3

or $1$4, which recovers real or complex weight vectors depending on the choice of $1$5 and $1$6 (Tan, 17 Apr 2026). In the purely quadratic case, the JKO minimizer admits the closed-form interpolation

$1$7

where $1$8 is the optimal transport map from $1$9 to $1$0. Projecting to the simplex yields

$1$1

which is described as an exact exponential moving average, and if $1$2 is generated by a mirror-descent or softmax update, the projected observable dynamics recover the usual mirror-descent-in-the-simplex rule exactly (Tan, 17 Apr 2026).

The framework also attributes synaptic competition and pruning to simplex constraints, provides a geometric account of phase alignment when $1$3, 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 (Tan, 17 Apr 2026). 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

$1$4

In this usage, a Wasserstein neuron is therefore a variational dynamical system with a latent measure-valued state and a projected external weight update (Tan, 17 Apr 2026).

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 (Pham et al., 2022). The target is a map

$1$5

Two architectures are proposed.

The bin-density architecture approximates a continuous density $1$6 on a rectangular box $1$7 by a piecewise-constant density over a uniform grid of $1$8 bins of volume $1$9, with coefficients

$\mathcal N(0,1)$0

When $\mathcal N(0,1)$1 is unknown, $\mathcal N(0,1)$2 is estimated from samples by

$\mathcal N(0,1)$3

The pair $\mathcal N(0,1)$4 is then fed into a standard feed-forward network with output $\mathcal N(0,1)$5 (Pham et al., 2022).

The cylindrical architecture instead computes empirical moments of a learned test-function network $\mathcal N(0,1)$6,

$\mathcal N(0,1)$7

approximated in practice by Monte Carlo, and then outputs $\mathcal N(0,1)$8 (Pham et al., 2022). Theoretical support is provided by two universal-approximation theorems: Theorem 2.1 for the bin-density construction under continuity in $\mathcal N(0,1)$9 and $\mu \in \mathcal P_2(Y)$0 on compactly supported measures with continuous density, and Theorem 2.2 for the cylindrical construction in $\mu \in \mathcal P_2(Y)$1 for a training law $\mu \in \mathcal P_2(Y)$2 on measures (Pham et al., 2022).

Training minimizes an empirical $\mu \in \mathcal P_2(Y)$3-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 (Pham et al., 2022). In $\mu \in \mathcal P_2(Y)$4, $\mu \in \mathcal P_2(Y)$5 experiments with $\mu \in \mathcal P_2(Y)$6 training measures and up to $\mu \in \mathcal P_2(Y)$7 samples per measure, the paper reports that bin networks work well for smooth $\mu \in \mathcal P_2(Y)$8 but require large $\mu \in \mathcal P_2(Y)$9 for density estimation, cylindrical networks reach the same training error with smaller $W_{\text{up}}$00, and both methods struggle on non-smooth $W_{\text{up}}$01 for distributions with near-Dirac peaks (Pham et al., 2022). 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 $W_{\text{up}}$02 at $W_{\text{up}}$03 (Pham et al., 2022).

A separate but conceptually allied construction appears in Wasserstein barycenter-based model fusion. There, each layer $W_{\text{up}}$04 is equipped with the uniform discrete measure

$W_{\text{up}}$05

and each node $W_{\text{up}}$06 is interpreted as a pair

$W_{\text{up}}$07

so that a neuron becomes a point in a transport–$W_{\text{up}}$08 space $W_{\text{up}}$09 (Akash et al., 2022). 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 (Akash et al., 2022).

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 $W_{\text{up}}$10, yielding a sparse entangled subnetwork in transformers (Kong et al., 29 Sep 2025). Second, there is a dynamical-geometric layer, where a neuron is fundamentally a probability measure updated by JKO steps in $W_{\text{up}}$11 and projected to observable weights (Tan, 17 Apr 2026). 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 (Pham et al., 2022, Akash et al., 2022).

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 (Kong et al., 29 Sep 2025). 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 (Tan, 17 Apr 2026). The mean-field work targets learning mappings on Wasserstein space with universal approximation guarantees and applications to time-dependent mean-field problems (Pham et al., 2022). The barycenter-fusion work targets layer-wise model fusion and evidence for permutation-corrected linear mode connectivity (Akash et al., 2022).

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.