Clustering Market Regimes Using the Wasserstein Distance

TL;DR

Markets switch between “regimes” (calm bull runs, turbulent crashes), and detecting those switches matters for trading, risk, and — relevant to anyone shipping ML in production — knowing when a deployed model is operating on data it was never trained for. This paper reframes regime detection as a clustering problem on the space of probability distributions: chop the return series into overlapping windows, treat each window as an empirical distribution, and run k-means where “distance” is the Wasserstein (optimal-transport) distance and “average” is the Wasserstein barycenter. In 1D these reduce to dirt-simple operations on sorted samples (sort, subtract, average — no measure theory needed at runtime). On real S&P data the method tags every historically known crisis including subtle ones (Eurozone 2010, US downgrade 2011, China 2015) that the moment-based baseline misses; on synthetic data with known regime switches it hits ~91% accuracy on non-Gaussian (jump-diffusion) returns where the moment baseline drops to 67% and the HMM collapses to detecting nothing. The headline: a small, principled metric swap buys you robustness to fat tails and outliers essentially for free.

Problem & Motivation

Financial returns are not a stationary stream. They cluster into periods of similar behavior — low-volatility uptrends, high-volatility selloffs — and then abruptly shift. The practical pain of detecting these shifts shows up in three places the authors call out: (1) making trading/allocation decisions, (2) classical risk management, and (3) model governance for ML systems — a regime shift is the signal that your pricing/hedging/generation model is now seeing out-of-distribution data and needs retraining. That third point is the one to internalize if you run AI systems: regime detection is drift detection wearing a finance hat.

The existing toolkit falls short:

• Technical analysis (moving-average crossovers, support breaks) is ad hoc and asset-specific.
• Hidden Markov Models (HMMs), the classical workhorse since Hamilton (1989), bake in two strong assumptions: the hidden regime is Markovian, and returns given a regime follow a parametric (usually Gaussian) distribution. Real returns are famously non-Gaussian — fat tails, skew, jumps — so a Gaussian-likelihood HMM is fighting the data.
• Moment-based clustering (summarize each window by its first few moments — mean, variance, skew, kurtosis — then cluster those vectors) throws away the shape of the distribution and gets dominated by a handful of extreme moments.

The core insight: each window of returns is a distribution. So why compress it into 2-4 numbers (moments) or force it through a Gaussian (HMM) before clustering? Cluster the distributions themselves, with a metric that actually respects distributional geometry.

What’s New (Core Contribution)

This is a method paper. The novelty is precise and modest-sounding but consequential:

• Reframe regime detection as clustering on P_p(R), the space of probability measures with finite p-th moment — not on Euclidean feature vectors. Before: cluster points in R^d (moments) or fit a parametric latent-state model (HMM). Now: each data point is an empirical distribution, and clustering happens directly in distribution space.
• Use the p-Wasserstein distance as the clustering metric. Before: Euclidean distance between moment vectors, or KL/Kolmogorov-Smirnov (the authors argue both are inadequate — KS lacks sensitivity, KL needs density estimation and lacks a tractable averaging operation). Now: optimal-transport distance, which “metrizes weak convergence” (distributions close in Wasserstein are genuinely close in the everyday sense) and, crucially, has a natural averaging operation.
• Use the Wasserstein barycenter as the centroid-update step. This is the part that makes k-means work in distribution space — you need a way to “average” a cluster of distributions into a single representative one. The barycenter provides it, and in 1D it’s just the coordinate-wise median of sorted samples.
• Show the whole thing is computationally trivial in 1D. The 1D Wasserstein distance between two N-atom empirical measures has a closed form: sort both, take the average of |α_i − β_i|^p. That’s O(N log N) (dominated by the sort), no optimization solver required.

What is not new: k-means itself, the Wasserstein distance, the barycenter, and using OT for clustering in other domains (images, documents). The contribution is the combination and the application to model-free financial regime detection, plus a careful validation methodology using the Maximum Mean Discrepancy (MMD).

How It Works (Technically)

The method, called WK-means (Wasserstein k-means), is ordinary k-means with two components swapped out: the distance function and the centroid-averaging function. Everything else — assign-to-nearest, recompute centroids, repeat until convergence — is unchanged. The genius is entirely in those two swaps, so let’s demystify them.

Step 0 — From a price path to a bag of distributions. Take a price series S = (s_0, ..., s_N). Compute log-returns r_i = log(s_{i+1}) − log(s_i). Now slice the return series into overlapping windows. Two hyperparameters control this: h_1 = window length (how many returns per window), h_2 = step/overlap offset. With (h_1, h_2) = (35, 28) on hourly data, each window is ~a trading week and consecutive windows are offset by ~a day, so they heavily overlap. Each window of h_1 returns becomes an empirical measure μ_i — literally just “the set of those return values, treated as a distribution that puts mass 1/N on each value.” After lifting you have a family K = {μ_1, ..., μ_M} of M distributions. This family is what gets clustered.

Step 1 — The distance: p-Wasserstein in 1D. The general definition (their Eq. 14) looks intimidating — it’s the minimum over all “transport plans” of the cost of moving mass from one distribution to the other:

W_p(μ,ν)^p = min over couplings P of ∫ d(x,y)^p dP(x,y)

In plain English: imagine μ is a pile of dirt and ν is a hole shaped differently. The Wasserstein distance is the minimum total work (mass × distance moved) to reshape the pile into the hole. That’s why it’s called “earth mover’s distance.” It respects geometry: moving mass a little costs a little, moving it far costs a lot — unlike KL divergence, which only cares about overlap of support.

Here’s the part that makes it usable: in one dimension, that minimization has a closed-form answer. For two empirical measures with the same number of atoms N, sort each one’s values ascending into (α_1 ≤ ... ≤ α_N) and (β_1 ≤ ... ≤ β_N). Then (their Eq. 21):

W_p(μ,ν)^p = (1/N) Σ_i |α_i − β_i|^p

That’s it. Sort both samples, pair them up in order, average the p-th power of the gaps. No optimization, no solver. The optimal transport plan in 1D is always “match smallest-to-smallest, largest-to-largest” (monotone matching), which is why sorting solves it. For p=1 it’s mean absolute difference of sorted samples; for p=2 it’s RMS of sorted-sample gaps. The whole cost is the O(N log N) sort.

Step 2 — The centroid: Wasserstein barycenter in 1D. k-means needs to “average” the distributions assigned to a cluster into one representative distribution. The Wasserstein barycenter (their Eq. 16) is the distribution μ minimizing total Wasserstein distance to all cluster members. Sounds hard. In 1D with p=1 and equal atom counts it collapses (their Prop. 2.6) to: stack all cluster members’ sorted samples and take the coordinate-wise median:

a_j = median(α_j^1, α_j^2, ..., α_j^M) for each rank position j

So the barycenter’s j-th smallest atom is the median of all the j-th smallest atoms across the cluster. Sort each member, line them up by rank, take medians down the columns. (For p=2 you’d take means instead of medians.) Using the median is what makes WK-means robust to outlier windows — a single crazy crash window can’t drag the centroid the way a mean would.

Step 3 — The loop. With those two pieces, k-means runs as usual:

flowchart LR
  A[Price path S] --> B[Log-returns r_i]
  B --> C[Lift: overlapping windows<br/>h1=window, h2=offset]
  C --> D[Family K of empirical<br/>distributions mu_i]
  D --> E[Init k centroids<br/>sample from K]
  E --> F{Assign each mu_j to<br/>nearest centroid<br/>via 1D Wasserstein}
  F --> G[Update each centroid =<br/>Wasserstein barycenter<br/>coordinate-wise median]
  G --> H{loss = sum W_p of<br/>old vs new centroids<br/>< tolerance?}
  H -- no --> F
  H -- yes --> I[k regime clusters<br/>+ centroid distributions]

The loss (their Eq. 23) is just the total Wasserstein movement of the centroids between iterations; when centroids stop moving (< ε), stop. Output: each window is labeled with a cluster (regime), and each cluster has a centroid distribution describing that regime.

Interactive: two 1D empirical distributions and the optimal transport that turns one into the other. Drag the slider to morph the right distribution; the bars show the sorted-sample matching and the running 1-Wasserstein cost (the average gap between paired sorted samples). This is the entire distance computation WK-means uses — schematic, built to teach the monotone-matching trick.

Validation — the MMD self-similarity score. On real data you don’t know the “true” regimes, so how do you prove the clustering is good? The authors use a second distributional metric, Maximum Mean Discrepancy (MMD), from the two-sample-test literature. MMD measures how distinguishable two samples are by embedding them through a kernel (Gaussian kernel here) and comparing means in that feature space — if two clusters are truly different regimes, the between-cluster MMD should be large; if a cluster is internally coherent, the within-cluster MMD (their “self-similarity score,” the median MMD over pairwise sub-samples) should be small. They report WK-means within-cluster scores of 0.0395 / 0.2304 vs the moment method’s 0.0631 / 1.1961 — the moment method’s second cluster is wildly heterogeneous because it just dumps outliers together.

The algorithm, simplified

import numpy as np

def wasserstein1d(mu, nu, p=1):
    # mu, nu: 1D arrays of N return values (an empirical distribution each)
    a, b = np.sort(mu), np.sort(nu)              # monotone matching IS the optimal plan in 1D
    return (np.mean(np.abs(a - b) ** p)) ** (1 / p)

def barycenter(cluster, p=1):
    # cluster: list of equal-length return windows; centroid = "average distribution"
    sorted_members = np.sort(np.array(cluster), axis=1)   # sort each window ascending
    agg = np.median if p == 1 else np.mean               # median (p=1) -> robust to outliers
    return agg(sorted_members, axis=0)                    # column-wise: j-th atom from j-th atoms

def wk_means(windows, k=2, tol=1e-6, p=1, max_iter=100):
    # windows: M empirical distributions (each a window of log-returns)
    centroids = [windows[i] for i in np.random.choice(len(windows), k, replace=False)]
    for _ in range(max_iter):
        # assign each window to nearest centroid under 1D Wasserstein
        labels = [min(range(k), key=lambda c: wasserstein1d(w, centroids[c], p))
                  for w in windows]
        new = [barycenter([windows[i] for i in range(len(windows)) if labels[i] == c], p)
               for c in range(k)]
        loss = sum(wasserstein1d(centroids[c], new[c], p) for c in range(k))  # Eq. 23
        centroids = new
        if loss < tol:                                    # centroids stopped moving
            break
    return labels, centroids

Read it once: the only “exotic” lines are np.sort (that’s the distance) and np.median(..., axis=0) (that’s the barycenter). Everything else is vanilla k-means.

Built on Prior Work

Results & Evidence

Real data (SPY hourly, 2005–2020, k=2, window≈1 week). Both WK-means and the moment baseline flag the 2008 Global Financial Crisis and the 2020 COVID crash. But only WK-means additionally isolates the subtler stress periods: the 2010 Eurozone/Greek debt crisis onset, the 2011 US credit-rating downgrade, and the 2015/16 China crash. In mean-variance scatter plots WK-means cleanly separates by variance (the natural risk axis); the moment method mostly just isolates outliers. The MMD self-similarity scores (Table 2) confirm WK-means clusters are far more internally coherent.

Synthetic — Gaussian (geometric Brownian motion), 50 runs:

When the data is genuinely Gaussian, the moment method has a slight edge on total/regime-off accuracy (expected — a Gaussian is fully described by mean+variance, so moments lose nothing). But even here WK-means is markedly better at the part that matters, detecting the regime change (regime-on: 87% vs 75%). The HMM is the worst, failing to resolve changes at this parameter granularity.

Synthetic — non-Gaussian (Merton jump diffusion), 50 runs:

This is the money table. With jumps and fat tails — i.e., real market behavior — WK-means barely degrades (91% vs 91%), while the moment method collapses (regime-on 27%) and the HMM’s “75% total” is a mirage: it scores 0.66% on regime-on because it never detects the alternate regime at all and dumps everything into one bucket (hence ~100% regime-off). Runtimes are all ~1 second for the whole experiment, so WK-means buys robustness at no real compute cost.

What the evidence does NOT establish. This is a controlled study, and you should read it as such: (1) all real-data experiments use a single asset (SPY) and k=2 for simplicity — no multi-asset or multi-regime validation on real data; (2) on real data there’s no ground truth, so “it found the crises” is qualitative pattern-matching, not a hard accuracy number; (3) the synthetic generators (gBm, Merton) are the authors’ choice and are favorable to a distribution-shape method; (4) k=2 is assumed, not learned — choosing the number of regimes is left to future work; (5) the method is a posteriori (it clusters history). The lag/reactivity in a live streaming setting is discussed only via the hyperparameter h_1 (window length) trade-off, not rigorously benchmarked. None of this is damning — it’s an honest, well-scoped method paper — but “91% accurate” is on their synthetic data, not a promise about your live book.

How You’d Use It

For an AI services company, the most durable framing is regime detection = distribution-drift detection, and this paper hands you a clean, dependency-light, model-free detector. Concrete slots:

• Model-governance / drift monitoring as a productized capability. Any client running an ML model in production (forecasting, pricing, recommendation, fraud) has the same problem the paper names: when does the live data distribution shift enough that the model is now extrapolating? Wrap WK-means around your feature stream: window it, cluster the windows’ distributions, and alarm when new windows start landing in a different cluster (or far from all centroids). This is a far more principled “is my model still valid?” trigger than monitoring raw accuracy (which you often can’t compute live without labels). The MMD self-similarity score doubles as a quantitative drift signal.
• A regime overlay for any quant/trading client. It’s ~50 lines, no training, no GPU, runs in milliseconds. You can offer it as a feature that conditions strategy parameters (size down in the high-variance cluster, up in the calm one) or as a risk dashboard that colors the equity curve by regime.
• A labeling/segmentation primitive in a larger agentic pipeline. If you have an agent that reasons over time-series (a “market analyst” agent, a “system-health analyst” agent), WK-means gives it a cheap, explainable tool call: detect_regimes(series) -> labeled_segments + regime_centroids. The centroid distributions are interpretable artifacts an LLM can summarize (“regime 2 is high-variance, negatively-skewed”), which beats a black-box latent state.
• Generalizes beyond finance. Nothing here is finance-specific. Any sensor/telemetry/log stream where “the distribution of values shifted” matters — server latency, IoT, manufacturing QC, API traffic — is a candidate. That’s the real moat: it’s a general distributional change detector with a one-page implementation.

The build-vs-buy call is firmly build — there’s no product to buy and the code is trivial; the value you add for a client is the integration, the windowing/k tuning, and the alerting glue, not the algorithm.

Build Your Own (Minimal Recipe)

You can stand up an 80%-of-the-value version in an afternoon. Components, in build order:

1. Windowing. Turn your series into log-returns (or just z-scored values for non-financial data) and slice into overlapping windows. Start with h_1 ≈ “one natural period” and h_2 ≈ 0.75·h_1. This is the only part with judgment in it.
2. 1D Wasserstein distance. The four lines above (np.sort both, mean of |a−b|^p). Or use scipy.stats.wasserstein_distance for p=1 if window sizes differ.
3. Barycenter update. Sort each cluster member, np.median(..., axis=0) for p=1. This is the one piece people get wrong — don’t average the raw windows, average them after sorting by rank.
4. k-means loop. Assign-nearest / update / check-centroid-movement. ~20 lines.
5. Validation. Add an MMD between/within-cluster score (Gaussian kernel, σ tuned to your data scale, e.g. 0.1) so you can measure whether clusters are coherent rather than eyeballing plots. Use scipy or a 10-line kernel-MMD.

The two genuinely hard parts: (a) choosing h_1 — too long and you lag/miss short regimes, too short and noise dominates; the authors call it “more art than science,” and the paper’s own sweep (Fig. 15) shows accuracy converges once h_1 is “large enough,” so err slightly long. (b) choosing k — the paper punts to k=2; for production use a silhouette/MMD sweep over k, or switch to a method that doesn’t fix k (see below). Libraries to reach for: numpy/scipy (everything), POT (Python Optimal Transport) if you ever go multivariate, scikit-learn only for the silhouette/Davies-Bouldin diagnostics.

How to Improve It

The paper’s own “future work” plus some leverage points:

1. Learn k instead of fixing it. Run an MMD- or silhouette-based sweep, or swap k-means for hierarchical or fuzzy c-means clustering (the authors flag both) so you don’t need the number of regimes up front and you allow soft membership (a window can be 70% calm / 30% turbulent during a transition — financially realistic).
2. Make it streaming / online. The paper is a posteriori. Build an online variant: maintain centroids, and for each new window compute its distance to existing centroids; if it’s far from all of them, spawn a candidate new regime. This turns it into a real-time drift alarm — the highest-value version for productionized ML monitoring.
3. Go multivariate properly. The 1D closed form is what makes it cheap; in d > 1 exact Wasserstein is expensive. Use the sliced Wasserstein distance (random 1D projections, average the 1D Wasserstein over them) the authors mention — it keeps the O(N log N)-per-slice cheapness and unlocks multi-asset / multi-feature regimes.
4. Tune p to your risk appetite. p=1 (median barycenter) is robust/sluggish; p=2 (mean barycenter) is more reactive but outlier-sensitive. Expose p as a knob — the paper hints the reactivity/robustness trade-off is application-dependent but never sweeps p itself. Easy, testable experiment.
5. Validate on real, multi-asset, multi-regime data with a proper backtest. The biggest gap is that real-data evidence is single-asset, k=2, qualitative. A genuine contribution would be a labeled (or weak-labeled via VIX/known-crisis dates) backtest across asset classes measuring detection lag, not just “it found 2008.”
6. Combine with the LLM-agent angle. Feed the centroid distributions (mean, variance, skew, kurtosis of each regime) to an LLM to auto-generate human-readable regime descriptions and transition narratives — turning a clustering output into an explainable monitoring report. Cheap to build, demos extremely well to non-technical clients.

Glossary

• Market regime — a contiguous period where returns behave as if drawn from one distribution (e.g. calm bull vs. volatile bear); a “regime change” is a switch between these.
• MRCP (Market Regime Clustering Problem) — the task of partitioning a return series into homogeneous regime segments.
• Empirical measure / distribution — the distribution that puts equal mass 1/N on each of your N observed values; “treat this window of returns as a distribution.”
• Wasserstein distance (W_p) — optimal-transport / “earth-mover’s” distance: the minimum work to reshape one distribution into another. In 1D = average gap between sorted samples.
• Wasserstein barycenter — the “average” of a set of distributions under Wasserstein distance; in 1D with p=1 it’s the coordinate-wise median of sorted samples.
• Optimal transport plan — the cheapest way to move probability mass from one distribution to another; in 1D it’s always smallest-to-smallest matching (monotone).
• k-means — unsupervised clustering: assign points to nearest of k centroids, recompute centroids as cluster averages, repeat until stable.
• MMD (Maximum Mean Discrepancy) — a kernel-based two-sample test statistic; how distinguishable two samples are after embedding through a kernel. Used here to score cluster quality.
• Kernel / RKHS / Gaussian kernel — a kernel κ(x,y) measures similarity; the Gaussian kernel exp(−‖x−y‖²/2σ²) is the standard “characteristic” choice that makes MMD a true metric.
• HMM (Hidden Markov Model) — classical regime model assuming a Markov latent state emitting (usually Gaussian) returns; the main baseline the paper beats on non-Gaussian data.
• Moments — mean (1st), variance (2nd), skew (3rd), kurtosis (4th)… summary numbers of a distribution; the moment-baseline clusters these vectors instead of the full distribution.
• gBm (geometric Brownian motion) — the standard Gaussian-log-return stock model used as the “easy” synthetic test.
• Merton jump diffusion — gBm plus random Poisson jumps; produces fat-tailed, skewed (non-Gaussian) returns — the “hard,” realistic synthetic test.
• Stylized facts — empirical regularities of financial returns (volatility clustering, fat tails, non-stationarity) that any honest model must contend with.
• Self-similarity score — the median within-cluster MMD; low = cluster members are genuinely alike (good clustering).
• h1 / h2 (window length / overlap) — the two hyperparameters: how many returns per window, and how far consecutive windows are offset.