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
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 L2-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
Figure 1: Schematic of linearized Wasserstein dictionary learning. We define an embedding Eρ from the nonlinear Wasserstein space into L2(ρ) , viewed as the tangent space at some base measure ρ . In the same linearized space, a set of dictionary elements {Sj}j=1m are learned; their sparse nonnegative combinations approximate the embeddings, which may then be synthesized into reconstructed measures.
Variant
Eρ(μ)
Setting
LWDL-OT
Tρ→μ
exact, 1D
LWDL-EOT
Tρ→μϵ
general d
LWDL-KP
fρ→μ
high d
Table 1: LWDL embeddings used in this work.
λ(μi;θa,θc)←LISTAEncoder(Eρ(μi),Sθa;θc).
Algorithm 1 Generic LWDL training algorithm
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 D . Original generating atoms are included as reference candidates and outlined in bold.
Data Set
Method
n
Mean W22
Epochs
Pavia
WDL ( γ=.5 )
100
4.6×10−3
117
WDL ( γ=2 )
100
4.9×10−3
84
LWDL-OT
100
5.0×10−5
240438
WDL ( γ=.5 )
1000
5.7×10−3
7∗
WDL ( γ=2 )
1000
6.0×10−3
7∗
LWDL-OT
1000
9.3×10−5
65183
Table 2: Comparison with WDL under a 1000s time budget. Mean W22 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 ϵ=0.025 and c=5×10−5 .
Figure 3: Top : example spectrum reconstruction. Bottom : atoms learned from LWDL-OT and NMF.
Method
Mean L2
Mean W2
LWDL-OT
2.72×10−6
4.71×10−5
NMF
2.98×10−7
6.69×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.
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.
Table 5: Summary of related methods and comparison with the proposed LWDL framework.
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.
ph(x)=Mh1∑j=1M2π1exp(−21(hx−yj)2).
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.
Figure 7: (Left) Projection of subset of data with r≤75 to principal component of TopK40 SAE. (Right) Projection of subset of data with r>75 to principal component of TopK40 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.
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.
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.
Boya Ma, Abram Magner, Maxwell McNeil +1
University at Albany – SUNY, Computer Science, Albany, NY, USA
Dictionary learning has long been studied from both optimization and probabilistic perspectives. While formulations with element-wise sparsity regularization (e.g., L1-based sparse coding) admit well-established probabilistic interpretations, many structured variants that impose global constraints lack a clear and tractable generative view. In this paper, we revisit a class of practically effective yet theoretically under-explored dictionary learning methods that impose a simple global regularization on the number of activated dictionary atoms, which we term parsimoniously activated dictionary learning (PADL). We show that PADL admits an equivalent formulation as maximum a posteriori estimation under a structured generative model, with auxiliary latent variables that govern global activation patterns. This formulation allows us to derive generalization guarantees that are difficult to obtain under the original formulation. More importantly, it yields an analytical characterization of the tradeoff between sparsity, storage cost, and reconstruction accuracy, enabling data-driven estimation of optimal hyperparameters. Based on this connection, we develop an efficient and interpretable PADL algorithm that eliminates manual hyperparameter tuning, achieving improved reconstruction performance under comparable sparsity levels on visual benchmarks. We further demonstrate its practical utility in accelerating inference for vision-language models.
Zihui Zhao, Yuanbo Tang, Yang Li
Tsinghua University, Institute of Data and Information Shenzhen Key Laboratory of Ubiquitous Data Enabling
Sparse autoencoders are widely used to uncover interpretable features in neural networks, yet reliable recovery remains difficult when features overlap or activate infrequently. These challenges involve both inferring which features explain an input and learning the dictionary that represents them. Here, we unify inference and dictionary learning as natural-gradient flows on a shared variational free energy. We instantiate this framework as BeFOND, an encoder-free sparse coding model with closed-form inference and learning dynamics. We show how recurrent explaining away reduces interference between overlapping features, while Fisher preconditioning can compensate for the slow learning of rare features. On synthetic data, BeFOND improves dictionary recovery and rare-feature detection, with a growing advantage over amortized baselines as superposition increases. On language-model activations, it improves single-feature concept detection and selective intervention, outperforming pretrained reference SAEs with substantially less training data. Its feature quality continues to improve with dictionary width, whereas the evaluated baselines largely plateau. Together, these results show how improving inference and learning within a unified probabilistic framework can make better use of data and dictionary capacity to interpret and intervene on neural representations.