Thompson Sampling for the Price of an SGD Step

Tue, 13 Oct 2026 00:20:00 +0000

Exploration at Scale · Part 1 of 2. Part 2 shows why d log T particles are enough and what carries over to neural networks. Based on Ensemble++ (NeurIPS 2025) and HyperAgent (ICML 2024).

Thompson sampling sends the learner where plausible models disagree, and stops exploring on its own once they agree.
Its price is a fresh posterior sample every round, which needs a matrix square root even in a linear model and has no closed form in a network.
Build the randomness into the model: a random index as a second input, so that a sample is one forward pass.
In a linear model the index weights M stored particles, and one incremental step per data point keeps the cloud on the posterior. Memory is the particles, nothing else.

The figures step through like slides. Use → or ▶, and step back with ←.

1. Thompson sampling turns uncertainty into exploration

Sample one plausible model, act on it, observe, update: the agent goes where the models disagree, as often as they say it might pay off.

value

options x

value

options x

value

options x

value

options x

12 / 12

After 40 rounds

Greedy stays at the known peak · value 0.68 per round

Undirected noise: epsilon-greedy random trials everywhere · value 0.79

Thompson sampling finds the best option, then stays · value 1.15

It found the best option and stopped exploring on its own. The price: a fresh sample from an updated posterior, every round.

A greedy learner plays the best option of the average model and stays at the peak it already knows. Undirected noise, such as ε-greedy, a sampling temperature or an entropy bonus, spends trials everywhere, including on options already known to be bad, and keeps paying for them for as long as it runs. Thompson sampling needed no rule to stop exploring: as the posterior narrowed, the samples stopped disagreeing. Over 200 runs it earns 1.15 per round, against 0.79 for ε-greedy and 0.68 for greedy.

2. With a model, the posterior sample becomes the bottleneck

Every round changes the posterior, so every round needs a new sample, and that sample is the expensive step.

Step	Linear model	Neural network
1 · Sample one plausible model	Closed form, but each sample needs a matrix square root, cubic in the dimension	No closed-form posterior to sample from
2 · Act on it	Best option under the sample	Best option under the sample
3 · Observe	A noisy response	A noisy response
4 · Update	One recursive step per data point	One SGD step per data point

The update is cheap in both cases; the sample is not. What I want is a posterior sample for the price of one update step.

3. The idea: a model that maps simple noise to posterior samples

Give the model a second input, a random index ξ, and train it so that simple noise is pushed forward onto the posterior.

6 / 6

A sample is one forward pass with a fresh index; an update is one step on the weights, and it moves every sample at once. The work is to train ν so that the noise lands on the posterior.

The model is $f_\nu(x, \xi)$, with one set of weights $\nu$ for every index. Sampling then costs one forward pass with a fresh index, and one update step moves every sample at once. In a linear model, where the exact posterior is known and every step can be checked, each sample is a point in parameter space:

functions of x: one curve per sample

parameter space

ellipse: the posterior over the parameter. Black dot: its mean μ

6 / 6

The model maps simple noise ξ ~ N(0, I) to a parameter μ + Aξ, which gives a function ⟨φ(x), θ̃(ξ)⟩. Its weights ν are the mean μ and M stored directions, the columns of A.

$$ \tilde\theta(\xi) = \mu + \mathbf{A}\,\xi, \qquad \xi \sim N(0, I_M), \qquad f_\nu(x,\xi) = \big\langle \phi(x),\ \tilde\theta(\xi)\big\rangle, \qquad \nu = (\mu, \mathbf{A}). $$

The weights are a mean and M stored directions. The rest of this post builds them.

4. A cloud of particles moves with every data point

Store M particles; one cheap step per data point, with noise we inject ourselves, keeps the cloud on the exact posterior.

prior

Before any data: the prior

prior: posterior before any data

M prior draws, kept: the particles

the arrows, side by side

the prior draws, as a table: column m is particle m

prior

Column m: the two coordinates of particle m.

prior

One random index, one sample

the index ξ: one weight per arrow

simple noise, drawn fresh for every sample

many indices, many samples: they fill the prior

the prior draws, as a table: column m is particle m

prior

Ensemble++: M stored vectors, one random index per sample.

prior

Why it is a prior sample

since ξ ~ N(0, I) is Gaussian

so this is all we need

same spread along a direction v

true in every direction: dashed ≈ solid

the prior draws, as a table: column m is particle m

prior

A random projection that preserves lengths. M of order d is enough: Johnson–Lindenstrauss.

posterior

Data: a linear-Gaussian model

the truth: a point in this plane (cross)

the goal: the cloud follows the posterior

the arrows at time t, from the mean

the prior draws, as a table: column m is particle m

prior

Wanted: one cheap step per data point keeps it true.

A data point arrives

a data point: input (arrow) and response

parameters that fit it exactly (dashed)

environment noise: the gap to the truth

input given here; later the agent picks it

the prior draws, as a table: column m is particle m

prior

Inject noise

injected noise: ours, one draw per particle

particle m’s own copy (thin lines)

our own noise: we add it, and we know it

prior draws, then injected noise: row t is a data point, column m is a particle

prior

New row: the 16 thin lines.

One step each

each particle steps toward its own line
the cloud now matches the exact posterior
the injected noise gives it the right width

prior draws, then injected noise: row t is a data point, column m is a particle

prior

Column m: everything that makes particle m different.

A second data point

a new row of injected noise
earlier rows are never redrawn
no old data kept or revisited, no matrix factored

prior draws, then injected noise: row t is a data point, column m is a particle

prior

One more row. Earlier rows are untouched.

enlarged ×2.5

After ten

collapsed around the truth
spread still matches the posterior (inset)
nothing reweighted, nothing resampled

prior draws, then injected noise: row t is a data point, column m is a particle

prior

Memory: the M particles, nothing else. Sampling has become incremental learning.

18 / 18

After ten: the cloud has collapsed around the truth and its spread still matches the posterior. Memory is the M particles, nothing else.

Two moves make it work. Sampling is a random combination: a Gaussian index weighting M stored arrows is exactly Gaussian, so the only thing to get right is the covariance of the cloud. Updating is a perturbed fit: each particle regresses on its own noisy copy of every response, anchored at its own prior draw:

$$ \tilde\theta_{t,m} = \arg\min_{\theta}\ \sum_{s\le t}\Big(Y_s + Z_{s,m} - \langle\theta, X_s\rangle\Big)^2 + \big\|\theta - \tilde\theta_{0,m}\big\|^2_{\Sigma_0^{-1}}, \qquad Z_{s,m}\sim N(0,1),\quad \tilde\theta_{0,m}\sim N(\mu_0,\Sigma_0). $$

When a data point arrives, the objective gains one term, so the update is one recursive step. Nothing is reweighted or resampled, which avoids the weight degeneracy of a particle filter (the idea goes back to the ensemble Kalman filter), and no data point is stored.

Why each particle is an exact posterior sample when the data are fixed

The minimizer is linear in the Gaussian perturbations $Z_{s,m}$ and $\tilde\theta_{0,m}$, so it is Gaussian. Its mean is the ridge solution, which is the posterior mean $\mu_t$. Its covariance is $\Sigma_t\big(\Sigma_0^{-1} + \sum_s X_s X_s^{\top}\big)\Sigma_t = \Sigma_t$. Hence $\tilde\theta_{t,m} \sim N(\mu_t, \Sigma_t)$, and the cloud is M posterior draws.

5. At any time, the same recipe gives a fresh sample

Mean plus stored arrows times a random index: a posterior sample with no square root anywhere.

after three data points, zoomed in

The posterior sample

weighted arrows

the same simple noise as for the prior

since ξ ~ N(0, I), as for the prior

what must hold: dashed ≈ solid

A posterior sample to act on, if the cloud keeps the right covariance. Does it, and with how few particles?

6 / 6

A posterior sample to act on, with no square root anywhere, if the cloud keeps the right covariance. Does it, and with how few particles?

$$ \tilde\theta(\xi) = \mu_t + \mathbf{A}_t\,\xi = \mu_t + \tfrac{1}{\sqrt M}\sum_{m}\xi_m\big(\tilde\theta_{t,m}-\mu_t\big), \qquad \text{a posterior sample if } \mathbf{A}_t\mathbf{A}_t^{\top} \approx \Sigma_t . $$

This is Thompson sampling at the cost of one incremental step per data point and one index per action, with memory equal to the M particles. It rests on one property: the cloud keeps the right covariance as data arrive. Part 2 shows when that holds, why it is harder when the agent picks its own data, how few particles it takes (d log T), and what carries over to a network.

References

Y. Li, J. Xu, B. Wang, Z.-Q. Luo. Scalable Exploration via Ensemble++. NeurIPS 2025.
Y. Li, J. Xu, L. Han, Z.-Q. Luo. Q-Star Meets Scalable Posterior Sampling: Bridging Theory and Practice via HyperAgent. ICML 2024.
Z. Li, Y. Li, Y. Zhang, T. Zhang, Z.-Q. Luo. HyperDQN: A Randomized Exploration Method for Deep Reinforcement Learning. ICLR 2022.
D. Russo, B. Van Roy, A. Kazerouni, I. Osband, Z. Wen. A Tutorial on Thompson Sampling. 2018.

Citation

@inproceedings{li2025ensemble,
 title = {Scalable Exploration via Ensemble++},
 author = {Li, Yingru and Xu, Jiawei and Wang, Baoxiang and Luo, Zhi-Quan},
 booktitle = {Advances in Neural Information Processing Systems},
 year = {2025}
}

Uncertainty | Yingru Li

Thompson Sampling for the Price of an SGD Step

1. Thompson sampling turns uncertainty into exploration

2. With a model, the posterior sample becomes the bottleneck

3. The idea: a model that maps simple noise to posterior samples

4. A cloud of particles moves with every data point

5. At any time, the same recipe gives a fresh sample

References

Citation