arXiv subjects

Viktor Stein

Publications and source records attributed to Viktor Stein.

10 recordsLinked to original sources

arXiv2026

SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds

Transformers owe much of their empirical success in natural language processing to the self-attention blocks. Recent perspectives interpret attention blocks as interacting particle systems, whose mean-field limits correspond to gradient flows of interaction energy functionals on probability density spaces equipped with Wasserstein-$2$-type metrics. We extend this viewpoint by introducing accelerated attention blocks derived from inertial Nesterov-type dynamics on density spaces. In our proposed architecture, tokens carry both spatial (feature) and velocity variables. The time discretization and the approximation of accelerated density dynamics yield Hamiltonian momentum attention blocks, which constitute the proposed accelerated attention architectures. In particular, for linear self-attention, we show that the attention blocks approximate a Stein variational gradient flow, using a bilinear kernel, of a potential energy. In this setting, we prove that elliptically contoured probability distributions are preserved by the accelerated attention blocks. We present implementable particle-based algorithms and demonstrate that the proposed accelerated attention blocks converge faster than the classical attention blocks while preserving the number of oracle calls.

cs.LG↗

arXiv2026

Well-Posed KL-Regularized Control via Wasserstein and Kalman-Wasserstein KL Divergences

Kullback-Leibler (KL) divergence regularization is widely used in reinforcement learning, but it becomes infinite under support mismatch and can degenerate in low-noise regimes. Using a unified information-geometric framework, we introduce KL analogs by replacing the Fisher-Rao geometry in the dynamical formulation of the KL with transport-based geometries, and derive closed-form expressions for common distribution families. Between elliptic distributions, these divergences remain finite for degenerating equal covariances and yield a geometric interpretation of regularization heuristics used in Kalman ensemble methods. We demonstrate the utility of these divergences in KL-regularized optimal control. In the fully tractable setting of linear time-invariant systems with Gaussian process noise, the classical KL reduces to a quadratic control penalty that becomes singular as process noise vanishes. Our variants remove this singularity and yield well-posed problems. In both the double integrator and cart-pole examples, the resulting controls preserve nontrivial feedback and achieve better closed-loop performance.

math.OC↗

arXiv2026

Approximating $f$-Divergences with Rank Statistics

We introduce a rank-statistic approximation of $f$-divergences that avoids explicit density-ratio estimation by working directly with the distribution of ranks. For a resolution parameter $K$, we map the mismatch between two univariate distributions $\mu$ and $\nu$ to a rank histogram on $\{ 0, \ldots, K\}$ and measure its deviation from uniformity via a discrete $f$-divergence, yielding a rank-statistic divergence estimator. We prove that the resulting estimator of the divergence is monotone in $K$, is always a lower bound of the true $f$-divergence, and we establish quantitative convergence rates for $K\to\infty$ under mild regularity of the quantile-domain density ratio. To handle high-dimensional data, we define the sliced rank-statistic $f$-divergence by averaging the univariate construction over random projections, and we provide convergence results for the sliced limit as well. We also derive finite-sample deviation bounds along with asymptotic normality results for the estimator. Finally, we empirically validate the approach by benchmarking against neural baselines and illustrating its use as a learning objective in generative modeling experiments.

stat.ML↗

arXiv2025

Biological Engineering: What does it mean? Where does it (need to) go?

Biological engineering, the convergence between engineering and biology, is at the forefront of significant advances in healthcare, agriculture, and environmental sustainability, making it highly relevant to current scientific and societal challenges. We take a comprehensive look at this broad and interdisciplinary domain, structure it into three main areas - bioinspired, biological and biohybrid approaches - and dissect inherent and fundamental challenges along with opportunities, highlighting specific examples. We describe how data-driven discovery and design, in conjunction with artificial intelligence, can mitigate the absence of reductionist models in these areas. Additionally, we address the education of a new generation of biological engineers, emphasizing mathematical, technical, and artificial intelligence frameworks.

q-bio.OT↗

arXiv2025

Towards understanding Accelerated Stein Variational Gradient Flow -- Analysis of Generalized Bilinear Kernels for Gaussian target distributions

Stein variational gradient descent (SVGD) is a kernel-based and non-parametric particle method for sampling from a target distribution, such as in Bayesian inference and other machine learning tasks. Different from other particle methods, SVGD does not require estimating the score, which is the gradient of the log-density. However, in practice, SVGD can be slow compared to score-estimation-based sampling algorithms. To design a fast and efficient high-dimensional sampling algorithm with the advantages of SVGD, we introduce accelerated SVGD (ASVGD), based on an accelerated gradient flow in a metric space of probability densities following Nesterov's method. We then derive a momentum-based discrete-time sampling algorithm, which evolves a set of particles deterministically. To stabilize the particles' position update, we also include a Wasserstein metric regularization. This paper extends the conference version \cite{SL2025}. For the bilinear kernel and Gaussian target distributions, we study the kernel parameter and damping parameters with an optimal convergence rate of the proposed dynamics. This is achieved by analyzing the linearized accelerated gradient flows at the equilibrium. Interestingly, the optimal parameter is a constant, which does not depend on the covariance of the target distribution. For the generalized kernel functions, such as the Gaussian kernel, numerical examples with varied target distributions demonstrate the effectiveness of ASVGD compared to SVGD and other popular sampling methods. Furthermore, we show that in the setting of Bayesian neural networks, ASVGD outperforms SVGD significantly in terms of log-likelihood and total iteration times.

math.OC↗

arXiv2025

Accelerated Stein Variational Gradient Flow

Stein variational gradient descent (SVGD) is a kernel-based particle method for sampling from a target distribution, e.g., in generative modeling and Bayesian inference. SVGD does not require estimating the gradient of the log-density, which is called score estimation. In practice, SVGD can be slow compared to score-estimation based sampling algorithms. To design fast and efficient high-dimensional sampling algorithms, we introduce ASVGD, an accelerated SVGD, based on an accelerated gradient flow in a metric space of probability densities following Nesterov's method. We then derive a momentum-based discrete-time sampling algorithm, which evolves a set of particles deterministically. To stabilize the particles' momentum update, we also study a Wasserstein metric regularization. For the generalized bilinear kernel and the Gaussian kernel, toy numerical examples with varied target distributions demonstrate the effectiveness of ASVGD compared to SVGD and other popular sampling methods.

stat.ML↗

arXiv2024

Wasserstein Gradient Flows of MMD Functionals with Distance Kernels under Sobolev Regularization

We consider Wasserstein gradient flows of maximum mean discrepancy (MMD) functionals $\text{MMD}_K^2(\cdot, \nu)$ for positive and negative distance kernels $K(x,y) := \pm |x-y|$ and given target measures $\nu$ on $\mathbb{R}$. Since in one dimension the Wasserstein space can be isometrically embedded into the cone $\mathcal C(0,1) \subset L_2(0,1)$ of quantile functions, Wasserstein gradient flows can be characterized by the solution of an associated Cauchy problem on $L_2(0,1)$. While for the negative kernel, the MMD functional is geodesically convex, this is not the case for the positive kernel, which needs to be handled to ensure the existence of the flow. We propose to add a regularizing Sobolev term $|\cdot|^2_{H^1(0,1)}$ corresponding to the Laplacian with Neumann boundary conditions to the Cauchy problem of quantile functions. Indeed, this ensures the existence of a generalized minimizing movement for the positive kernel. Furthermore, for the negative kernel, we demonstrate by numerical examples how the Laplacian rectifies a "dissipation-of-mass" defect of the MMD gradient flow.

math.AP↗

arXiv2024

Wasserstein Gradient Flows of MMD Functionals with Distance Kernel and Cauchy Problems on Quantile Functions

We give a comprehensive description of Wasserstein gradient flows of maximum mean discrepancy (MMD) functionals $\mathcal F_\nu := \text{MMD}_K^2(\cdot, \nu)$ towards given target measures $\nu$ on the real line, where we focus on the negative distance kernel $K(x,y) := -|x-y|$. In one dimension, the Wasserstein-2 space can be isometrically embedded into the cone $\mathcal C(0,1) \subset L_2(0,1)$ of quantile functions leading to a characterization of Wasserstein gradient flows via the solution of an associated Cauchy problem on $L_2(0,1)$. Based on the construction of an appropriate counterpart of $\mathcal F_\nu$ on $L_2(0,1)$ and its subdifferential, we provide a solution of the Cauchy problem. For discrete target measures $\nu$, this results in a piecewise linear solution formula. We prove invariance and smoothing properties of the flow on subsets of $\mathcal C(0,1)$. For certain $\mathcal F_\nu$-flows this implies that initial point measures instantly become absolutely continuous, and stay so over time. Finally, we illustrate the behavior of the flow by various numerical examples using an implicit Euler scheme, which is easily computable by a bisection algorithm. For continuous targets $\nu$, also the explicit Euler scheme can be employed, although with limited convergence guarantees.

math.AP↗

arXiv2024

Interpolating between Optimal Transport and KL regularized Optimal Transport using R\'enyi Divergences

Regularized optimal transport (OT) has received much attention in recent years starting from Cuturi's introduction of Kullback-Leibler (KL) divergence regularized OT. In this paper, we propose regularizing the OT problem using the family of $\alpha$-R\'enyi divergences for $\alpha \in (0, 1)$. R\'enyi divergences are neither $f$-divergences nor Bregman distances, but they recover the KL divergence in the limit $\alpha \nearrow 1$. The advantage of introducing the additional parameter $\alpha$ is that for $\alpha \searrow 0$ we obtain convergence to the unregularized OT problem. For the KL regularized OT problem, this was achieved by letting the regularization parameter $\varepsilon$ tend to zero, which causes numerical instabilities. We present two different ways to obtain premetrics on probability measures, namely by R\'enyi divergence constraints and by penalization. The latter premetric interpolates between the unregularized and the KL regularized OT problem with weak convergence of the unique minimizer, generalizing the interpolation property of KL regularized OT. We use a nested mirror descent algorithm to solve the primal formulation. Both on real and synthetic data sets R\'enyi regularized OT plans outperform their KL and Tsallis counterparts in terms of being closer to the unregularized transport plans and recovering the ground truth in inference tasks better.

math.OC↗

arXiv2024

Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces

Commonly used $f$-divergences of measures, e.g., the Kullback-Leibler divergence, are subject to limitations regarding the support of the involved measures. A remedy is regularizing the $f$-divergence by a squared maximum mean discrepancy (MMD) associated with a characteristic kernel $K$. We use the kernel mean embedding to show that this regularization can be rewritten as the Moreau envelope of some function on the associated reproducing kernel Hilbert space. Then, we exploit well-known results on Moreau envelopes in Hilbert spaces to analyze the MMD-regularized $f$-divergences, particularly their gradients. Subsequently, we use our findings to analyze Wasserstein gradient flows of MMD-regularized $f$-divergences. We provide proof-of-the-concept numerical examples for flows starting from empirical measures. Here, we cover $f$-divergences with infinite and finite recession constants. Lastly, we extend our results to the tight variational formulation of $f$-divergences and numerically compare the resulting flows.

stat.ML↗