Hessian Ascent for the SK Model IV
TVD sampling for the SK model via Hessian dynamics, Jarzynski's equality and localized sampling
Continuing on the program that Jonathan and I started with David in PHA 1 and PHA 2, with Ewan and Holden, we put out a PHA 3 preprint on the arXiv (Potential Hessian Ascent III: Sampling the Sherrington-Kirkpatrick Model at \(\beta\) < 1/2, DLSS26). Shortly after PHA 3, Holden, Jonathan and I put out a follow-up, creatively termed (Potential Hessian Ascent IV: Sampling the Sherrington-Kirkpatrick Model at \(\beta\) < 1, LSS26) that builds on the cavity interpolation theory and free probability toolkit developed in PHA 3 to introduce a “local” SDE error analysis which only uses (and proves) “local” regularity properties for the driver and drift terms underlying the matricial functions of certain quantities that drive the SDEs.
Together, these works yield a \(o_n(1)\) TVD sampler for the SK model up to the replica-symmetric threshold \(\beta < 1\), above which sampling from the Gibbs measured is conjectured to be hard. The algorithm combines algorithmic stochastic localization (ASL) with rejection sampling over path-space via Jarzynski’s equality (JE). The analysis uses local regularity properties of the TAP Hessian and the PHA framework, developing new cavity interpolation theory and an extension to the free-probability toolkit introduced in PHA 1, Section 4, which is then combined with delicate SDE error analysis and the ability to efficiently sample from sufficiently “localized” distributions, i.e., those with large enough external field induced by a stochastic localization (SL) process that’s been run long enough.
I have decided to write, with guest contributions from Holden, a 6-part blog post explaining the conceptual background and development of the algorithm, the main proof skeletons, and the main technical innovations.
-
In this first blog post, Holden and I will explain the connection between SL and Hessian dynamics (HD), introduce Jarzynski’s equality, discuss the over-determined system of ASL, TAP and PHD SDEs, and the natural development of the ‘‘desiderata’’ that must be proved for the algorithm to succeed. I will end with three consequences (and perspectives) of the result – one for theoretical physicists, another for analysts/probabilists, and a final one for theoretical computer scientists.
-
In the second blog post, I will spend time developing the cavity interpolation theory that gives exact moment estimates and ‘‘better-than-naive’’’ overlap-concentration along the SL process. This is necessary to prove the desideratum that the algorithmic covariance (Hessian of the TAP free energy) stays close to the true covariance along the SL path. It will turn out to be the case that this theory holds for \(\beta < 1\) even though this was not explicitly written down in PHA 3, and a couple of final details for this were not fleshed out until PHA 4. The key technical contribution here is to run a cavity interpolation for the planted SL model with SL tilt and compute “self-stability” estimates on doing Taylor expansions for bulk-deviation quantities (such as overlaps and magnetizations) to get exact moment estimates.
-
In the third blog post, I will spend time explaining the extension to the analysis of the non-commutative interpolation developed in PHA 1, Section 4 that controls the diagonal entries of the (squared) algorithmic covariance, which is accomplished in PHA 3, Appendix B and PHA 4, Section 5 – this will then imply various regularity estimates and show that the ASL-TAP process closely tracks the PHD process. We will first do this for \(\beta < 1/2\), and then I will then show that, in fact, one can reason about a “regularized” resolvent which agrees with the original resolvent only over a (possibly non-convex) subset of \((-1,1)^n\) and, simultaneously, significantly simplify the free interpolation analysis in PHA 3 by using a result of Bandeira, Boedihardjo and van-Handel, 2023. This allows one to reason about the concentration and free limits of the diagonal sub-algebra of the TAP Hessian (a resolvent) all the way to \(\beta < 1\) where it can be singular in a subset of \((-1,1)^n\).
-
In the fourth blog post, Holden will show how the localized distribution concentrates on a wedge of the hypercube (provided it is run for sufficiently long) and how a 2-stage decomposition on this wedge allows one to show the entropy contraction property, which is the final desideratum that needs to be proved – this is the approach adopted in PHA 3. After that, he will briefly outline how a result of Kumar et al., 2026, which allows one to sample from a Gibbs distribution with quadratic potential and sufficiently strong external field on a wedge of \(\lbrace -1,1\rbrace^n\) using an annealing-type procedure, is used in PHA 4 to directly sample from the localized distribution after invoking the fact that SL run for large time will cause the distribution to be heavily concentrated inside a wedge of sufficient size.
-
In the fifth blog post, I will talk about how to access a “safe set” of points \(S_{A}(c) \subset \lbrace -1,1\rbrace^n\) when one fixes a “good” event over the input \(A\) so that, with high probability over the SL paths, the TAP Hessian evaluated at any \(m \in S_{A}(c)\) is \(c\)-strongly convex. This is the place where, perhaps unsurpringly, we rely on a result about AMP state evolution. Essentially, we combine the fact that posterior means of the SL tilted measures are “stable” in time for most SL paths with the \(c\)-strong convexity of Hessians in small balls around AMP iterates, to build “tubes” around good SL paths that we then require to be contained in \(S_{A}(c)\). From this, one straightforwardly obtains that, with high probability over \(A\) and the SL paths, uniformly in time \(t \in [0,T]\), small enough balls around the SL-tilted means \(\lbrace m_t\rbrace_{0 \le t \le T}\) lie in \(S_{A}(c)\). I believe that extending the ability to get the this sort of “joint” state-evolution statements for PHA itertes independently of invoking results about AMP is an important open direction, as it will further elucidate the capabilities and limitations of PHD.
-
In the final blog-post, I will talk about the PHA program in general, mentioning the two main current projects underway as well a more longer-term set of two projects that are also being worked on (but I will discuss these far more vaguely and briefly, for reasons that I hope are obvious given the current circumstances). I will, in a related way, (briefly) discuss the next “phase” of the program. When the program was originally conceived, spin-glasses on the hypercube were seen merely as a “test-bed” for the program that had historical importance, sufficient technical depth, and structural similarity to models that are the long-term target of the program.
Table of Contents
Designing the algorithm and desiderata
We begin by introducing the precise equivalence between stochastic localization (SL) and Hessian dynamics (HD), and the “surrogate” covariance we use to “algorithmize” the SL/HD procedures as ASL/PHD. From this, it will be clear that the error of the produced sample with respect to the true distribution will be entirely a function of how well one (algorithmic) family of SDEs tracks another (ideal) family of SDEs. This error can be systematically bounded by building a list of desiderata that start from asserting the very natural \(O(1)\) Wasserstein-2 distance based on an elementary optimal-coupling argument with a covariance error bound, to \(1-\Omega(1)\) TVD error by boosting this covariance closeness with Girsanov’s theorem (requiring some extra regularity properties), to finally achieving \(o_n(1)\) TVD error by utilizing rejection sampling over path-space.
Stochastic localization and Hessian dynamics
Given a measurable space \((\Omega, \mathcal{B}(\Omega))\) a localization process is a stochastic process \(\lbrace\nu_t(\cdot)\rbrace_{t\ge 0}\) over the space of probability measures on \((\Omega, \mathcal{B}(\Omega))\). The process has two distinct properties: 1) for any \(A \in \mathcal{B}(\Omega)\), \(\lim_{t\to\infty} \nu_t(A) \in \lbrace0,1\rbrace\), and 2) \(\nu_t(\cdot)\) is a martingale.
The first property tells us that the stochastic process “localizes” at one particular event \(A\) in the measure space, and the fact that it is a martingale implies that the process is “stationary” upon averaging. The latter property allows us to think of the localization process as a convex decomposition of the measure \(\nu_0(\cdot)\), weighted by the “tilted” measures along the way. For us, there are two important consequences of this:
-
If we can algorithmically simulate the process by “following” the tilts that generate the sequence of tilted measures for a very large (but constant) time \(T\), then we will come close to a sample from the target measure, and
-
We can choose any measure-valued process that localizes and is a martingale based on its suitability for efficient simulation, and analyzing its paths to prove our desiderata.
From prior works, a good choice is the “linear-tilt” localization scheme, which considers an unknown sample (signal) \(x_0\) drawn from the target measure and consists of noisy observations of \(x_0\) that are increasingly informative as time \(t\) evolves; the posterior distribution of \(x_0\) then localizes towards Dirac measure \(\delta_{x_0}\)12. This scheme starts at \(y_0 = 0^n\) and reveals itself at time \(t\) as
\[y_t = tx_0 + B_t\, ,\]where \(x_0 \sim \nu_0\) and \(B_t \sim \mathcal{N}(0,t)\). The measure \(\nu_t(\cdot)\) induced by this scheme is
\[\nu_t(\sigma) \propto e^{\langle y_t,\sigma \rangle}\nu_0(\sigma)\,,\]and it is easy to see that \(\lim_{t \to\infty} \nu_t(\cdot) \to \delta_{x_0}\) almost-surely and that \(\lbrace\nu_t(\cdot)\rbrace_{t\ge 0}\) is a martingale process since they form a Doob martingale. At this point, using the many known equivalent characterizations of stochastic localization (see [Sections 1 & 2, Shi, Tian, and Zhang, 2026]) one can also rewrite the linear-tilt process as
\[dy_t = m_t dt + dB_t\,,\]where \(m_t = \mathbb{E}_{x\sim\nu_t}\left[x\right]\). A main conceptual insight in our work is that there is yet another equivalent rewrite for a stochastic localization process, and this rewrite tracks the evolution of the magnetizations/averages \(\lbrace m_t\rbrace_{t\ge 0}\) of the “tilted” measures \(\nu_t(\cdot)\) under the linear-tilt localization scheme \(\lbrace y_t\rbrace_{t\ge 0}\). This process is
\[dm_t = \mathsf{Cov}(\nu_t)dB_t\,\qquad \text{Hessian dynamics (HD)}\,,\]where \(\mathsf{Cov}(\nu_t)\) is the covariance matrix for the measure \(\nu_t\). It is not difficult to see that the magnetization process obeys this SDE, from combining its definition as the mean of the tilted measures, an application of Ito’s lemma, and the fact that the magnetizations form a martingale. It will turn out to be the case (for convenience more than anything else) that, though the actual algorithm will not run an algorithmic version of the “HD” version of SL, it will still be very useful to analyze HD to prove some of the desiderata3. The algorithmic version of this process, which we call Potential Hessian Dynamics (PHD), is
\[d\hat{m}_t = \hat{Q}(\hat{m}_t)dB_t\,.\]A critical part of proving our desiderata and writing down the final algorithm relies on gaining access to an efficiently computable sequence of matrix-valued functions \(\hat{Q} : [-1,1]^n \to M_n(\mathbb{R})_{\text{sa}}\) – this is exactly where the contiguity to a planted model and the TAP free energy will be of assistance.
Algorithmic surrogates via the TAP free energy
Before we introduce the TAP free energy for the SK model, as is relevant in the high-temperature regime, let us quickly remind ourselves of the (random) Gibbs measure associated with it. The target/Gibbs measure for the SK model is given as
\[\mu_{\beta A}(\sigma) := \frac{e^{\frac{1}{2}\left\langle\sigma,\beta A\sigma\right\rangle}}{Z}\,,\]where \(A\) is a \(\mathsf{GOE}(n)\) matrix, \(Z\) is the normalizing constant called the partition function, and the density is defined for every \(\sigma \in \Sigma_n := \lbrace -1,1 \rbrace^n\).
The SK model has a particularly nice form for its TAP free energy in the high-temperature regime (\(\beta < 1\)) whose validity was rigorously established in a series of papers – see, for instance [Chen and Panchenko, 2019]. The free energy is equivalent to the largest value of the TAP functional evaluated as a supremum over all possible magnetizations. The TAP functional for the SK model at some magnetization \(m \in [-1,1]^n\) is given as
\[\mathcal{F}_{\mathsf{TAP}}(m) = -\frac{\beta}{2}\langle \sigma, A \sigma \rangle - \sum_{i=1}^n h(m_i) - n\left(\frac{\beta^2}{4}\left(1-\frac{\|m\|^2_2}{n}\right)\right)\,,\]where \(h(\cdot) : [-1,1] \to [0,1]\) is the entropy functional for a Bernoulli random variable (see PHA 3, (2.2)). This structure of the TAP free energy is closely related to taking a Legendre transform of the standard definition of the free energy, and consists of adding the final term as a “correction term” to the Gibbs variational principle over product measures. For our purposes, we don’t just need the TAP free energy of the standard SK model, but we need some analogue of it for any sequence of tilts \(y_t\) generated by the linear-tilt localization scheme – this is easily accomplished by adding an external field term. Namely, setting \(\nu_0 = \mu_{\beta A}\) and using the fact that \(\nu_t(m) \propto e^{\langle y_t, m \rangle}\mu_{\beta A}\) gives us the following modification to the TAP free energy, given a fixed tilt \(y \in \mathbb{R}^n\):
\[\mathcal{F}_{\mathsf{TAP}}(y,m) := -\frac{\beta}{2}\langle \sigma, A \sigma \rangle - \langle y, m \rangle - \sum_{i=1}^n h(m_i) - n\left(\frac{\beta^2}{4}\left(1-\frac{\|m\|^2_2}{n}\right)\right)\,.\]We will be interested in viewing this as a Fenchel legendre transform, and given some algorithmic estimate \(\hat{y}_t\) of \(y_t\) (with sufficiently good approximation) we will be able to compute \(\hat{m}_t\) using the strong convexity of this TAP free energy (since that will guarantee the uniqueness of \(\hat{m}_t\)), which will in turn allow us to compute the next tilt \(\hat{y}_{t+\delta t}\).
This is fantastic for the algorithmic procedure, and it also suggests a natural covariance \(\hat{Q}(m)\) that can be defined everywhere4 in \((-1,1)^n\) – the inverse of the Hessian of the TAP free energy. This is so since \(y_t\) and \(m_t\) are dual to each other, and the Crouzeix identity tells us that the second-derivative with respect to one variable (\(y_t\)) is the functional inverse of the other (\(m_t\)) and this gives
\[\hat{Q}(m) := \left(\nabla^2 \mathcal{F}_{\mathsf{TAP}}(y,m)\right)^{-1} = \left(\beta^2\mathsf{tr}_n[D^{-1}(m)]I_n -\beta A + D(m)-\frac{2\beta^2}{n}mm^T\right)^{-1}\,,\]where \(D(m)\) is the diagonal matrix with entries \(\left\lbrace \frac{1}{1-(m_i)^2}\right\rbrace_{i=1}^n\). A pleasant consequence of the choice of \(\hat{Q}(m)\) is that it is actually independent of the tilt \(y\). Another immediate consequence is that, under the Loewner order sandwich on \(\hat{Q}(\cdot)\) (and consequently \(D(\cdot)\hat{Q}(\cdot)\)) as will be required by the first desideratum for the covariance matrix, one can easily obtain that
\[\begin{aligned} \|\hat{Q}(a) -\hat{Q}(b)\|^2_F &= \|\hat{Q}(a)D(a)\left(D^{-1}(a)-D^{-1}(b)\right)D(b)\hat{Q}(b) -\hat{Q}(a)(R(a)-R(b))\hat{Q}(b)\|^2_F \\ &\le \|D(a)\hat{Q}(a)\|_\infty^2\|D^{-1}(a) - D^{-1}(b)\|_F^2 + \|\hat{Q}(a)\|^2_\infty\|R(a) - R(b)\|^2_F \\ &\le C(\beta)\|a-b\|^2_2 + C'(\beta)\|R(a)-R(b)\|^2_F\, , \end{aligned}\]where \(R(a) = \left(\beta^2\mathsf{tr}_n[D^{-1}(a)]I_n-\beta A - \frac{2\beta^2}{n}aa^T\right)^{-1}\). Then, conditioning on the event that \(\|A\|_\infty \le 2+\delta_\beta\) and doing some elementary estimates using the fact that \(a,b \in (-1,1)^n\), yields that \(\|R(a)-R(b)\|^2_F \le C(\beta)\|a-b\|^2_2\). Substituting this into the bound above shows that \(\hat{Q}(\cdot)\) is Lipschitz with respect to \(\|\cdot\|_F\) and proves one of the four regularity properties in the desiderata (simply as a consequence of the first desideratum and surrogate choice of covariance).
I mentioned before that while the algorithmic version of HD, called potential Hessian dynamics (PHD), which we can now see will start at \(\hat{m}_0 = 0^n\) with updates
\[d\hat{m}_t = \left(\beta^2\mathsf{tr}_n[D^{-1}(\hat{m}_t)]I_n -\beta A + D(\hat{m}_t)-\frac{2\beta^2}{n}\hat{m}_t(\hat{m}_t)^T\right)^{-1}dB_t\,,\]will not be used directly, it will be useful in the analysis. One can ask, then, how so? One answer is the reason we can’t just boost the \(O_{T,\beta}(1)\) \(W_2\) bound we can assume (given our desiderata, which we will systematically prove) to a \(O(1)\)-KL divergence bound via Girsanov’s theorem – we now focus on why this is, and how to address it.
SDE error analysis
Let us now reason about how to bound the error between two continuous diffusion processes run for large constant time. Think of one process as representing an “ideal” process that, if run for infinite time, will localize on a point in the support of the Gibbs measure, and think of the other as one that uses a proxy covariance that can be algorithmically computed (but incurs approximation errors).
As a warm-up, consider HD and PHD. Denote the ideal process as \(dm_t = Q(m_t)dB_t\) and the algorithmic process as \(d\hat{m}_t = \hat{Q}(\hat{m}_t)dB_t\). Then, by a simple argument using the fact that the Wasserstein-\(2\) (\(W_2\)) distance between two Gaussians is upper bounded by the Frobenius norm of the difference between their covariances, we have
\[W_2(\mathsf{dist}(m_t),\mathsf{dist}(\hat{m}_t)) \le \mathbb{E}\|\hat{Q}(\hat{m}_t) - Q(m_t)\|^2_F\, .\]This already hints at why one of the most crucial desiderata in the entire program is going to be an error estimate tracking how well a “surrogate” covariance \(\hat{Q}(\cdot)\) tracks the “ideal” covariance \(Q(\cdot)\) across SL paths \(\lbrace m_t \rbrace_t\).
Now we introduce the two sets of Ito processes that we will actually compare. The ideal processes will be a set of two coupled processes, and the algorithmic processes will be similar (though not necessarily coupled, unless we enforce that manually). These processes are SDEs that also have drift terms. In addition to the diffusion process before, the ideal process will also include \(dy_t = m_tdt + dB_t\) and the algorithmic estimate will likewise include \(d\hat{y}_t = \hat{f}_t(\hat{m}_t)dt + dB_t\). Now, we set \(y_0 = m_0 = \hat{y}_0 = \hat{m}_0 = 0^n\) (which is the relevant initialization for us) and, assuming that the solutions exist and are well-posed, obtain the following representation for the ideal and algorithmic processes
\[d\begin{pmatrix} y_t \\ m_t \end{pmatrix} = \begin{pmatrix} m_t \\ 0^n \end{pmatrix} dt + \begin{pmatrix} I_n \\ Q(m_t) \end{pmatrix}dB_t\,,\]and
\[d\begin{pmatrix} \hat{y}_t \\ \hat{m}_t \end{pmatrix} = \begin{pmatrix} \hat{f}_t(\hat{m}_t) \\ 0^n \end{pmatrix} dt + \begin{pmatrix} I_n \\ \hat{Q}(\hat{m}_t) \end{pmatrix}dB_t\,.\]We can now upper-bound the \(W_2\) distance between the ideal and algorithmic processes using the cumulative difference process \(\text{err}_t = \begin{pmatrix} y_t - \hat{y}_t \\ m_t - \hat{f}_t(\hat{m}_t) \end{pmatrix}\). Indeed, a simple application of Ito’s lemma tells us that the rate of the average instantaneous error between the two processes is
\[\frac{d}{dt}\mathbb{E}\|\text{err}_t\|^2_2 = 2\mathbb{E}\langle\text{err}_t, (\hat{f}_t(\hat{m}_t),0^n)-(m_t,0^n)\rangle + \mathbb{E}\|Q(m_t) - \hat{Q}(\hat{m}_t)\|^2_F\,,\]whereupon an application of a \(c\)-weighted AM-GM inequality on the first term followed by some triangle inequalities and the fact that \((a+b)^2 \le 2a^2+2b^2\) tells us that
\[\frac{d}{dt}\mathbb{E}\|\text{err}_t\|^2_2 \le c\mathbb{E}\|\text{err}_t\|^2_2 + \frac{2}{c}\left(\underbrace{\mathbb{E}\|\hat{f}_t(\hat{m}_t)-\hat{f}_t(m_t)\|^2_2}_{\hat{f}_t \text{ Lipschitz error}} + \underbrace{\mathbb{E}\|\hat{f}_t(m_t)-m_t\|^2_2}_{\text{PHD-TAP drift error}}\right) + 2\underbrace{\mathbb{E}\|\hat{Q}(\hat{m}_t)-\hat{Q}(m_t)\|^2_F}_{\hat{Q}(\cdot)\text{ Lipschitz error}} + 2\underbrace{\mathbb{E}\|\hat{Q}(m_t)-Q(m_t)\|^2_F}_{\text{covariance estimtate error}}\,.\]At this point, it is clear that if we are only to run our algorithmic processes for finite time \(T\) and the final four terms in the bound above are of order \(O_{t,\beta}(1)\) at every \(0 \le t \le T\), a Gronwall’s inequality bound will immediately give that \(\mathbb{E}\|\text{err}_T\|^2_2 = O_{T,\beta}(1)\).
Now, if we stop our algorithmic processes at some time \(T\), we must either assert that some deterministic and efficient “rounding” function \(h(\hat{m}_T,\hat{y}_T)\) outputs a sample with \(o_n(1)\) TVD error, or we can use the pair \((\hat{m}_T,\hat{y}_T)\) as a “warm start” to another efficient algorithm which outputs a sample \(\sigma \in \lbrace -1,1\rbrace^n\) that is \(o_n(1)\) close in TVD error to the target measure. Since we stop at a large constant time \(T\), we are forced to adopt the latter approach. We will approximately sample from the localized distribution at time \(T\), which is the Gibbs measure tilted by \(\hat{y}_T\).
Between the bounds on the four final four quantities for the cumulative error process, the required sampler with the warm start mentioned above, and the fact that the algorithmic surrogate \(\hat{Q}(\cdot)\) must be a valid covariance, we already have a list of desiderata that the algorithm requires, at least to have \(W_2\) error that is \(O_{T,\beta}(1)\)5:
-
\(\hat{Q}(\cdot)\) is a regular and valid covariance, namely \(c(\beta)I_n \preceq \hat{Q}(m) \preceq C(\beta)I_n\) at all \(m \in \lbrace -1,1\rbrace^n\)6.
-
The covariance errors are small, that is \(\mathbb{E}\|\hat{Q}(m_t) - Q(m_t)\|^2_F \le O_{t,\beta}(1)\).
-
The functions \(\hat{f}_t(\cdot)\) are \(C_{\beta,t}\)-Lipschitz with respect to \(\|\cdot\|_2\) inside the solid cube.
-
The error PHD-TAP drift error for magnetization is bounded, that is \(\mathbb{E}\|\hat{f}_t(m)-m\|^2_2 \le O_{t,\beta}(1)\) for every \(m \in \lbrace -1,1\rbrace^n\)7.
-
Will be filled in after the below discussion to obtain \(o(1)\) error.
-
After running the (discretized) algorithmic processes for \((\hat{m}_t,\hat{y}_t)\) for time \(T\), there is a sampler that samples from a certain simpler (“stochastically localized”) distribution with \(o_n(1)\) TVD error in polynomial time.
While it is not clear right now, the covariance and PHD-TAP drift error estimates will only permit error of the desired order when \(\beta < 1\) – this will be discussed in the second and third blog post. Additionally, the PSD-ness property of the covariance will hold pointwise inside the cube only for \(\beta < 1/2\). To make the PSD-ness property, and the PHD-TAP drift error estimate (which relies on it in an indirect way), work for \(\beta < 1\), we will need to refine these desiderata to hold only at a certain “safe” set \(S_A(c) \subset \lbrace -1, 1\rbrace^n\) – the PHD-TAP error part of this will be the subject of the third blog post, and obtaining PSD-ness over the “safe” set will be established in the final blog post. Lastly, the Lipschitz error for \(\hat{Q}(\cdot)\) will be a consequence of the Loewner regularity of \(\hat{Q}(\cdot)\) combined with the resolvent structure it has based on the explicit choice we use – see (1.2).
We are now two steps away from getting a \(o_n(1)\)-TVD sampler, if we can show the desiderata outlined above. First, we need to reason about one more source of \(L_2\)-error to go from \(O_{T,\beta}(1)\)-\(W_2\) error to \(O_{T,\beta}(1)\)-KL divergence error via an application of Girsanov’s theorem. At that point, Pinsker’s inequality yields a \(O_{T,\beta}(1)\)-TVD error, but that is still not \(o_n(1)\). The second step, which achieves this, is to use rejection sampling (but over path space) to suppress the error further to \(o_{n}(1)\)-TVD error – this is where Jarzynski’s equality enters the picture. We will now introduce the ideal processes (SL/HD), then define the TAP free energy and use it to derive the algorithmic surrogates. At that point, we will be able to complete the ASL-TAP and PHD \(L_2\)-error bound needed to apply Girsanov’s theorem. We will then move on to defining Jarzynski’s equality, see how it allows us to do rejection sampling, which further suppresses the TVD error. Doing the last step will incur one more desideratum, at which point we will conclude with our final list of desiderata.
ASL-TAP and boosting to a KL divergence bound
In a nutshell, the main reason why the \(O_{t,\beta}(1)\) \(W_2\) error can’t be boosted to a KL-divergence error via Girsanov’s theorem is that the algorithmic processes for \(\hat{y}_t\) and \(\hat{m}_t\) are not coupled in any way, but they should satisfy the constraint that \(\nabla\mathcal{F}_{\mathsf{TAP}}(\hat{m}_t,\hat{y}_t) = 0\). This requirement is the algorithmic analogue of the fact that \(y_t\) and \(m_t\) satisfy the relationship \(m_t = \mathbb{E}_{x\sim\nu_t}[x]\). If not, \(\hat{m}_t\) and \(\hat{y}_t\) are simply not coupled in a way the ideal processes are, and this is what stops us from applying Girsanov’s theorem to boost the \(W_2\) error to a KL-divergence error. It turns out that we cannot satisfy all three algorithmic processes at the same time – the algorithmic processes based on the TAP surrogacy are over-determined. We need to satisfy the TAP fixed-point, so we can either choose PHD-TAP
\[d\hat{m}_t = \hat{Q}(\hat{m}_t)dB_t\,,\qquad\qquad d\hat{y}_t = \hat{f}(\hat{m}_t)dt + dB_t\]or ASL-TAP
\[d\hat{y}_t = \hat{m}_tdt + dB_t\,,\qquad\qquad d\hat{m}_t=\hat{Q}(\hat{m}_t)\left(\hat{m}_t-\hat{f}(\hat{m}_t)\right)dt + \hat{Q}(\hat{m}_t)dB_t\,,\]where
\[\hat{f}(z) := \left(D(z)E_{\mathsf{D}_n}\left[\hat{Q}(z)^2\right]D(z) - \beta^2\left(\mathsf{tr}_n\left[\hat{Q}(z)^2\right]I_n + 2\hat{Q}(z)^2\right)\right)z\]can be explicitly derived via an application of Ito’s lemma to the function that gives a precise relationship between \(\hat{y}_t\) and \(\hat{m}_t\), which itself can be straightforwardly derived from the TAP fixed-point equation.
For the algorithm, using the strong convexity of \(\hat{Q}(\cdot)\), we choose to satisfy the TAP fixed-point constraint and run the (discretized) ASL process \(d\hat{y}_t = \hat{m}_tdt + dB_t\). For the analysis, we actually show that, under the desiderata mentioned before, PHD-TAP \((\hat{m}_t, \hat{y}_t)\) stays close to the ideal processes \((y_t,m_t)\) in the sense that
\[\mathbb{E}\|\hat{m}(y_t) - m_t(y_t)\|^2_2 = O_{t,\beta}(1)\,.\]This follows from the calculation
\[\begin{align*} \mathbb{E} \|\hat{m}(y_t) - m_t(y_t)\|^2_2 &\lesssim \mathbb{E}[\|\hat m(y_t) - \hat m (\hat y_t)\|^2] + \mathbb{E}[\|\hat m_t(\hat y_t)- m_t\|^2]\\ &\lesssim \mathbb{E} \|\hat y_t - y_t\|^2 + \mathbb{E} \|\hat m_t-m_t\|^2 =O(1). \end{align*}\]where the first line uses the triangle inequality and the second line uses the Lipschitz-ness of \(\hat{m}\) (which can be straightforwardly concluded from the fact that its Jacobian \(\hat{Q}\) has bounded operator norm by the first desideratum) and the last step follows from the Gronwall argument described earlier.
After this, one can apply Girsanov’s theorem (after an application of the DPI) to show that ASL-TAP stays close to the ideal SL process for \(y_t\).
Consequently, we must account for one additional source of error, and that is the expected squared error between the ASL-TAP and PHD process. The final desiderata to obtain \(O(1)\) KL divergence error will come from this part of the argument.
Reducing to negligible TVD error via rejection sampling over path space
A natural idea to reduce \(O(1)\) error (in KL divergence) to \(o(1)\) error is rejection sampling. Rejection sampling gives a way to sample from a distribution \(P\) given the ability to sample from a proposal distribution \(\hat P\), which is efficient if we can bound the ratio \(\frac{dP}{d\hat P}\) with high probability under \(P\). It is natural to use the algorithmic process at time \(T\) as the proposal distribution; however, we do not have a formula for the density of this distribution.
Jarzynski’s equality
This is where Jarzynski’s equality enters. It is an identity from nonequilibrium statistical physics based on the Feynman-Kac formula, which yields a method to sample or compute normalizing constants (free energy) for a target distribution, given an annealing sequence \(\rho_t\propto e^{-U_t}\). Under basic regularity conditions, it allows us to correct the distribution of a SDE with any given drift towards a target distribution using importance weights as calculated below. If the proposal \(\hat\rho_t\) is the distribution of \(x_t\) given by
\[\begin{align} dx_t & = \left[-\frac12\nabla U_t(x_t) + b_t(x_t)\right]dt + dB_t, &x_0&\sim \rho_0\\ dw_t &= \left[\nabla \cdot b_t(x_t) - \left\langle\nabla U_t(x_t), b_t(x_t)\right\rangle - \partial_t U_t(x_t)\right]dt, & w_0&=0, \end{align}\]then8
\[\frac{d\rho_t}{d\hat p_t}(x) = \frac{\mathbb E[e^{w_t}|x_t=x]}{\mathbb E[e^{w_t}]}.\]Moreover, the ratio of normalizing constants can be computed as \(\frac{\int e^{-U_t(x)}dx}{\int e^{-U_0(x)}dx} = \mathbb E[e^{w_t}]\).
Two special cases are worth nothing. First, annealed Langevin dynamics is the case where \(b_t\equiv0\): the SDE for \(x_t\) at time \(t\) is exactly Langevin diffusion for \(\rho_t\), and the correction is needed because the stationary distribution would be attained only after running the time-\(t\) Langevin for infinite time. More relevant for us, if \(\hat p_t = p_t\) is the SL distribution at time \(t\), then choosing \(b_t(y) = -\frac 12 \nabla U_t(y) + \frac yt = \frac 12 (m(y) + \frac yt)\) recovers the SL process, and \(dw_t\) is a constant.
Back to our problem, what should we choose for the reference distribution \(\rho_t\)? We claim that a natural choice is
\[\rho_t(y) \propto e^{-\mathcal F_{\mathrm{TAP}}(\hat m(y), y) - \frac{\|y\|^2}{2t}},\]which can be seen in 2 ways:
-
We expect \(e^{-\mathcal F_{\mathrm{TAP}}(\hat m(y), y)}\approx Z_{\beta A, y}\), so using this approximation,
\[\begin{align*} p_t(y) & = \sum_{\sigma\in \{\pm 1\}^n} \mu_{\beta A}(\sigma) p(y_t=y|\sigma)\propto \sum_{\sigma\in \{\pm 1\}^n} e^{\frac\beta 2 \langle\sigma,A\sigma\rangle}e^{-\frac{\|y-t\sigma\|_2^2}{2t}}\\ &= \sum_{\sigma\in \{\pm1\}^n} e^{\frac\beta 2 \langle{\sigma,A\sigma}\rangle + \langle{y,\sigma}\rangle} e^{-\frac{\|y\|_2^2}{2t}} = {Z_{\beta A,y}} e^{-\frac{\|y\|_2^2}{2t}} \approx e^{-\mathcal F_{\mathrm{TAP}}(\hat m(y),y)}-\frac{\|y\|^2}{2t} \propto \rho_t(y). \end{align*}\] -
The score function involves \(\hat m(y)\):
\[\begin{align*} -\nabla_y \mathcal F_{\mathrm{TAP}}(\hat m(y),y) &= \hat m(y)\\ \implies \nabla_y \log \rho_t(y) &= \hat m(y) - \frac yt, \end{align*}\]while \(\nabla_y \log p_t(y)= m(y)- \frac yt\), corresponding to the drift for the algorithmic SDE being \(\hat m(y)\) rather than \(m(y)\).
Applying this to our situation gives
\[\begin{align*} d\hat w_t &= \omega(\hat m_t), \\ \text{where } \omega(m):&= \frac{1}{2}\left(\mathrm{Tr}[\hat{Q}(m)] + \|m\|^2_2\right) . \end{align*}\]Combined with the (ASL-TAP) process, we refer to the combined process as the (ASL-TAP-JE) process, and this is the main backbone of the sampling algorithm.
We remark that there has been renewed interest in Jarzynski’s equality recently, as a way to steer deep generative models to output samples from a given probability distribution. To our knowledge, this is the first use of Jarzynski’s equality in a theoretical sampling problem.
From Jarzynski’s equality to rejection sampling
Recall that our actual target in the first stage is the SL distribution \(p_T\); the reference is to give us something to compare to. Our rejection sampling ratio is hence
\[\frac{dp_T}{d\hat p_T} = \frac{d\rho_T}{d\hat p_T} \cdot \frac{dp_T}{d\rho_T} \propto \mathbb E \left[e^{\hat w_T} | \hat y_T\right] \cdot \frac{Z_{\beta A, \hat y_T}}{e^{-\mathcal F_{\mathrm{TAP}}(\hat m_T,\hat y_T)}}.\]We can estimate if we can estimate the partition function \(Z_{\beta A, \hat y_T}\); by a standard argument with simulated annealing, this is possible if we can sample from \(\mu_{\beta'A, \hat y_T}\), \(0\le \beta'\le \beta\). The efficiency of this derives from 2 separate estimates: first, that \(p_T\) is close to \(\hat p_T\) (\(\mathrm{KL}(p_T \| \hat p_T)= O(1)\)), and that second, that most of the mass of the expectation \(\mathbb E e^{\hat w_T}\) comes from when the relative size of \(e^{\hat w_T}\) is bounded, which we formalize as a desiderata below.
We note some subtleties in operationalizing this idea into an algorithmic guarantee. First, because the rejection sampling ratio is not normalized, we first draw multiple samples and normalize using an appropriate quantile. Second, we only know that we can sample from \(\mu_{\beta A,y}\) and hence estimate \(Z_{\beta A, y}\) over typical draws of \(y\) from the ideal SL distribution \(y_T\sim p_T\), whereas we are applying this to draws \(\hat y_T\sim \hat p_T\). This will not be a problem as long as we do not overestimate the partition function in the possibly “fictitious” \(y\)’s. Fortunately, we can ensure this using a functional inequality on a subset of the hypercube (a wedge) which does not depend on the tilt \(y\).
Final desideratum
Because the JE weights are exponentiated, a natural desideratum is that the moment generating function is bounded.
(d5). Let \((\hat y_t, \hat m_t,\hat w_t)\) evolve as in (ASL-TAP-JE). Then for all \(\lambda \in \mathbb R\), there is some finite \(C_{\mathrm{JE},T,\lambda} > 0\) so that
\[\mathbb E \left[\exp\left(\lambda\hat{w}_T - \mathbb E [\lambda\hat{w}_T]\right)\right] \le C_{\mathrm{JE},T,\lambda}.\]In fact, this follows from a more basic fact.
(d5’). The function \(\omega(m)\) in the ODE for \(\hat w_T\) is Lipschitz.
Calculation shows that \(\nabla \omega (m) = m - f(m)\), so this is in fact exactly the same drift error that is considered in desideratum 4.
Exponential moments (d5) may seem like a strong condition, but it is satisfied for Lipschitz functions of Gaussian random variables. The key insight is that Lipschitzness of \(\omega\) implies that \(\hat w_T\) is Lipschitz in a Gaussian space, not a finite one but an infinite-dimensional path space. We will show the implication between (d5’) and (d5) in a subsequent post. Upon adding this final desideraturm to the previously existing list as number 5), the list of required desiderata are complete and the next few blog posts will focus on how one proves these in the case of the SK model, highlighting reusable techniques and frameworks in the approach where possible.
Conclusions
We end this blog post by focusing on three consequences about the PHA paradigm for sampling.
The first is that sampling in TVD allows one to conclude that the distribution of overlaps of two samples produced by two independent runs of the algorithm on the same instance is certifiably close to that of the overlap distribution of the target Gibbs measure – this is simply not possible to conclude with \(o(n)\)-\(W_2\) error. For spin-glasses, the distribution of overlaps under the Gibbs is the order parameter of the system and so being able to compute it gives access to many statistics of the system.
To introduce the second consequence, we note that the literature has historically focused on getting efficient TVD sampling by proving functional inequalities for a Markov chain whose invariant measure is the target Gibbs measure. Unfortunately, proving functional inequalities straight-up for the Gibbs measure of spin-glasses seems to be difficult for various technical reasons. The PHA paradigm comes to the rescue by allowing one to sample even without access to such a functional inequality.
In fact, this leads to the third point: the PHA paradigm (in conjunction with Jarzynski’s equality) does actually help prove functional inequalities! It does so by providing access to certain regularity conditions that can be then “bootstrapped” into proving (weak) functional inequalities. More generally, the PHA paradigm can be seen as a framework about resolvent-driven SDEs whose dynamics can be analyzed systematically using tools from non-commutative analysis.
Approximating the overlap distribution
A question that I was asked on twitter/bluesky by Lenka Zdeborova was:
What is the statistic that the TVD sampler can compute which the Wasserstein-2 samplers cannot?
The answer to this question is that a sampler close in Wasserstein-2 distance, when run twice (with independent internal randomness) on the same instance, produces two samples whose normalized overlap distribution is not close to the distribution of the overlaps of two independent replicas sampled from the true Gibbs. However, a sampler that gets \(o_n(1)\) close in TVD error does have this property. First, let’s fix \(0<\beta <1\) and an interaction matrix \(A\) which is in the “good event” for which the negligible TVD sampler gets \(o_n(1)\) close to the true Gibbs on that sampler. Let the true Gibbs law be \(\mu_{\beta A}\) and let the output law of one sample of the algorithm be \(\nu_{\beta A}\). Then, on the good event, \(d_{TV}(\mu_{\beta A},\nu_{\beta A}) \le o_n(1)\) immediately implies that, given \(\sigma^1,\sigma^2 \sim \mu_{\beta A}^{\otimes 2}\) and \(\tau^1,\tau^2 \sim \nu_{\beta A}^{\otimes 2}\),
\[d_{TV}\left(\frac{\langle \sigma^1,\sigma^2\rangle}{n}, \frac{\langle \tau^1, \tau^2 \rangle}{n}\right) \le 2o_n(1)\,,\]since the probability that the two pairs disagree is \(\le 2o_n(1) - 2 o_n(1)^2\) by a simple coupling argument. In fact, by choosing the correct centering \(q^*\) and re-scaling \(\frac{1}{\sqrt{n}}\) for the overlaps (since they exhibit sub-Gaussian tails for the target measure in the high-temperature regime) one can use the TVD-closeness to assert that
\[d_{TV}\left(\sqrt{n}\frac{\langle \sigma^1,\sigma^2\rangle}{n}, \sqrt{n}\frac{\langle\tau^1,\tau^2\rangle}{n}\right) \le 2o_n(1)\,,\]and since a result of Guerra implies that the distribution of the overlaps of the true Gibbs measure converges for most instances of \(A\) (the disorder) to \(\mathcal{N}\left(0,\frac{1}{1-\beta^2}\right)\), this means that the overlap law induced by the algorithm’s law \(\nu_{\beta A}\) is close in TVD to this. Unfortunately, one cannot reach this conclusion with a (normalized by \(n\)) negligible \(o_n(1)\) Wasserstein-2 error guarantee alone, and one would need quantiative control on how small the Wasserstein-2 error is (and \(o(n^{1/4})\) would suffice). Additionally, any measurable function of the overlaps that shrinks with size \(n\) cannot be well-approximated from negligible Wasserstein-2 distance, but can be shown to be well-approximated by samples from a sampler with negligible TVD error.
Sampling without functional inequalities
Much of the literature on efficient sampling algorithms proceeds by proving a functional inequality (Poincaré or log-Sobolev), which establishes rapid mixing of the associated Markov chain (Glauber dynamics in discrete settings and Langevin dynamics in continuous settings). Our framework provides an alternative path, obtaining the strong guarantee of negligble TV distance without needing to prove the standard (stronger) functional inequalities. In the context of sampling algorithms, functional inequalities are the strongest results one can prove and hence desirable when achievable–however, for spin glass models, they may be exceedingly difficult to prove, if indeed they do hold, which is by no means clear.
Weak functional inequalities and resolvent driven SDEs
Subsequent work has upgraded the guarantees to functional inequalities for the SK model by building directly build on this algorithmically-inspired line of work – in our follow-up work [WPI via ASL] we establish a weak Poincaré inequality by combining the desiderata shown here with regularity of the SL process. Rather than implying exponential mixing from worst-case initialization, weak Poincaré inequalities imply mixing from a warm-start with a specified mixing rate. This broadens the possible landscape of sampling results, where different strengths of desiderata can imply different strengths of functional inequality and mixing results.
We find it remarkable that although the tools introduced here–ASL and Jarzynski’s equality–are algorithmic in nature and intended as such, they can (1) function to prove a purely mathematical result, and (2) be absorbed into the analysis of an algorithm that does not rely on simulating these processes. Thus, an algorithmic study of these models is fruitful for both computational and purely mathematical understanding.
We end by discussing a key mathematical perspective of the broader program. Towards establishing the regularity properties, one of the pre-meditated crimes of the PHA program is the underlying theme of resolvent-driven SDEs and the ability to analyze various (matricial) functions of its dynamics. By this, I mean consider a general vector-valued SDE of the form
\[d\hat{m}_t = \hat{Q}(t,\hat{m}_t)dB_t\]where \(\hat{Q}(t,\hat{m_t}) = \left(z(\hat{m}_t)\mathsf{Id} - \left(\mathcal{F}(x_1,\dots,x_m,\hat{m}_t) + M(\hat{m}_t)\right)\right)^{-1}\). Here, \((x_1,\dots,x_m)\) is an encoding of the input and \(\hat{m}_t \in \mathcal{D}\) (reflective of the set of all possible points in the SDE’s trajectory) with \(m = \mathrm{poly}(n)\) and \(\mathcal{F}: \mathsf{dom}(\mathcal{I}) \times \mathcal{D} \mapsto \mathbb{M}_{sa}(n)\) maps these points to a self-adjoint matrix. The matrix \(M: \mathcal{D} \mapsto \mathbb{M}_{sa}(n)\) can be seen as an “entropic” correction (encoding the geometry of the set of terminal points) to an “energy” term provided by \(\mathcal{F}\), and the whole term itself may be interpreted as the Hessian of a free-energy. Here is the main theme:
Utilize non-commutative “independence”s between \(\mathcal{F}\) and \(M\) and use analytic subordination to project onto sub-algebras.
This leads us to be able to control the projection of (reasonable) functions of \(\hat{Q}\) onto the algebra generated by \(M(\hat{m}_t)\). Specifically,
\[E_{\mathcal{M}(\hat{m}_t)}\left[\hat{Q}(t,\hat{m}_t)\right] = \left(f_{\mathcal{F},M}(z(\hat{m}_t))\mathsf{Id}-M(\hat{m}_t)\right)^{-1}\,,\]by a result of Biane, and this allows one to control the projection onto the sub-algebra \(\mathcal{M}(\hat{m}_t)\) of continuous functions \(g\) that act on the covariance. This follows simply by realizing that if \(z(\hat{m}_t)\) is not in the support of the free-sum (so that the resolvent is well-behaved) then one can use an application of the Weierstrass theorem in conjunction with the derivatives of a resolvent under Biane’s subordination to show that
\[\left\|E_{\mathcal{M}(\hat{m}_t)}\left[g\left(\hat{Q}_\infty(t,\hat{m}_t)\right)\right] - P\left(\left(f_{\mathcal{F},M}(z(\hat{m}_t))1_{\mathcal{A}} - M_\infty(\hat{m}_t) \right)^{-1}\right)\right\|_{\mathrm{op}} \le \text{small}\,,\]where \(P(A(z)) := c_0\mathbf{1} + \sum_{k=1}^n\frac{c_k(-1)^{k-1}}{(k-1)!}\frac{d^{k-1}}{dw^{k-1}}(A(w))^{-1}\mid_{w=z}\) and the coefficients \(\{c_k\}_k\) come from the Bernstein polynomial that witnesses the approximation to the function \(g\) in the Weierstrass theorem.
Note that the approximation above is for the ideal, infinite-dimensional objects in a free-algebra generated by the operators \(\{\hat{Q}_\infty(t,\hat{m}_t),M_\infty(\hat{m}_t)\}\) with the same spectral distributions that the finite-dimensional matrices \(\hat{Q}(t,\hat{m}_t)\) and \(M(\hat{m}_t)\) converge to. Of course, we want a finite dimensional version of this, and by recent advances in free probability theory, we can replace the terms with their finite-dimensional counterparts without incurring too much error. A systematic way to do this using strong-convergence results that are based on Gaussian interpolation along with a way to make the resolvent “sufficiently nice” is something that has been systematically developed in PHA 1, Section 4, PHA 3, Appendix B and PHA 4, Section 5. One can then use this control to prove various regularity statements by using the regularity of the corresponding free objects (see PHA 4, Section 6) or prove convergence of various Lipschitz functions of the iterates to one-dimensional SDEs (see PHA 1, Section 5).
The particular structure of the sum inside the resolvent is based on the Hessian of a free-energy – the free energy has an internal-energy term and an entropy term that which encodes the underlying target geometry. Therefore, the \(\mathcal{F}\) matrix can be thought of as a Hessian of the energy and the matrix \(M\) is the Hessian of an entropy term. The projection onto the algebra generated by \(M\) asks for the purely entropic contribution of the dynamics at any time \(t\) and point \(\hat{m}_t \in \mathcal{D}\).
FOOTNOTES
-
This seems cyclical, since we are starting at a sample drawn from the measure we wish to eventually sample from. However, it will turn out that the dependency between the “tilted” measures along the localization process at every time \(t\) and the input randomness can be decoupled – this is because, at high-temperature (\(\beta < 1\)), our measure turns out to be contiguous with respect to drawing the initial sample \(x_0 \sim \mathsf{Unif}\left(\lbrace-1,1\rbrace^n\right)\) and then running the process, albeit with a “plant” term that independently gets added to the input randomness. This obviously changes the structure of our surrogate covariance \(\hat{Q}(\cdot)\), but the added technical burden can be dealt with and is substantially easier that dealing with a localization process where the tilted measures cannot be decoupled (made independent) from the randomness of the instance. See [Section 2, EAMS’22] for more details about the planted model. ↩
-
This is in fact a reparameterization of diffusion models, as worked out in Montanari, 2023. ↩
-
For the particular linear-tilt scheme that we are using as our localization process, we will be able to assert that the algorithmic process we use allows us to efficiently estimate \(\hat{m}_t\) from \(\hat{y}_t\). ↩
-
Once again, this is only meaningfully true when the TAP free energy is strongly convex everywhere inside \([-1,1]^n\). This ceases to be true when \(1/\sqrt{2} < \beta < 1\) and we will need to “regularize” our choice of \(\hat{Q}\) for that situation – this is one of the key conceptual upgrades in PHA 4. We will discuss this in the final blog post. ↩
-
Note that even having a sampler with \(W_2\)-error of order \(O_{T,\beta}(1)\) is an improvement over the prior state of the art result, if it works for the entire replica-symmetric regime of \(0 <\beta< 1\). As we shall see, one only needs to argue that the PHD process stays close to the ASL-TAP process, which is derived in (1.3), to obtain \(O_{T,\beta}(1)\) error. Jarzynski’s equality only enters as a final step, to further decrease the error to \(o(1)\). ↩
-
Note that the fact that \(\hat{Q}(\cdot)\) must be \(C\)-Lispchitz is not immediately implied by the upper bound on the Loewner order, but requires using the resolvent identity and definition of the exact choice of \(\hat{Q}(\cdot)\) used in the algorithmic process. ↩
-
It will shortly become clear why I am calling this a “PHD-TAP drift error”. ↩
-
Jarzynski’s equality actually gives a Radon-Nikodym derivative of path space measures. Since we only care about the marginal distribution at time \(T\), we take the expectation conditional on the endpoint. ↩