Technical Summary: Gradient Flows for Empirical Bayes in High-Dimensional Linear Models

Problem Statement

The paper addresses the challenge of Empirical Bayes (EB) estimation in high-dimensional linear models where latent parameters interact through complex designs, a setting where theoretical and algorithmic understanding lags behind the well-developed sequence models (e.g., $y_i = \theta_i + \epsilon_i$). Specifically, the authors consider the linear model:
$$y = X\theta + \epsilon, \quad \theta_j \overset{iid}{\sim} g^*, \quad \epsilon_i \overset{iid}{\sim} N(0, \sigma^2)$$
where $n$ (sample size) and $p$ (dimension) are both large. The goal is to estimate the unknown prior distribution $g^*$ of the regression coefficients from the data $(X, y)$.

The primary difficulties in this setting are:

1. Statistical: Establishing consistency of the Nonparametric Maximum Likelihood Estimator (NPMLE) for general regression designs $X$ (where $X$ is not the identity matrix) as $n, p \to \infty$. While recent work [MSS23] established consistency under strong assumptions ($n \ge p$ and well-conditioned $X$), conditions for general designs (including $n < p$ and random designs) remained open.
2. Computational: The marginal log-likelihood $\bar{F}_n(g)$ involves an intractable posterior integral over high-dimensional $\theta$, making direct optimization difficult. Existing approximate methods (e.g., Mean-Field Variational Inference, MCEM) lack complete theoretical convergence guarantees, particularly in non-convex settings or when $n < p$.

Methodology

The authors propose a novel framework combining gradient flows on probability measures with adaptive Langevin dynamics.

1. Reparametrization and Variational Representation

To address the lack of regularity in the posterior of $\theta$ (which may become discrete), the authors introduce a smoothed latent variable $\phi = \theta + z$, where $z \sim N(0, \tau^2 I)$. This transforms the model into:
$$y = X\phi + \tilde{\epsilon}, \quad \phi_j \overset{iid}{\sim} N_\tau * g^*$$
where $N_\tau$ is a Gaussian kernel. This reparametrization ensures the posterior of $\phi$ admits a smooth, positive density, facilitating sampling.

The marginal log-likelihood is expressed via a Gibbs variational representation:
$$\bar{F}_n(g) = \min_{q \in \mathcal{P}^*(\mathbb{R}^p)} F_n(q, g)$$
where $F_n(q, g)$ is the Gibbs free energy. The optimization is performed jointly over the posterior density $q$ and the prior $g$.

2. Bivariate Gradient Flow System

The authors propose a system of simultaneous gradient flow equations to minimize $F_n(q, g)$:
$$\begin{cases} \frac{d}{dt} q_t(\phi) = -p \cdot \text{grad}_{W_2}^q F_n(q_t, g_t)[\phi] \\ \frac{d}{dt} g_t(\theta) = -\alpha \cdot \text{grad}_{FR}^g F_n(q_t, g_t)[\theta] \end{cases}$$

• $q$-flow (Wasserstein-2): The evolution of $q_t$ follows a Fokker-Planck equation corresponding to a Langevin diffusion with a time-dependent drift determined by the current prior $g_t$.
• $g$-flow (Fisher-Rao): The evolution of $g_t$ follows a Fisher-Rao gradient flow, which naturally preserves the positivity and bounded support constraints of the prior.

3. The EBflow Algorithm

Since the exact solution to the gradient flow is intractable, the authors derive a discrete-time algorithm, EBflow:

• E-step (Approximate): A single step of an Unadjusted Langevin Algorithm (ULA) updates the sample $\phi_t$ based on the current prior $g_t$.
• M-step (Approximate): The prior $g_t$ is updated using a single step of Fisher-Rao gradient descent. Crucially, the high-dimensional density $q_t$ is approximated by the empirical distribution of the coordinates of the single Langevin trajectory $\phi_t$, avoiding the need for multiple parallel chains.

Key Contributions

1. Statistical Consistency of Near-NPMLEs

The authors establish the statistical consistency of any "near-NPMLE" (an approximate optimizer of the marginal log-likelihood) under deterministic conditions on the design matrix $X$.

• Assumption: The design matrix must allow for the existence of test vectors $\{z_j\}$ that "approximately orthogonalize" columns of $X$ against others (reminiscent of the debiased Lasso). This extends recent results [MSS23] which required $n \ge p$ and well-conditioned $X$.
• Result: As $n, p \to \infty$, any sequence of priors $\hat{g}_n$ satisfying $\bar{F}_n(\hat{g}_n) \le \bar{F}_n(g^*) + o(1)$ converges weakly to the true prior $g^*$. This holds even in undersampled regimes ($n < p$) and for random designs with subgaussian rows.

2. Convergence of Gradient Flows

The paper provides a theoretical analysis of the continuous-time gradient flow system:

• Univariate Flows:
  • For fixed $q$, the Fisher-Rao flow for $g$ converges globally to the sequence model NPMLE at a rate of $O(1/t)$.
  • For fixed $g$, the Wasserstein flow for $q$ converges exponentially fast to the posterior $P_g(\phi|y)$ in KL-divergence.
• Joint Flow: Under a dimension-free Log-Sobolev Inequality (LSI) for the posterior laws (established for sufficiently large noise variance $\sigma^2$), the joint flow converges to a near-NPMLE.
• Time Complexity: If the marginal log-likelihood is convex on a sub-level set defined by the initialization, the algorithm reaches a near-NPMLE in time linear in $(n, p)$.

3. Algorithmic Implementation

The paper introduces EBflow, a practical algorithm that:

• Uses a single Langevin chain to estimate the average marginal distribution, leveraging a "single-chain propagation-of-chaos" intuition in high dimensions.
• Employs a reparametrization by $\phi$ to ensure the regularity required for Langevin sampling, even if the true prior $g^*$ is discrete.
• Demonstrates competitive performance against MCEM and Variational Inference baselines in simulations.

Results

Theoretical Guarantees

• Consistency: Theorem 3.2 proves that near-optimizers of the marginal likelihood are consistent estimators of $g^*$ under Assumption 3.1 (which covers random designs and $n < p$), relaxing the stronger conditions required in prior work [MSS23].
• Convergence Rates: Theorem 3.10 establishes that the joint gradient flow reaches a density satisfying an averaged first-order optimality condition within a time horizon $T = O(\epsilon^{-3.01}(n+p))$.
• LSI: Theorem 3.7 proves a dimension-free LSI for the posterior of the smoothed variable $\phi$ when $\sigma^2 > M^2 \|X\|_{op}^2$, a condition compatible with the consistency regime.

Empirical Performance

In numerical experiments comparing EBflow to Langevin-MCEM, Gibbs-MCEM, and Coordinate Ascent Variational Inference (CAVI):

• Accuracy: EBflow generally achieves lower Total Variation (TV) error in estimating the prior density compared to CAVI and is competitive with or superior to MCEM methods, particularly in settings with block-correlated or genotype designs.
• Efficiency: EBflow demonstrates faster runtimes than MCEM methods, which require expensive convex optimization (NPMLE) in every M-step.
• Robustness: The method performs well across various prior shapes (Gaussian, skewed, bimodal, scale-mixture) and design matrices (identity, i.i.d., block-correlated, genetic).

Significance and Claims

The paper claims to bridge a gap in Empirical Bayes theory for high-dimensional linear models by:

1. Extending consistency guarantees for NPMLE to general linear models (beyond the sequence model) under deterministic design assumptions that encompass random designs and $n < p$, relaxing the stronger conditions ($n \ge p$, well-conditioned $X$) of recent work [MSS23].
2. Developing a new computational paradigm using bivariate gradient flows and adaptive Langevin dynamics, offering a theoretically grounded alternative to variational inference and MCEM.
3. Establishing dimension-free convergence properties via new Log-Sobolev inequalities for high-dimensional posterior laws, which are of independent interest for Bayesian sampling.

The authors note that while convergence guarantees are stated for the continuous-time flow, they provide insights for the discrete EBflow algorithm. They acknowledge that the global landscape of the objective may be non-convex, and their theoretical guarantees apply to local convergence from specific initializations, leaving the global geometry as an open question.