# Gradient flows for empirical Bayes in high-dimensional linear models

Zhou Fan, Leying Guan, Yandi Shen11footnotemark: 1, Yihong Wu11footnotemark: 1
Department of Statistics and Data Science, Yale UniversityDepartment of
Biostatistics, Yale University
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
This work is supported in part by NSF DMS-2142476, NSF DMS-2310836, NSF
CCF-1900507.

###### Abstract

Empirical Bayes provides a powerful approach to learning and adapting to
latent structure in data. Theory and algorithms for empirical Bayes have a
rich literature for sequence models, but are less understood in settings where
latent variables and data interact through more complex designs.

In this work, we study empirical Bayes estimation of an i.i.d. prior in
Bayesian linear models, via the nonparametric maximum likelihood estimator
(NPMLE). We introduce and study a system of gradient flow equations for
optimizing the marginal log-likelihood, jointly over the prior and posterior
measures in its Gibbs variational representation using a smoothed
reparametrization of the regression coefficients. A diffusion-based
implementation yields a Langevin dynamics MCEM algorithm, where the prior law
evolves continuously over time to optimize a sequence-model log-likelihood
defined by the coordinates of the current Langevin iterate.

We show consistency of the NPMLE as $n,p\to\infty$ under mild conditions,
including settings of random sub-Gaussian designs when $n\asymp p$. In high
noise, we prove a uniform log-Sobolev inequality for the mixing of Langevin
dynamics, for possibly misspecified priors and non-log-concave posteriors. We
then establish polynomial-time convergence of the joint gradient flow to a
near-NPMLE if the marginal negative log-likelihood is convex in a sub-level
set of the initialization.

###### Contents

  1. 1 Introduction
    1. 1.1 Related work
    2. 1.2 Notational conventions
  2. 2 Method
    1. 2.1 Fisher-Rao gradient flow for the Gaussian sequence model: a review
    2. 2.2 Variable reparametrization and joint gradient flow
    3. 2.3 Discretization and particle implementation
    4. 2.4 Posterior inference
    5. 2.5 Smoothness regularization
  3. 3 Theoretical guarantees
    1. 3.1 Consistency of near-NPMLEs
    2. 3.2 Convergence of gradient flows
      1. 3.2.1 Univariate $g$-flow
      2. 3.2.2 Univariate $q$-flow
      3. 3.2.3 Joint flow
      4. 3.2.4 Local convexity
  4. 4 Experiments
    1. 4.1 Parameter choices
    2. 4.2 Method comparison
      1. 5 Discussion
        1. A Proof of consistency
          1. A.1 TV lower bound
          2. A.2 Covering net
          3. A.3 Completing the proof
          4. A.4 Assumption A for random designs
        2. B Analysis of algorithm dynamics
          1. B.1 Fisher-Rao flow for the sequence model
          2. B.2 LSI for Langevin dynamics in high noise
          3. B.3 Analysis of the joint gradient flow
          4. B.4 Proof of (3.17)
          5. B.5 Proof of Proposition 3.8
        3. C Solution properties of gradient flows
          1. C.1 Solution properties of the g-flow
          2. C.2 Solution properties of the q-flow
          3. C.3 Completing the proof
          4. C.4 Proof of Lemma B.4
        4. D Derivations of gradient flows
          1. D.1 General background
          2. D.2 Derivatives of $\bar{F}_{n}(g)$ and $F_{n}(q,g)$
          3. D.3 Non-convexity
        5. E Details of algorithms

##  1 Introduction

Introduced by Robbins [Rob51, Rob56] in the 1950s, empirical Bayes (EB) is a
powerful framework for large-scale inference that learns and adapts to latent
structure in data. In the most commonly studied setting, one or several
independent samples are observed corresponding to each latent parameter, for
example modeled via a Gaussian sequence model

$\displaystyle y_{i}=\theta_{i}+\varepsilon_{i}\qquad\text{ for
}i=1,\ldots,n.$ (1.1)

Treating $\\{\theta_{i}\\}_{i=1}^{n}$ as i.i.d. draws from a prior
distribution $g_{\ast}$, empirical Bayes methods proceed by estimating
$g_{\ast}$ from the noisy observations $\\{y_{i}\\}_{i=1}^{n}$ and using the
estimated prior $g_{\ast}$ to perform downstream Bayesian inference. Such
methods have found fruitful applications in numerous areas of statistics and
the applied sciences, and we refer to the review articles [Cas85, Zha03,
Efr21] and the monographs [ML89, CL96, Efr10b] for systematic treatments of
this broad subject.

In many applications, one may instead wish to model the data
$\mathbf{y}\in\mathbb{R}^{n}$ as having a more complex joint dependence on
latent parameters ${\bm{\theta}}\in\mathbb{R}^{p}$. In this paper, we consider
a canonical setting of a linear model with an i.i.d. prior,

$\displaystyle\mathbf{y}=\mathbf{X}{\bm{\theta}}+{\bm{\varepsilon}},\qquad\theta_{j}\overset{iid}{\sim}g_{\ast}\text{
for }j=1,\ldots,p.$ (1.2)

Assuming that the sample size $n$ and dimension $p$ are both large, we study
the analogous question of estimating $g_{\ast}$ from the data
$(\mathbf{X},\mathbf{y})$. One specific application in which this problem
arises is in estimating effect size distributions and genetic architectures of
complex traits, and we refer to [ZQPC18, O’C21, ZZ21, SSAAP22, MCW+23] for
examples of methods developed in this context. From the perspective of
statistical theory, basic questions such as conditions allowing for consistent
estimation of $g_{\ast}$ are, at present, only partially understood [MSS23].

We adopt a similar setting as in the recent work [MSS23], and focus on
nonparametric estimation of a prior $g_{\ast}$ supported on a fixed and known
bounded domain $[-M,M]\subset\mathbb{R}$, via the nonparametric maximum
likelihood estimator (NPMLE)

$\displaystyle\widehat{g}=\arg\,\min_{g\in\mathcal{P}(M)}\bar{F}_{n}(g)$ (1.3)

or its regularized variants. Here $\mathcal{P}(M)$ is the space of all
probability distributions on $[-M,M]$, and

$\displaystyle\bar{F}_{n}(g):=-\frac{1}{p}\log
P_{g}(\mathbf{y})=-\frac{1}{p}\log\int_{[-M,M]^{p}}\frac{1}{(2\pi\sigma^{2})^{n/2}}\exp\left(-\frac{\|\mathbf{y}-\mathbf{X}{\bm{\theta}}\|^{2}}{2\sigma^{2}}\right)\prod_{j=1}^{p}\mathrm{d}g(\theta_{j})$
(1.4)

is the (normalized) negative marginal log-likelihood of $\mathbf{y}$ under the
prior distribution $g$. The NPMLE was introduced in [Rob50, KW56] and has been
extensively studied in sequence models of the form (1.1). For the linear model
(1.2), in contrast to the sequence model setting of (1.1), $\bar{F}_{n}(g)$
does not decompose across coordinates $j=1,\ldots,p$ and is non-convex in the
prior $g$ for general regression designs. This renders optimization (or even
the evaluation) of the log-likelihood $\bar{F}_{n}(g)$ a substantially more
challenging task.

Inspired by and building upon a line of work on optimizing functionals of
probability measures using gradient flows [JKO98, CB18, MMN18, LCB+22, YWR23],
in this work we propose and study a system of bivariate gradient flow
equations for optimizing $\bar{F}_{n}(g)$. Briefly, our approach consists of
reparametrizing111 This reparameterization is introduced for multiple reasons,
such as ensuring the ${\bm{\varphi}}$ variable has a smooth density supported
on all of $\mathbb{R}^{p}$ which lends itself to Langevin-type sampling
algorithms of the posterior. See Section 2.3 for a detailed discussion. the
linear model (1.2) using an auxiliary variable
${\bm{\varphi}}\in\mathbb{R}^{p}$ as

$\mathbf{y}=\mathbf{X}{\bm{\varphi}}+\widetilde{\bm{\varepsilon}},$ (1.5)

where ${\bm{\varphi}}$ has i.i.d. coordinates with a prior
$\mathcal{N}_{\tau}*g$ that is a Gaussian convolution of $g$, and
$\widetilde{\bm{\varepsilon}}$ is a (correlated) Gaussian residual defined so
that (1.5) expresses the same model as (1.2). We write the Gibbs variational
representation of $\bar{F}_{n}(g)$ in this parametrization as

$\bar{F}_{n}(g)=\min_{q}F_{n}(q,g)=\min_{q}{-}\frac{1}{p}\int\Big{[}\log
P_{g}(\mathbf{y},{\bm{\varphi}})-\log
q({\bm{\varphi}})\Big{]}\mathrm{d}q({\bm{\varphi}})$ (1.6)

where the minimization is over all probability distributions $q$ on
$\mathbb{R}^{p}$. We then minimize $F_{n}(q,g)$ jointly over $(q,g)$ via the
simultaneous gradient flow equations

$\displaystyle\begin{cases}\frac{\mathrm{d}}{\mathrm{d}t}q_{t}({\bm{\varphi}})&={-}p\cdot\operatorname{grad}^{W_{2}}_{q}F_{n}(q_{t},g_{t})[{\bm{\varphi}}],\\\
\frac{\mathrm{d}}{\mathrm{d}t}g_{t}(\theta)&={-}\alpha\cdot\operatorname{grad}^{{\operatorname{FR}}}_{g}F_{n}(q_{t},g_{t})[\theta].\end{cases}$
(1.7)

Here, $\operatorname{grad}^{W_{2}}_{q}$ denotes a gradient with respect to the
Wasserstein-2 geometry in $q$, $\operatorname{grad}^{{\operatorname{FR}}}_{g}$
denotes a gradient with respect to the Fisher-Rao geometry in $g$, and
$\alpha>0$ is a parameter for the relative learning rate of $g$ versus $q$.
Our use of the Fisher-Rao geometry for $g$ is inspired by [YWR23], who studied
also hybrid flows under combined Fisher-Rao and Wasserstein-2 geometries for
computing the NPMLE in the sequence model (1.1).

Following the pioneering insights of [JKO98], the $q$-flow in (1.7) may be
simulated via a sample ${\bm{\varphi}}_{t}\sim q_{t}$ evolving in
$\mathbb{R}^{p}$ according to a Langevin diffusion process, in our setting
with time-evolving drift whose form depends on $g_{t}$. The Fisher-Rao
$g$-flow of (1.7) can be seen as a continuous-time version of the M-step in an
expectation-maximization (EM) algorithm. Consequently, the resulting method
may be understood as a continuous-time variant of a Monte-Carlo expectation-
maximization (MCEM) procedure [WT90] with latent parameter ${\bm{\varphi}}$.
We provide a more detailed description and discussion of this method in
Section 2. We discuss some issues of practical implementation and parameter
tuning in Section 4, where we also present numerical comparisons with
alternative MCEM and variational inference approaches. We note that related
methods have also been proposed recently for general latent variable models
with parametric priors in [KLJ23, ACG+23], and analyzed in contexts where the
full-data log-likelihood is jointly convex in the latent variable and prior
parameters. This joint convexity does not hold in our studied setting of
linear models, requiring the development of new ideas to analyze the gradient
flow dynamics. We discuss this distinction and related literature further in
Section 1.1.

Our work takes a step towards understanding the theoretical properties of the
NPMLE in the linear model, as well as of the gradient flow procedure (1.7) for
its optimization. We summarize these contributions as follows:

  1. 1.

In an asymptotic setting where $n,p\to\infty$, and under a certain
deterministic condition on the design $\mathbf{X}$, we establish consistency
(in the sense of weak convergence) for any approximate minimizer $\widehat{g}$
of the negative log-likelihood $\bar{F}_{n}(g)$.

Our condition for $\mathbf{X}$ requires the existence of test vectors
$\\{\mathbf{z}_{j}\\}_{j=1}^{p}$ that, roughly speaking, can decorrelate a
sufficient number of variables $\\{\mathbf{x}_{j}\\}_{j=1}^{p}$ of the design.
This builds upon and extends a recent result of [MSS23] that showed
consistency for well-conditioned design matrices $\mathbf{X}$ when $n\geq p$.
Our assumption is sufficiently general to encompass also settings where $n<p$,
and in particular, settings where $n$ and $p$ are of a similar order and
$\mathbf{X}$ is a random design having i.i.d. sub-Gaussian rows with well-
conditioned population covariance.

  2. 2.

For suitable initial conditions $(q_{0},g_{0})$, we show existence,
uniqueness, and smoothness for the solution $\\{(q_{t},g_{t})\\}_{t\geq 0}$ of
the coupled differential equations corresponding to (1.7), where $g_{t}$
evolves in the space of densities on $[-M,M]$ and $q_{t}$ describes the law of
a Langevin-type diffusion for a sample ${\bm{\varphi}}_{t}\in\mathbb{R}^{p}$.

  3. 3.

For fixed $q_{t}=q$ (or its approximation via a sample
${\bm{\varphi}}_{t}={\bm{\varphi}}$), the “univariate” evolution of $g_{t}$ in
(1.7) is a Fisher-Rao gradient flow for optimizing the marginal log-likelihood
in a sequence model (1.1). For any initialization $g_{0}$ fully supported on
$[-M,M]$, we prove the global convergence of this univariate flow
$\\{g_{t}\\}_{t\geq 0}$ as $t\to\infty$ to the sequence model NPMLE, at the
rate $O(1/t)$.

  4. 4.

For fixed $g_{t}=g$, the “univariate” evolution of $q_{t}$ in (1.7) describes
a standard Langevin diffusion with fixed drift. We show that for sufficiently
large (but dimension-independent) noise variance $\sigma^{2}>0$, its
stationary posterior law satisfies a log-Sobolev inequality (LSI), uniformly
over all priors $g$ supported on $[-M,M]$. Such a result is known to imply an
exponential rate of convergence of $\\{q_{t}\\}_{t\geq 0}$, in KL-divergence
as $t\to\infty$, to its stationary distribution.

Our proof of this result leverages the breakthrough insight of [BB19], who
established a similar log-Sobolev inequality for high-temperature spin
systems.

  5. 5.

Finally, for the joint evolution of $\\{(q_{t},g_{t})\\}_{t\geq 0}$ following
(1.7), under an assumption of a uniform LSI as guaranteed above, we show that
if the marginal objective $\bar{F}_{n}(g)$ is convex in the sub-level set
defined by the initial value $F_{n}(q_{0},g_{0})$, then $\\{g_{t}\\}_{t\geq
0}$ reaches, in polynomial time, an approximate NPMLE that has the same
statistical consistency properties as $\widehat{g}$.

These results are presented and discussed in Section 3.

We note that while the convergence guarantees in this paper are stated for the
continuous-time flow (1.7), we believe that they also yield useful insights
that pave the way to a discrete-time analysis. We discuss this and a few other
interesting directions for future research in Section 5.

###  1.1 Related work

Empirical Bayes. The NPMLE was proposed in [Rob50, KW56] for nonparametric
estimation of the unknown prior in sequence-type models. The monograph [Lin95]
provides a systematic treatment of structural and geometric properties, and we
refer to [PW20] for more recent results. Early consistency results on the
NPMLE include [Jew82, Lin83b, LT84, Pfa88]; see also [Che17] for a recent
survey. We refer to [GW00, GvdV01, Zha09, JZ09, SG20, PW21] for work on
related problems of marginal density estimation and EB-compound estimation.
Various algorithms have been proposed and studied for computing the NPMLE in
sequence models, including EM-type methods [DLR77, Lai78, Cov84, VL93, JZ09],
general constrained convex optimization techniques (e.g. the interior-point
method) [LZ07, KM14], and vertex direction methods [LK92, Böh95, Wan07,
GJW08]. We also mention the recent work [ZCST22] which proposed an EM-type
algorithm for the sequence model (1.1) under general dimensions without a
fixed support grid.

EB estimation in linear models is closely related to that in the sequence
model with correlated errors, which was discussed in [Efr10a, Sch10, SS18] in
a context of multiple testing and false discovery rate control. For designs
$\mathbf{X}$ having $n\geq p$ and full column rank, the sufficient z-score
statistics follow such a correlated-errors model,222In this case,
$(\mathbf{X}^{\top}\mathbf{X})^{-1}\mathbf{X}^{\top}\mathbf{y}={\bm{\theta}}+{\bm{\varepsilon}}^{\prime}$,
where
${\bm{\varepsilon}}^{\prime}\sim\mathcal{N}(0,\sigma^{2}(\mathbf{X}^{\top}\mathbf{X})^{-1})$.
and the two models are equivalent. They are no longer equivalent in settings
of $n<p$, where observations $\mathbf{y}$ depend on ${\bm{\theta}}$ only via
its projection onto the row span of $\mathbf{X}$. Some early and
representative works studying EB in Bayesian linear models include [NS86,
GF00, YL05] for parametric priors. A body of work in statistical genetics has
developed methods in more semiparametric and nonparametric contexts, for
example based on EM or hierarchical Bayes sampling algorithms [ZQPC18, ZZ21],
mean-field variational approximations [SSAAP22, MCW+23], or moment-matching
procedures in the Fourier domain [O’C21]. Nonparametric EB via a discretized
computation of the NPMLE was also suggested in [WZ21].

Theoretical results on nonparametric EB in the linear model seem more scarce,
and we are only aware of [JG21] who studied a mixture of (low-dimensional)
regression models that is more akin to the setting of (1.1), and the recent
work [MSS23] that analyzed the NPMLE and variational approximations in high-
dimensional regression models. This latter work [MSS23] showed that the NPMLE
is consistent as $n,p\to\infty$ when $n\geq p$ and $\mathbf{X}/\sigma$ has all
singular values of constant order, and established conditions under which
consistency holds also for the optimizer of a mean-field variational
approximation to the log-likelihood objective. Our current work was motivated
by applications where strong correlations between variables of $\mathbf{X}$
may render mean-field variational methods inaccurate, and by the goal of
designing a computationally efficient method for estimating the prior $g_{*}$
in such contexts.

Gradient flows and Bayesian inference. Gradient flows of probability measures
play a central role in our work, and we refer to [San17] for an introduction
and [AGS08, Vil09] for fuller and systematic treatments. The connection
between gradient-based optimization of entropy in the Wasserstein-2 geometry
and Langevin diffusions was recognized in [JKO98, Ott01], and has enabled a
fruitful line of work developing new analyses of Langevin dynamics as well as
new Bayesian sampling algorithms [Per16, Wib18, DMM19, AC21, MCC+21].
Wasserstein gradient flows have also been applied to understand training
dynamics of shallow neural networks [NS17, MMN18, SS20, RVE18, CB18] and to
develop new methods for variational inference [LCB+22, YY22].

Our diffusion-based implementation of the gradient flow equations leads to a
continuous-time Monte Carlo EM method [WT90, BH99, LC01], where an intractable
E-step is simulated using an Unadjusted Langevin Algorithm [RT96, Dal17,
DM17]. Classical theoretical treatments of MCEM include [CL95, FM03, Nea13],
but typically assume an asymptotically growing number of MCMC iterations
between M-step updates. Our analyses of the gradient flow equations may be
understood as studying MCEM convergence in an opposite continuous-time regime,
where the prior is continuously updated before the equilibration of the MCMC
samples to their stationary measure.

Closest to our current work are the recent papers [KLJ23, ACG+23], which also
proposed and studied EB methods for general latent variable models based on
joint optimization or sampling of the prior and posterior in a Gibbs
variational representation of the log-likelihood. These works focus on a
setting where the prior $g$ belongs to a fixed parametric family
$\\{g_{\bm{\alpha}}\\}$, and the analyses therein assume that the full-data
log-likelihood $\log P({\bm{\theta}},\mathbf{y}\mid{\bm{\alpha}})$ of the
latent variable ${\bm{\theta}}$ and observation $\mathbf{y}$ is jointly
strongly concave in $({\bm{\theta}},{\bm{\alpha}})$. Such a condition is
sufficient to ensure that the marginal log-likelihood $\log
P(\mathbf{y}\mid{\bm{\alpha}})$ has a unique local maximizer
$\widehat{\bm{\alpha}}$, and that furthermore the Gibbs variational objective
is Wasserstein-2 convex in $q$ for each fixed ${\bm{\alpha}}$. It is then
shown in [KLJ23] that a bivariate gradient flow in ${\bm{\alpha}}$ and $q$
converges at an exponential rate to $\widehat{\bm{\alpha}}$ and its
corresponding posterior. It may be checked that this joint convexity does not
hold in our problem, and furthermore the Gibbs variational objective may not
be Wasserstein-2 convex in $q$ under priors $g$ where the posterior is non-
log-concave. These issues motivate our choice of writing the variational
objective using the reparametrized regression variable in (1.5), and our
development of different ideas to analyze the gradient flow dynamics. We
mention also the related works [Chi22, NWS22] that established quantitative
convergence of McVean-Vlasov Langevin diffusions for optimizing entropy-
regularized functionals $F(q)$ under uniform LSI assumptions. However, these
works also assume that $F(q)$ is globally convex in the linear geometry. Such
an assumption again does not hold in our setting, where the “entropy sandwich”
arguments of [Chi22, NWS22] are not available.

For performing nonparametric optimization over the prior $g$, we will use (for
simplicity) the Fisher-Rao geometry. The ideas of the Fisher-Rao geometry date
back to the classical work [Hel09, Hot30, RR45, Kak48]. More recently,
together with the Wasserstein geometry, it has played an important role in the
study of unbalanced optimal transport and the associated Wasserstein-Fisher-
Rao flow [CPSV18b, CPSV18a, KMV16, LMS18, LMS16, YWR23]. For sampling, recent
work [LLN19, LSW23, DEP23] showed that dynamics resulting from the combined
Wasserstein-Fisher-Rao gradient flow can accelerate the convergence of
Langevin dynamics, especially when the target prior is not log-concave.

###  1.2 Notational conventions

$\mathcal{P}(M),\mathcal{P}(\mathbb{R}),\mathcal{P}(\mathbb{R}^{p})$ are the
spaces of all probability distributions on $[-M,M]$, $\mathbb{R}$, and
$\mathbb{R}^{p}$, and
$\mathcal{P}_{*}(M),\mathcal{P}_{*}(\mathbb{R}),\mathcal{P}_{*}(\mathbb{R}^{p})$
are the spaces of probability densities on these domains with respect to
Lebesgue measure.

For two measures $P,Q$ such that $P$ is absolutely continuous to $Q$, the
Kullback-Leibler divergence is
$D_{\mathrm{KL}}(P\|Q)=\int\log\frac{\mathrm{d}P}{\mathrm{d}Q}\mathrm{d}P$.
The total variation distance is
$d_{\mathrm{TV}}(P,Q)=\frac{1}{2}\int|\mathrm{d}P-\mathrm{d}Q|$.

We use $\mathcal{N}_{\sigma}$ as shorthand for the Gaussian law
$\mathcal{N}(0,\sigma^{2})$, and $\mathcal{N}_{\sigma}(x)$ for its density
function at $x\in\mathbb{R}$. We use $P*Q$ for the convolution of two
measures, so $[\mathcal{N}_{\sigma}*g](x)$ is the density of the convolution
of $\mathcal{N}_{\sigma}$ with $g$.

Let $\nabla f$, $\Delta f$, and $\nabla\cdot f$ denote the gradient,
Laplacian, and divergence of a function $f:\mathbb{R}^{d}\to\mathbb{R}$. A
function $f:\Omega\to\mathbb{R}$ on an open domain $\Omega$ belongs to
$C^{r}(\Omega)$ (resp. $C^{\infty}(\Omega)$) if it is $r$-times (resp.
infinitely-times) continuously differentiable on $\Omega$. We use
$\lVert\mathbf{A}\rVert_{{\operatorname{op}}}$ for the operator norm of a
matrix $\mathbf{A}$, and $\lambda_{\min}(\mathbf{A})$ (resp.
$\lambda_{\max}(\mathbf{A})$) for its smallest (resp. largest) eigenvalue when
$\mathbf{A}$ is symmetric.

##  2 Method

In our exposition, we treat $\mathbf{X}$ as deterministic, i.e. all results
are conditional on $\mathbf{X}$. We assume throughout that $g_{\ast}$ is a
prior distribution supported on a fixed and known interval $[-M,M]$, and we
seek to estimate $g_{\ast}$ from $(\mathbf{X},\mathbf{y})$ following the
linear model (1.2).

###  2.1 Fisher-Rao gradient flow for the Gaussian sequence model: a review

Consider first the sequence model (1.1) with Gaussian errors of variance
$\sigma^{2}$, corresponding to the identity design matrix
$\mathbf{X}=\operatorname{Id}$ with $p=n$. The NPMLE for priors supported on
$[-M,M]$ is

$\displaystyle\widehat{g}=\arg\,\min_{g\in\mathcal{P}(M)}\bar{F}_{n}(g),\quad\text{where
}\bar{F}_{n}(g)=-\frac{1}{n}\sum_{i=1}^{n}\log[\mathcal{N}_{\sigma}\ast
g](y_{i})$ (2.1)

and $\mathcal{N}_{\sigma}*g$ denotes the convolution of
$\mathcal{N}_{\sigma}\equiv\mathcal{N}(0,\sigma^{2})$ with $g$.

Gradient flows for optimizing $\bar{F}_{n}(g)$ were studied in [YWR23], and we
review here a simple such flow operating in the Fisher-Rao geometry for $g$:
For any functional $F(g)$ of a positive density $g$ on $[-M,M]$, we define its
first variation $\delta F[g]:[-M,M]\rightarrow\mathbb{R}$ as the unique (up to
an additive constant) function for which

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}\varepsilon}\Big{|}_{\varepsilon=0}F(g+\varepsilon\chi)=\int_{-M}^{M}\delta
F[g](x)\chi(x)\mathrm{d}x$ (2.2)

for every bounded function $\chi:[-M,M]\rightarrow\mathbb{R}$ satisfying
$\int_{-M}^{M}\chi(x)\mathrm{d}x=0$. In particular, the first variation of
$\bar{F}_{n}(g)$ in (2.1) may be computed to be

$\delta\bar{F}_{n}[g](\theta)=-\frac{1}{n}\sum_{i=1}^{n}\frac{\mathcal{N}_{\sigma}(y_{i}-\theta)}{\mathcal{N}_{\sigma}*g(y_{i})}+1,$
(2.3)

where we choose the constant $+1$ so that it has mean 0 under $g$. The Fisher-
Rao gradient flow is then

$\frac{\mathrm{d}}{\mathrm{d}t}\,g_{t}(\theta)={-}g_{t}(\theta)\cdot\delta\bar{F}_{n}[g_{t}](\theta)=g_{t}(\theta)\left[\frac{1}{n}\sum_{i=1}^{n}\frac{\mathcal{N}_{\sigma}(y_{i}-\theta)}{\mathcal{N}_{\sigma}*g_{t}(y_{i})}-1\right].$
(2.4)

This may be understood as an infinitesimal version of the widely studied EM
algorithm [DLR77]

$\displaystyle
g_{t+1}(\theta)=\frac{1}{n}\sum_{i=1}^{n}\frac{g_{t}(\theta)\mathcal{N}_{\sigma}(y_{i}-\theta)}{\mathcal{N}_{\sigma}*g_{t}(y_{i})},$
(2.5)

and we refer to [YWR23] for other interpretations and a more detailed
discussion. Both the Fisher-Rao flow (2.4) and EM algorithm (2.5) may be
approximated numerically by restricting to
$g=\sum_{k=1}^{K}w_{k}\delta_{b_{k}}$ supported on a fixed grid
$b_{1},\ldots,b_{K}\in[-M,M]$ and updating the weights $(w_{1},\ldots,w_{K})$.

For more general designs $\mathbf{X}$, the Fisher-Rao gradient flow may be
written as (c.f. Appendix D.2 for a derivation)

$\frac{\mathrm{d}}{\mathrm{d}t}g_{t}(\theta)=\bar{\mu}[g_{t}](\theta)-g_{t}(\theta)$
(2.6)

where

$\bar{\mu}[g](x)=\frac{1}{p}\sum_{j=1}^{p}\int_{\mathbb{R}^{p-1}}P_{g}(\theta_{1},\ldots,\theta_{j-1},x,\theta_{j+1},\ldots,\theta_{p}\mid\mathbf{y})\prod_{i:i\neq
j}\mathrm{d}\theta_{i}$ (2.7)

is the average marginal density of coordinates $\theta_{1},\ldots,\theta_{p}$
under the posterior law
$\mu[g]({\bm{\theta}}):=P_{g}({\bm{\theta}}\mid\mathbf{y})$. Exact computation
of $\bar{\mu}[g]$ is intractable for general regression designs. As mentioned
in Section 1, we will address this challenge by performing the gradient flow
instead on a Gibbs variational representation (1.6) of $\bar{F}_{n}(g)$,
coupled with a gradient flow over the space of high-dimensional posterior
measures on $\mathbb{R}^{p}$.

###  2.2 Variable reparametrization and joint gradient flow

Consider now the linear model with i.i.d. prior $g_{\ast}$ and noise variance
$\sigma^{2}>0$, which we restate here for clarity:

$\mathbf{y}=\mathbf{X}{\bm{\theta}}+{\bm{\varepsilon}},\qquad\theta_{j}\overset{iid}{\sim}g_{\ast}\text{
for
}j=1,\ldots,p,\qquad{\bm{\varepsilon}}\sim\mathcal{N}(0,\sigma^{2}\operatorname{Id})\text{
independent of }{\bm{\theta}}.$ (2.8)

To describe the reparametrization (1.5), fix any parameter $\tau^{2}>0$
satisfying $\tau^{2}<\sigma^{2}/\lambda_{\max}(\mathbf{X}\mathbf{X}^{\top})$,
where $\lambda_{\max}(\cdot)$ denotes the largest eigenvalue. Decompose the
Gaussian noise ${\bm{\varepsilon}}$ as
${\bm{\varepsilon}}=\widetilde{\bm{\varepsilon}}+\mathbf{X}\mathbf{z}$ where
$\mathbf{z}\sim\mathcal{N}(0,\tau^{2}\operatorname{Id})$ and
$\widetilde{\bm{\varepsilon}}\sim\mathcal{N}(0,\,{\bm{\Sigma}})$ with

${\bm{\Sigma}}=\sigma^{2}\operatorname{Id}-\tau^{2}\mathbf{X}\mathbf{X}^{\top}\succ
0.$

Setting ${\bm{\varphi}}={\bm{\theta}}+\mathbf{z}$, the marginal law of
$\mathbf{y}$ is equivalently expressed as

$\displaystyle\mathbf{y}=\mathbf{X}{\bm{\varphi}}+\widetilde{\bm{\varepsilon}},\qquad\varphi_{j}\overset{iid}{\sim}\mathcal{N}_{\tau}*g_{\ast},\qquad\widetilde{\bm{\varepsilon}}\sim\mathcal{N}(0,\,{\bm{\Sigma}})\text{
independent of }{\bm{\varphi}},$ (2.9)

where we recall that $\mathcal{N}_{\tau}*g$ denotes the convolution of the
$\mathcal{N}(0,\tau^{2})$ density with $g$. Then the negative log-likelihood
$\bar{F}_{n}(g)={-}\frac{1}{p}\log P_{g}(\mathbf{y})$ takes the form
(equivalent to (1.4))

$\bar{F}_{n}(g)={-}\frac{1}{p}\log\int_{\mathbb{R}^{p}}\frac{1}{(2\pi)^{n/2}(\det{\bm{\Sigma}})^{1/2}}\exp\left({-}\frac{1}{2}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})\right)\prod_{j=1}^{p}[\mathcal{N}_{\tau}*g](\varphi_{j})\mathrm{d}\varphi_{j}.$
(2.10)

This functional $\bar{F}_{n}(g)$ has the Gibbs variational
representation,333This representation follows from rearranging the identity
$0=\min_{q}D_{\mathrm{KL}}(q({\bm{\varphi}})\|P({\bm{\varphi}}\mid\mathbf{y}))$,
noting that the minimization may be restricted to probability densities
$q(\cdot)$ since the posterior law $P({\bm{\varphi}}\mid\mathbf{y})$ admits a
density on $\mathbb{R}^{p}$. up to an additive constant not depending on $g$,

$\displaystyle\bar{F}_{n}(g)=$
$\displaystyle\min_{q\in\mathcal{P}_{*}(\mathbb{R}^{p})}F_{n}(q,g),$
$\displaystyle F_{n}(q,g):=$
$\displaystyle\frac{1}{p}\int\bigg{[}\frac{1}{2}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})-\sum_{j=1}^{p}\log[\mathcal{N}_{\tau}*g](\varphi_{j})+\log
q({\bm{\varphi}})\bigg{]}q({\bm{\varphi}})\mathrm{d}{\bm{\varphi}}$
$\displaystyle\hskip
72.26999pt+\frac{n}{2p}\log(2\pi)+\frac{1}{2p}\log\det{\bm{\Sigma}}.$ (2.11)

The minimization is over the space of all probability densities on
$\mathbb{R}^{p}$, and the minimizer for each fixed prior $g$ is given by the
posterior law

$\nu[g]({\bm{\varphi}}):=P_{g}({\bm{\varphi}}\mid\mathbf{y}).$ (2.12)

Thus the NPMLE $\widehat{g}:=\arg\,\min_{g}\bar{F}_{n}(g)$ may be computed by
joint minimization of $F_{n}(q,g)$. This representation is also the basis of
empirical Bayes learning using variational inference approaches [BBB+03],
where the domain of the candidate posterior law $q(\cdot)$ is restricted to a
variational sub-class of densities on $\mathbb{R}^{p}$, as well as using
expectation-maximization (EM) algorithms which may be understood as optimizing
$F_{n}(q,g)$ using alternating updates of $q$ and $g$ [NH98].

Here, we study optimization of $F_{n}(q,g)$ over densities $q$ on
$\mathbb{R}^{p}$ and $g$ on $[-M,M]$ via the simultaneous gradient flows
(1.7), under a Wasserstein-2 geometry for $q$ [JKO98, Vil09] and a Fisher-Rao
geometry for $g$ [KMV16, YWR23]. The first variations of $F_{n}(q,g)$ are (up
to additive constants)

$\displaystyle\delta_{q}F_{n}(q,g)[{\bm{\varphi}}]$
$\displaystyle=\frac{1}{p}\bigg{[}\frac{1}{2}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})-\sum_{j=1}^{p}\log[\mathcal{N}_{\tau}*g](\varphi_{j})+\log
q({\bm{\varphi}})\bigg{]},$ (2.13) $\displaystyle\delta_{g}F_{n}(q,g)[\theta]$
$\displaystyle=-\left[\mathcal{N}_{\tau}*\frac{\bar{q}}{\mathcal{N}_{\tau}*g}\right](\theta)+1,$

where, corresponding to any density $q(\cdot)$ on $\mathbb{R}^{p}$, we define
its coordinate-wise average marginal density $\bar{q}(\cdot)$ on $\mathbb{R}$
by

$\displaystyle\bar{q}(x)=\frac{1}{p}\sum_{j=1}^{p}\int_{\mathbb{R}^{p-1}}q(\varphi_{1},\ldots,\varphi_{j-1},x,\varphi_{j+1},\ldots,\varphi_{p})\prod_{i:i\neq
j}\mathrm{d}\varphi_{i}.$ (2.14)

The gradient flow equations (1.7) then take the explicit form of the coupled
system of differential equations

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}q_{t}({\bm{\varphi}})$
$\displaystyle=\nabla\cdot\left[q_{t}({\bm{\varphi}})\left(\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}(\mathbf{X}{\bm{\varphi}}-\mathbf{y})-\left(\frac{[\mathcal{N}_{\tau}*g_{t}]^{\prime}(\varphi_{j})}{[\mathcal{N}_{\tau}*g_{t}](\varphi_{j})}\right)_{j=1}^{p}\right)\right]+\Delta
q_{t}({\bm{\varphi}}),$ (2.15)
$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}g_{t}(\theta)$
$\displaystyle=\alpha\,g_{t}(\theta)\left(\left[\mathcal{N}_{\tau}*\frac{\bar{q}_{t}}{\mathcal{N}_{\tau}*g_{t}}\right](\theta)-1\right),$
(2.16)

where $\nabla\cdot f({\bm{\varphi}})$ and $\Delta f({\bm{\varphi}})$ denote
the divergence and Laplacian in ${\bm{\varphi}}$. We provide a heuristic
derivation of (2.13) and (2.15–2.16) in Appendix D.2.444We will henceforth
treat the coupled system (2.15–2.16) as the definition of the gradient flow
equations (1.7), and not formally justify their equivalence with a theory of
gradient flows on metric spaces as developed in [AGS08]. We emphasize that the
Fisher-Rao gradient in (2.16) is computed with respect to the unsmoothed prior
law $g$ for ${\bm{\theta}}$, not the smoothed prior law $\mathcal{N}_{\tau}*g$
for ${\bm{\varphi}}$, as performing a gradient flow in the latter would not
directly yield an estimate for $g_{\ast}$ and, more importantly, would require
a method of constraining the density to the space of Gaussian convolutions.

###  2.3 Discretization and particle implementation

As first recognized in [JKO98], the evolution (2.15) resulting from the
Wasserstein-2 gradient of $F_{n}(q,g)$ in $q$ is a Fokker-Planck equation with
$q_{t}$ being the density of a sample ${\bm{\varphi}}_{t}\in\mathbb{R}^{p}$
evolving according to the following Langevin diffusion:

$\mathrm{d}{\bm{\varphi}}_{t}=\left({-}\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}(\mathbf{X}{\bm{\varphi}}_{t}-\mathbf{y})+\left(\frac{[\mathcal{N}_{\tau}*g_{t}]^{\prime}(\varphi_{t,j})}{[\mathcal{N}_{\tau}*g_{t}](\varphi_{t,j})}\right)_{j=1}^{p}\right)\mathrm{d}t+\sqrt{2}\,\mathrm{d}\mathbf{B}_{t},$
(2.17)

where $\\{\mathbf{B}_{t}\\}_{t\geq 0}$ is a standard Brownian motion in
$\mathbb{R}^{p}$. In our setting, this diffusion is a McVean-Vlasov process in
the sense that the first drift term is coupled to the evolution (2.16) for
$g_{t}$, which in turn depends on the law $q_{t}$ of its solution
${\bm{\varphi}}_{t}$ via the average marginal density $\bar{q}_{t}$.

We remark that estimating the high-dimensional density $q_{t}$ on
$\mathbb{R}^{p}$ may be challenging, but estimating its average marginal
$\bar{q}_{t}$ on $\mathbb{R}$ is conceivably easier. In this work, we will
simulate the coupled system (2.16) and (2.17) using an estimate of
$\bar{q}_{t}$ in (2.16) by the empirical distribution of coordinates of
${\bm{\varphi}}_{t}$,

$\bar{q}_{t}\approx\frac{1}{p}\sum_{j=1}^{p}\delta_{\varphi_{t,j}}.$ (2.18)

To reduce the influence of Monte Carlo variance of this estimate on the prior
$g_{t}$, we will set a small relative learning rate parameter $\alpha$ in
(1.7) in our numerical implementations, so that $g_{t}$ depends in aggregate
on estimates of $\\{\bar{q}_{t}\\}$ over a longer effective history of past
time. We conjecture that for high dimensions $p$ and design matrices
$\mathbf{X}$ that have limited long-range dependence across variables, the
estimate on the right side of (2.18) may concentrate well around $\bar{q}_{t}$
without needing to average across multiple parallel samples of
${\bm{\varphi}}_{t}$. We discuss this further in Section 5, and leave the
theoretical exploration of this to future work.

For simplicity, we opt for a discretization of (2.16) and (2.17) in time using
a forward Euler scheme, and also discretize the prior distribution $g$ in
space as

$g=\sum_{k=1}^{K}w_{k}\delta_{b_{k}},$ (2.19)

where we fix an equally-spaced grid of $K$ support points
$b_{1},\ldots,b_{K}\in[-M,M]$ with corresponding probability weights
$w=(w_{1},\ldots,w_{K})$. Together with the approximation (2.18), this yields

$\displaystyle\left[\mathcal{N}_{\tau}*\frac{\bar{q}_{t}}{\mathcal{N}_{\tau}*g_{t}}\right](\theta)$
$\displaystyle=\int\frac{\mathcal{N}_{\tau}(\theta-\varphi)}{\sum_{k=1}^{K}w_{k}\,\mathcal{N}_{\tau}(b_{k}-\varphi)}\bar{q}_{t}(\varphi)\mathrm{d}\varphi\approx\frac{1}{p}\sum_{j=1}^{p}\frac{\mathcal{N}_{\tau}(\theta-\varphi_{t,j})}{\sum_{k=1}^{K}w_{k}\,\mathcal{N}_{\tau}(b_{k}-\varphi_{t,j})},$
$\displaystyle\frac{[\mathcal{N}_{\tau}*g]^{\prime}(\varphi)}{[\mathcal{N}_{\tau}*g](\varphi)}$
$\displaystyle={-}\frac{1}{\tau^{2}}\cdot\frac{\sum_{k=1}^{K}(\varphi-
b_{k})w_{k}\,\mathcal{N}_{\tau}(\varphi-
b_{k})}{\sum_{k=1}^{K}w_{k}\,\mathcal{N}_{\tau}(\varphi-b_{k})}.$

Applying these discretizations to (2.16) and (2.17) leads to the discretized
“EBflow” algorithm

$\displaystyle{\bm{\varphi}}_{t+1}$
$\displaystyle={\bm{\varphi}}_{t}-\eta^{\varphi}_{t}\Bigg{[}\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}(\mathbf{X}{\bm{\varphi}}_{t}-\mathbf{y})+\frac{1}{\tau^{2}}\bigg{(}\frac{\sum_{k=1}^{K}(\varphi_{t,j}-b_{k})w_{t,k}\,\mathcal{N}_{\tau}(\varphi_{t,j}-b_{k})}{\sum_{k=1}^{K}w_{t,k}\,\mathcal{N}_{\tau}(\varphi_{t,j}-b_{k})}\bigg{)}_{j=1}^{p}\Bigg{]}+\sqrt{2\eta^{\varphi}_{t}}\,\mathbf{b}_{t},$
(2.20) $\displaystyle w_{t+1,k}$
$\displaystyle=w_{t,k}+\eta^{w}_{t}w_{t,k}\Bigg{[}\frac{1}{p}\sum_{j=1}^{p}\frac{\mathcal{N}_{\tau}(b_{k}-\varphi_{t+1,j})}{\sum_{i=1}^{K}w_{t,i}\,\mathcal{N}_{\tau}(b_{i}-\varphi_{t+1,j})}-1\Bigg{]}\text{
for }k=1,\ldots,K,$ (2.21)

where $\mathbf{b}_{t}\sim\mathcal{N}(0,\operatorname{Id})$ is a standard
Gaussian vector in $\mathbb{R}^{p}$ generated independently in each iteration,
$w_{t}:=(w_{t,1},\ldots,w_{t,K})$ are the probability weights representing
$g_{t}=\sum_{k=1}^{K}w_{t,k}\delta_{b_{k}}$, and
$\eta_{t}^{\varphi},\eta_{t}^{w}>0$ are step sizes for the updates of
${\bm{\varphi}}_{t}$ and $w_{t}$ respectively, with
$\eta_{t}^{w}/\eta_{t}^{\varphi}$ representing the relative learning rate
$\alpha>0$ in (1.7) and (2.16). (In our experiments, we will allow
$\eta_{t}^{w}$ and $\eta_{t}^{\varphi}$ to vary with the iteration $t$, fixing
their ratio $\eta_{t}^{w}/\eta_{t}^{\varphi}$.)

Note that the iterations (2.20) represent steps of an Unadjusted Langevin
Algorithm (ULA) [RT96] with time-varying prior. The prior is updated via
(2.21), which has the equivalent form

$w_{t+1,k}=(1-\eta^{w}_{t})w_{t,k}+\eta_{t}^{w}\cdot\frac{1}{p}\sum_{j=1}^{p}\frac{w_{t,k}\,\mathcal{N}_{\tau}(b_{k}-\varphi_{t+1,j})}{\sum_{i=1}^{K}w_{t,i}\,\mathcal{N}_{\tau}(b_{i}-\varphi_{t+1,j})}.$

As such, as long as $\eta_{t}^{w}\leq 1$, we have $w_{t}\geq 0$ for all
iterations $t\geq 1$, an appealing property resulting from the form of the
Fisher-Rao gradient $\operatorname{grad}_{g}^{\operatorname{FR}}F_{n}(q,q)$.

Let us now briefly discuss the role of the initial variable reparametrization
${\bm{\theta}}\mapsto{\bm{\varphi}}$ in (2.9), and the interpretation of this
algorithm as a Monte-Carlo EM (MCEM) procedure. If the true prior $g_{*}$ has
disconnected intervals of support, then the support of its posterior
$\mu[g_{*}]({\bm{\theta}})$ would be similarly disconnected, and it is clear
that a Langevin diffusion for sampling from this posterior would not be
ergodic. Even if the true prior $g_{*}$ were such that this Langevin diffusion
is rapidly mixing, such a property may not hold under the priors $\\{g_{t}\\}$
evolving according to the gradient flow (2.15–2.16). Indeed, for identity
design (i.e. the sequence model (1.1)) the exact NPMLE $\widehat{g}$ is known
to be supported on a finite set of atoms [Lin83a, PW20], and we expect for
more general designs that $\\{g_{t}\\}$ may also converge to a discrete
distribution in the limit $t\to\infty$, if no additional regularization is
imposed. In contrast, the prior $\mathcal{N}_{\tau}*g_{t}$ induces a
continuous posterior distribution $\nu[g_{t}]({\bm{\varphi}})$ for
${\bm{\varphi}}$ even if $g_{t}$ is discrete. For sufficiently large noise
variance $\sigma^{2}>0$, we will further verify in Section 3 that this
posterior $\nu[g_{t}]({\bm{\varphi}})$ satisfies a log-Sobolev inequality,
which holds irrespective of the connectedness of the support of $g_{t}$ or a
lower bound for its density. For the discretized algorithm (2.20–2.21),
reparametrization by ${\bm{\varphi}}$ likewise enables posterior sampling
using a continuous Langevin diffusion even if the support of $g_{t}(\cdot)$ is
discretized to a grid.

The two steps (2.20–2.21) of EBflow may be understood as “incremental”
versions of the E-step and M-step in an EM algorithm with latent variable
${\bm{\varphi}}$. Indeed, an exact EM algorithm in this parametrization would
minimize $F_{n}(q,g)$ via the following alternating updates of $q$ and $g$:

(E-step) $\displaystyle q_{t}$ $\displaystyle=\nu[g_{t}]\qquad\text{(the
posterior law of ${\bm{\varphi}}$ under the prior
$\mathcal{N}_{\tau}*g_{t}$)},$ (2.22) (M-step) $\displaystyle g_{t+1}$
$\displaystyle=\arg\,\min_{g}\int{-}\log[\mathcal{N}_{\tau}*g](\varphi)\bar{q}_{t}(\varphi)\mathrm{d}\varphi.$

Both steps have computational difficulties in practice: The E-step requires
computation of the posterior $\nu[g_{t}]$, which is intractable for general
design matrices $\mathbf{X}$. The update (2.20) of EBflow replaces this E-step
by a single step of ULA, and then approximates $\bar{q}_{t}$ as needed in the
M-step by the empirical distribution of coordinates of the sampled
${\bm{\varphi}}_{t}=(\varphi_{t,j})_{j=1}^{p}$ cf. (2.18). Under this
approximation, $g_{t+1}$ in the M-step is precisely the solution of an NPMLE
problem in the Gaussian sequence model (1.1) with observations
$(\varphi_{t,1},\ldots,\varphi_{t,p})$. Computing $g_{t+1}$ exactly may be
expensive to carry out in every iteration, and the update (2.21) of EBflow
replaces this M-step by a single step of Fisher-Rao gradient descent for
solving this sequence model NPMLE problem.

###  2.4 Posterior inference

For downstream tasks of posterior inference, one may perform Monte Carlo
approximation using samples ${\bm{\varphi}}_{t}\in\mathbb{R}^{p}$ of
(2.20–2.21) for a set of final iterations $t\in\\{T+1,\ldots,T+T^{\prime}\\}$
after fixing the final estimated weights $w_{T}=(w_{T,1},\ldots,w_{T,K})$.
Typically one is interested in computing expectations under the posterior law
of the original parameter ${\bm{\theta}}$, rather than under the
reparametrized variable ${\bm{\varphi}}$. Such expectations may often be
estimated using samples $\\{{\bm{\varphi}}_{t}\\}$ via the relation

$\displaystyle\mathbb{E}[f({\bm{\theta}})\mid\mathbf{y}]$
$\displaystyle=\int_{\mathbb{R}^{p}}\int_{[-M,M]^{p}}f({\bm{\theta}})P({\bm{\theta}}\mid{\bm{\varphi}})P({\bm{\varphi}}\mid\mathbf{y})\mathrm{d}{\bm{\theta}}\mathrm{d}{\bm{\varphi}}$
$\displaystyle\approx\frac{1}{T^{\prime}}\sum_{t=T+1}^{T+T^{\prime}}\int_{[-M,M]^{p}}f({\bm{\theta}})P({\bm{\theta}}\mid{\bm{\varphi}}_{t})\mathrm{d}{\bm{\theta}}.$

Here, coordinates of ${\bm{\theta}}$ are independent conditional on
${\bm{\varphi}}$, with
$P({\bm{\theta}}\mid{\bm{\varphi}})=\prod_{j=1}^{p}P(\theta_{j}\mid\varphi_{j})\propto\prod_{j=1}^{p}g(\theta_{j})\mathcal{N}_{\tau}(\varphi_{j}-\theta_{j})$.
This enables Monte Carlo approximation of
$\mathbb{E}[f({\bm{\theta}})\mid\mathbf{y}]$ for functions $f$ that depend
only on low-dimensional marginals of ${\bm{\theta}}$.

For example, the posterior mean
$\mathbb{E}[{\bm{\theta}}\mid\mathbf{y}]\in\mathbb{R}^{p}$ may be estimated as

$\displaystyle\mathbb{E}[{\bm{\theta}}\mid\mathbf{y}]$
$\displaystyle\approx\frac{1}{T^{\prime}}\sum_{t=T+1}^{T+T^{\prime}}\int_{[-M,M]^{p}}{\bm{\theta}}P({\bm{\theta}}\mid{\bm{\varphi}}_{t})\mathrm{d}{\bm{\theta}}$
$\displaystyle=\frac{1}{T^{\prime}}\sum_{t=T+1}^{T+T^{\prime}}\left(\frac{\int_{-M}^{M}\theta\cdot
g(\theta)\,\mathcal{N}_{\tau}(\theta-\varphi_{t,j})\mathrm{d}\theta}{\int_{-M}^{M}g(\theta)\mathcal{N}_{\tau}(\theta-\varphi_{t,j})\mathrm{d}\theta}\right)_{j=1}^{p}$
$\displaystyle=\frac{1}{T^{\prime}}\sum_{t=T+1}^{T+T^{\prime}}\left(\frac{\sum_{k=1}^{K}b_{k}w_{T,k}\,\mathcal{N}_{\tau}(b_{k}-\varphi_{t,j})}{\sum_{k=1}^{K}w_{T,k}\,\mathcal{N}_{\tau}(b_{k}-\varphi_{t,j})}\right)_{j=1}^{p},$
(2.23)

where the last step holds under the space discretization (2.19) of $g$. This
estimate (2.23) for $\mathbb{E}[{\bm{\theta}}\mid\mathbf{y}]$ may then be used
for downstream tasks such as predicting the response for new samples.

###  2.5 Smoothness regularization

The perspective of minimizing $F_{n}(q,g)$ using a joint gradient flow enables
easy extensions and variations upon this basic algorithm, for example to use
stochastic or preconditioned gradient updates for ${\bm{\varphi}}$, or to
incorporate regularization for the prior $g$.

The unregularized NPMLE $\widehat{g}$ may be adequate for purposes of
prediction on new samples, but does not share the smoothness properties of
$g_{\ast}$ if the latter admits a smooth density. We describe here one simple
approach for smoothness regularization, via addition of a smoothing-spline
penalty to the log-likelihood objective:

$\bar{F}_{n,\lambda}(g)=-\frac{1}{p}\log
P_{g}(\mathbf{y})+\frac{\lambda}{2}\int_{-M}^{M}g^{\prime\prime}(\theta)^{2}d\theta.$
(2.24)

On the discrete grid $b_{1},\ldots,b_{K}\in[-M,M]$ with equal spacing
$\Delta=b_{k+1}-b_{k}$, approximating $g$ with weights $\\{w_{k}\\}_{k=1}^{K}$
as in (2.19) and applying $g(b_{k})\approx w_{k}/\Delta$, we may then
approximate

$\frac{\lambda}{2}\int_{-M}^{M}g^{\prime\prime}(\theta)^{2}d\theta\approx\frac{\lambda\Delta}{2}\sum_{k=1}^{K}\left[D\left(\frac{w}{\Delta}\right)\right]_{k}^{2},$
(2.25)

where $D$ is a differential quadrature matrix approximating the second
derivative, for example $D\in\mathbb{R}^{(K-2)\times K}$ having entries
$D[i,i+1]=-2/\Delta^{2}$, $D[i,i]=D[i,i+2]=1/\Delta^{2}$, and remaining
entries 0. This leads to the discretized algorithm where (2.21) is replaced by

$\displaystyle w_{t+1,k}$
$\displaystyle=w_{t,k}+\eta_{t}^{w}w_{t,k}\Bigg{[}\frac{1}{p}\sum_{j=1}^{p}\frac{\mathcal{N}_{\tau}(b_{k}-\varphi_{t+1,j})}{\sum_{i=1}^{K}w_{t,i}\,\mathcal{N}_{\tau}(b_{i}-\varphi_{t+1,j})}$
(2.26) $\displaystyle\hskip
144.54pt-\lambda\left(\frac{D^{\top}Dw_{t}}{\Delta}\right)_{k}-1+\lambda\sum_{i=1}^{K}w_{t,i}\left(\frac{D^{\top}Dw_{t}}{\Delta}\right)_{i}\Bigg{]}.$

Here, $\lambda(D^{\top}Dw)/\Delta$ represents the gradient of the right side
of (2.25), and the last term centers the update so that $\sum_{k}w_{t,k}=1$
remains constant. Noting that $D^{\top}D[w/\Delta]$ is an approximation to the
fourth derivative $g^{(4)}(\theta)$, this term $\lambda(D^{\top}Dw)/\Delta$
may also be understood as a discrete approximation to the first variation of
the continuous penalty on the left side of (2.25).

##  3 Theoretical guarantees

###  3.1 Consistency of near-NPMLEs

Formalizing (1.3), we may define any (constrained) NPMLE over $\mathcal{P}(M)$
by

$\displaystyle\widehat{g}\in\arg\,\min_{g\in\mathcal{P}(M)}\bar{F}_{n}(g).$

We note that the existence of a minimizer $\widehat{g}$ follows from the weak
compactness of $\mathcal{P}(M)$ and weak continuity of
$g\mapsto\bar{F}_{n}(g)$, although to our knowledge, uniqueness is not
resolved except for identity designs [Lin83a]. As such, our consistency result
applies to any such $\widehat{g}$, and more generally, for any near-minimizer
of $\bar{F}_{n}(g)$.

We start with the following assumption on the design matrix $\mathbf{X}$.

######  Assumption A.

Let $\mathbf{x}_{1},\ldots,\mathbf{x}_{p}$ denote the columns of $\mathbf{X}$,
and let the noise variance $\sigma^{2}>0$ be a fixed constant. There are
constants $C_{\ast},c_{\ast},\alpha,\beta>0$ such that the following hold as
$n,p\to\infty$:

  1. 1.

$\|\mathbf{X}\|_{{\operatorname{op}}}\leq C_{\ast}$ and
$\operatorname{Tr}\mathbf{X}^{\top}\mathbf{X}\geq c_{\ast}p$.

  2. 2.

Let $\mathbf{1}\in\mathbb{R}^{p}$ be the all-1’s vector. Then
$\|\mathbf{X}\mathbf{1}\|_{2}^{2}\geq c_{\ast}p$.

  3. 3.

For a subset $\mathcal{S}\subseteq\\{1,\ldots,p\\}$ of size
$|\mathcal{S}|\geq\alpha p$ and for each $j\in\mathcal{S}$, there exists a
vector $\mathbf{z}_{j}\in\mathbb{R}^{n}$ (deterministic conditional on
$\mathbf{X}$) such that $c_{\ast}\leq\|\mathbf{z}_{j}\|_{2}^{2}\leq C_{\ast}$,
$|\mathbf{z}_{j}^{\top}\mathbf{X}\mathbf{1}|\leq C_{\ast}$, and

$\displaystyle\max_{j\in\mathcal{S}}|\mathbf{z}_{j}^{\top}\mathbf{x}_{j}-1|\to
0,\qquad\max_{j\in\mathcal{S}}\sum_{k:k\neq
j}|\mathbf{z}_{j}^{\top}\mathbf{x}_{k}|^{2+\beta}\to 0.$ (3.1)

  4. 4.

For any $\\{a_{j}:j\in\mathcal{S}\\}$ satisfying $|a_{j}|\leq 1$,

$\displaystyle\left\|\sum_{j\in\mathcal{S}}a_{j}\mathbf{z}_{j}\right\|_{2}^{2}\leq
C_{\ast}p.$

Let us give a brief discussion of Assumption A. Conditions (1) and (2) are
mild, where (1) fixes a scaling of $\mathbf{X}$ to have typical singular
values of constant order, and (2) rules out cases such as $\mathbf{X}\bm{1}=0$
where the mean of $g_{\ast}$ would not be identifiable. Conditions (3) and (4)
are key conditions which posit the existence of a sufficient number of “test
vectors” $\\{\mathbf{z}_{j}\\}_{j\in\mathcal{S}}$ such that:

  1. 1.

each $\mathbf{z}_{j}$ is sufficiently aligned with $\mathbf{x}_{j}$;

  2. 2.

$\mathbf{z}_{j}$ has small alignment with the remaining
$\\{\mathbf{x}_{k}\\}_{k\neq j}$.

The role of $\\{\mathbf{z}_{j}\\}$ is easy to understand: The distribution of
$\mathbf{z}_{j}^{\top}\mathbf{y}$ for each $j\in\mathcal{S}$ will “isolate”
the $j$-th coordinate $\theta_{j}$ of ${\bm{\theta}}$, and be close to an
observation from a Gaussian sequence model (1.1) with a larger error variance.
This assumption is reminiscent of the construction of the debiased Lasso
[ZZ14], but an important difference is that here
$\\{\mathbf{z}_{j}\\}_{j\in\mathcal{S}}$ is used only as a proof device for
analyzing the NPMLE, and plays no role in the eventual algorithm.

Under Assumption A, the following result establishes the consistency of near
optimizers of $\bar{F}_{n}$, including any exact NPMLE $\widehat{g}$. Its
proof is presented in Appendix A.

######  Theorem 3.1.

Let $g_{*}\in\mathcal{P}(M)$, and suppose Assumption A holds as
$n,p\to\infty$. Let $\widetilde{g}\in\mathcal{P}(M)$ be any estimator of
$g_{*}$ such that, almost surely,

$\displaystyle\bar{F}_{n}(g_{\ast})-\bar{F}_{n}(\widetilde{g})=\frac{1}{p}\Big{(}\log
P_{\widetilde{g}}(\mathbf{y})-\log
P_{g_{*}}(\mathbf{y})\Big{)}\geq-\delta_{n}$

for some $\delta_{n}\downarrow 0$ as $n,p\to\infty$. Then
$\widetilde{g}\Rightarrow g_{*}$ almost surely, in the sense of weak
convergence.

A related consistency result was shown in the recent work [MSS23] studying
variational inference, in settings where $p\leq n$ and $\mathbf{X}$ has
singular values satisfying
$c\leq\sigma_{\min}(\mathbf{X})\leq\sigma_{\max}(\mathbf{X})\leq C$. Under
such conditions, the ‘test vectors’ in Assumption A can be constructed as
$\mathbf{z}_{j}=\mathbf{X}(\mathbf{X}^{\top}\mathbf{X})^{-1}\mathbf{e}_{j}$
with $\mathbf{e}_{j}$ denoting the $j$-th basis vector in $\mathbb{R}^{p}$, so
that $\\{\mathbf{z}_{j}^{\top}\mathbf{y}\\}_{j=1}^{p}$ are the usual linear
regression z-scores. Our Assumption A can be seen as a relaxation of these
conditions in [MSS23], allowing for a more general $(n,p)$-ratio, including
settings where the row span of $\mathbf{X}$ is not all of $\mathbb{R}^{p}$.

The following result shows that Assumption A is satisfied for random sub-
Gaussian designs with well-conditioned population covariance. We present its
proof in Appendix A.4.

######  Proposition 3.2.

Suppose $n/p\geq\gamma$ for any constant $\gamma>0$. Suppose
$\mathbf{X}\in\mathbb{R}^{n\times p}$ has i.i.d. sub-Gaussian rows
$\\{\frac{1}{\sqrt{n}}x^{(i)}\\}_{i=1}^{n}$ where $x^{(i)}$ has mean 0,
covariance ${\bm{\Sigma}}_{X}\in\mathbb{R}^{p\times p}$ satisfying

$\displaystyle
c\leq\lambda_{\min}({\bm{\Sigma}}_{X})\leq\lambda_{\max}({\bm{\Sigma}}_{X})\leq
C,$

and sub-Gaussian norm555 The sub-Gaussian norm of $U\in\mathbb{R}^{p}$ is
defined as
$\|U\|_{\psi_{2}}\triangleq\sup\\{t>0:\sup_{a\in\mathbb{R}^{p}:\lVert
a\rVert_{2}\leq 1}\mathbb{E}[\exp(t(a^{\top}U)^{2})]<\infty$. We say $U$ is
$K$-sub-Gaussian if $\|U\|_{\psi_{2}}\leq K$. bounded above by a constant
$K>0$. Then almost surely as $n,p\to\infty$, Assumption A is satisfied with
$\alpha=1$, any $\beta>0$, some $c_{\ast},C_{\ast}$ depending only on
$(K,\gamma,C,c)$, and

$\displaystyle\mathbf{z}_{j}=\Pi\mathbf{X}{\bm{\Sigma}}_{X}^{-1}\mathbf{e}_{j},$

where $\Pi\in\mathbb{R}^{n\times n}$ is the projection matrix orthogonal to
$\mathbf{X}\mathbf{1}$, and $\mathbf{e}_{j}\in\mathbb{R}^{p}$ is the $j$-th
standard basis vector.

As mentioned above, the vectors $\mathbf{z}_{j}$ are used only as a proof
device to analyze the NPMLE, and they may depend on the covariance
${\bm{\Sigma}}_{X}$ even when ${\bm{\Sigma}}_{X}$ is unknown and cannot be
consistently estimated in operator norm.

###  3.2 Convergence of gradient flows

In Section 3.2.1, we study convergence of the Fisher-Rao flow $\\{g_{t}\\}$ in
(2.16) where $q_{t}$ is replaced by a fixed $q\in\mathcal{P}(\mathbb{R}^{p})$,
corresponding to a sequence model NPMLE problem. In Section 3.2.2, we study
the Langevin diffusion corresponding to a Wasserstein-2 flow $\\{q_{t}\\}$ in
(2.15) where $g_{t}$ is fixed at $g\in\mathcal{P}(M)$, showing a uniform log-
Sobolev inequality for its stationary posterior measure under large noise
variance $\sigma^{2}>0$. Lastly in Section 3.2.3, we combine the previous two
ingredients to study the joint gradient flow (1.7).

####  3.2.1 Univariate $g$-flow

First consider the univariate Fisher-Rao gradient flow

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\,g_{t}(\theta)=g_{t}(\theta)\left[\mathcal{N}_{\tau}*\frac{\bar{q}}{\mathcal{N}_{\tau}*g_{t}}(\theta)-1\right]:=g_{t}(\theta)\left[\int\frac{\mathcal{N}_{\tau}(\theta-\varphi)}{\mathcal{N}_{\tau}*g_{t}(\varphi)}\mathrm{d}\bar{q}(\varphi)-1\right]$
(3.2)

for $\theta\in[-M,M]$ and a fixed probability distribution
$\bar{q}\in\mathcal{P}(\mathbb{R})$. This corresponds to the evolution of
$g_{t}$ alone in (2.16), which is a gradient flow for minimizing the objective

$\bar{F}(g\mid\bar{q}):=\int{-}\log[\mathcal{N}_{\tau}*g](\varphi)\mathrm{d}\bar{q}(\varphi).$

If $\bar{q}=\frac{1}{n}\sum_{i=1}^{n}\delta_{y_{i}}$ is an empirical
distribution of $n$ samples, then this objective is exactly the marginal log-
likelihood $\bar{F}_{n}(g)$ in the Gaussian sequence model (1.1) with noise
variance $\tau^{2}$, and (3.2) is exactly the gradient flow (2.4) for
computing the sequence model NPMLE.

Corollary C.2 in Appendix C.1 shows that, for any
$\bar{q}\in\mathcal{P}(\mathbb{R})$ with finite moment-generating function, if
(3.2) is initialized at $g_{0}\in\mathcal{P}_{*}(M)$ having strictly positive
density on all of $[-M,M]$, then it has a unique solution $\\{g_{t}\\}_{t\geq
0}$ in the space of positive densities. The following (deterministic) result
establishes convergence of this solution as $t\to\infty$.

######  Theorem 3.3.

Let $\bar{q}\in\mathcal{P}(\mathbb{R})$ with $\int
e^{\lambda\varphi}\bar{q}(\varphi)\mathrm{d}\varphi<\infty$ for all
$\lambda\in\mathbb{R}$. Suppose that (3.2) is initialized from a density
$g_{0}\in\mathcal{P}_{*}(M)$ that is strictly positive on $[-M,M]$. Then the
solution $\\{g_{t}\\}_{t\geq 0}$ to (3.2) satisfies, for any
$h\in\mathcal{P}_{*}(M)$,

$\displaystyle\bar{F}(g_{t}\mid\bar{q})-\bar{F}(h\mid\bar{q})\leq\frac{D_{\mathrm{KL}}(h\|g_{0})}{t}.$
(3.3)

If $\bar{q}=\mathcal{N}_{\tau}*g_{\ast}$ is a Gaussian mixture for some true
prior with density $g_{\ast}\in\mathcal{P}_{*}(M)$, then the minimizer of
$\bar{F}(g\mid\bar{q})$ is exactly $g_{\ast}$, and applying Theorem 3.3 with
$h=g_{\ast}$ bounds directly the suboptimality gap of $g_{t}$. If instead
$\bar{q}$ is the empirical distribution of a sample $(y_{1},\ldots,y_{n})$ and
$\widehat{g}$ is the minimizer of $\bar{F}(g\mid\bar{q})=\bar{F}_{n}(g)$ given
by the NPMLE in the sequence model (1.1), then we may apply Theorem 3.3 with
$h$ being a smoothed approximation of $\widehat{g}$, to show

$\bar{F}_{n}(g_{t})-\bar{F}_{n}(\widehat{g})\leq\frac{D_{\mathrm{KL}}(h\|g_{0})}{t}+\bar{F}_{n}(h)-\bar{F}_{n}(\widehat{g}).$

In particular, taking $t\to\infty$ first, followed by
$h\Rightarrow\widehat{g}$ weakly, and noting that $\bar{F}_{n}(g)$ is weakly
continuous in $g$, we have
$\lim_{t\to\infty}\bar{F}_{n}(g_{t})=\bar{F}_{n}(\widehat{g})$, so $g_{t}$ is
a near-NPMLE for large $t$.

The proof of Theorem 3.3 can be found in Appendix B.1. Motivated by viewing
(3.2) as a mirror flow [NY83], the analysis exploits the convexity of
$\bar{F}(g)\equiv\bar{F}(g\mid\bar{q})$ in the standard linear geometry for
any fixed $\bar{q}$, namely

$\bar{F}(\lambda
g_{0}+(1-\lambda)g_{1})\leq\lambda\bar{F}(g_{0})+(1-\lambda)\bar{F}(g_{1}),$
(3.4)

to bound $\bar{F}(g_{t})-\bar{F}(h)$ by the integrated first variation
$\int_{-M}^{M}{-}\delta\bar{F}[g_{t}](\theta)h(\theta)\mathrm{d}\theta$. Some
related results include [LLN19, LSW23, DEP23] which established exponential
convergence of the Fisher-Rao gradient flow for optimizing a different
objective $g\mapsto D_{\mathrm{KL}}(g\|\pi)$ with $\pi$ denoting the measure
to sample from, and [YWR23, Theorem 4] which showed a conditional convergence
result for a discrete-time version of (3.2). Complementary to this result,
Theorem 3.3 provides a quantitative and unconditional convergence guarantee
for the continuous-time flow.

####  3.2.2 Univariate $q$-flow

Next, define the negative log-posterior-density for prior $g(\cdot)$ (up to a
normalizing constant),

$\displaystyle
U_{g}({\bm{\varphi}})=\frac{1}{2}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{y}-\mathbf{X}{\bm{\varphi}})-\sum_{j=1}^{p}\log\mathcal{N}_{\tau}\ast
g(\varphi_{j}).$ (3.5)

We consider the Langevin diffusion

$\mathrm{d}{\bm{\varphi}}_{t}={-}\nabla
U_{g}({\bm{\varphi}}_{t})+\sqrt{2}\,\mathrm{d}\mathbf{B}_{t}$ (3.6)

with a fixed prior $g\in\mathcal{P}(M)$, whose law $q_{t}$ satisfies the
Fokker-Planck evolution

$\frac{\mathrm{d}}{\mathrm{d}t}q_{t}({\bm{\varphi}})=\nabla\cdot\left[q_{t}({\bm{\varphi}})\nabla
U_{g}({\bm{\varphi}})\right]+\Delta q_{t}({\bm{\varphi}}).$ (3.7)

This corresponds to the evolution of $q_{t}$ alone in (2.15).

We prove a uniform log-Sobolev inequality (LSI) for the stationary measure
$\nu[g]({\bm{\varphi}})=P_{g}({\bm{\varphi}}\mid\mathbf{y})$ of the diffusion
(3.6), for any prior $g\in\mathcal{P}(M)$ and any noise level
$\sigma^{2}>M^{2}\|\mathbf{X}\|_{\operatorname{op}}^{2}$. As this type of
result may be of independent interest in other contexts, we also provide a LSI
for the posterior measure
$\mu[g]({\bm{\theta}})=P_{g}({\bm{\theta}}\mid\mathbf{y})$ in the original
${\bm{\theta}}$ parametrization when $g\in\mathcal{P}_{*}(M)$ has a density
bounded from above and below.

For any $g\in\mathcal{P}(M)$, let us denote the posterior density of
${\bm{\varphi}}$ given $(\mathbf{X},\mathbf{y})$, under its i.i.d. prior
$\mathcal{N}_{\tau}*g$, by

$\displaystyle\nu[g]({\bm{\varphi}}):=P_{g}({\bm{\varphi}}\mid\mathbf{y})=\frac{\exp(-U_{g}({\bm{\varphi}}))}{\int\exp(-U_{g}({\bm{\varphi}}^{\prime}))\mathrm{d}{\bm{\varphi}}^{\prime}}.$
(3.8)

When the prior has a density $g\in\mathcal{P}_{*}(M)$, we denote the posterior
density of ${\bm{\theta}}$ given $(\mathbf{X},\mathbf{y})$ in the original
parametrization of the linear model (2.8) by
$\mu[g]({\bm{\theta}})=P_{g}({\bm{\theta}}\mid\mathbf{y})$. Following [Led99,
Eq. (2.18)], we say that a distribution $\rho$ on an open domain
$\Omega\subseteq\mathbb{R}^{p}$ satisfies a log-Sobolev inequality (LSI) with
constant $C>0$ if,666The inequality (3.9) corresponds to “LSI($\lambda$)” in
the formulation of [Vil09] with $\lambda=1/C$. for any bounded, Lipschitz, and
continuously differentiable function $f:\Omega\rightarrow\mathbb{R}$,

$\displaystyle\operatorname{\text{ent}}_{\rho}(f^{2}):=\mathbb{E}_{\rho}f^{2}\log
f^{2}-\mathbb{E}_{\rho}f^{2}\cdot\log\mathbb{E}_{\rho}f^{2}\leq
2C\cdot\mathbb{E}_{\rho}\lVert\nabla f\rVert^{2}.$ (3.9)

The following (deterministic) result, whose proof is given in Section B.2,
establishes a LSI for the posterior distribution in both parametrizations.

######  Theorem 3.4.

Let $g\in\mathcal{P}(M)$ and fix a constant $\delta>0$. Suppose that
$\sigma^{2}/\lVert\mathbf{X}\rVert_{{\operatorname{op}}}^{2}>M^{2}+\delta$.
Then:

  1. 1.

If $\tau^{2}\in(0,\sigma^{2}/\|\mathbf{X}\|_{\operatorname{op}}^{2}-\delta)$,
then $\nu[g]({\bm{\varphi}})$ satisfies a LSI with constant
$C=C(M,\tau,\delta)>0$.

  2. 2.

Suppose that $g\in\mathcal{P}_{*}(M)$ and
$\kappa(g)\equiv\frac{\sup_{\theta\in[-M,M]}g(\theta)}{\inf_{\theta\in[-M,M]}g(\theta)}<\infty$.
Then $\mu[g]({\bm{\theta}})$ satisfies a LSI on $(-M,M)^{p}$ with constant
$C\cdot\kappa(g)^{2}$ for some $C=C(M,\delta)>0$.

A direct and well-known consequence of Theorem 3.4-(1) (see [Vil09, p. 288])
is that the Langevin diffusion (3.6) converges exponentially fast in KL-
divergence to $\nu[g]$.

######  Corollary 3.5.

Under the assumption of Theorem 3.4-(1), let
$q_{0}\in\mathcal{P}(\mathbb{R}^{p})$, and let $\\{q_{t}\\}_{t\geq 0}$ be the
marginal laws of $\\{{\bm{\varphi}}_{t}\\}_{t\geq 0}$ in (3.6). Then for any
$t\geq 0$,

$\displaystyle D_{\mathrm{KL}}(q_{t}\|\nu[g])\leq\exp(-2t/C)\cdot
D_{\mathrm{KL}}(q_{0}\|\nu[g]),$

where $C$ is the same constant as in Theorem 3.4-(1).

We mention in passing that Theorem 3.4-(2) likewise implies exponential
convergence in KL for a Langevin diffusion in ${\bm{\theta}}$ reflected at the
boundaries of $[-M,M]^{p}$ [BEL18].

The posterior laws
$\nu[g]({\bm{\varphi}})=P_{g}({\bm{\varphi}}\mid\mathbf{y})$ and
$\mu[g]({\bm{\theta}})=P_{g}({\bm{\theta}}\mid\mathbf{y})$ are, in general,
not log-concave and also not close to product measures, and a dimension-free
LSI cannot be obtained by a naive application of Holley-Stroock perturbation
in $\mathbb{R}^{p}$. Our proof of Theorem 3.4 uses instead the idea of [BB19]
to express these posterior laws as log-concave mixtures of product
distributions. In our setting of a Bayesian linear model, this idea has an
intuitive interpretation related to the reparametrization of ${\bm{\theta}}$
by ${\bm{\varphi}}$: For the posterior of ${\bm{\theta}}$, we may consider the
generative process
${\bm{\theta}}\mapsto{\bm{\varphi}}^{\prime}\mapsto\mathbf{y}$ where
${\bm{\varphi}}^{\prime}={\bm{\theta}}+\mathbf{z}$ is given as in (2.9) for
some ${\tau^{\prime}}^{2}>0$ to be specified, and write

$\displaystyle P_{g}({\bm{\theta}}\mid\mathbf{y})=\int
P_{g}({\bm{\theta}}\mid{\bm{\varphi}}^{\prime})P_{g}({\bm{\varphi}}^{\prime}\mid\mathbf{y})\mathrm{d}{\bm{\varphi}}^{\prime}.$

Then a LSI for $P_{g}({\bm{\theta}}\mid\mathbf{y})$ follows from those for
$P_{g}({\bm{\theta}}\mid{\bm{\varphi}}^{\prime})$ (uniform over
${\bm{\varphi}}^{\prime}$) and $P_{g}({\bm{\varphi}}^{\prime}\mid\mathbf{y})$
(uniform over $\mathbf{y}$), which are established as follows:

  * •

Since $P_{g}({\bm{\theta}}\mid{\bm{\varphi}}^{\prime})$ is a product measure,
by tensorization it suffices to establish a uniform LSI over its univariate
marginals. We do this using a Holley-Stroock perturbation argument and the
characterization of LSI for one-dimensional measures [BG99] (c.f. Lemma B.2).

  * •

The smoothed prior $\mathcal{N}_{\tau^{\prime}}\ast g$, and hence also the
posterior measure $P_{g}({\bm{\varphi}}^{\prime}\mid\mathbf{y})$, become
strongly log-concave when ${\tau^{\prime}}^{2}>M^{2}$, and a LSI follows from
the Bakry-Emery criterion [BE85].

The LSI for $P_{g}({\bm{\varphi}}\mid\mathbf{y})$ follows from a similar
argument, where we consider instead the generative process
${\bm{\theta}}\mapsto{\bm{\varphi}}\mapsto{\bm{\varphi}}^{\prime}\mapsto\mathbf{y}$
to represent

$\displaystyle P_{g}({\bm{\varphi}}\mid\mathbf{y})=\int
P_{g}({\bm{\varphi}}\mid{\bm{\varphi}}^{\prime})P_{g}({\bm{\varphi}}^{\prime}\mid\mathbf{y})\mathrm{d}{\bm{\varphi}}^{\prime}.$

We remark that if $\tau^{2}>M^{2}$ in our starting reparametrization by
${\bm{\varphi}}$, then we may simply take
${\bm{\varphi}}^{\prime}={\bm{\varphi}}$ and the LSI follows directly from
log-concavity. However, choosing a large value of $\tau^{2}$ leads to a
difficult sequence-model NPMLE problem to be solved by the evolution of
$\\{g_{t}\\}$ in the joint flow, which we observe in practice can lead to
unstable behavior of the discretized algorithm and slow convergence. Theorem
3.4 allows for any value $\tau^{2}>0$, for which the prior
$\mathcal{N}_{\tau}*g$ is not necessarily log-concave.

The mixing of Langevin diffusions for posterior sampling in Bayesian
regression models has been studied before, for example in [DT12] using a Meyn-
Tweedie approach and in [Dal17] for log-concave priors. We mention also the
recent results of [NW22] that show rapid mixing of Langevin diffusions in
settings where the posterior concentrates in a locally log-concave region of
the parameter space. However, to our knowledge, these results do not seem to
directly imply polynomial-time mixing in many settings of interest in our
application, such as that of Proposition 3.2 where $p$ scales linearly with
$n$ and the posterior $P_{g}({\bm{\varphi}}\mid\mathbf{y})$ may not be log-
concave or contract around the true regression coefficient parameter.

####  3.2.3 Joint flow

Equipped with the understanding of its two components, we now study the joint
flow (2.15–2.16), recorded here for the reader’s convenience:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}q_{t}({\bm{\varphi}})$
$\displaystyle=\nabla\cdot\left[q_{t}({\bm{\varphi}})\nabla
U_{g_{t}}({\bm{\varphi}})\right]+\Delta q_{t}({\bm{\varphi}})$ (3.10)
$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}g_{t}(\theta)$
$\displaystyle=\alpha\cdot
g_{t}(\theta)\left(\left[\mathcal{N}_{\tau}*\frac{\bar{q}_{t}}{\mathcal{N}_{\tau}*g_{t}}\right](\theta)-1\right).$

The following first verifies the existence and uniqueness of a solution to
this joint gradient flow in
$\mathcal{P}_{*}(\mathbb{R}^{p})\times\mathcal{P}_{*}(M)$.

######  Theorem 3.6.

Fix any $T>0$. Suppose the initialization
$(q_{0},g_{0})\in\mathcal{P}_{*}(\mathbb{R}^{p})\times\mathcal{P}_{*}(M)$
satisfies:

  1. (i)

The density $q_{0}$ is continuous and upper bounded, and for any
${\bm{\lambda}}\in\mathbb{R}^{p}$ we have $\mathbb{E}_{{\bm{\varphi}}\sim
q_{0}}e^{{\bm{\lambda}}^{\top}{\bm{\varphi}}}<\infty$.

  2. (ii)

The density $g_{0}$ is strictly positive on all of $[-M,M]$.

Then there exists a unique classical solution
$\\{(q_{t},g_{t})\\}_{t\in[0,T]}$ to (2.15–2.16) such that
$(q_{t},g_{t})\in\mathcal{P}_{*}(\mathbb{R}^{p})\times\mathcal{P}_{*}(M)$ for
all $t\in[0,T]$. This solution $\\{(q_{t},g_{t})\\}_{t\in[0,T]}$ satisfies:

  1. 1.

$t\mapsto g_{t}(\theta)$ is $C^{\infty}$ on $(0,T)$ for each fixed
$\theta\in[-M,M]$. Furthermore, $g_{t}(\theta)\geq g_{0}(\theta)e^{-t}$ for
each $\theta\in[-M,M]$ and $t\in[0,T]$

  2. 2.

$\\{q_{t}\\}_{t\in[0,T]}$ are the densities of the marginal laws of
$\\{{\bm{\varphi}}_{t}\\}_{t\in[0,T]}$ solving the SDE (2.17) with
initialization ${\bm{\varphi}}_{0}\sim q_{0}$. For any
${\bm{\lambda}}\in\mathbb{R}^{p}$,

$\displaystyle\sup_{t\in[0,T]}\mathbb{E}_{{\bm{\varphi}}\sim
q_{t}}e^{{\bm{\lambda}}^{\top}{\bm{\varphi}}}<\infty.$ (3.11)

Furthermore, $(t,{\bm{\varphi}})\mapsto q_{t}({\bm{\varphi}})$ is $C^{\infty}$
on $(0,T)\times\mathbb{R}^{p}$.

We prove Theorem 3.6 in Appendix C.3, by identifying (2.15–2.16) as the fixed
point of a contractive mapping from a suitably defined Banach space to itself.

The following (deterministic) result shows, under the joint gradient flow
(2.15–2.16), that $\\{g_{t}\\}$ reaches a near-NPMLE in polynomial time if the
marginal log-likelihood objective $\bar{F}_{n}$ is convex in a sub-level set
defined by the initial objective value. In the absence of such a local
convexity guarantee, we show more generally that $\\{g_{t}\\}$ reaches a
density that approximately satisfies an averaged first-order optimality
condition for $\bar{F}_{n}(g)$, defined in terms of its (centered) first
variation integrated against any fixed test density $h\in\mathcal{P}_{*}(M)$.
For any density $g\in\mathcal{P}_{*}(M)$ with full support, this first
variation takes the form (c.f. Appendix D.2)

$\delta\bar{F}_{n}[g](\theta)={-}\frac{\bar{\mu}[g]}{g}(\theta)+1={-}\mathcal{N}_{\tau}*\frac{\bar{\nu}[g]}{\mathcal{N}_{\tau}*g}(\theta)+1$
(3.12)

where $\bar{\mu}[g]$ (resp. $\bar{\nu}[g]$) is the average marginal density of
$\mu[g]$ (resp. $\nu[g]$) as defined by (2.7) and (2.14). (As explained in
Appendix D.2, the first variation does not depend on $\tau$, and the last
expression in (3.12) holds for any $\tau>0$ and more generally for
$g\in\mathcal{P}(M)$.)

######  Theorem 3.7.

Suppose that $\nu[g]({\bm{\varphi}})$ satisfies a LSI with some constant
$C=C(M,\tau)>0$, uniformly over all priors $g\in\mathcal{P}(M)$. Suppose that
$(q_{0},g_{0})\in\mathcal{P}_{*}(\mathbb{R}^{p})\times\mathcal{P}_{*}(M)$
satisfy the conditions of Theorem 3.6, and let $\\{(q_{t},g_{t})\\}_{t\geq 0}$
be the solution to (2.15–2.16). Let
$F_{n,\ast}=\min_{g}\bar{F}_{n}(g)=\min_{(q,g)}F_{n}(q,g)$, and define
$\delta\bar{F}_{n}[g]$ by (3.12).

Then there exist $(M,\tau)$-dependent constants $C,c>0$ such that for any
density $h\in\mathcal{P}_{*}(M)$, any
$\varepsilon\in(0,\min(c,\frac{1}{16\alpha}))$, and some time

$\displaystyle t\leq
T(h,\varepsilon):=\frac{C}{\varepsilon}D_{\mathrm{KL}}(h\|g_{0})+\frac{Cp\cdot
e^{C\sqrt{\log(1/\varepsilon)}}}{\varepsilon^{2}}(F_{n}(q_{0},g_{0})-F_{n,*}),$
(3.13)

it holds that

$\displaystyle\int_{-M}^{M}-\delta\bar{F}_{n}[g_{t}](\theta)h(\theta)\mathrm{d}\theta\leq\varepsilon.$
(3.14)

If $\bar{F}_{n}$ is convex (in the usual sense of (3.4)) over the sub-level
set $\mathcal{K}_{n}\equiv\\{g:\bar{F}_{n}(g)\leq F_{n}(q_{0},g_{0})\\}$, and
if $h\in\mathcal{K}_{n}$, then there exists some $t\leq T(h,\varepsilon)$ such
that

$\displaystyle\bar{F}_{n}(g_{t})-\bar{F}_{n}(h)\leq\varepsilon.$ (3.15)

Let us make several remarks regarding the conditions and interpretation of
this theorem.

  1. 1.

The required uniform LSI condition is provided, for example, by Theorem 3.4 in
high-noise settings where
$\sigma^{2}>M^{2}\|\mathbf{X}\|_{\operatorname{op}}^{2}$. Note that under
Assumption A, this represents a constant and dimension-independent requirement
for $\sigma^{2}$, for which consistency of the NPMLE as guaranteed by Theorem
3.1 can still hold. Reparametrization by ${\bm{\varphi}}$ is important for
this LSI to hold for the posterior laws $\nu[g_{t}]$ even if $\\{g_{t}\\}$
converges to a discrete measure.

  2. 2.

The resulting guarantee (3.14) may be understood as an averaged version of a
first-order optimality condition for $g_{t}$ to locally minimize
$\bar{F}_{n}(g)$. Indeed, any local minimizer $\widetilde{g}$ of
$\bar{F}_{n}(g)$ must satisfy
$\frac{\mathrm{d}}{\mathrm{d}t}\bar{F}_{n}\big{(}(1-t)\widetilde{g}+t\delta_{\theta}\big{)}|_{t=0+}\geq
0$ for all $\theta\in[-M,M]$, which implies that

$\displaystyle-\delta\bar{F}_{n}(\widetilde{g})(\theta)\leq 0\quad\text{for
all }\theta\in[-M,M].$ (3.16)

The statement (3.14) averages this inequality with respect to a fixed
(arbitrary) reference density $h$. If $\bar{F}_{n}$ is convex on the sub-level
set $\mathcal{K}_{n}$ containing $h$, then

$\displaystyle\bar{F}_{n}(h)-\bar{F}_{n}(g)$
$\displaystyle\geq\int_{-M}^{M}\delta\bar{F}_{n}[g](\theta)\big{(}h(\theta)-g(\theta)\big{)}\mathrm{d}\theta=\int_{-M}^{M}\delta\bar{F}_{n}[g](\theta)h(\theta)\mathrm{d}\theta,$

so in particular, the statement (3.14) is sufficient to imply the bound (3.15)
on the sub-optimality gap.

  3. 3.

Under Assumption A and any choice
$\tau^{2}<\sigma^{2}/\|\mathbf{X}\|_{\operatorname{op}}^{2}-\delta$, one may
check (see Appendix B.4) that on an event of probability approaching 1,

$F_{n}(q_{0},g_{0})-F_{n,*}\leq
C\Big{(}1+\frac{n+\mathbb{E}_{q_{0}}[\|{\bm{\varphi}}\|_{2}^{2}]}{p}\Big{)}+\frac{1}{p}D_{\mathrm{KL}}\big{(}q_{0}\big{\|}(\mathcal{N}_{\tau}*g_{0})^{\otimes
p}\big{)}$ (3.17)

for a constant $C:=C(\delta,\sigma^{2},M)>0$. If
$g_{\ast}\in\mathcal{P}_{*}(M)$ admits a density on $[-M,M]$, then choosing
$h=g_{\ast}$ in (3.15) verifies that, under the local convexity assumption of
the theorem and the mild initial conditions
$D_{\mathrm{KL}}(g_{*}\|g_{0})=O(1)$,
$D_{\mathrm{KL}}(q_{0}\|(\mathcal{N}_{\tau}*g_{0})^{\otimes p})=O(p)$, and
$\mathbb{E}_{q_{0}}[\|{\bm{\varphi}}\|_{2}^{2}]=O(p)$, the gradient flow
reaches in time $t=O(\varepsilon^{-2+o(1)}(n+p))$ a near-NPMLE $g_{t}$ (in the
sense of Theorem 3.1) satisfying
$\bar{F}_{n}(g_{t})\leq\bar{F}_{n}(g_{*})+\varepsilon$.

If $g_{\ast}$ does not admit a density, then the theorem may be applied with
$h$ being a smoothed approximation of $g_{\ast}$ similar to the discussion
after Theorem 3.3 for the univariate $g$-flow.

  4. 4.

We emphasize that the local convexity assumption needed for (3.15) pertains to
the marginal objective $\bar{F}_{n}$, and the theorem does not require
$F_{n}(q,g)$ to be jointly convex on the sub-level set
$\\{(q,g):F_{n}(q,g)\leq F_{n}(q_{0},g_{0})\\}$.

Fixing any initial prior $g_{0}$ and any $\delta>0$, under the assumed LSI,
one may first run a Langevin diffusion for a dimension-independent time
$T_{0}(\delta)$ to obtain an initialization $q_{0}$ that satisfies
$F_{n}(q_{0},g_{0})<F_{n}(\nu[g_{0}],g_{0})+\delta=\bar{F}_{n}(g_{0})+\delta$.
The theorem may then be applied with this $(q_{0},g_{0})$, under the condition
that $\bar{F}_{n}$ is convex over
$\\{g:\bar{F}_{n}(g)<\bar{F}_{n}(g_{0})+\delta\\}$.

In certain asymptotic settings as $n,p\to\infty$, we expect $\bar{F}_{n}(g)$
and its functional derivatives to converge to those of a deterministic and
dimension-independent functional $\bar{F}(g)$. In such settings,
$\bar{F}_{n}(g)$ for all large $n,p$ would be convex in a fixed sub-level set
$\mathcal{K}$ of $\bar{F}(g)$ containing $g_{\ast}$, whose definition is
dimension-free. We formalize such a statement in Section 3.2.4 below.

The proof of Theorem 3.7 is given in Appendix B.3. On a high level, the proof
analyzes the integrated first variation
$\int_{-M}^{M}{-}\delta\bar{F}_{n}[g_{t}](\theta)h(\theta)\mathrm{d}\theta$ as
in Theorem 3.3 on convergence of the $g$-flow, and shows that under the
assumed LSI, replacing the posterior law $\bar{\nu}[g_{t}]$ in (D.2) by
$\bar{q}_{t}$ can inflate the value of
$\int_{-M}^{M}{-}\delta\bar{F}_{n}[g_{t}](\theta)h(\theta)\mathrm{d}\theta$ by
at most a constant factor plus a quantity that is controllable in terms of
$D_{\mathrm{KL}}(q_{t}\|\nu[g_{t}])$. The divergence
$D_{\mathrm{KL}}(q_{t}\|\nu[g_{t}])$ may in turn be bounded by the time-
derivative of the objective value $F_{n}(q_{t},g_{t})$ via a second
application of the assumed LSI.

As mentioned in the introduction, convergence of MCEM methods to fixed points
of the EM algorithm has been classically studied in [CL95, FM03, Nea13], in
settings of asymptotically growing numbers of EM iterations between discrete
M-step updates. Theorem 3.7 may be understood as a convergence result for a
continuous-time version of MCEM, in a regime where the prior $g_{t}$ is
continuously updated before the law $q_{t}$ of the posterior sample reaches
the stationary law $\nu[g_{t}]$.

####  3.2.4 Local convexity

Finally, let us formalize an asymptotic condition under which the sub-level
set $\mathcal{K}_{n}$ of Theorem 3.7 may be taken to contain a fixed,
dimension-free neighborhood of the true prior $g_{\ast}$, for all large $n,p$.

For any strictly positive density $g\in\mathcal{P}_{*}(M)$ and bounded
function $\chi:[-M,M]\to\mathbb{R}$ satisfying
$\int_{-M}^{M}\chi(\theta)\mathrm{d}\theta=0$, we may define a functional
Hessian of $\bar{F}_{n}$ in the direction of $\chi$ by

$\bar{H}_{n}[g](\chi)=\frac{\mathrm{d}^{2}}{\mathrm{d}^{2}\varepsilon}\Big{|}_{\varepsilon=0}\bar{F}_{n}(g+\varepsilon\chi).$

Explicitly, writing $\left\langle f\right\rangle_{g}=\int
f({\bm{\varphi}})\nu[g]({\bm{\varphi}})\mathrm{d}{\bm{\varphi}}$ as shorthand
for the mean of $f({\bm{\varphi}})$ under the posterior law
$\nu[g]({\bm{\varphi}})=P_{g}({\bm{\varphi}}\mid\mathbf{y})$, we may check by
direct differentiation of the form (2.10) for $\bar{F}_{n}(g)$ (c.f. Lemma
B.5) that

$\displaystyle\bar{H}_{n}[g](\chi)=\frac{1}{p}\Bigg{[}\left\langle\sum_{j=1}^{p}\Big{(}\frac{\mathcal{N}_{\tau}\ast\chi}{\mathcal{N}_{\tau}\ast
g}(\varphi_{j})\Big{)}^{2}\right\rangle_{g}-\left\langle\Big{(}\sum_{j=1}^{p}\frac{\mathcal{N}_{\tau}\ast\chi}{\mathcal{N}_{\tau}\ast
g}(\varphi_{j})\Big{)}^{2}\right\rangle_{g}+\left\langle\sum_{j=1}^{p}\frac{\mathcal{N}_{\tau}\ast\chi}{\mathcal{N}_{\tau}\ast
g}(\varphi_{j})\right\rangle^{2}_{g}\Bigg{]}.$ (3.18)

We may extend this definition (3.18) of $\bar{H}_{n}[g](\chi)$ to all
probability measures $g\in\mathcal{P}(M)$ and signed measures $\chi$ on
$[-M,M]$ such that $\int_{-M}^{M}\chi(\mathrm{d}\theta)=0$.

We consider a general setting where there exist limiting functionals
$\bar{F}(g)$ and $\bar{H}[g](\chi)$ to which $\bar{F}_{n}(g)$ and
$\bar{H}_{n}[g](\chi)$ converge uniformly, as $n,p\to\infty$. Let
$\lVert\chi\rVert_{\operatorname{\text{TV}}}=\int_{-M}^{M}|\chi(\mathrm{d}\theta)|$
be the total variation norm of the signed measure $\chi$, and define

$\displaystyle\Gamma=\Big{\\{}\chi:\chi\text{ is a signed measure supported on
}[-M,M],\;\int_{-M}^{M}\chi(\mathrm{d}\theta)=0,\;\lVert\chi\rVert_{\operatorname{\text{TV}}}=1\Big{\\}}.$

######  Proposition 3.8.

Define $\bar{H}_{n}[g]$ by (3.18), and suppose that there exist deterministic
functionals $\bar{F}:\mathcal{P}(M)\to\mathbb{R}$ and
$\bar{H}:\mathcal{P}(M)\times\Gamma\to\mathbb{R}$ (not depending on $n,p$)
such that, as $n,p\to\infty$,

$\displaystyle\sup_{g\in\mathcal{P}(M)}\Big{|}\bar{F}_{n}(g)-\bar{F}(g)\Big{|}\overset{P}{\to}0,\qquad\sup_{g\in\mathcal{P}(M),\,\chi\in\Gamma}\Big{|}\bar{H}_{n}[g](\chi)-\bar{H}[g](\chi)\Big{|}\overset{P}{\to}0.$
(3.19)

Suppose furthermore that the true prior $g_{\ast}$ is the unique minimizer of
$\bar{F}$, and $\bar{H}$ satisfies

$\displaystyle\inf_{\chi\in\Gamma}\bar{H}[g_{\ast}](\chi)>0.$ (3.20)

Let $F_{n,*}=\min_{g}\bar{F}_{n}(g)$. Then there exist a weak neighborhood
$\mathcal{K}$ of $g_{\ast}$ and a constant $c_{0}>0$, both independent of
$(n,p)$, such that with probability approaching 1, $\bar{F}_{n}(g)$ is convex
over the sub-level set $\mathcal{K}_{n}=\\{g:\bar{F}_{n}(g)\leq
F_{n,*}+c_{0}\\}$ which contains $\mathcal{K}$.

The proof of Proposition 3.8 is given in Appendix B.5. One simple type of
design under which these asymptotic conditions may be verified is a block
design where $\mathbf{X}$ is block-diagonal with $N$ equal blocks
$\bar{\mathbf{X}}\in\mathbb{R}^{B_{1}\times B_{2}}$ independent of $(n,p)$,
and $(n,p)=(NB_{1},NB_{2})$. Then we may simply take
$\bar{F}=\mathbb{E}\bar{F}_{n}$ and $\bar{H}=\mathbb{E}\bar{H}_{n}$, noting
that these functionals are the same for all $N\geq 1$. We verify in
Proposition B.8 of Appendix B.5 that the conditions of Proposition 3.8 indeed
hold for this example.

We believe that the conditions of Proposition 3.8 may hold also for random
designs $\mathbf{X}$ with i.i.d. entries, and more generally for random sub-
Gaussian designs of the form of Proposition 3.2 under certain asymptotic
conditions for the design covariance ${\bm{\Sigma}}_{X}$, in the proportional
asymptotics limit as $n,p\to\infty$ with $n/p\to\gamma\in(0,\infty)$. Here,
$\bar{F}(g)$ would be the limit evidence or free energy of the model under a
misspecified prior $g$, and $\bar{H}[g]$ its functional Hessian. Establishing
such a convergence result is outside the scope of the current paper, although
we refer to [GKKZ23] for recent related results in a different problem of low-
rank matrix estimation and [BCPS21] in a linear regression model where the
misspecified prior $g$ has a parametric Gaussian form.

##  4 Experiments

We compare the discretized method of Section 2—hereafter referred to as
EBflow—with several more “vanilla” versions of Monte-Carlo expectation
maximization (MCEM) and with an approach based on mean-field variational
inference using coordinate ascent updates (CAVI).

For MCEM, the most direct comparison is with an implementation using ULA to
approximate the E-step, with the same latent parameter
${\bm{\varphi}}={\bm{\theta}}+\mathbf{z}$ (with
$\mathbf{z}\sim\mathcal{N}(0,\tau^{2}\operatorname{Id})$) as in EBflow. We
perform $T_{\text{iter}}$ ULA iterations (2.20) with fixed step size
$\eta^{\varphi}$ and prior weights $w$ between M-steps, and update $w$ in each
M-step by computing the NPMLE in the sequence model (1.1) from the coordinates
$(\varphi_{t,j})_{j=1}^{p}$ in the past $T_{\text{iter}}$ iterations. This
update is a convex program, which we solve using CVXR. We refer to this method
as Langevin-MCEM. We also compare to a Gibbs-MCEM method and to CAVI in the
original parametrization by ${\bm{\theta}}$. For Gibbs-MCEM, we perform
$T_{\text{iter}}$ Gibbs sampling iterations between updates of prior weights
$w$, each iteration resampling $\\{\theta_{j}\\}_{j\in[p]}$ in sequence. For
CAVI, we optimize the variational representation of the marginal log-
likelihood over priors and product posteriors for $\bm{\theta}$, by updating
the posterior distribution for each coordinate $\\{\theta_{j}\\}_{j\in[p]}$ in
sequence, followed by the prior. As a short burn-in period, we perform the
first 200 iterations of EBflow, Langevin-MCEM, and Gibbs-MCEM without updating
prior weights $w$, under the initialization
$w=\operatorname{Uniform}(\\{b_{1},\ldots,b_{K}\\})$. Further details of the
methods are discussed in Appendix E.

In Section 4.1 we first discuss some theoretical considerations on parameter
tuning, such as step sizes and the reparametrization parameter $\tau$. We
focus on a simple setting with Gaussian prior and i.i.d. Gaussian design as a
benchmark to choose the parameters of EBflow and other methods, so that they
all have reasonable performance in this pilot example. In Section 4.2, we then
compare these methods across other priors and designs.

###  4.1 Parameter choices

Figure 1: TV distance $d_{\mathrm{TV}}(g_{t},g_{\ast})$ across 10000
iterations for estimating a standard Gaussian prior
$g_{\ast}=\mathcal{N}(0,1)$ via optimization of a smoothed log-likelihood
objective, using various parameter tunings for EBflow, Langevin-MCEM, and
Gibbs-MCEM (each with $10$ independent runs).

Variable reparametrization. It is understood that the approximation accuracy
of ULA for continuous-time Langevin dynamics improves with the smoothness of
the log-density of its stationary distribution [VW19, BCE+22]. In the
parametrization by ${\bm{\varphi}}$, the negative Hessian of the log-
posterior-density is (by Tweedie’s formula)

$-\nabla^{2}\log
P({\bm{\varphi}}\mid\mathbf{y})=\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}\mathbf{X}+\frac{1}{\tau^{2}}\operatorname{Id}-\frac{1}{\tau^{4}}\operatorname{\textrm{diag}}\Big{(}\operatorname{Var}[\theta_{j}\mid\varphi_{j}]\Big{)}_{j=1}^{p}$
(4.1)

where we recall
${\bm{\Sigma}}=(\sigma^{2}\operatorname{Id}-\tau^{2}\mathbf{X}\mathbf{X}^{\top})^{-1}$.
Ignoring the last ${\bm{\varphi}}$-dependent term, which becomes negligible
for bounded $\theta_{j}$ and large $\tau$, and treating the largest eigenvalue

$\lambda_{\max}:=\lambda_{\max}(\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}\mathbf{X}+\tau^{-2}\operatorname{Id})$

of the first two terms as a proxy for the smoothness of this log-posterior, we
will fix the choice

$\tau^{2}=0.5\,\sigma^{2}/\lambda_{\max}(\mathbf{X}\mathbf{X}^{\top})$ (4.2)

in the reparametrization of ${\bm{\theta}}$ by ${\bm{\varphi}}$, which is
easily checked to be the value that minimizes $\lambda_{\max}$. Theoretical
results of [VW19, BCE+22] on ULA suggest a step size inversely proportional to
the smoothness parameter. As such, we will rescale the step size
$\eta_{t}^{\varphi}$ in EBflow and Langevin-MCEM relative to
$1/\lambda_{\max}$, i.e. all further reference to $\eta_{t}^{\varphi}$ will be
in the following reparametrized form of (2.20):

${\bm{\varphi}}_{t+1}={\bm{\varphi}}_{t}-\frac{\eta^{\varphi}_{t}}{\lambda_{\max}}\Bigg{[}\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}(\mathbf{X}{\bm{\varphi}}_{t}-\mathbf{y})+\frac{1}{\tau^{2}}\bigg{(}\frac{\sum_{k=1}^{K}w_{t,k}(\varphi_{t,j}-b_{k})\mathcal{N}_{\tau}(\varphi-
b_{k})}{\sum_{k=1}^{K}w_{t,k}\,\mathcal{N}_{\tau}(\varphi_{t,j}-b_{k})}\bigg{)}_{j=1}^{p}\Bigg{]}+\sqrt{\frac{2\eta^{\varphi}_{t}}{\lambda_{\max}}}\,\mathbf{b}_{t}.$
(4.3)

Step sizes. A larger step size $\eta_{t}^{\varphi}$ in EBflow leads to faster
mixing, but can result in a higher discretization bias for ULA. The step size
$\eta_{t}^{w}$ determines an exponential weighting of past iterates when
updating the prior. An overly large $\eta_{t}^{w}$ can lead to high Monte
Carlo variance in the estimated prior, while an overly small $\eta^{w}_{t}$
can lead to slow convergence.

To illustrate these trade-offs, we generate data
$\mathbf{y}=\mathbf{X}{\bm{\theta}}+{\bm{\varepsilon}}$ with dimensions
$n=500$ and $p=1000$, with standard Gaussian prior
$g_{\ast}=\mathcal{N}(0,1)$, i.i.d. Gaussian design $\mathbf{X}$, and noise
variance $\sigma^{2}$ chosen777In our experiment, we first generate
$\mathbf{X}$ and ${\bm{\theta}}$, and then set $\sigma^{2}$ equal to the
sample variance of the entries of $\mathbf{X}\theta$. so that
${\bm{\varepsilon}}$ and $\mathbf{X}{\bm{\theta}}$ each explain 50% of the
variance of $\mathbf{y}$. Let $w_{\ast,1},\ldots,w_{\ast,K}$ denote weights
representing $g_{*}$ discretized to the grid $b_{1},\ldots,b_{K}\in[-3,3]$.
Figure 1(a) displays the total-variation (TV) error
$d_{\mathrm{TV}}(g_{t},g_{\ast})\approx\frac{1}{2}\sum_{k=1}^{K}|w_{t,k}-w_{\ast,k}|$
of the density estimate for 10,000 iterations of EBflow, over 10 independent
trials on the same data $(\mathbf{X},\mathbf{y})$. To produce a smooth density
estimate, we apply all methods to minimize the negative log-likelihood
objective with an additional smoothing spline regularizer as in (2.24), with
common fixed regularization $\lambda=0.003$.

In this example, fixing larger step sizes $\eta^{\varphi}_{t}=1.0$ and
$\eta^{w}_{t}=0.01$ results in reasonable performance, with smaller
$(\eta_{t}^{w},\eta_{t}^{\varphi})$ leading to slower convergence and larger
$\eta_{t}^{w}/\eta_{t}^{\varphi}$ leading to higher Monte Carlo variance in
the estimated prior. There is a small amount of ULA bias incurred by the large
step size choice $\eta^{\varphi}_{t}=1.0$, and Figure 1(b) displays the TV
error under an alternative log-linear step size decay over 10,000 iterations,

$\eta_{t}^{\varphi}=a\cdot b^{t}\text{ with }\eta_{t}^{\varphi}=1.0\text{ at
}t=1\text{ and }\eta_{t}^{\varphi}=0.1\text{ at
}t=10000,\qquad\eta_{t}^{w}=0.01\eta_{t}^{\varphi}.$ (4.4)

The reduction of ULA bias at the final step size $\eta_{t}^{\varphi}=0.1$
yields a reduction of the average TV error of the estimated density to 0.046,
from 0.075 for the fixed choice $\eta_{t}^{\varphi}=1.0$. We recommend using
this log-linear decay (4.4) as the default option for problems of this size,
with the possibility of increasing $\eta_{t}^{w}$ relative to
$\eta_{t}^{\varphi}$ for problems with higher dimension $p$.

Analogous trade-offs apply for the other MCEM approaches when determining the
number of iterations $T_{\text{iter}}$ between prior updates. Figure 1(c–d)
displays TV error for Langevin-MCEM with $\eta_{t}^{\varphi}\in\\{1.0,0.1\\}$
and $T_{\text{iter}}\in\\{100,1000\\}$, and for Gibbs-MCEM with
$T_{\text{iter}}\in\\{10,100,1000\\}$, both across 10 independent trials for
the same data $(\mathbf{X},\mathbf{y})$. We will henceforth fix
$T_{\text{iter}}=100$ for both methods to achieve a reasonable balance between
Monte Carlo variance and convergence speed.

###  4.2 Method comparison

Priors and designs. We generate data
$\mathbf{y}=\mathbf{X}{\bm{\theta}}+{\bm{\varepsilon}}$ where
$\\{\theta_{j}\\}_{j=1}^{p}\overset{iid}{\sim}g_{\ast}$, and the prior
distribution $g_{\ast}$ is one of four smooth densities:

  * •

(Gaussian) $g_{\ast}=\mathcal{N}(0,1)$

  * •

(Cauchy) $g_{\ast}=\operatorname{Cauchy}(0,0.6)$

  * •

(Skew)
$g_{\ast}=\frac{1}{3}\mathcal{N}(-2,0.5^{2})+\frac{1}{3}\mathcal{N}(-1.5,1^{2})+\frac{1}{3}\mathcal{N}(0,2^{2})$

  * •

(Bimodal)
$g_{\ast}=\frac{1}{2}\mathcal{N}(-1.5,0.5^{2})+\frac{1}{2}\mathcal{N}(1.5,0.5^{2})$

For consistency, each prior is truncated to $[-3,3]$ and discretized over a
grid of $K=61$ points $b_{1},\ldots,b_{K}$ with equal spacing 0.1. We generate
$\mathbf{X}$ to be one of the following designs:

  * •

(identity) $\mathbf{X}=\operatorname{Id}$ and $n=p$.

  * •

(iid) $\mathbf{X}$ has i.i.d. Gaussian entries.

  * •

(block02corr0.9) $\mathbf{X}$ has multivariate Gaussian rows, consisting of
$p/2$ independent pairs of variables with correlation 0.9.

  * •

(block10corr0.5) $\mathbf{X}$ has multivariate Gaussian rows, consisting of
$p/10$ independent blocks of variables with pairwise correlation 0.5.

We fix the dimension $p=1000$ for all designs. For the latter three random
designs, we vary the sample sizes $n\in\\{500,1000,2000\\}$. The noise
variance $\sigma^{2}$ is set so that ${\bm{\varepsilon}}$ and
$\mathbf{X}{\bm{\theta}}$ each explain 50% of the variance of $\mathbf{y}$ for
the Gaussian prior, and ${\bm{\varepsilon}}$ explains 20% of the variance of
$\mathbf{y}$ for the Cauchy, skew, and bimodal priors. To study the effects of
Monte Carlo variance of the MCEM methods on the estimated prior, we fix a
single generated instance $(\mathbf{X},\mathbf{y})$ for each prior and design,
and then evaluate 10 independent runs of each algorithm (except for CAVI which
is deterministic).

Experimental setup. We fix the prior support for all methods to be the grid
$b_{1},\ldots,b_{K}$ over $[-3,3]$. EBflow is run with $\tau^{2}$ set as in
(4.2), and the log-linear step size decay (4.4) over 10,000 iterations. We
tested also a preconditioned variant [RS02] of EBflow with the
${\bm{\varphi}}$ update in (4.3) replaced by

${\bm{\varphi}}_{t+1}={\bm{\varphi}}_{t}-\eta^{\varphi}_{t}\mathbf{Q}^{-1}\Bigg{[}\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}(\mathbf{X}{\bm{\varphi}}_{t}-\mathbf{y})+\frac{1}{\tau^{2}}\bigg{(}\frac{\sum_{k=1}^{K}w_{t,k}(\varphi_{t,j}-b_{k})\mathcal{N}_{\tau}(\varphi-
b_{k})}{\sum_{k=1}^{K}w_{t,k}\,\mathcal{N}_{\tau}(\varphi_{t,j}-b_{k})}\bigg{)}_{j=1}^{p}\Bigg{]}+\sqrt{2\eta^{\varphi}_{t}}\mathbf{Q}^{-1/2}\,\mathbf{g}_{t},$
(4.5)

fixing
$\mathbf{Q}=\mathbf{X}^{\top}{\bm{\Sigma}}^{-1}\mathbf{X}+\tau^{-2}\operatorname{Id}$
as a rough proxy for the Hessian (4.1). Related forms of preconditioning have
been suggested, for example, for posterior sampling in logistic regression
models in [Dal17] to accelerate mixing of Langevin dynamics with ill-
conditioned designs. One may check that this preconditioning corresponds to
the original gradient flow (1.7) with an alternative choice of inner-product
on $\mathbb{R}^{p}$ underlying the Wasserstein-2 geometry. Langevin-MCEM and
Gibbs-MCEM are run fixing $T_{\text{iter}}=100$ Monte Carlo iterations between
M-step updates, also over 10,000 total Monte Carlo iterations. CAVI is run for
1000 iterations, as it converges much faster than the MCEM methods.

We compare prior density estimates and computation time for all methods
applied to minimize the negative log-likelihood with a small smoothing spline
regularization (2.24), fixing $\lambda=0.003$ for the Gaussian prior and
$\lambda=0.001$ for the other three priors. We report also the achieved
marginal log-likelihood for identity design and mean-squared prediction error
for random designs, and verify that these are similar with and without
smoothness regularization.

Figure 2: Estimated prior densities across 6 methods, and 10 independent runs
for each method except CAVI. Priors left to right: Gaussian, Cauchy, skew,
bimodal. All settings have sample size $n=1000$ and dimension $p=1000$. Design
is block02corr0.9.

Results. Figure 2 displays the estimated priors in an illustrative example of
the block02corr0.9 design with $n=p=1000$, which highlights some of the main
trends across other priors and designs.

In general, the estimates of EBflow, Langevin-MCEM with large step size
$\eta_{t}^{\varphi}=1.0$, and Gibbs-MCEM are more accurate than the remaining
two methods. Preconditioning EBflow can yield more stable estimates, due to
faster mixing of the ULA chain. Langevin-MCEM at the smaller step size
$\eta_{t}^{\varphi}=0.1$ suffers from high Monte Carlo variance. CAVI produces
density estimates that tend to underestimate the prior variance, due to
inaccuracy of the variational approximation for this type of design that has
strong pairwise correlations between variables.

Design | $n/p$ | Method | TV | TV.sd | Time (iters) | Log-likelihood
---|---|---|---|---|---|---
identity | 1.0 | EBflow | 0.063 | 0.0016 | 500 | 0.511
EBflow.preconditioned | 0.063 | 0.0016 | 500 | 0.511
Langevin.eta1.0 | 0.068 | 0.0097 | 200 | 0.496
Langevin.eta0.1 | 0.081 | 0.0142 | 200 | 0.514
Gibbs | 0.084 | 0.0033 | 300 | 0.514
CAVI | 0.084 | NA | 2 | 0.514
Design | $n/p$ | Method | TV | TV.sd | Time (iters) | Prediction MSE
iid | 2.0 | EBflow | 0.052 | 0.0059 | 400 | 0.409
EBflow.preconditioned | 0.050 | 0.0038 | 400 | 0.407
Langevin.eta1.0 | 0.044 | 0.0063 | 200 | 0.405
Langevin.eta0.1 | 0.069 | 0.0149 | 200 | 0.408
Gibbs | 0.057 | 0.0019 | 300 | 0.408
CAVI | 0.076 | NA | 1 | 0.419
iid | 1.0 | EBflow | 0.044 | 0.0075 | 700 | 0.585
EBflow.preconditioned | 0.045 | 0.0052 | 600 | 0.583
Langevin.eta1.0 | 0.065 | 0.0085 | 400 | 0.580
Langevin.eta0.1 | 0.091 | 0.0191 | 700 | 0.586
Gibbs | 0.073 | 0.0038 | 500 | 0.584
CAVI | 0.078 | NA | 2 | 0.611
iid | 0.5 | EBflow | 0.046 | 0.0069 | 1000 | 0.789
EBflow.preconditioned | 0.052 | 0.0034 | 900 | 0.788
Langevin.eta1.0 | 0.075 | 0.0072 | 600 | 0.790
Langevin.eta0.1 | 0.066 | 0.0228 | 1100 | 0.787
Gibbs | 0.062 | 0.0065 | 800 | 0.786
CAVI | 0.118 | NA | 2 | 0.788
block02corr0.9 | 2.0 | EBflow | 0.036 | 0.0123 | 600 | 0.303
EBflow.preconditioned | 0.031 | 0.0093 | 600 | 0.302
Langevin.eta1.0 | 0.041 | 0.0140 | 200 | 0.301
Langevin.eta0.1 | 0.075 | 0.0228 | 200 | 0.303
Gibbs | 0.058 | 0.0041 | 400 | 0.303
CAVI | 0.135 | NA | 1 | 0.333
block02corr0.9 | 1.0 | EBflow | 0.064 | 0.0070 | 800 | 0.525
EBflow.preconditioned | 0.060 | 0.0062 | 700 | 0.524
Langevin.eta1.0 | 0.071 | 0.0067 | 400 | 0.528
Langevin.eta0.1 | 0.097 | 0.0219 | 1150 | 0.531
Gibbs | 0.083 | 0.0021 | 600 | 0.525
CAVI | 0.161 | NA | 1 | 0.545
block02corr0.9 | 0.5 | EBflow | 0.032 | 0.0120 | 1100 | 0.646
EBflow.preconditioned | 0.037 | 0.0104 | 950 | 0.646
Langevin.eta1.0 | 0.053 | 0.0099 | 650 | 0.647
Langevin.eta0.1 | 0.062 | 0.0171 | 1750 | 0.644
Gibbs | 0.038 | 0.0036 | 800 | 0.640
CAVI | 0.157 | NA | 2 | 0.670
block10corr0.5 | 2.0 | EBflow | 0.081 | 0.0141 | 600 | 0.304
EBflow.preconditioned | 0.079 | 0.0102 | 600 | 0.302
Langevin.eta1.0 | 0.074 | 0.0213 | 200 | 0.300
Langevin.eta0.1 | 0.109 | 0.0344 | 200 | 0.305
Gibbs | 0.088 | 0.0052 | 400 | 0.299
CAVI | 0.349 | NA | 1 | 0.449
block10corr0.5 | 1.0 | EBflow | 0.031 | 0.0072 | 900 | 0.478
EBflow.preconditioned | 0.026 | 0.0080 | 800 | 0.475
Langevin.eta1.0 | 0.049 | 0.0076 | 400 | 0.473
Langevin.eta0.1 | 0.077 | 0.0211 | 1650 | 0.477
Gibbs | 0.051 | 0.0054 | 600 | 0.471
CAVI | 0.283 | NA | 1 | 0.513
block10corr0.5 | 0.5 | EBflow | 0.076 | 0.0134 | 1500 | 0.650
EBflow.preconditioned | 0.075 | 0.0125 | 1300 | 0.645
Langevin.eta1.0 | 0.082 | 0.0123 | 800 | 0.642
Langevin.eta0.1 | 0.100 | 0.0377 | 2750 | 0.645
Gibbs | 0.068 | 0.0073 | 1000 | 0.636
CAVI | 0.156 | NA | 2 |

Table 1: Comparison of methods: Gaussian prior.

Table 2: Comparison of methods: Cauchy prior.

Table 3: Comparison of methods: Skew prior.

Table 4: Comparison of methods: Bimodal prior.

Tables 4.2–4.2 summarize estimation accuracy and runtime across all tested
methods, priors, and designs. We report the mean and standard deviation of the
TV error of the estimated density, across 10 independent runs for all methods
except CAVI. As rough proxies for convergence speed, we report the median
number of iterations (rounded to the nearest 100) used to reach a TV error
less than 0.2, or NA if the final TV error is larger than this value.888The
real compute times of Langevin-MCEM and Gibbs-MCEM are slower than EBflow in
our straightforward implementations of these methods. For Langevin-MCEM, this
is due to the complexity of solving the sequence-model NPMLE problem in each
M-step (which is not reflected in the iteration count). For Gibbs sampling, it
is due to a loop over coordinates as compared to a vectorized matrix-vector
computation in EBflow. In our experiments at this problem size, EBflow was
about 2–3 times faster than Gibbs sampling, and 5–10 times faster than
Langevin-MCEM, although these comparisons are necessarily implementation
dependent. We report also the attained marginal log-likelihood value
$\bar{F}_{n}$ (up to an additive constant) for identity design,999For general
design, it is difficult to evaluate the marginal log-likelihood which involves
a $p$-dimensional integral. and the relative mean-squared prediction error

$\text{MSE}=\frac{\|\mathbf{X}_{\text{new}}({\bm{\theta}}_{\ast}-\mathbb{E}[{\bm{\theta}}\mid\mathbf{X},\mathbf{y}])\|_{2}^{2}}{\|\mathbf{X}_{\text{new}}{\bm{\theta}}_{\ast}\|_{2}^{2}}$

across $n_{\text{new}}=1000$ test samples $\mathbf{X}_{\text{new}}$ for random
designs, where ${\bm{\theta}}_{\ast}$ is the true regression parameter. The
posterior mean $\mathbb{E}[{\bm{\theta}}\mid\mathbf{X},\mathbf{y}]$ is
estimated in EBflow and Langevin-MCEM from $T^{\prime}=50{,}000$ additional
iterates of ${\bm{\varphi}}_{t}$ using (2.23), $T^{\prime}=50{,}000$
additional iterates of ${\bm{\theta}}_{t}$ in Gibbs-MCEM, and the average of
the final posterior marginals estimated in CAVI.

Table 4.2 compares results with and without spline smoothing, for a single run
of each method at dimensions $n=p=1000$. The smoothing regularization (2.24)
is needed to obtain accurate prior density estimates, but we observe that
marginal log-likelihood and prediction MSE are nearly identical with and
without smoothing for all methods except CAVI (where the prediction MSE
without smoothing is sometimes worse).

Table 5: Comparison of results with (left) and without (right) spline
smoothing

We summarize here a few high-level trends and take-aways from these results,
many of which are also visually apparent in Figure 2:

  * •

Mean TV-accuracy of the prior density is similar for EBflow with and without
preconditioning. However, preconditioning can reduce Monte Carlo variance
across independent trials, especially for more correlated designs.

  * •

Mean TV-accuracy of the prior density for EBflow is competitive with the best
tunings of Langevin-MCEM and Gibbs-MCEM for the Gaussian, Cauchy, and skew
priors, and in many settings EBflow yields the lowest error.

  * •

For the bimodal prior and many of the tested designs, no method obtains an
accurate prior density estimate with TV error less than 0.2. In settings where
an accurate estimate is obtained, Gibbs-MCEM produces the best estimates,
likely because it more rapidly mixes between the modes of the posterior
distribution in comparison to Langevin dynamics, even when the latter is
carried out in the smoothed parameter ${\bm{\varphi}}$.

  * •

In terms of iteration count, EBflow has competitive runtime with Langevin-MCEM
and Gibbs-MCEM.

  * •

CAVI is the fastest method, and is accurate under the identity design
$\mathbf{X}=\operatorname{Id}$ (where the variational approximation is exact)
as well as several settings of i.i.d. Gaussian design. However, its accuracy
in terms of both the prior density estimate and prediction MSE can
substantially degrade under correlated designs.

  * •

All methods achieve similar marginal log-likelihood under identity design, and
the MCEM methods achieve similar prediction MSE under random designs. We do
note that Langevin-MCEM and EBflow with final step size
$\eta_{t}^{\varphi}=0.1$ have higher per-variable autocorrelation than with
step size $\eta_{t}^{\varphi}=1.0$ and than Gibbs sampling, resulting in the
need to use more posterior samples to obtain an accurate estimate of
$\mathbb{E}[{\bm{\theta}}\mid\mathbf{X},\mathbf{y}]$.

##  5 Discussion

We have studied empirical Bayes estimation of an unknown prior distribution
for the regression coefficients of a linear model, introducing a system of
bivariate gradient flow equations for optimizing the marginal log-likelihood
via its Gibbs variational representation. This gradient flow system may be
simulated using a Langevin diffusion, where the prior is continuously updated
based on an estimate of the average marginal distribution of coordinates of
the Langevin sample. Although the regression model and Langevin sample are
both high-dimensional, the one-dimensional nature of the estimand $g_{\ast}$
enables its estimation using a fully nonparametric approach. Similar ideas can
be explored in models beyond linear regression, where a high-dimensional
latent parameter is comprised of independent draws from an unknown univariate
or low-dimensional prior law.

We believe there are many interesting directions for future work, and we
discuss here two of these directions:

  1. 1.

(Global landscape analysis) For general regression designs, we believe the
optimization landscape of the marginal log-likelihood $\bar{F}_{n}(g)$ may
have multiple local minimizers, and we find empirically that the prior density
estimated by EBflow is sometimes sensitive to the choice of initialization. In
such settings, we have only provided a theoretical guarantee of convergence to
a prior $g$ that satisfies an approximate first-order optimality condition for
$\bar{F}_{n}(g)$.

It is an open question to obtain a better understanding of the global
landscape geometry of $\bar{F}_{n}(g)$, and to characterize possible
conditions under which the local minimizer of $\bar{F}_{n}(g)$ may be unique.
One interesting setting is that of random “mean-field” designs, in which
$\bar{F}_{n}(g)$ may asymptotically converge to a deterministic limit
functional $\bar{F}(g)$ in the sense of Proposition 3.8. Although simpler
algorithms may be available for empirical Bayes estimation in mean-field
contexts, we believe that nonetheless this is a natural setting to study
questions about the optimization landscape, as they may be studied via the
limit functional $\bar{F}(g)$.

  2. 2.

(Time discretization and propagation-of-chaos) In this work we have studied
the dynamics and convergence of the idealized continuous-time gradient flow,
but have not yet taken the steps to translate this to convergence guarantees
for the time-discretized EBflow algorithm. We believe that a careful analysis
of the EBflow algorithm may lead to improved understanding of issues related
to parameter tuning, such as the choice of variance $\tau^{2}$ in our
reparametrization by the smoothed regression variable ${\bm{\varphi}}$.

One particular question of interest is to understand conditions under which
the true law $\bar{q}_{t}$ may be consistently estimated by the empirical
average of coordinates (2.18) over a finite time horizon $[0,T]$, using only
one realization $\\{{\bm{\varphi}}_{t}\\}_{t\in[0,T]}$ of the Langevin
diffusion rather than an asymptotically growing number of independent parallel
chains. We note that this type of “single-chain propagation-of-chaos” result
is only possible in settings of growing dimensions $p\to\infty$, and we
believe that its study may lead to improved understanding of functional
inequalities and concentration of measure for the law
$\\{q_{t}\\}_{t\in[0,T]}$ as it evolves along the gradient flow.

Our numerical simulations suggest that EBflow is competitive in many settings
with alternative approaches such as MCEM with coordinate-wise Gibbs sampling,
in terms of both estimation accuracy and runtime. More importantly, we hope
that this EBflow approach can make available a broader toolbox of ideas
related to Langevin-based algorithms for empirical Bayes estimation, including
ideas around adaptive preconditioning [GC11] and underdamped/Hamiltonian
dynamics [Nea11, CCBJ18] to handle more ill-conditioned designs, and
stochastic gradient approximations [WT11, MCF15] to scale to larger-$(n,p)$
problems. In extensions to more complex models where latent parameters are
organized in multi-layer archictectures, such gradient-based ideas rather than
coordinate-based computations may become even more advantageous
computationally, due to the lower complexity of computing gradients using a
back-propagation procedure. We believe that such extensions to more complex
models and sampling algorithms would be interesting to explore also in future
work.

##  Appendix A Proof of consistency

In this appendix, we prove Theorem 3.1 and Proposition 3.2.

###  A.1 TV lower bound

For any $\eta>0$ and $g\in\mathcal{P}(M)$, we write as shorthand
$g_{\eta}=\mathcal{N}_{\eta}*g=\mathcal{N}(0,\eta^{2})*g$. Let

$\displaystyle
d_{\eta}(g,g^{\prime})=W_{1}\big{(}g_{\eta},g_{\eta}^{\prime}\big{)}$ (A.1)

be the smoothed Wasserstein-1 distance on $\mathcal{P}(M)$.

All constants in the lemmas and proofs may depend on the noise variance
$\sigma^{2}>0$, on $M$ defining the common support of $\mathcal{P}$, on the
bound $\gamma$ for $n/p$, and on $C_{\ast},c_{\ast},\alpha,\beta$ in
Assumption A.

######  Lemma A.1.

There exists a sufficiently large constant $v>0$, constants $c,c^{\prime}>0$,
and a sequence $\delta_{n}\to 0$ such that for all large $n,p$ and all
$g,g^{\prime}\in\mathcal{P}(M)$,

$\displaystyle
d_{\mathrm{TV}}(P_{g}(\mathbf{y}),P_{g^{\prime}}(\mathbf{y}))\geq
1-2e^{-cp(d_{v}(g,g^{\prime})-\delta_{n})_{+}^{2}}.$

###### Proof.

Fixing any constant $\eta>0$, the smoothed distribution $g_{\eta}$ satisfies
an LSI with uniformly bounded LSI constant for all $g\in\mathcal{P}(M)$
[BGMZ18, Theorem 1.1]. Note that convolving $g$ with $\mathcal{N}_{\eta}$
corresponds to adding an independent
$\mathcal{N}(0,\eta^{2}\,\mathbf{X}\mathbf{X}^{\top})$ noise vector to
$\mathbf{y}$. Hence by the data processing inequality,

$\displaystyle
d_{\mathrm{TV}}\big{(}P_{g}(\mathbf{y}),P_{g^{\prime}}(\mathbf{y})\big{)}\geq
d_{\mathrm{TV}}\big{(}P_{g_{\eta}}(\mathbf{y}),P_{g_{\eta}^{\prime}}(\mathbf{y})\big{)},$

and it suffices to lower bound the TV distance on the right side for priors
$g_{\eta}$ and $g_{\eta}^{\prime}$.

Noting that $g_{\eta}$ has variance at most $M^{2}+\eta^{2}$ for every
$g\in\mathcal{P}(M)$, we pick a constant $v>0$ satisfying

$v^{2}\geq\max_{j\in\mathcal{S}}\left(\eta^{2}+v_{j}\cdot(M^{2}+\eta^{2})+w_{j}\right),\qquad\text{
where }v_{j}=\sum_{k:k\neq j}(\mathbf{z}_{j}^{\top}\mathbf{x}_{k})^{2},\qquad
w_{j}=\sigma^{2}\|\mathbf{z}_{j}\|_{2}^{2},$ (A.2)

and $\\{\mathbf{z}_{j}\\}_{j\in\mathcal{S}}$ are given by Assumption A. We
will check below that $v_{j},w_{j}$ are bounded above by a constant uniformly
over $j\in\mathcal{S}$, so that we may indeed pick $v$ here as a constant
independent of $n,p$. Denote
$\Delta=d_{v}(g,g^{\prime})=W_{1}(g_{v},g_{v}^{\prime})$. To establish the
desired TV lower bound, it suffices to exhibit an explicit test statistic
$T(\mathbf{y})$ with rejection threshold
$t=(\mathbb{E}_{g_{\eta}}[T(\mathbf{y})]+\mathbb{E}_{g_{\eta}^{\prime}}[T(\mathbf{y})])/2$
such that

$\displaystyle\max(\mathbb{P}_{g_{\eta}}[T(\mathbf{y})<t],\mathbb{P}_{g_{\eta}^{\prime}}[T(\mathbf{y})\geq
t])\leq e^{-cp(\Delta-\delta_{n})_{+}^{2}}.$ (A.3)

Indeed, this immediately implies the conclusion of the lemma,

$d_{\mathrm{TV}}(P_{g_{\eta}}(\mathbf{y}),P_{g_{\eta}^{\prime}}(\mathbf{y}))\geq\Big{|}\mathbb{P}_{g_{\eta}}[T(\mathbf{y})\geq
t]-\mathbb{P}_{g_{\eta}^{\prime}}[T(\mathbf{y})\geq t]\Big{|}\geq
1-2e^{-cp(\Delta-\delta_{n})_{+}^{2}}.$

Let $\kappa_{1},\kappa_{2}$ and $\kappa_{1}^{\prime},\kappa_{2}^{\prime}$
denote the means and variances of $g_{\eta},g_{\eta}^{\prime}$ respectively.
To construct $T(\mathbf{y})$ satisfying (A.3), we consider three cases. The
constants $C,c>0$ below will change from instance to instance.

Case 1: $|\kappa_{1}-\kappa_{1}^{\prime}|>c_{1}\Delta$ for a constant
$c_{1}>0$ to be specified. Suppose without loss of generality
$\kappa_{1}>\kappa_{1}^{\prime}$, and consider the test statistic

$\displaystyle
T(\mathbf{y})=\mathbf{1}^{\top}\mathbf{X}^{\top}\mathbf{y}=\mathbf{1}^{\top}\mathbf{X}^{\top}\mathbf{X}{\bm{\theta}}+\mathbf{1}^{\top}\mathbf{X}^{\top}{\bm{\varepsilon}}.$

Under Assumption A.2, we have

$\mathbb{E}_{g_{\eta}}[T(\mathbf{y})]-\mathbb{E}_{g_{\eta}^{\prime}}[T(\mathbf{y})]=\mathbf{1}^{\top}\mathbf{X}^{\top}\mathbf{X}\mathbf{1}\cdot(\kappa_{1}-\kappa_{1}^{\prime})\geq
c\Delta p.$ (A.4)

Computing the gradient of $T(\mathbf{y})$ in
$({\bm{\theta}},{\bm{\varepsilon}})$, we observe that $T(\mathbf{y})$ is
$L$-Lipschitz in $({\bm{\theta}},{\bm{\varepsilon}})$ with respect to the
$\ell_{2}$-norm, where
$L^{2}=\mathbf{1}^{\top}(\mathbf{X}^{\top}\mathbf{X})^{2}\mathbf{1}+\mathbf{1}^{\top}\mathbf{X}^{\top}\mathbf{X}\mathbf{1}\leq
Cp$ under Assumption A.1–2. Then by the LSI for
$({\bm{\theta}},{\bm{\varepsilon}})$ under both $g_{\eta}$ and
$g_{\eta}^{\prime}$, for any $s>0$, $T(\mathbf{y})$ has the sub-Gaussian
concentration arounds its mean

$\displaystyle\max\Big{(}\mathbb{P}_{g_{\eta}}[T(\mathbf{y})-\mathbb{E}_{g_{\eta}}T(\mathbf{y})\leq-s\Delta
p],\;\mathbb{P}_{g_{\eta}^{\prime}}[T(\mathbf{y})-\mathbb{E}_{g_{\eta}^{\prime}}T(\mathbf{y})\geq
s\Delta p]\Big{)}\leq e^{-cs^{2}\Delta^{2}p}.$

Choosing a constant $s>0$ small enough and combining with (A.4), this shows
that (A.3) holds with $\delta_{n}=0$.

Case 2: $|\kappa_{2}-\kappa_{2}^{\prime}|\geq c_{2}\max(\Delta,\Delta^{2})$
for a constant $c_{2}>0$ to be specified. Suppose without loss of generality
$\kappa_{2}>\kappa_{2}^{\prime}$, let $\Pi$ denote the projection orthogonal
to $\mathbf{X}\mathbf{1}$, and consider the test statistic

$\displaystyle
T(\mathbf{y})=\|\mathbf{X}^{\top}\Pi\mathbf{y}\|^{2}=\|\mathbf{X}^{\top}\Pi(\mathbf{X}{\bm{\theta}}+{\bm{\varepsilon}})\|^{2}=\begin{pmatrix}{\bm{\theta}}\\\
{\bm{\varepsilon}}\end{pmatrix}^{\top}\underbrace{\begin{pmatrix}\mathbf{X}^{\top}\Pi\mathbf{X}\mathbf{X}^{\top}\Pi\mathbf{X}&\mathbf{X}^{\top}\Pi\mathbf{X}\mathbf{X}^{\top}\Pi\\\
\Pi\mathbf{X}\mathbf{X}^{\top}\Pi\mathbf{X}&\Pi\mathbf{X}\mathbf{X}^{\top}\Pi\end{pmatrix}}_{:=M(\mathbf{X},\Pi)}\begin{pmatrix}{\bm{\theta}}\\\
{\bm{\varepsilon}}\end{pmatrix}$

Under Assumptions A.1–2, the matrix $\mathbf{X}^{\top}\Pi\mathbf{X}$ has at
least $cp$ eigenvalues greater than $c$, for some sufficiently small constant
$c>0$. Then, irrespective of the means $\kappa_{1},\kappa_{1}^{\prime}$, we
have

$\displaystyle\mathbb{E}_{G_{\eta}}[T(\mathbf{y})]-\mathbb{E}_{G_{\eta}^{\prime}}[T(\mathbf{y})]=\operatorname{Tr}\mathbf{X}^{\top}\Pi\mathbf{X}\mathbf{X}^{\top}\Pi\mathbf{X}\cdot(\kappa_{2}-\kappa_{2}^{\prime})\geq
c\max(\Delta,\Delta^{2})p.$

We have also $\|M(\mathbf{X},\Pi)\|_{\operatorname{op}}\leq C$ and
$\|M(\mathbf{X},\Pi)\|_{\mathrm{F}}^{2}\leq Cp$ for the above matrix
$M(\mathbf{X},\Pi)$, the latter holding because each block of
$M(\mathbf{X},\Pi)$ has rank at most $p$. Then by the Hanson-Wright inequality
[RV13, Theorem 1.1], viewing $T(\mathbf{y})$ as a quadratic form in
$({\bm{\theta}},{\bm{\varepsilon}})$ where ${\bm{\theta}}$ is sub-Gaussian
under both $g_{\eta}$ and $g_{\eta}^{\prime}$, for any $s>0$,

$\displaystyle\mathbb{P}_{g_{\eta}}[T(\mathbf{y})-\mathbb{E}_{g_{\eta}}T(\mathbf{y})$
$\displaystyle\leq-s\max(\Delta,\Delta^{2})p]\leq
e^{-c\min(s,s^{2})\Delta^{2}p}$
$\displaystyle\mathbb{P}_{g_{\eta}^{\prime}}[T(\mathbf{y})-\mathbb{E}_{g_{\eta}^{\prime}}T(\mathbf{y})$
$\displaystyle\geq s\max(\Delta,\Delta^{2})p]\leq
e^{-c\min(s,s^{2})\Delta^{2}p}.$

Then (A.3) holds also in this case with $\delta_{n}=0$.

Case 3: $|\kappa_{1}-\kappa_{1}^{\prime}|<c_{1}\Delta$ and
$|\kappa_{2}-\kappa_{2}^{\prime}|<c_{2}\max(\Delta,\Delta^{2})$. Let
$\mathcal{S}\subseteq\\{1,\ldots,p\\}$ and
$\\{\mathbf{z}_{j}:j\in\mathcal{S}\\}$ be as in Assumption A. For each
$j\in\mathcal{S}$, consider

$\displaystyle
s_{j}=\mathbf{z}_{j}^{\top}\mathbf{y}=\underbrace{\mathbf{z}_{j}^{\top}\mathbf{x}_{j}\cdot\theta_{j}}_{\equiv
a_{j}}+\underbrace{\sum_{k:k\neq
j}\mathbf{z}_{j}^{\top}\mathbf{x}_{k}\cdot\theta_{k}}_{\equiv
b_{j}}+\underbrace{\mathbf{z}_{j}^{\top}{\bm{\varepsilon}}}_{\equiv e_{j}},$

where these three terms $a_{j},b_{j},e_{j}$ are independent. Let
$\mathcal{L}(f({\bm{\theta}}))$ denote the law of $f({\bm{\theta}})$ under
$g_{\eta}$. Then by Assumption A.3,

$\displaystyle
W_{1}(\mathcal{L}(a_{j}),g_{\eta})=W_{1}(\mathcal{L}(a_{j}),\mathcal{L}(\theta_{j}))\leq|\mathbf{z}_{j}^{\top}\mathbf{x}_{j}-1|\cdot\mathbb{E}_{g_{\eta}}[|\theta_{j}|]\to
0.$

For $b_{j}$, we have

$\displaystyle\mathbb{E}_{g_{\eta}}[b_{j}]=m_{j}\kappa_{1},\qquad\operatorname{Var}_{g_{\eta}}[b_{j}]=v_{j}\kappa_{2},\qquad\text{where
}m_{j}=\sum_{k:k\neq j}\mathbf{z}_{j}^{\top}\mathbf{x}_{k},\qquad
v_{j}=\sum_{k:k\neq j}(\mathbf{z}_{j}^{\top}\mathbf{x}_{k})^{2}.$

The Wasserstein-$(2+\beta)$ CLT of [Rio09, Eq. (4.1)] shows that for some
$c=c(\beta)>0$,

$\displaystyle
W_{2+\beta}\Big{(}\mathcal{L}(b_{j}),\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2})\Big{)}^{2+\beta}\leq
c\sum_{k:k\neq
j}|\mathbf{z}_{j}^{\top}\mathbf{x}_{k}|^{2+\beta}\cdot\mathbb{E}_{g_{\eta}}[|\theta_{k}|^{2+\beta}].$

Assumption A.3 provides the Lyapunov condition ensuring that the right side
converges to 0 as $n,p\to\infty$, uniformly over $j\in\mathcal{S}$. Then also

$\displaystyle
W_{1}\Big{(}\mathcal{L}(b_{j}),\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2})\Big{)}\leq
W_{2+\beta}\Big{(}\mathcal{L}(b_{j}),\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2})\Big{)}\to
0.$

Finally, $e_{j}\sim\mathcal{N}(0,w_{j})$ with variance
$w_{j}=\sigma^{2}\|\mathbf{z}_{j}\|_{2}^{2}$. Putting this together,

$W_{1}\Big{(}\mathcal{L}(s_{j}),\,g_{\eta}*\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2}+w_{j})\Big{)}\leq\delta_{n}/4$
(A.5)

for some $\delta_{n}\to 0$. Similarly, writing $\mathcal{L}^{\prime}(\cdot)$
for the law under $g_{\eta}^{\prime}$,

$W_{1}\Big{(}\mathcal{L}^{\prime}(s_{j}),\,g_{\eta}^{\prime}*\mathcal{N}(m_{j}\kappa_{1}^{\prime},v_{j}\kappa_{2}^{\prime}+w_{j})\Big{)}\leq\delta_{n}/4.$
(A.6)

We use $|\kappa_{1}-\kappa_{1}^{\prime}|<c_{1}\Delta$ and
$|\kappa_{2}-\kappa_{2}^{\prime}|<c_{2}\max(\Delta,\Delta^{2})$ for some
sufficiently small $c_{1},c_{2}>0$ to compare these normal distributions. Note
that for any two univariate normal laws, by the coupling $(\mu+\sigma
Z,\mu^{\prime}+\sigma^{\prime}Z)$ for $Z\sim\mathcal{N}(0,1)$,

$\displaystyle
W_{1}(\mathcal{N}(\mu,\sigma^{2}),\mathcal{N}(\mu^{\prime},{\sigma^{\prime}}^{2}))$
$\displaystyle\leq|\mu-\mu^{\prime}|+|\sigma-\sigma^{\prime}|\cdot\mathbb{E}[|Z|]$
$\displaystyle\leq|\mu-\mu^{\prime}|+C\min\Big{(}|\sigma^{2}-{\sigma^{\prime}}^{2}|,|\sigma^{2}-{\sigma^{\prime}}^{2}|^{1/2}\Big{)},$

the second inequality holding when $\sigma^{2},{\sigma^{\prime}}^{2}$ are
lower bounded by a constant. By the definitions of $m_{j},v_{j}$ and under
Assumption A.3, we have
$|m_{j}|=|\mathbf{z}_{j}^{\top}\mathbf{X}\mathbf{1}-\mathbf{z}_{j}^{\top}\mathbf{x}_{j}|\leq
C$,
$v_{j}\leq\mathbf{z}_{j}^{\top}\mathbf{X}\mathbf{X}^{\top}\mathbf{z}_{j}\leq
C$, and $w_{j}=\sigma^{2}\|\mathbf{z}_{j}\|_{2}^{2}\in[c,C]$ for some
constants $C,c>0$. Then for a sufficiently small choice of constants
$c_{1},c_{2}>0$, we have

$\displaystyle
W_{1}(\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2}+w_{j}),\,\mathcal{N}(m_{j}\kappa_{1}^{\prime},v_{j}\kappa_{2}^{\prime}+w_{j}))<\Delta/2.$

Combining with (A.6), this implies

$W_{1}\Big{(}\mathcal{L}^{\prime}(s_{j}),\,g_{\eta}^{\prime}*\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2}+w_{j})\Big{)}\leq\delta_{n}/4+\Delta/2.$
(A.7)

Then by (A.5), (A.7), and the triangle inequality and translation invariance
for $W_{1}$,

$\displaystyle W_{1}(\mathcal{L}(s_{j}),\mathcal{L}^{\prime}(s_{j}))$
$\displaystyle\geq
W_{1}\Big{(}g_{\eta}*\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2}+w_{j}),\,g_{\eta}^{\prime}*\mathcal{N}(m_{j}\kappa_{1},v_{j}\kappa_{2}+w_{j})\Big{)}-\delta_{n}/2-\Delta/2$
$\displaystyle=W_{1}\Big{(}g_{\sqrt{\eta^{2}+v_{j}\kappa_{2}+w_{j}}},\,g_{\sqrt{\eta^{2}+v_{j}\kappa_{2}+w_{j}}}^{\prime}\Big{)}-\delta_{n}/2-\Delta/2.$

Recalling the choice of $v$ in (A.2) and applying $\kappa_{2}\leq
M^{2}+\eta^{2}$, we have $\eta^{2}+v_{j}\kappa_{2}+w_{j}\leq v^{2}$. Then by
monotonicity of the smoothed Wasserstein distance in the smoothing parameter
[GG20, Theorem 3],
$W_{1}(g_{\sqrt{\eta^{2}+v_{j}\kappa_{2}+w_{j}}},\,g_{\sqrt{\eta^{2}+v_{j}\kappa_{2}+w_{j}}}^{\prime})\geq
W_{1}(g_{v},g_{v}^{\prime})=\Delta$, so

$\displaystyle
W_{1}(\mathcal{L}(s_{j}),\mathcal{L}^{\prime}(s_{j}))\geq\Delta/2-\delta_{n}/2.$

Finally, this implies by the Kantorovich duality for $W_{1}$ that there exists
for each $j\in\mathcal{S}$ a smooth 1-Lipschitz test function
$f_{j}:\mathbb{R}\to\mathbb{R}$ such that
$\mathbb{E}_{g_{\eta}}[f_{j}(s_{j})]-\mathbb{E}_{g_{\eta}^{\prime}}[f_{j}(s_{j})]\geq(\Delta-\delta_{n})_{+}/3$.
We consider the test statistic

$\displaystyle
T(\mathbf{y})=\sum_{j\in\mathcal{S}}f_{j}(s_{j})=\sum_{j\in\mathcal{S}}f_{j}(\mathbf{z}_{j}^{\top}(\mathbf{X}{\bm{\theta}}+{\bm{\varepsilon}})),$

which then satisfies
$\mathbb{E}_{g_{\eta}}[T(\mathbf{y})]-\mathbb{E}_{g_{\eta}^{\prime}}[T(\mathbf{y})]\geq
cp(\Delta-\delta_{n})_{+}$ because $|\mathcal{S}|\geq\alpha p$. Computing the
gradient in $({\bm{\theta}},{\bm{\varepsilon}})$, we see that $T(\mathbf{y})$
is $L$-Lipschitz in $({\bm{\theta}},{\bm{\varepsilon}})$ for

$\displaystyle
L^{2}=\bigg{\lVert}\mathbf{X}^{\top}\sum_{j\in\mathcal{S}}\mathbf{z}_{j}\cdot
f_{j}^{\prime}(s_{j})\bigg{\rVert}^{2}+\bigg{\lVert}\sum_{j\in\mathcal{S}}\mathbf{z}_{j}\cdot
f_{j}^{\prime}(s_{j})\bigg{\rVert}^{2}.$

Applying $|f_{j}^{\prime}(s_{j})|\leq 1$ for all $j\in\mathcal{S}$,
$\|\mathbf{X}\|_{{\operatorname{op}}}\leq C_{*}$, and the final condition of
Assumption A.4, this gives $L^{2}\leq Cp$. Then (A.3) holds by the LSI for
$({\bm{\theta}},{\bm{\varepsilon}})$, similarly to Case 1. ∎

A direct consequence of Lemma A.1 is the following pointwise comparison of the
likelihood objective.

######  Corollary A.2.

Let $v>0$ and the sequence $\delta_{n}\downarrow 0$ be as in Lemma A.1. Fix
any constant $\delta>0$. Then there exists some constants $c=c(\delta)$ and
$\tau=\tau(\delta)$ such that for all large $n,p$ and any $g\in\mathcal{P}(M)$
with $d_{v}(g,g_{*})\geq\delta$,

$\displaystyle\mathbb{P}_{g_{*}}\left[\frac{1}{p}\Big{(}\log
P_{g}(\mathbf{y})-\log P_{g_{*}}(\mathbf{y})\Big{)}\geq-\tau\right]\leq
e^{-cp}.$

###### Proof.

Writing as shorthand $P_{g}=P_{g}(\mathbf{y})$, Lemma A.1 shows that for some
constant $c>0$, all large $n,p$, and any $g\in\mathcal{P}(M)$ with
$d_{v}(g,g_{*})\geq\delta$, we have

$\displaystyle e^{-cp\delta^{2}}$ $\displaystyle\geq
1-d_{\mathrm{TV}}(P_{g_{*}},P_{g})=\int\min(P_{g_{*}},P_{g})\mathrm{d}\mathbf{y}$
$\displaystyle\geq\int e^{-p\tau}P_{g_{*}}\mathbf{1}\\{P_{g}\geq
e^{-p\tau}P_{g_{*}}\\}\mathrm{d}\mathbf{y}=e^{-p\tau}\mathbb{P}_{g_{*}}\left[p^{-1}\left(\log
P_{g}-\log P_{g_{*}}\right)\geq{-}\tau\right]$

for any $\tau>0$. The corollary follows upon choosing $\tau=c\delta^{2}/2$ and
adjusting constants. ∎

###  A.2 Covering net

We now define the following pseudo-metric $\widetilde{d}_{\eta,B}$ over
$\mathcal{P}(M)$: Let $G_{\eta}(\cdot),G_{\eta}^{\prime}(\cdot)$ denote the
CDFs of $g_{\eta}=\mathcal{N}_{\eta}*g$ and
$g_{\eta}^{\prime}=\mathcal{N}_{\eta}*g^{\prime}$. (Here $G_{\eta}^{\prime}$
is not to be confused with a derivative of $G_{\eta}$.) We set

$\displaystyle\widetilde{d}_{\eta,B}(g,g^{\prime})=\max_{x:|x|\leq
B}|G_{\eta}(x)-G_{\eta}^{\prime}(x)|.$

######  Lemma A.3.

Fix any constants $C_{0},\iota>0$. Let $\Pi_{X}$ be the projection onto the
column span of $\mathbf{X}$. Then there exist constants $\eta,B,L>0$ such
that, on the event where $\|\Pi_{X}\mathbf{y}\|_{2}^{2}\leq C_{0}p$, for any
$g,g^{\prime}\in\mathcal{P}(M)$ we have

$\displaystyle\frac{1}{p}\Big{|}\log P_{g}(\mathbf{y})-\log
P_{g^{\prime}}(\mathbf{y})\Big{|}\leq
L\,\widetilde{d}_{\eta,B}(g,g^{\prime})+\iota.$

###### Proof.

Let $[\mathbf{V}_{X}\mid\mathbf{W}_{X}]\in\mathbb{R}^{n\times n}$ be an
orthogonal matrix where $\Pi_{X}=\mathbf{V}_{X}\mathbf{V}_{X}^{\top}$ and the
columns of $\mathbf{V}_{X}$ form an orthonormal basis for the column span of
$\mathbf{X}$. Write $\widetilde{\mathbf{y}}=\mathbf{V}_{X}^{\top}\mathbf{y}$
and $\widetilde{\mathbf{X}}=\mathbf{V}_{X}^{\top}\mathbf{X}$. Then
$\widetilde{\mathbf{y}}=\widetilde{\mathbf{X}}{\bm{\theta}}+\mathbf{V}_{X}^{\top}{\bm{\varepsilon}}$
is independent of
$\mathbf{W}_{X}^{\top}\mathbf{y}=\mathbf{W}_{X}^{\top}{\bm{\varepsilon}}$, and
the law of the latter does not depend on the prior $g$. Thus

$\log P_{g}(\mathbf{y})=\log P_{g}(\widetilde{\mathbf{y}})+\text{constant}$

for a constant independent of $g$, where $P_{g}(\widetilde{\mathbf{y}})$
denotes the marginal density of $\widetilde{\mathbf{y}}$ under $g$. Observe
that under Assumption A.1,
$\|\widetilde{\mathbf{X}}\widetilde{\mathbf{X}}^{\top}\|_{{\operatorname{op}}}=\|\mathbf{X}\mathbf{X}^{\top}\|_{{\operatorname{op}}}<C$
for a constant $C>0$. Then we may choose a constant $\eta>0$ small enough such
that
$\eta\cdot\|\widetilde{\mathbf{X}}\widetilde{\mathbf{X}}^{\top}\|_{\operatorname{op}}<\sigma^{2}/2$,
and reparametrize the law of $\widetilde{\mathbf{y}}$ equivalently as

$\displaystyle\widetilde{\mathbf{y}}=\widetilde{\mathbf{X}}{\bm{\varphi}}+\widetilde{\bm{\varepsilon}},\qquad\varphi_{j}\overset{iid}{\sim}g_{\eta},\qquad\widetilde{\bm{\varepsilon}}\sim\mathcal{N}(0,{\bm{\Sigma}}),$

where
${\bm{\Sigma}}=\sigma^{2}\operatorname{Id}-\eta\,\widetilde{\mathbf{X}}\widetilde{\mathbf{X}}^{\top}$
satisfies
$(\sigma^{2}/2)\operatorname{Id}\preceq{\bm{\Sigma}}\preceq\sigma^{2}\operatorname{Id}$.
The marginal log-likelihood $\log P_{g}(\mathbf{y})$ then has the equivalent
form

$\log
P_{g}(\mathbf{y})=\log\int\exp\left(-\frac{1}{2}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})\right)\prod_{j=1}^{p}g_{\eta}(\varphi_{j})\mathrm{d}{\bm{\varphi}}+\text{constant}$
(A.8)

for a constant independent of $g$.

We now perform a change-of-measure from $G_{\eta}$ to $G_{\eta}^{\prime}$ by
introducing $F_{\eta}(x)={G_{\eta}^{\prime}}^{-1}(G_{\eta}(x))$ where
${G_{\eta}^{\prime}}^{-1}(\cdot)$ is the functional inverse of the CDF
$G_{\eta}^{\prime}(\cdot)$. (This is the univariate optimal transport map from
$G_{\eta}$ to $G_{\eta}^{\prime}$.) Let us write $F_{\eta}({\bm{\varphi}})$
for the coordinate-wise application of $F_{\eta}$ to ${\bm{\varphi}}$. Then
the marginal log-likelihood $\log P_{g^{\prime}}(\mathbf{y})$ may also be
written as

$\log
P_{g^{\prime}}(\mathbf{y})=\log\int\exp\left(-\frac{1}{2}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}}))^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}}))\right)\prod_{j=1}^{p}g_{\eta}(\varphi_{j})\mathrm{d}{\bm{\varphi}}+\text{constant}.$
(A.9)

Taking the difference of (A.8) and (A.9),

$\displaystyle\log P_{g^{\prime}}(\mathbf{y})-\log
P_{g}(\mathbf{y})=\log\left\langle
e^{-\frac{1}{2}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}}))^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}}))+\frac{1}{2}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})}\right\rangle,$

where we introduce the shorthand

$\langle f({\bm{\varphi}})\rangle=\frac{\int
f({\bm{\varphi}})\,e^{-H({\bm{\varphi}})}\mathrm{d}G_{\eta}({\bm{\varphi}})}{\int
e^{-H({\bm{\varphi}})}\mathrm{d}G_{\eta}({\bm{\varphi}})},\quad
H({\bm{\varphi}})=\frac{1}{2}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}}),\quad\mathrm{d}G_{\eta}({\bm{\varphi}})=\prod_{j=1}^{p}g_{\eta}(\varphi_{j})\mathrm{d}{\bm{\varphi}}$

for the expectation with respect to the posterior law of ${\bm{\varphi}}$
under $g_{\eta}$ given $(\widetilde{\mathbf{X}},\widetilde{\mathbf{y}})$.

Let $C,c>0$ denote constants changing from instance to instance, which may
depend on $\eta$. By Jensen’s inequality,

$\displaystyle\log P_{g^{\prime}}(\mathbf{y})-\log P_{g}(\mathbf{y})$
$\displaystyle\geq\frac{1}{2}\left\langle-(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}}))^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}}))+(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})^{\top}{\bm{\Sigma}}^{-1}(\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}})\right\rangle.$

Setting
$\mathbf{u}=\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}{\bm{\varphi}}$ and
$\mathbf{v}=\widetilde{\mathbf{y}}-\widetilde{\mathbf{X}}F_{\eta}({\bm{\varphi}})$,
we apply

$\displaystyle\langle{-}\mathbf{v}^{\top}{\bm{\Sigma}}^{-1}\mathbf{v}+\mathbf{u}^{\top}{\bm{\Sigma}}^{-1}\mathbf{u}\rangle$
$\displaystyle=\big{\langle}{-}(\mathbf{u}-\mathbf{v})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{u}-\mathbf{v})+2(\mathbf{u}-\mathbf{v})^{\top}{\bm{\Sigma}}^{-1}\mathbf{u}\big{\rangle}$
$\displaystyle\geq-\langle(\mathbf{u}-\mathbf{v})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{u}-\mathbf{v})\rangle-2\langle(\mathbf{u}-\mathbf{v})^{\top}{\bm{\Sigma}}^{-1}(\mathbf{u}-\mathbf{v})\rangle^{1/2}\langle\mathbf{u}^{\top}{\bm{\Sigma}}^{-1}\mathbf{u}\rangle^{1/2}.$

Together with the conditions
${\bm{\Sigma}}^{-1}\preceq(2/\sigma^{2})\operatorname{Id}$,
$\|\widetilde{\mathbf{X}}\|_{\operatorname{op}}\leq C$ under Assumption A.1,
and Cauchy-Schwarz, this gives

$\log P_{g^{\prime}}(\mathbf{y})-\log
P_{g}(\mathbf{y})\geq-C\left(\big{\langle}\|F_{\eta}({\bm{\varphi}})-{\bm{\varphi}}\|_{2}^{2}\big{\rangle}+\big{\langle}\|F_{\eta}({\bm{\varphi}})-{\bm{\varphi}}\|_{2}^{2}\big{\rangle}^{1/2}\big{\langle}H({\bm{\varphi}})\big{\rangle}^{1/2}\right).$
(A.10)

To lower bound (A.10), we may first bound $\langle H({\bm{\varphi}})\rangle$
as follows: Introduce the partition function $Z=\int
e^{-H({\bm{\varphi}})}\mathrm{d}G_{\eta}({\bm{\varphi}})$. Applying
$H({\bm{\varphi}})\geq 0$, $e^{-H({\bm{\varphi}})}/Z\leq 1$ on the event
$e^{-H({\bm{\varphi}})}\leq Z$, and $H({\bm{\varphi}})<-\log Z$ on the
complementary event $e^{-H({\bm{\varphi}})}>Z$, we have

$\langle H({\bm{\varphi}})\rangle=\int
H({\bm{\varphi}})\frac{e^{-H({\bm{\varphi}})}}{Z}\mathrm{d}G_{\eta}({\bm{\varphi}}$
≤∫H(φ)dGη(φ)+max(0,-logZ)≤2∫H(φ)dGη(φ),

the last step applying the Jensen’s inequality lower bound $\log Z\geq-\int
H({\bm{\varphi}})\mathrm{d}G_{\eta}({\bm{\varphi}})$. Thus we have bounded the
mean of $H({\bm{\varphi}})$ under the posterior by its mean under the prior.
On the event
$\|\widetilde{\mathbf{y}}\|_{2}^{2}=\|\Pi_{X}\mathbf{y}\|_{2}^{2}\leq C_{0}p$,
applying ${\bm{\Sigma}}^{-1}\preceq(2/\sigma^{2})\operatorname{Id}$,
$\|\widetilde{\mathbf{X}}\|_{\operatorname{op}}\leq C$, and that $g_{\eta}$ is
sub-Gaussian (since $g$ has bounded support), we have $\int
H({\bm{\varphi}})\mathrm{d}G_{\eta}({\bm{\varphi}})\leq Cp$. This implies

$\langle H({\bm{\varphi}})\rangle\leq Cp.$ (A.11)

To bound $\langle\|F_{\eta}({\bm{\varphi}})-{\bm{\varphi}}\|_{2}^{2}\rangle$,
let us fix a small constant $t>0$ to be determined, pick $B=B(t)>0$
sufficiently large, and define
$N_{B}({\bm{\varphi}})=|\\{j:|\varphi_{j}|>B\\}|$. Applying again $\log
Z\geq-\int H({\bm{\varphi}})\mathrm{d}G_{\eta}({\bm{\varphi}})\geq-Cp$ and
$H({\bm{\varphi}})\geq 0$, we have

$\displaystyle\langle\mathbf{1}\\{N_{B}({\bm{\varphi}})>tp\\}\rangle=\int\mathbf{1}\\{N_{B}({\bm{\varphi}})>tp\\}\frac{e^{-H({\bm{\varphi}})}}{Z}\mathrm{d}G_{\eta}({\bm{\varphi}})\leq
e^{Cp}\int\mathbf{1}\\{N_{B}({\bm{\varphi}})>tp\\}\mathrm{d}G_{\eta}({\bm{\varphi}}).$

Under the product prior law $\prod_{j=1}^{p}g_{\eta}(\varphi_{j})$, this
quantity $N_{B}({\bm{\varphi}})$ is a binomial random variable whose success
probability is at most $q=2e^{-cB^{2}}$ for a constant $c:=c(M,\eta)$, by sub-
Gaussianity of $g_{\eta}$. For any fixed constant $t>0$, by choosing
$B=B(t)>0$ large enough, we have $q\leq t/2$, so by the Chernoff bound for the
binomial distribution [BLM13, pp. 24]

$\displaystyle\int\mathbf{1}\\{N_{B}({\bm{\varphi}})>tp\\}\mathrm{d}G_{\eta}({\bm{\varphi}})\leq\mathbb{P}_{g_{\eta}}\Big{(}\text{Bin}(p,q)-pq\geq\frac{t}{2}p\Big{)}\leq
e^{-ph_{q}(q+\frac{t}{2})},$

where $h_{q}(a)=\operatorname{\text{KL}}(\text{Ber}(a)||\text{Ber}(q))$. Since
$h_{q}(q+t/2)\rightarrow\infty$ as $q\downarrow 0$, we may then choose
$B=B(t)$ large enough such that