Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Sparse Transformer Architectures via Regularized Wasserstein Proximal Operator with $L_1$ Prior

Fuqun Han, Stanley Osher, Wuchen Li

TL;DR

The paper addresses the challenge of incorporating structural priors, notably sparsity, into high-dimensional generative modeling. It introduces a sparse transformer architecture built from the regularized Wasserstein proximal operator (RWPO), embedding a prior-driven proximal update into each layer and coupling it with a learnable transport potential to realize a probability-flow framework. Theoretical results show that a sparse L1 prior accelerates convergence in KL divergence and contracts mass more effectively under a viscous (diffusive) regularization, while numerical experiments across synthetic benchmarks, Bayesian inverse problems, and MNIST demonstrate improved accuracy and faster convergence relative to traditional neural ODE–based flows. This approach unifies optimal transport, mean-field dynamics, and transformer-style interactions, offering a scalable, prior-informed alternative for generative modeling and posterior estimation with practical impact for scientific computing and inverse problems.

Abstract

In this work, we propose a sparse transformer architecture that incorporates prior information about the underlying data distribution directly into the transformer structure of the neural network. The design of the model is motivated by a special optimal transport problem, namely the regularized Wasserstein proximal operator, which admits a closed-form solution and turns out to be a special representation of transformer architectures. Compared with classical flow-based models, the proposed approach improves the convexity properties of the optimization problem and promotes sparsity in the generated samples. Through both theoretical analysis and numerical experiments, including applications in generative modeling and Bayesian inverse problems, we demonstrate that the sparse transformer achieves higher accuracy and faster convergence to the target distribution than classical neural ODE-based methods.

Sparse Transformer Architectures via Regularized Wasserstein Proximal Operator with $L_1$ Prior

TL;DR

The paper addresses the challenge of incorporating structural priors, notably sparsity, into high-dimensional generative modeling. It introduces a sparse transformer architecture built from the regularized Wasserstein proximal operator (RWPO), embedding a prior-driven proximal update into each layer and coupling it with a learnable transport potential to realize a probability-flow framework. Theoretical results show that a sparse L1 prior accelerates convergence in KL divergence and contracts mass more effectively under a viscous (diffusive) regularization, while numerical experiments across synthetic benchmarks, Bayesian inverse problems, and MNIST demonstrate improved accuracy and faster convergence relative to traditional neural ODE–based flows. This approach unifies optimal transport, mean-field dynamics, and transformer-style interactions, offering a scalable, prior-informed alternative for generative modeling and posterior estimation with practical impact for scientific computing and inverse problems.

Abstract

In this work, we propose a sparse transformer architecture that incorporates prior information about the underlying data distribution directly into the transformer structure of the neural network. The design of the model is motivated by a special optimal transport problem, namely the regularized Wasserstein proximal operator, which admits a closed-form solution and turns out to be a special representation of transformer architectures. Compared with classical flow-based models, the proposed approach improves the convexity properties of the optimization problem and promotes sparsity in the generated samples. Through both theoretical analysis and numerical experiments, including applications in generative modeling and Bayesian inverse problems, we demonstrate that the sparse transformer achieves higher accuracy and faster convergence to the target distribution than classical neural ODE-based methods.
Paper Structure (16 sections, 5 theorems, 71 equations, 4 figures, 1 table, 1 algorithm)

This paper contains 16 sections, 5 theorems, 71 equations, 4 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Proposed network structure and PDE formulation
  3. Probability flow and particle updates
  4. Regularized Wasserstein proximal operator layer
  5. Sparse transformer with L1 prior
  6. Learning objective and training procedure
  7. Convergence and stability of flow with sparse transformer
  8. Convergence with the sparse prior
  9. Stability of the optimization problem
  10. Numerical experiments
  11. Benchmark examples
  12. Bayesian inverse problems
  13. Image generation for MNIST dataset.
  14. Summary and discussion
  15. Postponed derivation and proof
  16. ...and 1 more sections

Key Result

Lemma 1

Let $\psi:\mathbb{R}^d\to\mathbb{R}$ be a convex potential, and let $\rho_k \in \mathbb{P}^2(\mathbb{R}^d)$. Then the regularized Wasserstein proximal operator eqn:Wprox admits the integral representation where the kernel depends on $\psi$, the step size $h>0$, and the inverse temperature constant (diffusion strength) $\beta>0$.

Figures (4)

  • Figure 1: Network structure: the drift potential $\phi$ is approximated by a ResNet, followed by a sparse-transformer (RWPO) layer that enforces sparsity through proximal updates.
  • Figure 2: Comparison of generated samples across benchmark datasets. Each column corresponds to a different dataset. Top row: target density $\rho_0$. Middle row: generated samples. Bottom row: KDE-estimated distribution using a Gaussian kernel.
  • Figure 3: Top row: Moons dataset. Bottom row: 2-Spirals dataset. Columns show the target density and generated samples without and with the sparse transformer, respectively.
  • Figure 7: Digit interpolations on MNIST using the learned flow. The first example is generated without the sparse transformer, and the following three are generated with it.

Theorems & Definitions (9)

  • Lemma 1
  • Proposition 2
  • proof
  • Proposition 3: Log-density evolution along trajectories
  • Proposition 4
  • Proposition 5
  • proof : Proof of Lemma \ref{['lem:rwpo']}
  • proof : Proof of Proposition \ref{['prop_KL_L1']}
  • proof : Proof of Proposition \ref{['prop:moment-ordering']}