Implicit KL Proximal Descent (IKLPD)

• The paper introduces IKLPD as a unified framework that employs KL divergence as a proximal penalty to regulate optimization in both parameter and measure spaces.
• IKLPD recovers classical EM iterations and their accelerated variants, offering provable linear to superlinear convergence under standard regularity conditions.
• The framework extends to optimal transport and generative modeling through JKO flows and neural approximations, showcasing practical implementations via trust-region and flow-based methods.

Implicit KL Proximal Descent (IKLPD) is a unified framework for implicit optimization in statistical inference and measure space, where the Kullback–Leibler (KL) divergence plays the role of a proximal penalty. IKLPD encompasses classic EM iterations, their accelerated relaxations, and modern entropy-gradient flows in optimal transport and generative modeling, by formulating each iterate as a variational minimization with a KL or Wasserstein-KL composite objective. This algorithmic class has provable convergence guarantees—linear or superlinear, depending on the regime—and admits efficient implementation in both parameter and probability measure spaces (Chrétien et al., 2012, Lascu et al., 2024, Cheng et al., 2023).

1. Formulation of Implicit KL Proximal Descent

IKLPD addresses variational regularization by implicit minimization of a target functional augmented by a divergence-based proximity term. For finite-dimensional parameter inference—e.g., maximum-likelihood with incomplete data—IKLPD seeks to maximize the incomplete log-likelihood $\ell(\theta) = \log g(y;\theta)$ by iterating:

$$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$

where $k(x|y;\theta)$ is the complete-data posterior for $\theta$ and $\tau_k>0$ is the relaxation parameter (Chrétien et al., 2012).

In measure-theoretic formulations, particularly in Wasserstein space, the IKLPD step takes the form:

$$\rho_{k+1} = \arg\min_{\rho \in \mathcal{P}_2(\mathbb{R}^d)} \left\{ \mathrm{KL}(\rho \,\|\, \pi) + \frac{1}{2\tau} W_2^2(\rho, \rho_k) \right\}$$

where $\mathrm{KL}(\rho\,\|\,\pi)$ is the relative entropy and $W_2$ is the 2-Wasserstein distance (Lascu et al., 2024, Cheng et al., 2023).

2. Algorithmic Variants and Implementation

IKLPD offers a spectrum of practical procedures, determined chiefly by the relaxation parameter schedule:

• Parameter setting $\tau_k=1$ recovers the classical EM algorithm, using the surrogate $Q$-function in each M-step.
• Relaxed and over-relaxed EM ($$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$0) delivers accelerated (often superlinear) convergence without loss of monotonicity.
• Trust-region style adapts $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$1 by model-vs-actual likelihood increments.
• Wasserstein–KL (JKO flow) update for densities leverages optimal transport, widely implemented via neural ODEs or invertible residual networks for flow-based generative models (Cheng et al., 2023).

A typical step in the proximal-in-Wasserstein regime involves optimization over the parameterization of a transport map $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$2, using mini-batch SGD to minimize the composite of log-likelihood and transport cost objectives.

Pseudocode for classical parameter IKLPD (Chrétien et al., 2012):

$\theta$2

Pseudocode for JKO/flow-based IKLPD (Cheng et al., 2023):

$\theta$3

3. Theoretical Properties and Convergence

All IKLPD schemes share key guarantees:

• Monotonicity: Each step ensures nondecreasing (for maximization) or nonincreasing (for minimization) objective, e.g., $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$3 for log-likelihood (Chrétien et al., 2012), or $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$4 in measure space (Lascu et al., 2024).
• Global convergence: Under standard regularity (strict concavity/convexity and differentiability), accumulation points of the iterates solve the first-order stationarity condition.
• Superlinear convergence: For $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$5 with $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$6 functional and divergence, convergence is superlinear; with constant $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$7, the decay is linear/geometric.

On the Wasserstein space, under $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$8-strongly convex $$\theta^{k+1} = \arg\max_{\theta} \left\{ \ell(\theta) - \frac{1}{\tau_k} D_{\mathrm{KL}}(k(\cdot\,|\,\theta^k)\,\|\,k(\cdot\,|\,\theta)) \right\}$$9 (potential) and the logarithmic Sobolev inequality (LSI), the entropy minimization IKLPD scheme converges at a linear rate $k(x|y;\theta)$0 in KL—the decay is geometric in the number of steps (Lascu et al., 2024).

4. Relation to EM, Trust-Region, and Gradient Flow Methods

IKLPD unifies several classic and recent schemes:

• EM algorithm: $k(x|y;\theta)$1 yields the exact EM M-step. Over-relaxed schemes with $k(x|y;\theta)$2 correspond to "EM relaxation".
• Trust-region EM: The KL-proximal term delimits a region in parameter space or measure space within which the likelihood is optimized (Chrétien et al., 2012).
• Proximal Newton: As $k(x|y;\theta)$3, IKLPD approaches a proximal Newton method penalized by KL divergence.
• Exponentiated-gradient: In mixture learning, certain IKLPD updates coincide with exponentiated-gradient steps.
• JKO (Jordan–Kinderlehrer–Otto) flow: In the continuous-measure domain, the IKLPD step is the minimizing-movement scheme for the Wasserstein gradient flow of entropy, foundational to variational approaches in optimal transport and generative modeling (Lascu et al., 2024, Cheng et al., 2023).

5. Implementation Guidance and Practical Aspects

Efficient IKLPD implementation demands judicious tuning of the relaxation parameter and solver:

• Parameter scheduling: $k(x|y;\theta)$4 is a simple decreasing schedule; trust-region criteria compare observed to predicted improvement to adapt $k(x|y;\theta)$5.
• Numerical solvers: Closed-form updates are possible in exponential-family or finite-mixture settings; otherwise, quasi-Newton, limited-memory BFGS, or conjugate-gradient solvers are applied.
• Inexact solves: Monotonicity is preserved even with inexact minimization, enabling "warm-start" and approximate inner steps.
• Neural approximation: In flow-based models, proximal maps are approximated with small neural "residual blocks" (e.g., invertible ResNets or neural ODEs), and inversion consistency losses are sometimes enforced (Cheng et al., 2023).
• Extensions: The scheme generalizes to objective functionals beyond KL divergence (e.g., $k(x|y;\theta)$6-divergences) and can handle distributions without densities by initial mollification (e.g., short-time diffusion) (Cheng et al., 2023).

6. Extensions and Applications

IKLPD generalizes to a variety of settings:

• High-dimensional inverse problems: Accelerated maximum likelihood estimation in tomography and dynamic imaging leverages relaxed IKLPD (Chrétien et al., 2012).
• Probability measure optimization: Entropy minimization and related variational problems in Wasserstein space exploit the JKO variant of IKLPD (Lascu et al., 2024).
• Generative modeling: Normalizing flows and residual block architectures utilize block-wise proximal-in-Wasserstein steps for efficient data distribution generation, achieving exponential KL accuracy with $k(x|y;\theta)$7 steps (Cheng et al., 2023).
• Mean-field and metric measure spaces: The framework accommodates both geodesic and “flat” convexity, relevant for broader classes of energy functionals.

7. Summary Table of IKLPD Variants

Setting	Proximal Objective	Main Guarantee
Parameter (EM/relax-EM)	$k(x|y;\theta)$8	Superlinear or linear conv.
Measure (JKO flow/entropy)	$k(x|y;\theta)$9	Linear (geometric) conv.
Generative models	$\theta$0	$\theta$1 KL accuracy

All convergence and guarantee claims are subject to the regularity and convexity conditions detailed in (Chrétien et al., 2012, Lascu et al., 2024), and (Cheng et al., 2023).