variational approach to concentration

A technique for multivariate concentration. Let $(S_t)$ be some stochastic process, say in ${\mathbb{R}}^d$. For instance, $S_n={\sum }_{i=1}^n{X}_i$ for multivariate observations $X_i\in {\mathbb{R}}^d$. We are aiming to generate a high probability bound on

$$\left|S_n\right|=\underset{v:|v|=1}{\sup }⟨v,S_n⟩.$$

The idea is to use a PAC-Bayes approach (which is itself based on variational inequalities, hence the name), in order to simultaneously bound $⟨v,S_t⟩$ in each direction $v$. Recall that a PAC-Bayes bound has the form

$$\mathbb{P}(\forall \rho \in \mathcal{M}(\Theta ):\text{Something holds})\ge 1-\delta .$$

That is, a PAC-Bayes bound provides a high probability bound simultaneously over all posteriors. The variational approach to concentration translates this into a high probability bound over all directions.

This approach was pioneered by Catoni and Giulini (here and here), and has now been used by a few authors to prove bounds in a variety of settings:

• Zhivotovskiy for bounding the singular values of random matrices
• Nakakita et al. for bounding the mean of high-dimensional random matrices under heavy-tails;
• Giulini for estimating the Gram operator in Hilbert spaces.
• Myself and others for estimating the mean of random vectors.

Recall a very general PAC-Bayes inequality:

Master theorem

Let $M_t(\theta )$ be a nonnegative supermartingale with initial value 1 for all $\theta \in \Theta$. Let $\nu$ be a data-free prior. Then,

$$\mathbb{P}(\forall t,\forall \rho \in \mathcal{M}(\Theta ):{\mathbb{E}}_{\rho }\log M_t(\theta )\le D_{\text{KL}}(\rho \Vert \nu )+\log (1/\delta ))\ge 1-\delta .$$

This is the time-uniform version of the master theorem, but we can also state a master fixed-time version. This reads: Let $N(\theta )$ be nonnegative have expected value at most 1 (it is an e-value) for all $\theta \in \Theta$. Then

$$\mathbb{P}(\forall \rho \in \mathcal{M}(\Theta ):{\mathbb{E}}_{\rho }\log N(\theta )\le D_{\text{KL}}(\rho \Vert \nu )+\log (1/\delta ))\ge 1-\delta .$$Link to original

To see how this works, let’s consider two use cases.

Example 1: Sub-Gaussian random vectors

This comes from our paper on time-uniform confidence spheres. Consider $n$ iid copies $X_1,\dots ,X_n$ of a $\Sigma$-sub-Gaussian random vector $X\in {\mathbb{R}}^d$ (see sub-Gaussian distributions). That is,

$$\mathbb{E}\exp (\lambda ⟨\theta ,X⟩)\le \exp \left(\frac{{\lambda }^2}{2}⟨\theta ,\Sigma \theta ⟩\right),$$

for all $\lambda \in \mathbb{R}$ and $\theta \in {\mathbb{R}}^d$. This implies that

$$N(\theta )=\exp \left\{\lambda \sum _{i\le n}⟨\theta ,X_i⟩-\frac{n{\lambda }^2}{2}⟨\theta ,\Sigma \theta ⟩\right\},$$

has expectation at most 1 (i.e., it is an e-value). Let $\nu$ be a Gaussian with mean 0 and covariance ${\beta }^{-1}I$ for some $\beta >0$. Consider the family of distributions $\{{\rho }_u:\left|u\right|=1\}$ where ${\rho }_u$ is a Gaussian with mean $u$ and covariance ${\beta }^{-1}I$. Then the KL divergence between ${\rho }_u$ and $\nu$ is $D_{\text{KL}}({\rho }_u\Vert \nu )=\beta /2$. Using the master theorem above, we obtain that, with probability $1-\delta$, simultaneously for all distributions $\rho$,

$$\lambda \sum _{i\le n}{\mathbb{E}}_{\rho }⟨\theta ,X_i⟩\le \frac{n{\lambda }^2}{2}{\mathbb{E}}_{\rho }⟨\theta ,\Sigma \theta ⟩+\frac{\beta }{2}+\log (1/\delta ).$$

Now, for $\rho ={\rho }_u$, ${\mathbb{E}}_{\rho }⟨\theta ,X_i⟩=⟨u,X_i⟩$ and

$${\mathbb{E}}_{\rho }⟨\theta ,\Sigma \theta ⟩=⟨u,\Sigma u⟩+{\beta }^{-1}\text{Tr}(\Sigma )\le \left|\Sigma \right|+{\beta }^{-1}\text{Tr}(\Sigma ),$$

using basic properties of the expectation of quadratic forms under Gaussian distributions (see eg here), and definition of the operator norm as $\left|A\right|={\sup }_{u,v:|u|=|v|=1}⟨u,\Sigma v⟩$. Since this holds simultaneously for all ${\rho }_u$, we obtain that, with probability $1-\delta$,

$$\underset{u}{\sup }\lambda \sum _{i\le n}⟨u,X_i⟩\le \frac{n{\lambda }^2}{2}(\left|\Sigma \right|+{\beta }^{-1}\text{Tr}(\Sigma ))+\frac{\beta }{2}+\log (1/\delta ).$$

The left hand side is equal to $\lambda \left|{\sum }_{i\le n}{X}_i\right|$, which gives us our concentration result. One can then optimize $\lambda$ using some calculus.

Example 2: Random matrices with finite Orlicz-norm

This example is adapted from Zhivotovskiy (2024). Let $M_1,\dots ,M_n$ be iid copies of a zero-mean random matrix $M$ with finite sub-exponential Orlicz norm, in the sense that, for some $C>0$,

$${\left|⟨\theta ,M\varphi ⟩\right|}_{{\psi }_1}\le C⟨\theta ,\Sigma \varphi ⟩,$$

for all $\theta ,\varphi \in {\mathbb{R}}^d$ where $Q=\mathbb{E}M$.

Here we highlight the main ingredients that are used in Zhivotovskiy’s results. For the details, see his paper or my blog post.

We take our parameter space in the master theorem above to be $\Theta ={\mathbb{R}}^d\times {\mathbb{R}}^d$. Let $\nu$ again be Gaussian with mean 0 and covariance ${\beta }^{-1}\Sigma$ and let ${\mu }_u$ be a truncated Gaussian mean $u$, covariance ${\beta }^{-1}\Sigma$ and radius $r$. For a vector $u\in {\Sigma }^{1/2}{\mathbb{S}}^{d-1}$, the KL-divergence between a truncated normal ${\mu }_u$ and $\nu$ is $\log (1/Z)+\beta /2$, where $Z=\mathbb{P}(\left|\theta -u\right|\le r)$ where $\theta \sim {\rho }_u$. Equivalently, $Z=\mathbb{P}(\left|Y\right|\le r)$ where $Y$ is a normal with mean $0$ and covariance ${\beta }^{-1}\Sigma$. Hence $1-Z=\mathbb{P}(\left|Y\right|>r)\le \mathbb{E}{\left|Y\right|}^2/r^2={\beta }^{-1}\text{Tr}(\Sigma )/r^2$. Thus, taking $r=\sqrt{2{\beta }^{-1}\text{Tr}(\Sigma )}$ yields $Z\ge 1/2$, and we obtain

$$\begin{array}{r}{D}_{\text{KL}}({\mu }_u\Vert \nu )\le \log (2)+\frac{\beta }{2}.\end{array}$$

Now it remains to construct a relevant quantity to use in the PAC-Bayes theorem. Consider

$$N(\theta ,\varphi )=\exp \left\{\lambda \sum _{i\le n}⟨\theta ,{\Sigma }^{-1/2}{M}_i{\Sigma }^{-1/2}\varphi ⟩-n\log \mathbb{E}\exp (\lambda ⟨\theta ,{\Sigma }^{-1/2}M{\Sigma }^{-1/2}\varphi ⟩)\right\},$$

where the expectation is over $M$. It’s easy to see this has expectation at most 1 (it can be written as the product of terms each with expectation exactly one). Apply the master theorem with the product distribution ${\mu }_u\times {\mu }_v$ for $u,v$ in the unit sphere. Then, to bound the right hand side of the process, use one of the exponential inequalities. In particular, for a random variable $Y$,

$$\mathbb{E}[\exp (\lambda (Y-\mathbb{E}Y))]\le \exp (4{\lambda }^2{\left|Y-\mathbb{E}Y\right|}_{{\psi }_1}),\forall |\lambda |\le \frac{1}{2{\left|Y-\mathbb{E}Y\right|}_{{\psi }_1}}.$$

A bunch of algebra then gives the result.