SOL - Sliced Optimal Transport for Language

SOL measures the gap between two text corpora with a double-sliced Wasserstein distance over contextual token embeddings. Lower is closer; identical corpora score zero.

Paper · Website · API · Models · Changelog · Contributing

Install

pip install 'sol-metric[text]'    # compare texts (PyTorch + Transformers)
pip install 'sol-metric[torch]'   # score existing embeddings on a GPU
pip install sol-metric            # score existing embeddings with NumPy only

Python 3.10 or newer. For a specific CUDA version, install PyTorch first; add accelerate to load large checkpoints straight onto the GPU.

Use

from sol_metric import SOL

reference = ["The train arrived at noon.", "It rained throughout the night."]
generated = ["The afternoon train was late.", "Heavy rain fell overnight."]
print(SOL().from_texts(reference, generated))

The default encoder is Dream 7B, as in the paper: a 15 GB download that needs a GPU with about 16 GB. Any Hugging Face model can be used instead:

from sol_metric import SOL, TextEncoder

encoder = TextEncoder('gpt2-large')
metric = SOL()
print(metric.from_texts(reference, generated, encoder=encoder))

Features can be saved and reused, for example one reference against several generated corpora:

from tempfile import TemporaryDirectory
from sol_metric import Features

with TemporaryDirectory() as tmp:
    metric.featurize_texts(reference, encoder=encoder, path=f'{tmp}/reference').close()
    with Features.load(f'{tmp}/reference') as ref:
        with metric.featurize_texts(generated, encoder=encoder) as gen:
            print(metric.from_features(ref, gen))

Existing token embeddings need only NumPy. Stack the token states of all documents in one array and mark where each document starts and ends:

import numpy as np
from sol_metric import SOL

rng = np.random.default_rng(0)
x = rng.normal(size=(12, 16))  # corpus 1: 12 token states of dimension 16
x_offsets = [0, 4, 7, 12]      # 3 documents: rows 0-3, 4-6 and 7-11
y = rng.normal(size=(10, 16))  # corpus 2: 10 token states
y_offsets = [0, 5, 10]         # 2 documents: rows 0-4 and 5-9
print(SOL().from_tokens(x, x_offsets, y, y_offsets))

Good to know:

• Results are reproducible: the random directions and GP draws are seeded (seed, gp_seed, both 0 by default). Change them to measure the spread over seeds, as the paper does with nine seed pairs.
• The paper's encoders (Dream 7B, GPT-2 Large, OLMo 2 1B, Qwen3 0.6B) are pinned to the paper's checkpoint commits; pass revision='main' for the latest.
• Encoders run in FP16 on a GPU and FP32 on a CPU. Each document is encoded whole up to the model's limit (2,048 tokens for Dream, 1,024 for GPT-2 Large); longer documents are truncated with a warning.
• A GPU is used when PyTorch sees one; pass device='cpu' to force the CPU.
• metric.last_run records the score, encoder, revision and all settings.

Reproduce the paper

The repository bundles LM1B and OpenWebText comparisons, encoded with GPT-2 Large. The 5k OpenWebText pair is the paper's seed-0 subset for MDLM.

python examples/compare_corpora.py owt                             # 32 texts per corpus
python examples/compare_corpora.py owt --validate --device cuda:0  # 5,000 per corpus

--validate re-encodes both corpora and checks the score against the paper's evaluation code on the same texts; this package matches it to 1e-10. On an RTX 5090 the OpenWebText check takes about 2.5 minutes. The paper's table reports 2.167356 because it used GPT-2 token IDs instead of the decoded text. The 5k texts (20 MB) are in the Git repository, not the PyPI package; see the validation manifest.

Settings

Memory is handled automatically: features stay in RAM while they fit and go to temporary files otherwise, and batch sizes never change scores; if memory runs out, the error names the batch size to lower. All settings are in the API reference.

What SOL computes

Each document is the uniform measure μᵢ on its token states in ℝᵈ, and a corpus is the uniform distribution over its documents. A unit direction ξ and a Gaussian-process draw g map every document to one number, M(μᵢ) = ⟨g, Q((e_ξ)#μᵢ)⟩, the GP-weighted integral of its projected quantile function. SOL averages one-dimensional Wasserstein distances over both:

\mathrm{SOL}(\boldsymbol\mu,\boldsymbol\nu)^2
= d\;\mathbb E_{\xi\sim\mathcal U(\mathbb S^{d-1}),\,g\sim\gamma}
\Big[\mathrm W_2^2\big((M_{\xi,g})_\#\boldsymbol\mu,\ (M_{\xi,g})_\#\boldsymbol\nu\big)\Big]

with d the hidden dimension and γ a zero-mean GP with RBF covariance. The estimator uses n_directions × n_projections slices, exact quantile-bin integrals and exact transport for unequal corpus sizes.

Cite

@misc{kornhardt2026solmeasuringgapstext,
      title={SOL: Measuring Gaps between Text Distributions by Double Sliced Wasserstein Metrics},
      author={Gregor Kornhardt and Moritz Piening and Jannis Chemseddine and Gabriele Steidl},
      year={2026},
      eprint={2610.06513},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2610.06513},
}

License

Apache-2.0. Model checkpoints and corpora keep their own licenses.