Add documentation for Stochastic Gradient Samplers #629
Conversation
AoifeHughes
commented
Aug 4, 2025
AoifeHughes
commented
- Docs on SGHMC / SGLD? Turing.jl#2270 adds docs to support this
|
Preview the changes: https://turinglang.org/docs/pr-previews/629 |
There was a problem hiding this comment.
I can comment on the clarity of explanations, but I can't comment on some of the content, most importantly the Summary section, because I know nothing about these samplers. E.g. the recommendations for hyperparameters, I have no idea about them. @yebai, who would be a good reviewer for that?
| # Define a simple Gaussian model | ||
| @model function gaussian_model(x) | ||
| μ ~ Normal(0, 10) | ||
| σ ~ truncated(Normal(0, 5), 0, Inf) |
There was a problem hiding this comment.
| σ ~ truncated(Normal(0, 5), 0, Inf) | |
| σ ~ truncated(Normal(0, 5); lower=0) |
The Inf version causes trouble with AD, see JuliaStats/Distributions.jl#1910. We are trying to guide users towards the kwargs lower and upper.
|
|
||
| ```{julia} | ||
| #| output: false | ||
| setprogress!(false) |
There was a problem hiding this comment.
This needs to be moved up, or replaced with progress=false in the sample call. Currently the above cell still produces loads of lines of progress output that don't render nicely: https://turinglang.org/docs/pr-previews/629/usage/stochastic-gradient-samplers/
| ``` | ||
|
|
||
| ```{julia} | ||
| plot(chain_sgld) |
There was a problem hiding this comment.
The results on https://turinglang.org/docs/pr-previews/629/usage/stochastic-gradient-samplers/ don't look convincing to me, it looks like sampling hasn't converged. Can we increase sample counts without it taking too long? Or it could be a problem with some hyperparameters, I wouldn't know.
| ``` | ||
|
|
||
| ```{julia} | ||
| plot(chain_sghmc) |
| summarystats(chain_hmc) | ||
| ``` | ||
|
|
||
| Compare the trace plots: |
There was a problem hiding this comment.
Could we comment on the conclusions from this, what do we learn from this comparison? Also, the first trace plot looks weird.
|
|
||
| ### When to Use Stochastic Gradient Samplers | ||
|
|
||
| - **Large datasets**: When full gradient computation is prohibitively expensive |
There was a problem hiding this comment.
Isn't this in contradiction with the statement below that with Turing full gradients are computed anyway, and noise is added?
| Pkg.instantiate(); | ||
| ``` | ||
|
|
||
| Turing.jl provides stochastic gradient-based MCMC samplers that are designed for large-scale datasets where computing full gradients is computationally expensive. The two main stochastic gradient samplers are **Stochastic Gradient Langevin Dynamics (SGLD)** and **Stochastic Gradient Hamiltonian Monte Carlo (SGHMC)**. |
There was a problem hiding this comment.
The first sentence seems to be immediately undermined by the next paragraph that says that you can't actually use them for this purpose. Maybe better to lead with what they are currently useful for and then comment on possible future uses on if we ever get to implementing these better, rather than the other way around.
| @@ -0,0 +1,219 @@ | |||
| --- | |||
There was a problem hiding this comment.
This is a general comment, not related to the line it's attached to: The navigation bar on the left needs a new link to this page, I think currently there's no way to navigate to it without knowing the URL.
| model = gaussian_model(data) | ||
| ``` | ||
|
|
||
| SGLD requires very small step sizes to ensure stability. We use a `PolynomialStepsize` that decreases over time: |
There was a problem hiding this comment.
Do we have other options for stepsize in Turing, other than PolynomialStepsize?
|
|
||
| ## Automatic Differentiation Backends | ||
|
|
||
| Both samplers support different AD backends: |
There was a problem hiding this comment.
This could link to the AD page in our docs for more information.
* Updated theme-colors to match main site for consistency * fixed search results color in dark mode * fix copy button css in dark mode * search bar background udpdate * removed current default footer and added custom one * Add custom footer and update styles to match TuringLang/turinglang.github.io#119 * Update styles to match original site * cleanup code * Added SCSS styles to match main site * Add all icons in navbar + match few tweaks with main PR * Enable Open Graph and Twitter Cards for SEO * fix corrupted png * remove old styles --------- Co-authored-by: Penelope Yong <[email protected]>
* Fix external sampler docs * Remove MCHMC as a dep * update * Explain docs in more detail * Bump to 0.39.9
…to add-stoc-docs
|
https://turinglang.org/docs/pr-previews/629/usage/stochastic-gradient-samplers/ - renders okay at least. Looking into the convergence things atm |
|
At the extent of my knowledge on these, some visual things I dont understand in the final figure, not sure why it's not converging properly. Happy to make changes if someone can direct what is needed |
|
I don't know why HMC is having such trouble with this quite simple model, the numerical integration errors just blow up quite often. Seems that you can fix that with a decent initial value for the chain though. Try adding the keyword argument We should set the same |
|
Thanks @AoifeHughes and @mhauru. Quick comments:
In conclusion, it might be better if we close this PR and revisit in the future. |
|
I wish this had come up earlier so @AoifeHughes could have avoided putting in effort trying to make the examples in this PR work. I would still advocate for merging a version of this, even if it only explains that we have implementations of SGLD and SGHMC in principle included in the codebase, but that they don't actually do anything useful at the moment because we can't compute stochastic gradients / don't have minibatching. Hence this is only useful for research needs until it gets improved in the future. Some poor user might otherwise go through the same exercise that Aoife and I have been through here, of slowly understanding that there is little point to our current SGLD/SGHMC implementations, and waste time. |
|
Sorry for not being able to provide feedback earlier.
EDIT: I could have misremembered a conversation with Aoife, in which she asked me a question about this issue. |
|
Maybe we could put some of these conclusions in the Turing docstrings? |
