Wasserstein Dictionary Learning: Optimal Transport-based unsupervised non-linear dictionary learning

1.1. Previous works.

• Dictionary learning customizes a dictionary and codes to factorize a dataset, while related variants impose sparsity, independence, positivity, or alternative fitting losses.The introduction contrasts such linear approaches with nonlinear reconstructions based on neural networks or geodesic interpolation.Source: Linear dictionary learning · pp. 2–3, Nonlinear dictionary learning · p. 3

Computational optimal transport.

• Entropic regularization makes OT strictly convex with a unique solution computable by Sinkhorn iterations using elementwise operations and matrix-vector products.For large images, separable convolution can replace some matrix-vector products and avoid handling a full cost matrix.Source: Computational optimal transport · p. 4

Wasserstein barycenters.

• Wasserstein barycenters provide the geometric building block, while exact computation for empirical measures involves a multimarginal transport problem whose size grows exponentially with support size.Regularized barycenter algorithms and generalized Sinkhorn methods offer computational alternatives.Source: Wasserstein barycenters · p. 4

Applications to image processing.

• The introduction connects OT image processing to color-distribution retrieval, color grading, image harmonization, and transport-based warping of grayscale images.It also reviews Wasserstein distances used as fitting losses in PCA-like methods, NMF, dictionary learning, and supervised learning.Source: Applications to image processing · p. 5, Wasserstein loss and fidelity · p. 5

1.2. Contributions.

• Wasserstein dictionary learning reconstructs probability histograms with Wasserstein barycenters of learned atoms rather than the usual linear dictionary product.The fitting loss can be any chosen criterion on the simplex, not necessarily Wasserstein distance.Source: §1.2 · p. 6
• The proposed method learns both atoms and reconstruction weights by differentiating the iterative generalized Sinkhorn scheme for Wasserstein barycenters.The resulting gradients can be used in a solver of choice; this work uses simultaneous quasi-Newton updates.Source: Wasserstein barycentric coordinates · pp. 4–5, §1.2 · p. 6
• On five slightly translated discretized Gaussians, linear generative models fail to capture the geometric changes that the OT approach captures.The learned atoms are discrete measures, unlike the PCA and NMF components.Source: §1.2 · p. 6
• The paper previews log-domain computation for sharp reconstructions, separability-based cost reduction, Sinkhorn warm starts or momentum, and an extension to unbalanced transport.These variants extend an earlier conference version that learned only two atoms and did not include the listed improvements.Source: §1.2 · p. 6

2.1. OT distances.

• For discrete histograms, an admissible transport plan is a matrix whose rows and columns sum to the two histograms, and its cost defines their OT distance.The section restricts attention to histograms on a fixed grid, represented as vectors in the probability simplex, with transportation costs stored in a cost matrix.Source: §2.1 · p. 8
• Entropic regularization makes the transport problem strongly convex and amenable to Sinkhorn scaling iterations.As γ approaches zero, the regularized distance converges to the original Wasserstein distance; larger γ promotes more diffuse transport matrices.Source: §2.1 · p. 9
• On a uniform grid with squared Euclidean costs, the kernel operation can be approximated by separable computation and is Gaussian with standard deviation γ/2.When γ is small, kernel values may become negligible and scaling vectors may become numerically unstable.Source: §2.1 · p. 9

2.2. Wasserstein barycenter.

• A Wasserstein barycenter minimizes the weighted sum of squared Wasserstein distances from a variable histogram to the input measures.With entropic regularization, the barycenter is uniquely defined and can be approximated with an efficient iterative scheme.Source: §2.2 · pp. 9–10

3. Wasserstein dictionary learning.

• The learned dictionary atoms and per-example weights reconstruct each histogram through a Wasserstein barycenter chosen to best fit it under a selected loss.The data and atoms are histograms in the probability simplex.Source: §3.1 · p. 10

3.1. Overview.

• Unlike linear dictionary learning's matrix product DΛ, reconstruction uses the nonlinear Wasserstein barycenter operator P(D, Λ).The paper describes this formulation as allowing nonlinear relationships between atoms and reconstructed data.Source: §3.1 · p. 10
• The method differentiates the finite Sinkhorn iteration scheme for approximate barycenters instead of relying on analytical Jacobians that require prohibitively large linear systems.The gradients are used in a descent method; the objective is not guaranteed to be convex jointly in atoms and weights or separately.Source: §3.1 · pp. 10–11

3.2. Backward recursive differentiation.

• A forward barycenter loop is followed by recursive backward differentiation to compute gradients with respect to atoms and weights.The procedure stores two additional vectors of size SNL; automatic differentiation can have memory costs from retaining the computation graph.Source: §3.1 · pp. 11–12

3.1. Overview.

• The fitting loss can vary, with examples including total variation, quadratic, Kullback-Leibler, and Wasserstein criteria.The loss gradient depends on the chosen fitting criterion.Source: §3.1 · pp. 11–12, Table 1 · p. 12
• Atoms and weights are optimized simultaneously with a quasi-Newton method after a logistic change of variables.The applications use an L-BFGS solver; rescaling weight gradients can help continue optimization when convergence is reached for one variable group before the other.Source: §3.1 · p. 14

4.1.1. Stabilization.

• Log-domain Sinkhorn updates enable barycenter computation at arbitrarily low γ, addressing numerical problems that arise when sharp reconstructions require very low entropy.Extending stabilization through the backward loop could require storing or recomputing many nonseparable kernels, making memory or operation costs too high at large scale.Source: §4.1.1 · p. 16

4.1.2. Separable log kernel.

• For a d-dimensional histogram, separable kernel application costs O(nd+1) rather than O(N^2), motivating a stable, separable log-domain kernel.The computation transfers inputs to the log-domain and shifts values by their maximum.Source: §4.1.2 · p. 17

4.3. Sinkhorn heavyball.

• Negative τ introduces extrapolation akin to a momentum term and can accelerate Sinkhorn; automatic differentiation feeds gradients from this barycenter loop into dictionary learning.The acceleration scheme is used within the method by differentiating the forward Sinkhorn loop.Source: §4.3 · p. 19

4.2. Warm start.

• In Figure 4a, warm start with L=2 yielded better results than L=100 without warm start, with a reported 36 gain factor in time.Across the two datasets, warm start improved reconstructions at equal Sinkhorn iterations, though the second dataset did not match L=100 with only two iterations.Source: §4.2 · pp. 18–19, Figure 4 · p. 20

4.4. Unbalanced.

• KL-penalized marginal mismatch relaxes mass conservation and can yield barycenters retaining more input structure; ρ controls departure from balanced OT, with ρ=∞ recovering usual OT.Automatic differentiation supplies dictionary and weight gradients for the unbalanced-barycenter representation.Source: §4.4 · pp. 20–21

5.1. Comparison with Wasserstein principal geodesics.

• For digits 2, 3, and 4, learned two-atom geodesics are visually extremely close to first WPGs; the 1s case does not recover the first WPG.The authors note that the learned span need not contain the input isobarycenter, and non-convex optimization may reach another local minimum.Source: Figure 5 · p. 21, §5.1 · p. 21
• With three atoms, the learned simplex combines captured warps, including loop and vertical-size changes.Intermediate weights allow combinations of the captured warpings, including a loopless, vertically shrunk 2.Source: §5.1 · pp. 21–22, Figure 6 · p. 22

5.2. Point spread functions.

• 1.71 × 10−3 average normalized mean square error for our method versus 2.62 × 10−3 for NMF, with middle-wavelength warping captured more fully.The learned atoms are visually close to the PSFs at the spectrum’s extremes.Source: Figure 8 · p. 23, §5.2 · p. 24

5.3. Cardiac sequences.

• The four-atom barycentric path forms a cycle, and reconstructed frames resemble inputs, indicating that learned key frames represent the sequence.Figure 12 plots the reconstructed frames by their barycentric coordinates in the four-atom basis.Source: §5.3 · p. 24, Figure 12 · p. 26

5.4. Wasserstein faces.

• Color inversion improved face results in most cases, atom interpolation supports emotion editing, and the fitting loss affects learned-atom appearance.The editing starts at the atoms’ isobarycenter and interpolates weights toward a selected atom.Source: §5.4 · p. 27, Figure 15 · p. 29

5.5. Literature learning.

• Five learned atoms group 20 novels largely by author; Dickens shares vocabulary too broadly to receive a distinct atom.The authors used only words seven to eight letters long in this experiment.Source: §5.5 · p. 27

5.4. Wasserstein faces.

• 33.8 mean PSNR for our method versus 33.6 for PCA, 33.5 for NMF, and 33.6 for K-SVD.The comparison uses five atoms for each method and 20 images of one person performing five expressions.Source: §5.4 · p. 27, Figure 13 · p. 28

5.6. Multimodal distributions.

• For data with disjoint modes, unbalanced reconstructions preserve the right mode count and grid regions, whereas balanced reconstructions can add or omit mass.The comparison illustrates a setting where the authors suggest the unbalanced version may be preferable.Source: §5.6 · pp. 27, 30

6. Conclusion.

• The method fits data to Wasserstein barycenters of learned atoms and offers entropically regularized computation schemes and variants.The paper illustrates the resulting representation on several applications.Source: §6 · p. 30
• Unbalanced barycenters place reconstructed modes at the right grid positions, while balanced barycenters can add spurious mass or omit original modes.The balanced examples also include mass in areas empty across all datapoints.Source: Figure 17 · p. 31
• Dictionary and barycentric-weight gradients are derived by differentiating the barycenter computation with respect to each.The conclusion points to faster Sinkhorn routines, including Greenkhorn [4] and multiscale methods [71], as possible integrations to speed the algorithm.Source: §6 · p. 31, Appendix A · p. 32

Appendix A. Proof of Proposition 3.1.

• Differentiating (15) supplies the next step in the proof.The passage gives no further readable detail about this differentiation.Source: Appendix A · p. 32
• The proof substitutes (27) and (29) into (11), using definitions (20)–(21).These substitutions are used in the subsequent derivation.Source: Appendix A · p. 32
• A constant initialization of b(0) establishes (18), and differentiation with respect to λ leads to (19).The derivation details for the differentials are deferred to Appendix C.Source: Appendix A · p. 32

Appendix B. Stabilized backward loop.

• Algorithm 4 takes data, dictionary atoms, and current barycentric weights as inputs.Its log-domain variables are marked with a tilde.Source: Appendix B · p. 33
• The algorithm runs a forward Sinkhorn loop followed by backward loops for the weights and dictionary gradients.The displayed fragments specify separate backward passes for the two gradient types.Source: Appendix B · p. 33
• Its outputs are the approximate barycenter and gradients with respect to the dictionary and weights.The output notation names these as P(L)(D, λ), ∇DE(L), and ∇λE(L).Source: Appendix B · p. 33

C.1. Computation of ∂bϕ.

• The passages state that taking the logarithm of (16) yields a result, and display ∂ϕ(bS,dS) in a fragment that also contains ∂b2.Source: §C.1 · p. 34, §C.2 · p. 34

C.4. Computation of Φb.

• C.4 computes Φb and notes that dependencies on an atom di pass through Ψ for the other indexed terms.The passage states this dependency for s ≠ i.Source: §C.4 · p. 35, §C.5 · p. 36

Appendix D. Generalized barycenters.

• Appendix D gives accelerated HeavyballSinkhorn and GeneralizedSinkhorn algorithms for approximate and unbalanced barycenters, respectively.The unbalanced algorithm also takes a KL parameter ρ > 0.Source: Appendix D · p. 37

• The appendix adds figures for the application in subsection 5.1.Source: §E.1 · p. 38

E.1. MNIST and Wasserstein geodesics.

• The two-atom dictionary span uses weights (1 −t, t) for images of digits 1, 2, 3, and 4.The caption refers readers to the first columns of [74, Figure 5] for comparison with first WPGs.Source: Figure E.1 · p. 38, §E.1 · p. 38
• Additional figures show results when training on images of digits 1, 3, and 4.They are described as matching Figure 6; the appendix also includes additional figures for subsection 5.2.Source: Figure E.2 · p. 39, Figure E.3 · p. 39, Figure E.4 · p. 40

E.2. Point spread functions.

• Extreme-wavelength dataset PSFs are shown alongside NMF-learned atoms, with the method’s learned atoms in Figure 9.Source: Figure E.5 · p. 40

E.3. Wasserstein faces.

• Mean PSNR rankings for atoms learned with different loss functions are 33.81, 33.72, 32.95, and 32.34.The comparison is described as similar to Figure 14.Source: Figure E.6 · p. 41
• The method is compared with Eigenfaces [81], NMF, and K-SVD for representing faces in a low-dimensional space.The comparison is described as similar to Figure 13; the appendix also adds figures for subsection 5.4.Source: §E.3 · p. 41, Figure E.7 · p. 42