optimal-transport spectroscopy deep-learning research

OptimalTransportforSpectralMatching:WhyWassersteinBeatsMSE

March 22, 2026 · 14 min read

When you compare two spectra, the standard loss function has a blind spot. Consider a single Gaussian peak at 1720 cm⁻¹ — the C=O stretch. Shift it by 1 cm⁻¹ (calibration drift) and shift it by 100 cm⁻¹ (a completely different functional group). Under mean squared error, both shifts produce massive pointwise loss. MSE doesn’t know that 1 cm⁻¹ and 100 cm⁻¹ are different amounts of “wrong.” It sees mismatched bins and penalizes them equally.

This matters because spectra are not just vectors of numbers — they are distributions on a metric space. Each bin sits at a specific wavenumber, and the distance between wavenumbers has physical meaning. The right loss function should respect this geometry. Optimal transport does exactly that, and Sinkhorn’s algorithm makes it fast enough to use inside a training loop.

The Problem with MSE

Mean squared error computes pointwise differences:

$\text{MSE}(\mu, \nu) = \frac{1}{N}\sum_{i=1}^{N}(\mu_i - \nu_i)^2$

Every bin is treated identically. A mismatch at bin $i$ is penalized the same regardless of whether the nearest matching peak is 1 bin away or 500 bins away. For spectra, this is pathological: a peak shifted by a single wavenumber creates loss at both the source position (where intensity dropped) and the destination (where it appeared), with no credit given for the shift being small.

mse_vs_shift.py

1

2

3

4

5

6

7

8

9

10

11

12

The gradient problem is the real issue. MSE’s gradient at the original peak position says “increase intensity here” and at the shifted position says “decrease intensity there.” It never says “slide the peak left by 5 cm⁻¹.” The optimizer sees two independent problems instead of one geometric correction. In practice, this means MSE-trained models produce blurry, averaged-out reconstructions — they learn to hedge by spreading intensity across nearby bins rather than committing to a precise peak position.

Spectra as Probability Distributions

Vibrational spectra have a natural interpretation as distributions. Intensity values are non-negative (absorbance or scattering intensity cannot be negative) and can be area-normalized to integrate to 1. The wavenumber axis is a metric space with the standard Euclidean distance: the “cost” of moving a unit of intensity from wavenumber $w_i$ to wavenumber $w_j$ is $|w_i - w_j|$.

This puts spectra in the Wasserstein space: the space of probability measures on a metric ground space. In this framing, comparing two spectra is equivalent to asking: what is the cheapest way to rearrange the intensity of one spectrum to match the other, where “cost” is proportional to how far each unit of intensity has to move?

This is the Monge-Kantorovich formulation of optimal transport — one of the deepest ideas in probability theory, and exactly the right mathematical framework for spectral comparison.

Wasserstein Distance: The Right Loss

The Wasserstein-$p$ distance between two distributions $\mu$ and $\nu$ on a metric space is:

$W_p(\mu, \nu) = \left(\inf_{\pi \in \Pi(\mu,\nu)} \int c(x,y)^p \, d\pi(x,y)\right)^{1/p}$

For discrete spectra over $N$ wavenumber bins, this becomes a linear program. Using the squared Euclidean cost $c(w_i, w_j) = |w_i - w_j|^2$, the 1-Wasserstein distance is:

$W_1(\mu, \nu) = \min_{\pi \in \Pi(\mu,\nu)} \langle C, \pi \rangle$

where $C_{ij} = |w_i - w_j|^2$ is the cost matrix and $\Pi(\mu,\nu) = \{\pi \geq 0 : \pi \mathbf{1} = \mu, \pi^T \mathbf{1} = \nu\}$ is the set of all valid transport plans — joint distributions with marginals $\mu$ and $\nu$.

The transport plan $\pi$ tells you exactly how much intensity to move from each source bin to each target bin. For a spectrum with a single shifted peak, the optimal plan is trivially simple: move all the mass from the old position to the new position, paying cost proportional to the shift distance.

wasserstein_1d.py

1

2

3

4

5

6

7

8

9

10

11

The Sinkhorn Algorithm

The 1D closed-form is elegant but limited. In practice we need Wasserstein distance as a differentiable loss function inside a training loop, often in a latent space where the “spectrum” is a $d$-dimensional embedding rather than a literal 1D signal. This means we need the general $N \times N$ optimal transport computation, which is an $O(N^3)$ linear program — far too slow for backpropagation.

The solution is entropic regularization. Instead of the exact OT problem, we solve:

$W_\varepsilon(\mu,\nu) = \min_{\pi \in \Pi(\mu,\nu)} \langle C, \pi \rangle - \varepsilon H(\pi)$

where $H(\pi) = -\sum_{ij} \pi_{ij} \log \pi_{ij}$ is the entropy of the transport plan and $\varepsilon > 0$ controls the regularization strength. The entropy term makes the problem strictly convex, and the optimal plan takes the form:

$\pi^* = \text{diag}(\mathbf{u}) \, K \, \text{diag}(\mathbf{v})$

where $K_{ij} = e^{-C_{ij}/\varepsilon}$ is the Gibbs kernel. The scaling vectors $\mathbf{u}$ and $\mathbf{v}$ are found by alternating normalization — the Sinkhorn iterations:

$\mathbf{u}^{(k+1)} = \frac{\mu}{K \mathbf{v}^{(k)}}, \qquad \mathbf{v}^{(k+1)} = \frac{\nu}{K^T \mathbf{u}^{(k+1)}}$

Each iteration is a matrix-vector product — fully parallelizable on GPU. Convergence typically takes 50–100 iterations for well-chosen $\varepsilon$.

sinkhorn_explorer.pyinteractive

ε1.00iters50shift5 cm⁻¹

WassersteinMSE

Wasserstein cost scales smoothly with shift distance.

The Epsilon Problem

The regularization parameter $\varepsilon$ controls the sharpness of the transport plan. This is where theory meets numerical reality:

• $\varepsilon \to 0$: the solution approaches exact OT. But the Gibbs kernel entries $K_{ij} = e^{-C_{ij}/\varepsilon}$ become exponentially small for distant bins. With $C_{ij} = (i - j)^2$ and $\varepsilon = 0.05$, a distance of just $|i - j| = 4$ gives $K_{ij} = e^{-16/0.05} = e^{-320} \approx 0$. The entire kernel matrix becomes numerically zero except on the diagonal.

• $\varepsilon \to \infty$: the entropy term dominates and the optimal plan approaches the product measure $\mu \otimes \nu$ — a uniform smear with no geometric information. The gradients become uninformative.

In Spektron’s training, we started with $\varepsilon = 0.05$ following the literature. In float64, this worked. When we enabled mixed-precision training (AMP) for speed, the Gibbs kernel entries fell below float16’s minimum representable value ($\sim 6 \times 10^{-8}$), producing NaN within 100 training steps.

The fix: increase $\varepsilon$ to 1.0 for our 128-dimensional embeddings, and use log-domain Sinkhorn for extra numerical safety:

$\log u_i^{(k+1)} = \log \mu_i - \text{logsumexp}_j\!\left(-\frac{C_{ij}}{\varepsilon} + \log v_j^{(k)}\right)$

This replaces $e^{-320}$ with $-320$ in log space — no underflow possible.

epsilon tuning

ε = 0.05 (float64 → AMP)

ε = 1.0 + log-domain

sinkhorn_stability.py

1

2

3

4

5

6

7

8

9

10

OT as a Reconstruction Loss

In Spektron’s pretraining objective, the model reconstructs masked spectral patches. The reconstruction loss needs three properties: differentiability (for backpropagation), geometric awareness (shifted peaks cost less than wrong peaks), and numerical stability under AMP.

Sinkhorn OT satisfies all three. The OT loss for a single patch:

$\mathcal{L}_{\text{OT}} = W_\varepsilon\!\left(\hat{s}_{\text{patch}},\, s_{\text{patch}}\right)$

In practice, pure OT converges slower than MSE because its gradients are smoother (a feature for geometry, a bug for speed). The hybrid loss combines both:

$\mathcal{L} = (1 - \alpha)\,\text{MSE} + \alpha\,W_\varepsilon$

Ablation over $\alpha$ shows $\alpha = 0.3$ balances convergence speed with geometric sensitivity.

loss_ablation.py

1

2

3

4

5

6

7

8

9

10

11

12

The peak position error reduction is the key result. MSE-trained models reconstruct approximate peak shapes but systematically blur the positions. Adding the Sinkhorn OT term sharpens position accuracy by nearly 3×, because the loss function now explicitly rewards moving predicted peaks toward their true positions rather than averaging over possible locations.

Beyond Loss Functions: Spectral Comparison

Optimal transport is not limited to training objectives. As a spectral similarity metric, it has immediate applications:

Library search. Given an unknown spectrum, find the closest match in a reference database. Wasserstein distance between the query and each reference spectrum produces more chemically meaningful rankings than cosine similarity, because it respects the fact that a near-miss in peak position is more informative than a near-miss in peak intensity.

Quality control. In process analytical technology (PAT), spectra are continuously monitored for deviations from a reference. OT-based control charts detect peak shifts at lower signal-to-noise ratios than MSE-based charts, because the shift signal is not diluted across all bins.

Calibration transfer. Different spectrometers produce systematically shifted spectra. The optimal transport plan between spectra from instrument A and instrument B directly encodes the calibration function — no parametric model required.

Related

This post is part of a series on the mathematical foundations behind Spektron. The spectral inverse problem post covers the group-theoretic constraints on spectral inversion. The state-space models post explains why sequence architectures outperform CNNs for spectral processing. The spectral identifiability post derives the information-theoretic limits of structure determination from spectra. Next: masked pretraining for scientific spectra — how BERT-style self-supervised learning creates general-purpose spectral representations, using OT as part of the reconstruction loss.