Sparse dictionary learning for measure-valued data based on linearized optimal transport: efficient algorithms and sample complexity

How to perform scalable sparse dictionary learning for probability measures in Wasserstein-2 space and analyze its generalization.

• Scenario: The authors state that learning and inference on probability measure data heavily use the nonlinear geometry of Wasserstein-2. Wasserstein dictionary learning (WDL) learns a dictionary of measures representing the data, enabling sample synthesis from latent codes, as well as dimensionality reduction, sparse representation, and compression.
• Limitations: The authors state that WDL suffers from computational and statistical difficulties for scalable learning in this space; synthesis often requires fixed-point iterations and repeated computation of W2, and data, atoms, and reconstructions are often on a fixed finite support. The authors state that it remains unclear which data and tasks are suitable for WDL.
• Idea: Fix an absolutely continuous reference measure ρ, use linearized optimal transport (LOT) to embed measures into the Hilbert space L2(ρ), then learn atoms in this flat geometry whose sparse nonnegative combinations can approximate the embeddings.
• Proposed object: Linearized Wasserstein dictionary learning (LWDL) with a sparsity penalty, where embeddings include exact transport maps, entropic transport maps, and Kantorovich potentials, parameterized by a sparse autoencoder. Also provide generalization bounds expressed in terms of the covering number of the atom class, and explicit rates for dictionaries of OT and entropic transport maps.

WDL in nonlinear geometry

• Schmitz et al. (2018): minimize squared W2 reconstruction on a finite deterministic training set, coefficients without regularization.
• Mueller et al. (2023): add geometric regularization so that the learned reference measure is close to the data points it represents. The authors state that such methods are slow in nonlinear W2 geometry, synthesis relies on fixed-point iteration and multiple W2 computations; the support is fixed; this geometry also makes theoretical analysis difficult.

Fixed reference measure, learning only coefficients

• Bonneel et al. (2016), Mallery et al. (2025), Werenski et al. (2022): represent observations as coefficients of Wasserstein barycenters or entropic regularized barycenters. The authors state that they work in fully nonlinear geometry and do not learn the reference measure.
• Gunsilius et al. (2024), Werenski et al. (2025): learn coefficients under the LOT metric for a fixed reference measure. The authors state that this corresponds to not optimizing over the atom set in objective (1) and g being identically 0.
• Do et al. (2025): use a greedy algorithm for a fixed reference measure and add an explicit sparsity penalty on the coefficients.

Baseline methods and benchmarks/datasets

• Baseline methods: WDL (Schmitz et al., 2018) with γ = 0.5 and 2; non-negative matrix factorization (NMF) on hyperspectral data, and a non-negative sparse autoencoder (N-SAE) trained directly on normalized spectra.
• Benchmarks/datasets: Pavia, Botswana, Salinas A; MNIST transport maps; ModelNet-10; synthetic text corruption on Pythia-410M hidden states and The Pile.

Table 5 in Appendix D marks this paper as simultaneously learning the measure and coefficients, with variable support, using linearization and sparsity promotion; the other methods in the table lack at least one of these. The authors state that the SAE parameterization can learn atoms faster than methods based on nonlinear W2.

Embed probability measures into L2(ρ) at the reference measure ρ, then use the shared encoder of LWDL with method-specific atoms to jointly learn sparse nonnegative codes. The three embeddings are the exact transport map, the entropic transport displacement field, and the centered Kantorovich potential.

Problem Setting

• Embedding: Data are measures in P2(R^d). Eρ maps into L2(ρ). ρ-LOT is the L2(ρ) norm of the difference of transport maps from ρ; the authors view it as linearizing W2 at ρ.
• Objective: LMDL learns m atoms such that the embedding can be approximated by a sparse linear combination with nonnegative coefficients. The penalty g is a sparsity-inducing term, with example g(λ)=c||λ||1, c>0.
• Three Variants: In Table 1, LWDL-OT takes the exact map, implemented only in one dimension; LWDL-EOT takes the entropic map, used for dimensions above one; LWDL-KP uses scalar potentials, aimed at high-dimensional coding. The authors say the three are different linearizations of WDL (Appendix B).
• Synthesis: OT and EOT can generate measures by pushing forward the learned map. The atoms of KP are scalar functions; the authors say one cannot directly synthesize measures in ambient space and must first differentiate the potentials.

Atom Parameterization

• LWDL-OT: Restricted to one-dimensional measures on [0,1]. Atoms are normalized by positive increments of a decoding matrix, are monotone on the grid of ρ, with left endpoint 0 and right endpoint 1 (Algorithm 2). The authors' stated reason is that the Brenier map on an interval is exactly monotone nondecreasing; in higher dimensions, forcing a vector field to be the gradient of a convex potential is harder.
• LWDL-EOT: ρ has finite support, and a target point set Y is additionally fixed so that its convex hull almost contains the data support. Atoms are first parameterized as finite-dimensional entropic maps; in experiments, the displacement field is used instead, and the embedding also subtracts the identity map (Algorithm 3). The authors say the displacement field empirically outperforms the map itself. Applicable for any d≥1.
• LWDL-KP: The embedding is the Kantorovich potential minus its ρ-mean; atoms are scalar functions on the support points, centered by the same convention (Algorithm 4). The authors say this replaces vector-valued maps with scalar functions, suitable for high-dimensional coding.

Encoder and Training

• LISTA-type Encoder: First compute the L2(ρ) inner products of each atom with the embedding, perform L layers of affine updates from the zero vector, then apply an activation mapping to the nonnegative orthant, outputting λ(L) (Algorithm 5, Eq. (2)). The authors say it retains the iterative affine-threshold form, but the weights and activation are not tied to a fixed ISTA update. The activation can be ReLU, JumpReLU, or a k-sparse activation.
• Loss: Eq. (3) averages one-half the squared reconstruction over samples, plus c times the ℓ1 of the coefficients. Atom parameters and encoder parameters are trained end to end.
• Optimization and Discretization: Algorithm 1 uses minibatch SGD, AdamW, and a cosine learning rate. In the numerical implementation, ρ has finite support, embeddings and atoms are represented by their values on the support points, and L2(ρ) becomes a weighted finite-dimensional norm. The output dimension of map embeddings is d, and that of potential embeddings is 1.

Generalization Bounds

• Setting: The pointwise loss uses the entropic map (ε>0) or the OT map when it exists (ε=0), with nonnegative combinations of unit-norm atoms plus g. The data second moment has an upper bound; g is lower semicontinuous, is 0 at 0, and is no less than c times ℓ1; the atom class is compact.
• Theorem 6.1: The generalization gap between the empirical minimizer dictionary and the population minimizer dictionary is controlled by an integral of the covering number of the atom class in L2(ρ), and depends on the moment bound and regularization strength. The authors say the main idea in the proof is known, but this exact statement has not appeared in the literature.
• Theorem 6.2: On a product interval with density having positive finite bounds, the rate for the OT map dictionary is O(√(m/n)) for d=1 and O(m^{1/d} n^{-2/(3d)}) for d>1. When the base measure and data are σ-subgaussian, the rate for the EOT map dictionary is O(√(m/n)). The constants depend on D, d, β, ϑ, c, or on σ, d, ε, ϑ, c, respectively.

Experimental Setup

• Method variants: LWDL-OT, LWDL-EOT, LWDL-KP. The main text §4.1 says to run Algorithm 5; Appendix G says to run Algorithm 3 and gives the hyperparameters for that run.
• Data: Synthetic MNIST transport maps 60,000 samples, 54,000 for training, 6,000 for testing, ρ support M=400; when compared with WDL, Pavia is one-dimensional with 102 support points, MNIST is two-dimensional, with n taken as 100 and 1000. Hyperspectral datasets are Pavia Centre, Botswana, Salinas A. ModelNet-10 takes six classes with at least 300 instances, totaling 1,800 measures, each cloud with 1,024 points. Text corruption 19,200 samples; The Pile about 80,000 documents, 100 to 2,000 tokens.
• Baselines: WDL, with γ 0.5 and 2, m=10, time budget 1,000 seconds; on hyperspectral data, NMF and N-SAE. Dictionary size is set to the number of non-background classes plus one.
• Metrics: L2(ρ) reconstruction error, Mean W2^2, aligned coefficient MSE, clustering accuracy, stability (least-squares slope of mean accuracy versus corruption fraction k; the authors state that closer to 0 is more robust), linear classification or logistic regression accuracy, source purity.
• Evaluation and repetition: For hyperspectral data, the dictionary is trained on clean spectra and then fixed, corrupted spectra are encoded, and after clustering, labels are matched. Table 4 averages over five corruption seeds; LWDL-OT and N-SAE are additionally averaged over five training seeds. The paper does not use LLM review and does not report significance tests.
• Feature extraction: Text experiments use intermediate-layer token hidden states from Pythia-410M, dmodel=1024, 24 layers.

Main Results

According to Table 2, final Mean W2^2 under a 1,000-second budget or upon numerical failure. * indicates early termination due to numerical instability.

Swipe horizontally to read the table

Dataset	Method	n	Mean W2^2	Epochs
Pavia	WDL (γ=.5)	100	4.6×10^{-3}	117
Pavia	LWDL-OT	100	5.0×10^{-5}	240438
Pavia	WDL (γ=2)	1000	6.0×10^{-3}	7*
Pavia	LWDL-OT	1000	9.3×10^{-5}	65183
MNIST	WDL (γ=2)	100	1.0×10^{-2}	101
MNIST	LWDL-EOT	100	2.6×10^{-3}	50978
MNIST	WDL (γ=.5)	1000	1.3×10^{-2}	10
MNIST	LWDL-EOT	1000	2.7×10^{-3}	12965

• Omitted rows: On Pavia, WDL (γ=2, n=100) is 4.9×10^{-3} (84 epochs); WDL (γ=.5, n=1000) is 5.7×10^{-3} (7*); on MNIST, WDL (γ=.5, n=100) is 1.2×10^{-2} (101 epochs), and WDL (γ=2, n=1000) is 1.3×10^{-2} (11 epochs). The largest error in the table is 1.3×10^{-2}.
• The authors state that within this budget, LWDL's Mean W2^2 is roughly 4 to 100 times lower, partly because it completes far more epochs than WDL; they also state that epochs are not directly comparable, but the gap indicates that LWDL has lower per-epoch cost.
• The authors state that WDL's time is mainly spent on dense kernel matrices, which is typical for Sinkhorn-type barycenters; the linearized method computes the OT or EOT map only once, followed by autoencoder optimization. Appendix H states that the reported time includes the computation of the initial map embedding.

Synthetic Dictionary Recovery

• What was measured: Transport map atoms for five MNIST digits (0, 2, 3, 7, 8), with coefficients having exactly two nonzeros and lying on the simplex; evaluation includes test-set reconstruction, atom matching, and coefficient recovery.
• Results: Mean squared L2(ρ) error is 9.22×10^{-4}; atom matching error is 1.74×10^{-2}; aligned coefficient MSE is 8.50×10^{-3}. Of the 25 nearest neighbors shown, 24 digits match the matched atom, and the highest nearest neighbor for all five atoms is the matching digit.
• Authors' interpretation: The authors state that the data maps can be reconstructed to a high degree, but the learned atoms have weak visual similarity to the digits; the nearest-neighbor results are described by the authors as evidence that at least the true atom classes can be identified. Despite residual in the atoms, the authors state that the coefficients remain close to the generating coefficients. In Appendix G, the authors suggest that fine-scale jitter may come from ρ being more diffuse than the data.

Hyperspectral Reconstruction and Corrupted Clustering

• Reconstruction: In Table 3 on Pavia spectra, LWDL-OT has Mean L2 of 2.72×10^{-6} and Mean W2 of 4.71×10^{-5}; NMF has 2.98×10^{-7} and 6.69×10^{-5}, respectively. The authors state that LWDL-OT is lower than NMF in W2 but higher in L2([0,1]); transport atoms often concentrate mass at the endpoints of the wavelength interval.
• DropUniform: In Table 4, LWDL-OT stability is −0.16 for Pavia, −0.59 for Botswana, and −0.57 for Salinas A, and on all three datasets it is closer to 0 than NMF and N-SAE. The authors state that it has the highest accuracy on 13 of the 15 nonzero corruption levels. Counterexamples in the table are Pavia at k=0.1 and 0.2: NMF is .790 and .785, LWDL-OT is .778 and .754.
• Clean data: The authors state that at k=0, each method is highest on one dataset. In Table 4, Pavia is LWDL-OT .763, Botswana is N-SAE .633 (LWDL-OT is .556), and Salinas A is NMF .853.

Point Clouds and Text Hidden States

• Point clouds: On coefficients for six object classes, linear classifier held-out accuracy is 96.1%. Appendix F states that after 2,000 epochs, mean squared L2(ρ) error is 0.00173. The authors state that atoms form interpretable parts and that Figure 4(b) further shows class structure in the codes; whether the clusters in the figure are separated is not described in the text with verifiable coordinates.
• Synthetic corruption: 19,200 variants, four corruption types, twelve intensities. The authors state that atom usage strongly depends on class, and this pattern persists across JumpReLU and TopK, and with 10 and 40 atoms. Over 10 random splits, logistic regression on the first two principal components exceeds 90%, and on the first six about 95%.
• The Pile: 512 atoms. In Table 7, the first 19 rows have purity 1.00. The authors state that f469, f327, f132, and f44 each have over 99% of documents from Wikipedia, but the topics of the top 20 differ; they also write that f98 and f469 have nearly identical top documents and state that there is redundancy among high-purity atoms.

Other Ablations and Analyses

• Appendix E Table 6: A set of N-SAE (m=15, c=10^{-3}) on Salinas A gives 0.2853 for all k in DropUniform.
• Same table LogDiffeo: accuracy does not decrease monotonically with k; Pavia's N-SAE (m=10, c=10^{-4}) goes from 0.7295 at k=None to 0.8199 at k=1000. Figure 6 caption states that performance on the other two corruption types varies with dataset and method.
• Appendix J: After changing the base measure, pushed-forward atoms inherit the qualitative properties of their respective base measures; the authors explain this from the perspective of the tangent space at ρ.
• Table 2 note: LWDL uses B=100 for n=100 and B=256 for n=1000; LWDL-EOT uses ε=0.025 and c=5×10^{-5}.
• Appendix G: m=5, ε=0.025, c=10^{-4}, 20-layer encoder, 1,000 epochs, reporting the checkpoint at epoch 862.
• Appendix I.2: E(x^2, 407)=133.82, while this quantity is 0 for j=134 or 263; the authors state that the corresponding documents are dominated by Math StackExchange posts.

Limitations and future directions noted by the authors

• Atom identifiability: Section 7 states that it is necessary in the future to guarantee atom identifiability under the LWDL generative model.
• Base measure: The same section states that it is necessary to understand the influence of the base measure and possibly introduce the base measure itself into the learning process.

Experimental coverage

• Embedding variants are LWDL-OT (exact one-dimensional maps), LWDL-EOT (entropic maps and displacement fields), and LWDL-KP (centered potentials) (Table 1, Section 3).
• Comparisons with WDL are on Pavia (support of 102 points) and two-dimensional MNIST, with n equal to 100 and 1000, a budget of 1000 seconds, m=10, and γ equal to 0.5 and 2 (Section 4.2, Table 2). The paper does not report the number of repetitions for this comparison.
• Hyperspectral data are Pavia Centre, Botswana, Salinas A; corruptions are DropUniform, DropContiguous, and LogWarp/LogDiffeo; baselines are NMF and N-SAE. Table 4 averages over five corruption seeds, and LWDL-OT and N-SAE additionally average over five training seeds (Section 5.1, Appendix E).
• Point clouds are six classes from ModelNet-10 with 1800 measures; text uses intermediate hidden states from Pythia-410M, with 19,200 synthetic corruptions and about 80,000 documents from The Pile, with 512 atoms (Sections 5.2–5.3). The paper does not report the test set size or number of repetitions for point cloud linear classification.
• Statistical results are Theorem 6.1 and Theorem 6.2. Recovery experiments with known dictionaries are in Section 4.1. The paper does not report significance tests, nor does it report the reviewing model or agreement with human judgment.

LWDL embeds probability measures into a Hilbert space, then performs sparse dictionary learning.

• Problem: The authors state that dictionary learning on the nonlinear W2 is computationally slow and theoretically difficult, and common implementations fix the support on a finite set of points.
• Approach: At a reference measure ρ, use exact transport maps, entropic transport displacement fields, or centered Kantorovich potentials to embed into L2(ρ), and jointly learn atoms and nonnegative codes with a LISTA-type encoder and ℓ1 penalty. OT and EOT can compose synthetic measures via pushforward; KP targets high-dimensional coding.
• Comparison with WDL: Within the 1000-second budget of Table 2, on Pavia with n=100, LWDL-OT has Mean W2^2 of 5.0×10^{-5}, while WDL (γ=.5) has 4.6×10^{-3}. With n=1000, WDL stops at epoch 7 due to numerical instability, while LWDL-OT achieves 9.3×10^{-5}. On MNIST with n=1000, LWDL-EOT achieves 2.7×10^{-3}, while WDL achieves 1.3×10^{-2}. The authors state the gap is roughly 4 to 100 times, and relate it to lower per-epoch cost.
• Recovery and applications: The squared L2(ρ) error on the synthetic test set is 9.22×10^{-4}, the alignment coefficient MSE is 8.50×10^{-3}, and the atom matching error is 1.74×10^{-2}; the authors state that the visual similarity between atoms and digits is weak, but nearest neighbors mostly belong to the matching digit. On Pavia, LWDL-OT has lower Mean W2 than NMF but higher Mean L2 (Table 3). Across the 15 nonzero levels of DropUniform, the authors state that LWDL-OT has the highest accuracy on 13, and the three stability measures are also closer to 0. On ModelNet-10, the held-out linear classification accuracy is 96.1%. On the first two principal components of text corruption codes, logistic regression exceeds 90%.
• Theory and future work: Theorem 6.2 gives rates of the form O(√(m/n)) for one-dimensional OT and for EOT in the sub-Gaussian case, while the rate for higher-dimensional OT contains n^{-2/(3d)}. The authors state that empirical results show better computational efficiency relative to existing WDL, and that atoms and reconstructions have utility and interpretability; future work concerns atom identifiability under generative models and the role of the base measure.