Theory Digest — Mar 27, 2026

1. Motivation (2–3 sentences): Learning optimal transport (OT) maps between probability measures is crucial in machine learning, statistics, and computational geometry, but suffers from the curse of dimensionality. Existing methods either lack theoretical guarantees or are computationally intractable in high dimensions.

2. Mathematical setup: Let $\varrho_0 \in \mathcal{P}_{ac}^2(\mathbb{R}^d)$ be an absolutely continuous source probability measure and $\gamma \in \mathcal{P}_2(\mathbb{R}^d)$ be a target probability measure. Define the cone of gradients of convex functions:

   \[K_{\varrho_0} := \{\nabla\phi \mid \phi \in \dot{H}^1_{\varrho_0}(\mathbb{R}^d, \mathbb{R}) \text{ is convex}\} \subset L^2_{\varrho_0}(\mathbb{R}^d, \mathbb{R}^d)\]

   By Brenier’s theorem, the optimal transport map $T_{\varrho_0}^\gamma$ between $\varrho_0$ and $\gamma$ exists uniquely and belongs to $K_{\varrho_0}$. For divergence $D: \mathcal{P}_2(\mathbb{R}^d) \to \mathbb{R}$, define the lifted functional $F: T \mapsto D(T_*\varrho_0)$.

   Assumptions:

   1. $\varrho_0$ has density with respect to Lebesgue measure
   2. $D$ is lower semicontinuous with respect to weak topology on $\mathcal{P}_2(\mathbb{R}^d)$
   3. $D$ is Wasserstein differentiable
   4. $D$ is $\lambda$-convex along generalized geodesics

3. Toy example: When $d=2$, $\varrho_0 = \mathcal{N}(0, I_2)$, and $\gamma = \mathcal{N}(\mu, I_2)$, the optimal transport map is simply $T(x) = x + \mu$. Using relative entropy $D(\varrho) = H(\varrho

   \gamma)$, the lifted functional becomes $F(T) = H(T_*\varrho_0

   \gamma)$, and the constrained gradient flow seeks to minimize this over $K_{\varrho_0}$.

4. Formal objective:

   \[T_{\varrho_0}^\gamma \in \arg\min_{T \in K_{\varrho_0}} D(T_*\varrho_0 | \gamma)\]

The method defines a constrained gradient flow in the space $L^2_{\varrho_0}(\mathbb{R}^d, \mathbb{R}^d)$ constrained to the cone $K_{\varrho_0}$ of optimal transport maps:

\[\partial_t T_t = -\text{proj}_{\text{Tan}_{T_t} K_{\varrho_0}}(\nabla_W D(T_{t*}\varrho_0) \circ T_t)\]

with $T_0 = \text{id}$, where $\text{proj}$ denotes projection onto the tangent cone.

Key steps:

1. Compute Wasserstein gradient $\nabla_W D(T_{t*}\varrho_0)$
2. Compose with current map: $\nabla_W D(T_{t*}\varrho_0) \circ T_t$
3. Project onto tangent cone of $K_{\varrho_0}$ at $T_t$

4. Update: $T_{t+dt} = T_t - dt \cdot \text{projection}$

   For practical implementation, two discrete schemes are proposed:

   Explicit scheme:

   \[\hat{T}_{k+1}^{\tau} \in \arg\min_{T \in K_{\varrho_0}} \int_{\mathbb{R}^d} \left\|\frac{-\nabla_W D(\hat{T}_k^{\tau*}\varrho_0) \circ \hat{T}_k^{\tau} - (T - \hat{T}_k^{\tau})}{\tau}\right\|^2 d\varrho_0\]

   Implicit scheme:

   \[\hat{T}_{k+1}^{\tau} \in \arg\min_{T \in K_{\varrho_0}} D(T_*\varrho_0) + \frac{1}{2\tau}\|T - \hat{T}_k^{\tau}\|_{L^2_{\varrho_0}}^2\]

   Applied to toy example: When $D$ is relative entropy and $\gamma = \mathcal{N}(\mu, I_2)$, $\nabla_W H(\cdot

   \gamma) = \nabla\log\varrho + \nabla V$ where $V(x) = |x-\mu|^2/2$. The flow evolves $T_t$ toward the translation $x \mapsto x + \mu$.

Step 1 — Prior work:

1. “Geometry and geodesics in spaces of probability measures” (Modin, 2016) suggested lifting divergences to transport map spaces but only studied finite-dimensional Gaussian case
2. “Neural optimal transport” (Korotin et al., 2021) used neural networks for OT but without theoretical convergence guarantees

3. “Gradient flows in Wasserstein space” (Ambrosio et al., 2008) established theory for Wasserstein gradient flows but in probability space, not transport map space

   Step 2 — Delta: This paper provides the first rigorous theoretical analysis of constrained gradient flows in the space of transport maps. New contributions include:

4. existence theorem for constrained gradient flows
5. exponential convergence guarantees to optimal transport map
6. connection to natural gradient descent

7. practical discrete schemes with convergence analysis.

   Step 3 — Theory-specific assessment:

   • Main theorem is moderately surprising: while gradient flows in convex sets are well-understood, the transport map constraint creates a non-trivial geometric setting
   • Proof technique is largely routine, assembling known results from convex analysis and optimal transport, but the constraint geometry requires careful handling
   • Convergence rates are dimension-independent, which is notable for high-dimensional OT problems
   • No lower bounds are established; optimality of rates unknown

   The theoretical guarantees for constrained flows to transport maps represent a clear advance, but the techniques are incremental applications of established convex optimization theory.

   Verdict: INCREMENTAL — solid theoretical analysis extending known gradient flow techniques to transport map constraints, but uses standard proof methods without fundamental innovations.

The main proof strategy follows standard convex analysis in Hilbert spaces, adapted to the transport map constraint:

1. Convexity and closedness of $K_{\varrho_0}$ (Proposition 2.1): Uses continuity of pushforward mapping and uniqueness of optimal transport to show that strong convergence in $L^2_{\varrho_0}$ preserves membership in $K_{\varrho_0}$.

2. Existence via minimizing movements (Theorem 2.11): Applies generalized minimizing movement scheme. Key inequality for proximal step well-posedness:

   \[J_\tau(T) = F(T) + \mathbf{1}_{K_{\varrho_0}}(T) + \frac{1}{2\tau}\|T - \hat{T}_k^\tau\|_{L^2_{\varrho_0}}^2\]

   For small enough $\tau$, $J_\tau$ is $(\tau^{-1} + \lambda)$-convex, ensuring unique minimizers.

3. Convergence analysis (Theorem 2.13): Uses projection characterization and convexity. Central energy estimate:

   \[\frac{d}{dt}[F(T_t) - F(T_{\varrho_0}^\gamma)] = \langle\nabla F(T_t), \partial_t T_t\rangle_{L^2_{\varrho_0}} = -\|w_t\|_{L^2_{\varrho_0}}^2 \leq 0\]

   For $\lambda$-convex case, star-convexity around minimizer gives:

   \[F(T_t) - F(T_{\varrho_0}^\gamma) \leq \langle T_t - T_{\varrho_0}^\gamma, -w_t\rangle_{L^2_{\varrho_0}} - \frac{\lambda}{2}\|T_t - T_{\varrho_0}^\gamma\|_{L^2_{\varrho_0}}^2\]

   Combined with Young’s inequality yields Polyak-Łojasiewicz condition, leading to exponential convergence via Grönwall’s lemma.

4. Lifted functional analysis (Appendix B): Shows that convexity and differentiability properties transfer from $D$ on $\mathcal{P}_2(\mathbb{R}^d)$ to lifted functional $F$ on $L^2_{\varrho_0}$. Gradient relationship:

   \[\nabla F(T) = \nabla_W D(T_*\varrho_0) \circ T\]

Purely theoretical with limited numerical illustration. The paper provides a proof-of-concept using Input Convex Neural Networks (ICNNs) to parameterize $K_{\varrho_0}$, comparing explicit vs implicit schemes against standard Euclidean gradient descent and Adam optimizer.

Empirical validation would require:

1. comprehensive comparison with state-of-the-art OT solvers on standard benchmarks
2. scalability analysis across dimensions
3. statistical analysis of approximation quality
4. computational cost comparisons. The current experiments only demonstrate feasibility rather than practical advantages.

Limitations:

1. RESTRICTIVE: Requires source measure $\varrho_0$ to be absolutely continuous with respect to Lebesgue measure
2. RESTRICTIVE: Target measure $\gamma$ must be log-concave for strongest results (exponential convergence)
3. TECHNICAL: Power-type growth condition (pg) needed for convergence rates in merely convex case
4. RESTRICTIVE: John domain assumption for relative entropy convergence rates significantly limits applicability

5. NATURAL: Convexity constraint parameterization (ICNNs) standard in the field

   Open problems:

6. Statistical analysis: Develop finite-sample convergence guarantees and generalization bounds when working with empirical measures
7. Computational complexity: Establish iteration complexity bounds for discrete schemes and compare with existing OT solvers in high dimensions

1. Motivation: Stochastic differential equation-based Bayesian neural networks (SDE-BNNs) provide strong theoretical foundations for modeling uncertainty in deep learning, but they require excessive computational cost due to large numbers of function evaluations (NFEs) during numerical SDE solving. This limits their practical applicability in real-world scenarios requiring efficient training and inference.

2. Mathematical setup: Consider a joint SDE system for weights $w_t \in \mathbb{R}^d$ and hidden states $h_t \in \mathbb{R}^p$ on time interval $[0,T]$:

   \[d\begin{pmatrix} w_t \\ h_t \end{pmatrix} = \begin{pmatrix} f_\phi(w_t, t) \\ f_{w_t}(h_t, t) \end{pmatrix} dt + \begin{pmatrix} g_\theta(w_t, t) \\ 0 \end{pmatrix} dB_t\]

   where $B_t$ is Brownian motion, $f_{w_t}$ represents the dynamics of network activations parameterized by weights $w_t$, and $f_\phi$ is a learned drift for the weight posterior. The training objective is the ELBO:

   \[L_{\text{ELBO}}(\phi) = \mathbb{E}_{q_\phi(w)}\left[\log p(y|x,w) - \int_0^T \frac{1}{2}\|u(w_t,t,\phi)\|_2^2 dt\right]\]

   where $u(w_t,t,\phi) = g(w_t,t)^{-1}[f_q(w_t,t,\phi) - f_p(w_t,t)]$.

   Assumptions:

   1. Prior over weights follows Ornstein-Uhlenbeck process with drift $f_p(w_t,t) = -w_t$
   2. Diffusion term is constant: $g(w_t,t) = \sigma I_d$
   3. Lipschitz continuity of drift functions for numerical stability

3. Toy example: For $d=2$ weights and scalar hidden state, the system becomes a 3-dimensional SDE. Standard solvers require 200-400 NFEs per forward pass, making training prohibitively expensive.

4. Formal objective: Minimize NFEs while maintaining or improving predictive accuracy:

   \[\min_\phi \left\{ \text{NFEs}, -L_{\text{ELBO}}(\phi) \right\}\]

The method integrates Nesterov’s accelerated gradient (NAG) into SDE-BNN by extending to second-order dynamics. The joint state becomes $H_t = (x_t, m_t, w_t)$ where $x_t$ is auxiliary variable and $m_t$ is momentum:

\[h_t = \sigma_f\left(t^{-3/2} e^{t/2}\right) x_t\] \[dH_t = \begin{pmatrix} \sigma_f(m_t) \\ -m_t - \sigma_f(f_{w_t}(h_t,t) + \epsilon\xi h_{\text{temp}}) \\ f_\phi(w_t,t) \end{pmatrix} dt + \begin{pmatrix} 0 \\ 0 \\ g(w_t,t) \end{pmatrix} dB_t\]

The key innovation is the NFE-dependent residual skip connection:

\[\epsilon = \begin{cases} 1 & \text{if NFE is odd} \\ 0 & \text{if NFE is even} \end{cases}\] \[h_{\text{temp}} = h_t \text{ if NFE is even}\]

Training procedure:

1. Initialize $w_0$, drift parameters $\phi$
2. For each batch: sample weight trajectory from posterior
3. Solve the augmented SDE system using adaptive solvers

4. Compute ELBO loss and update parameters via gradient descent

   Toy example application: For $d=2$, the method caches $h_t$ at odd NFEs and injects it at even NFEs, creating skip connections every two solver steps. This reduces required NFEs from ~400 to ~240 while maintaining solution quality.

Step 1 — Prior work:

• “Neural Ordinary Differential Equations” (Chen et al., 2018): introduced continuous-depth networks via ODEs
• “Improving Neural ODEs with Nesterov’s Accelerated Gradient Method” (Nguyen et al., 2022): applied NAG to deterministic Neural ODEs
• “Infinitely Deep Bayesian Neural Networks with Stochastic Differential Equations” (Xu et al., 2022): combined SDEs with Bayesian inference for uncertainty quantification

Step 2 — Delta: This paper combines NAG acceleration with SDE-BNN framework and introduces NFE-dependent residual connections. The combination is new, but each component exists separately.

Step 3 — Theory-specific assessment:

• Main contribution is engineering-focused: reducing computational cost while maintaining performance
• No new theoretical insights about convergence rates, approximation error, or Bayesian properties
• The NFE-dependent residual mechanism is an ad-hoc heuristic without theoretical justification
• No analysis of how NAG affects the posterior approximation quality or ELBO convergence

The bounds and convergence properties inherit from existing NAG and SDE-BNN analyses. No lower bounds are established for the NFE reduction.

Verdict: INCREMENTAL — straightforward combination of existing techniques with empirical validation but no theoretical novelty.

This is primarily an applied paper without formal proofs. The method combines existing techniques:

1. NAG dynamics from Nguyen et al. (2022): extends first-order ODE

   \[\frac{dh_t}{dt} = f_\theta(h_t, t)\]

   to second-order system

   \[\frac{d^2h_t}{dt^2} + \frac{3}{t}\frac{dh_t}{dt} + f_\theta(h_t, t) = 0\]

2. SDE-BNN framework from Xu et al. (2022): joint dynamics for weights and activations with ELBO objective:

   \[L_{\text{ELBO}}(\phi) = \mathbb{E}_{q_\phi(w)}\left[\log p(y|x,w) - \int_0^T \frac{1}{2}\|u(w_t,t,\phi)\|_2^2 dt\right]\]

3. The main technical contribution is the NFE-dependent residual connection scheme, which is presented as an engineering heuristic without theoretical analysis.

   No concentration inequalities, convergence rate analysis, or approximation error bounds are provided. The paper relies on empirical validation rather than mathematical proofs.

Datasets: MNIST, CIFAR-10, Walker2D kinematic simulation, 1D toy regression

Baselines: Standard SDE-BNN with both fixed-step and adaptive solvers

Key numbers:

• MNIST: 99.04% vs 98.90% accuracy, 240 vs 400 NFEs
• CIFAR-10: 88.36% vs 87.56% accuracy, 170 vs 270 NFEs
• Walker2D: 145 vs 180 NFEs, faster convergence
• Consistent 20-40% reduction in NFEs across tasks

Implementation: PyTorch on NVIDIA A100, midpoint solver, various learning rates and architectures depending on task.

The experiments are reasonably comprehensive but limited to relatively simple datasets. No comparison with other acceleration methods or recent SDE variants beyond standard SDE-BNN.

Limitations:

1. NFE-dependent residual scheme lacks theoretical justification - TECHNICAL (heuristic design without analysis)
2. Limited to simple datasets (MNIST, CIFAR-10) - RESTRICTIVE (scalability unclear)
3. No analysis of how NAG affects Bayesian posterior quality - TECHNICAL (could be analyzed)
4. Comparison only against vanilla SDE-BNN, not other acceleration methods - TECHNICAL (broader baselines needed)

5. Hyperparameter sensitivity not thoroughly studied - NATURAL (standard limitation)

   Open problems:

6. Theoretical analysis of convergence rates and posterior approximation quality under NAG dynamics
7. Scalability evaluation on large-scale datasets and modern architectures like transformers

1. Motivation (2-3 sentences): Diffusion language models (DLMs) offer flexible generation order and parallel token generation, but their efficiency critically depends on the decoding strategy. Confidence-based decoding strategies have shown strong empirical performance but lack theoretical understanding.

2. Mathematical setup:
   • Discrete vocabulary: $X$ with mask symbol $M$, extended vocabulary $\bar{X} = X \cup {M}$
   • Target data distribution: $X_0 \in X^L$ sampled from $p_{\text{data}}$
   • Mask predictor: Product of conditional marginals $p_\theta = \prod_{i=1}^L p_i^\theta$ where
   \[p_i(x | z_S) := P[X_i^0 = x | X_S^0 = z_S]\]

   for any unmasked position set $S \subseteq [L] \setminus {i}$.

   Assumptions:

   1. Mask predictor is optimal: $p_\theta = \prod_{i=1}^L p_i$ (product of true conditional marginals)
   2. Entropy sum-based stopping: iteration terminates when cumulative entropy exceeds threshold $\eta$
   3. Random permutation order for token unmasking

3. Toy example: When $L=2$, $X_0 = (X_1^0, X_2^0)$ with entropy $H(X_0)$, the algorithm randomly orders positions as $(\Pi_1, \Pi_2)$ and unmasks $\Pi_1$ first if $H(p_{\Pi_1}^\theta(\cdot

   Y_0)) \leq \eta$, otherwise stops and unmasks both tokens.

4. Formal objective: Achieve $\varepsilon$-accurate sampling measured by expected KL divergence:

   \[E_\Pi[KL(p_{X_0} \| p_{Y_T|\Pi})] \leq \varepsilon\]

The entropy sum-based decoding strategy (Algorithm 1) works as follows:

1. Draw random permutation $\Pi$ of $[L]$ uniformly

2. At each iteration $t$, starting with unmasking set $D_t = \emptyset$ and entropy sum $S = 0$: - For next position $\Pi_k$ in permutation order, compute entropy:

   \[H(p_{\Pi_k}^\theta(\cdot | Y_{t-1}))\]

   - Sample token $Y\_{\Pi\_k}^t \sim p\_{\Pi\_k}^\theta(\cdot | Y\_{t-1})$
   - Add $\Pi\_k$ to $D\_t$ and update entropy sum $S$
   - Stop iteration when $S > \eta$ or all tokens generated
   

3. Stopping criterion: iteration ends at first index $k_t$ where:

   \[k_t := \min\left\{i > W_{t-1} : \sum_{k=W_{t-1}+1}^i H(p_{\Pi_k}^\theta(\cdot | Y_{t-1})) > \eta\right\} \wedge L\]

   Applied to toy example: With $L=2$, if $\Pi = (1,2)$ and $\eta = 0.5$:

   • If $H(p_1^\theta(\cdot

     (M,M))) = 0.3 < \eta$, unmask position 1

   • If $H(p_2^\theta(\cdot

     Y_1)) = 0.4$, then $S = 0.3 + 0.4 = 0.7 > \eta$, so stop and unmask position 2

   • Total: 1 iteration, 2 tokens unmasked

Step 1 — Prior work:

• Li and Cai (2025): Established $O(1/T)$ convergence for uniform decoding with balanced step-sizes, leading factor is sum of mutual information $\sum_i I(X_i; X_{\setminus i})$
• Chen et al. (2025b): Refined convergence using total correlation, developed optimal step-size schedules requiring prior knowledge
• Ben-Hamu et al. (2025): Proposed empirical entropy sum-based strategy with entropy sorting, showed strong empirical results

Step 2 — Delta: This paper provides the first theoretical analysis of confidence-based (specifically entropy-based) decoding for DLMs. Key contributions:

1. proves $\tilde{O}(H(X_0)/\varepsilon)$ iteration complexity for entropy sum-based strategy
2. shows automatic adaptation to data complexity without hyperparameter tuning

3. develops novel analysis framework for adaptive decoding dynamics.

   Step 3 — Theory-specific assessment:

   • Main theorem is moderately surprising: shows sublinear iteration complexity $\tilde{O}(H(X_0)/\varepsilon)$ vs. $\Theta(L)$ when $H(X_0) \ll L$, which wasn’t obvious from prior uniform decoding analysis
   • Proof technique is genuinely new: introduces novel dyadic size envelopes to handle adaptive unmasking sets, uses KL decomposition via mutual information with careful indicator function analysis
   • Bounds appear reasonably tight: iteration complexity scales with entropy $H(X_0)$ (natural complexity measure), though no matching lower bounds exist yet

   Verdict: SIGNIFICANT — provides first rigorous theory for practically important confidence-based decoding, with novel proof techniques and meaningful complexity improvement.

The proof uses five key technical steps:

1. KL decomposition along adaptive trajectory: Express sampling error as:

   \[E_\Pi[KL(p_{X_0} \| p_{Y_T|\Pi})] = E_{X,\Pi}\left[\sum_{t=1}^T KL\left(p_{D_t}(\cdot | X_{W_{t-1}}) \| \prod_{i \in D_t} p_i(\cdot | X_{W_{t-1}})\right)\right]\]

2. Mutual information decomposition (Lemma 1): For any joint distribution vs. product of marginals:

   \[KL\left(p_{Z_1,\ldots,Z_d} \| \prod_{i=1}^d p_{Z_i}\right) \leq \sum_{i \in [d] \setminus \{j\}} I(Z_i; Z_{[d] \setminus \{i\}})\]

3. Dyadic size envelopes: Define sequence $(D_t)_{t=1}^L$ with $D_1 = 0$ and:

   \[D_t := \begin{cases} D_{t-1}, & \text{if } |D_{t-1}| < D_{t-1} \\ 2D_{t-1}, & \text{if } |D_{t-1}| \geq D_{t-1} \end{cases}\]

4. Key averaging lemma (Lemma 2): After averaging over random token position:

   \[E[I(X_i; X_{D_{t_i} \setminus \{i\}} | \cdots) \mathbf{1}_{\text{entropy} \leq \eta}] \leq 2E[(H(X_i | \cdots) \wedge \eta) \mathbf{1}_{|D_{t_i}| \geq D_{t_i}}]\]

5. Entropy budget control: Stopping rule ensures at most $\log_2 L + 1$ iterations have $

   D_t

   \geq D_t$, and entropy sum per iteration is bounded by $2\eta$, yielding:

   \[E_\Pi[KL(p_{X_0} \| p_{Y_T|\Pi})] \leq 4\eta(\log_2 L + 1)\]

Purely theoretical. Empirical validation would require:

1. comparing iteration counts of entropy sum-based vs. uniform decoding on datasets with varying entropy $H(X_0)$
2. measuring actual KL divergence convergence rates
3. testing adaptation to different data complexities without hyperparameter tuning
4. verifying $\tilde{O}(H(X_0)/\varepsilon)$ scaling empirically.

Limitations:

1. Assumes optimal mask predictor $p_\theta = \prod_i p_i$ - RESTRICTIVE (ignores training error effects on confidence-based decoding)
2. Uses random permutation instead of deterministic entropy-sorted order - TECHNICAL (simplifies analysis but differs from practice)
3. Analysis limited to entropy-based confidence measures - TECHNICAL (other measures like probability mass, margin gaps not covered)

4. No lower bounds on iteration complexity - NATURAL (leaves optimality question open)

   Open problems:

5. Extend analysis to suboptimal mask predictors with training error
6. Establish matching lower bounds for confidence-based decoding strategies