Sparse Dictionary Learning for Measure-Valued Data via Linearized Optimal Transport: Efficient Algorithms and Sample Complexity

Brendan Mallery; Shuchin Aeron; James Murphy

stat.MLOct 9, 2026

Sparse Dictionary Learning for Measure-Valued Data via Linearized Optimal Transport: Efficient Algorithms and Sample Complexity

Authors: Brendan Mallery, Shuchin Aeron, James Murphy

Organizations: Independent Researcher · Department of Electrical and Computer Engineering, Tufts University · The NSF AI Institute for Artificial Intelligence and Fundamental Interactions (IAIFI) · Department of Mathematics, Tufts University

Abstract

We study dictionary learning for probability measure-valued data by embedding measures into Hilbert spaces using features derived from optimal transport geometry. We learn data-adapted dictionaries through sparse autoencoder parameterizations, producing fast, scalable representations of probability measures as sparse finite-dimensional codes. The resulting framework combines the rich, geometric structure of optimal transport with the interpretability and efficiency of sparse representation learning. We demonstrate its utility across diverse data modalities, including hyperspectral images, geometric point clouds, and document embeddings in large language model latent spaces. We also prove generalization bounds for sparse dictionary learning objectives on broad classes of $L^2$ -valued measure embeddings, expressed in terms of the metric complexity of the atom class. For optimal transport and entropic transport map dictionaries, we derive new covering number bounds that yield explicit rates. Together, our computational and theoretical results support a new, scalable model for representation learning in the space of probability measures via embedding into Hilbert spaces.

Figures & tables

Panel of Figure 1 — Figure 1: Schematic of linearized Wasserstein dictionary learning. We define an embedding $\mathcal{E}_{\rho}$ from the nonlinear Wasserstein space into $L^{2}(\rho)$ , viewed as the tangent space at some base measure $\rho$ . In the same linearized space, a set of dictionary elements $\{S_{j}\}_{j=1}^{m}$ are learned; their sparse nonnegative combinations approximate the embeddings, which may then be synthesized into reconstructed measures.

Variant	$\mathcal{E}_{\rho}(\mu)$	Setting
LWDL-OT	$T_{\rho\to\mu}$	exact, 1D
LWDL-EOT	$T^{\epsilon}_{\rho\to\mu}$	general $d$
LWDL-KP	$f_{\rho\to\mu}$	high $d$

Table 1: LWDL embeddings used in this work.

\lambda(\mu_{i};\theta_{a},\theta_{c})\leftarrow\operatorname{LISTAEncoder}\bigl(\mathcal{E}_{\rho}(\mu_{i}),\mathcal{S}^{\theta_{a}};\theta_{c}\bigr).

Algorithm 1 Generic LWDL training algorithm

Panel of Figure 2 — Figure 2: LWDL-EOT recovery on sparse convex combinations of five MNIST transport-map atoms, corresponding to digits $0,2,3,7,8$ . (Top left): Representative held-out convex-combination maps and their LWDL-EOT reconstructions. (Top right): Source points, original target map, and learned reconstruction map in a held-out example. (Bottom): Nearest-neighbor search for learned LWDL-EOT atoms against a held-out pool restricted to digits in $\mathcal{D}$ . Original generating atoms are included as reference candidates and outlined in bold.

Data Set	Method	$n$	Mean $W_{2}^{2}$	Epochs
Pavia	WDL ( $\gamma=.5$ )	100	$4.6\times 10^{-3}$	117
	WDL ( $\gamma=2$ )	100	$4.9\times 10^{-3}$	84
	LWDL-OT	100	$\mathbf{5.0\times 10^{-5}}$	240438
	WDL ( $\gamma=.5$ )	1000	$5.7\times 10^{-3}$	$7^{*}$
	WDL ( $\gamma=2$ )	1000	$6.0\times 10^{-3}$	$7^{*}$
	LWDL-OT	1000	$\mathbf{9.3\times 10^{-5}}$	65183

Table 2: Comparison with WDL under a 1000s time budget. Mean

W_{2}^{2}

is the final squared Wasserstein reconstruction error. ∗ denotes early termination due to numerical instability. LWDL uses

B=100

for

n=100

and

B=256

for

n=1000

; LWDL-EOT uses

\epsilon=0.025

and

c=5\times 10^{-5}

Panel of Figure 3 — Figure 3: Top : example spectrum reconstruction. Bottom : atoms learned from LWDL-OT and NMF.

Method	Mean $L^{2}$	Mean $W_{2}$
LWDL-OT	$2.72\times 10^{-6}$	$4.71\times 10^{-5}$
NMF	$2.98\times 10^{-7}$	$6.69\times 10^{-5}$

Table 3: Reconstruction errors on Pavia spectra; lower is better.

Data	Method	Corruption level $k$						Stab.
		$0$	$.1$	$.2$	$.3$	$.4$	$.5$
Pavia	NMF	.753	.790	.785	.616	.644	.536	$-.48$
	N-SAE	.730	.546	.503	.450	.461	.464	$-.47$
	LWDL-OT	.763	.778	.754	.735	.715	.695	$-.16$
Botswana	NMF	.620	.221	.162	.159	.146	.139	$-.75$
	N-SAE	.633	.312	.246	.204	.182	.168	$-.79$

Table 4: Accuracy under DropUniform corruption. Each row uses one fixed configuration across all corruption levels.

k

is the fraction of spectral bands dropped. Accuracies are averaged over five corruption seeds and, for LWDL-OT and N-SAE, five training seeds. Stability is the least-squares slope of these mean accuracies versus

k

; values closer to zero indicate greater robustness.

Panel of Figure 4 — Figure 4: Application diagnostics for LWDL-EOT and LWDL-KP. (a)–(b) Point-cloud reconstruction and code structure. (c)–(d) Transformer corruption-code structure and linear separability.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

\lambda(\mu_{i};\theta_{a},\theta_{c})\leftarrow\operatorname{LISTAEncoder}\bigl(\mathcal{E}_{\rho}(\mu_{i}),\mathcal{S}^{\theta_{a}};\theta_{c}\bigr).

Appendix

Algorithm 2 LWDL-OT in

\mathbb{R}

\zeta_{j,k}^{\epsilon,\theta_{a}}(x)\leftarrow\frac{\exp\!\left(-\frac{\|x-y_{k}\|^{2}+[W_{\mathrm{dec}}^{\theta_{a}}]_{j,k}}{2\epsilon}\right)}{\sum_{\ell=1}^{r}\exp\!\left(-\frac{\|x-y_{\ell}\|^{2}+[W_{\mathrm{dec}}^{\theta_{a}}]_{j,\ell}}{2\epsilon}\right)},\qquad k=1,\dots,r.

Appendix

Algorithm 3 LWDL-EOT

\lambda(\mu_{i};\theta_{a},\theta_{c})\leftarrow\operatorname{LISTAEncoder}\bigl(\mathcal{E}_{\rho}(\mu_{i}),\mathcal{S}^{\theta_{a}};\theta_{c}\bigr).

Appendix

Algorithm 4 LWDL-KP

h(\mu_{i},\theta_{a})\leftarrow\Bigl(\bigl\langle S_{1}^{\theta_{a}},\mathcal{E}_{\rho}(\mu_{i})\bigr\rangle_{L^{2}(\rho)},\dots,\bigl\langle S_{m}^{\theta_{a}},\mathcal{E}_{\rho}(\mu_{i})\bigr\rangle_{L^{2}(\rho)}\Bigr)\in\mathbb{R}^{m}.

Appendix

Algorithm 5 LISTA-type encoder for LWDL

Method	Learns measures	Learns codes	Free support	Linearized	Sparsity promoting
[ 44 , 38 ]	✓	✓	✗	✗	✗
[ 9 , 51 , 34 ]	✗	✓	✗	✗	✗
[ 24 , 52 ]	✗	✓	✗	✓	✗
[ 18 ]	✗	✓	✗	✗	✓
This paper	✓	✓	✓	✓	✓

Appendix

Table 5: Summary of related methods and comparison with the proposed LWDL framework.

Panel of Figure 5 — Figure 5: Hyperspectral image structure. A hyperspectral data cube records a spectrum at each spatial pixel: the spatial coordinates $(x,y)$ index image locations, the spectral axis indexes wavelength bands, and a selected pixel yields both a one-dimensional spectrum and corresponding intensities across individual bands.

\widetilde{p}_{h}(x)=\frac{1}{Mh}\sum_{j=1}^{M}\frac{1}{\sqrt{2\pi}}\exp\!\left(-\frac{1}{2}\Bigl(\frac{x-y_{j}}{h}\Bigr)^{2}\right).

Appendix

Algorithm 7 Kernel Density Estimate from Transport Map

Figure 6: Clustering accuracy of LWDL-OT, NMF, and the nonnegative sparse autoencoder (N-SAE) on the Pavia Centre, Botswana, and Salinas A hyperspectral data sets under increasing corruption strength. Results are shown for uniform band dropout, contiguous band dropout, and logarithmic spectral warping. LWDL-OT is particularly robust to uniform band dropout, while performance under the other corruptions varies across data sets and methods.

Table 6: Performance comparison for DropUniform across data sets and

k

values.

Panel of Figure 7 — Figure 7: (Left) Projection of subset of data with $r\leq 75$ to principal component of $\text{TopK}_{40}$ SAE. (Right) Projection of subset of data with $r>75$ to principal component of $\text{TopK}_{40}$ SAE.

Rank	Feature	Purity	Docs	Top source	Source distribution
1	83	1.00	100	Pile-CC	Pile-CC (100%)
2	263	1.00	100	StackExchange	StackExchange (100%)
3	82	1.00	100	Pile-CC	Pile-CC (100%)
4	324	1.00	100	USPTO Backgrounds	USPTO Backgrounds (100%)
5	314	1.00	100	PubMed Abstracts	PubMed Abstracts (100%)
6	74	1.00	100	PubMed Abstracts	PubMed Abstracts (100%)

Appendix

Table 7: Top learned features ranked by source purity. Purity is computed as the fraction of labeled top-activating documents belonging to the most common Pile source.

Panel of Figure 8 — Figure 8: Dictionary atoms learned by our method when trained on point-cloud data, illustrating the dependence on the base measure. Colors indicate the value of the functional embedding on the point cloud.

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 798 papers

Multi-Dictionary Learning for Low Rank Sparse Coding
Sep 12, 2025Boya Ma, Abram Magner, Maxwell McNeil +1Representation Learning Sample Complexity
University at Albany – SUNY, Computer Science, Albany, NY, USA
On the Sparsity-Storage-Accuracy Tradeoff in Parsimoniously Activated Dictionary Learning
Jun 21, 2026Zihui Zhao, Yuanbo Tang, Yang LiStructured Sparsity Dictionary Learning
Tsinghua University, Institute of Data and Information Shenzhen Key Laboratory of Ubiquitous Data Enabling
Inference and learning in sparse autoencoders as natural gradient flow
Oct 5, 2026Hadi Vafaii, Tejas Rao, David Chanin +6Gradient Descent Sparse Autoencoders
Astera Institute · UC Berkeley · UCL +2

Sep 12, 2025J/K move · Enter open · S save

Multi-Dictionary Learning for Low Rank Sparse Coding

Boya Ma, Abram Magner, Maxwell McNeil, Petko Bogdanov

University at Albany – SUNY, Computer Science, Albany, NY, USA

Representation Learning Sample Complexity Dictionary Learning Sparse Coding

Open paper AskPDF arXiv

Sparse dictionary coding represents signals as linear combinations of a few dictionary atoms. It has been applied to images, time series, graph signals and multi-way spatio-temporal data by jointly employing temporal and spatial dictionaries. Data-agnostic analytical dictionaries, such as the discrete Fourier transform, wavelets and graph Fourier, have seen wide adoption due to efficient implementations and good practical performance. On the other hand, dictionaries learned from data offer sparser and more accurate solutions but require learning of both the dictionaries and the coding coefficients. This becomes especially challenging for multi-dictionary scenarios since encoding coefficients correspond to all atom combinations from the dictionaries. To address this challenge, we propose a low-rank coding model for 2-dictionary scenarios and study its data complexity. Namely, we establish upper and lower bounds on the number of samples needed to learn dictionaries that generalize to unseen samples from the same distribution. We propose an alternating convex optimization solution, called AODL, which employs alternating optimization between the sparse coding matrices and the learned dictionaries. We demonstrate its quality for data reconstruction and missing value imputation in both synthetic and real-world datasets. For a fixed reconstruction quality, AODL learns up to

90\%

sparser solutions compared to non-low-rank and analytical (fixed) dictionary baselines. In addition, the learned dictionaries reveal interpretable insights into patterns from training samples.

Abstract

Representation Learning

Explore similar work

Multi-Dictionary Learning for Low Rank Sparse Coding

On the Sparsity-Storage-Accuracy Tradeoff in Parsimoniously Activated Dictionary Learning

Inference and learning in sparse autoencoders as natural gradient flow