Skip to content

Add Autodiff compatable emulator and Barker gradient-based sampler#339

Merged
odunbar merged 33 commits intomainfrom
orad/autodiff-mcmc
Jan 31, 2025
Merged

Add Autodiff compatable emulator and Barker gradient-based sampler#339
odunbar merged 33 commits intomainfrom
orad/autodiff-mcmc

Conversation

@odunbar
Copy link
Member

@odunbar odunbar commented Jan 16, 2025

Purpose

Closes #341
Closes #342
Closes #343
Closes #344
Closes #345

Content

From a private repos of @KotaOK-00 - and with additional changes to interface

  • Copy implementation of the gradient-based sampler proposals
  • Implemented the calculation of ESJD
  • Added esjd unit test.
  • Created a new system AutodiffProtocol containing derived types GradFreeProtocol ForwardDiffProtocol and ReverseDiffProtocol. autodiff options are created with the MCMCProtocol, e.g.,
BarkerSampling() # creates BarkerSampling{ForwardDiffProtocol}()
BarkerSampling{ReverseDiffProtocol}() # creates BarkerSampling{ReverseDiffProtocol}()
pCNMHSampling() # creates pCNMHSampling{GradFreeProtocol}()

Within the MCMC propose step, we make calls to autodiff_gradient and autodiff_hessian

autodiff_gradient(model, current_state.params, sampler)
autodiff_hessian(model, current_state.params, sampler)

that dispatch off the pre-specified autodiff protocol in the sampler. e.g. for ForwardDiffProtocol this runs

ForwardDiff.gradient(x -> AdvancedMH.logdensity(model, x), current_state.params)
Symmetric(ForwardDiff.hessian(x -> AdvancedMH.logdensity(model, x), current_state.params))

This will make adding new autodiff options easier

  • Created a new type AGPJL for AbstractGPs. When building an emulator with AGPJL, As there is no optimizer of such kernels, we allow users to train e.g. with GPJL, then do the following:
opt_params_per_model = Emulators.get_params(gp_optimized)
kernel_params = [
Dict(
       "log_rbf_len" => params[1:end-2],
       "log_std_sqexp" => params[end-1],
       "log_std_noise" => params[end],
) 
for params in opt_params_per_model]
agp = GaussianProcess(AGPJL(), ...)
Emulator( agp, ... ; ..., kernel_params = kernel_params) # builds

If the user does not provide kernel_params, a useful error message appears to direct them what to do.

  • Unit tests in GaussianProcesses/runtests.jl for AGP interface as above, and that GPJL vs AGP gives very similar emulator predicitions
  • Unit tests in MarkovChainMonteCarlo/runtests.jl for AGP on RWM and pCN sampling - give similar posterior means to GPJL (NB it is much slower, 5-10x)
  • Unit testing in MarkovChainMonteCarlo/runtests.jl for Barker update (ForwardDiff and ReverseDiff) with comparable results to RWM.

Emulators produce similar outputs.

[ Info: testing algorithm: RWMHSampling{GradFreeProtocol}
[ Info: ESJD = [0.03796417690519561, 0.21983371208939942, 0.4385500000002978]
[ Info: 1.4149419254957858
...
[ Info: testing algorithm: BarkerSampling{ForwardDiffProtocol}
[ Info: ESJD = [0.0202305530376406, 0.07430262173309042, 0.4598200000003191]
[ Info: 1.384449605016855
...
[ Info: testing algorithm: BarkerSampling{ReverseDiffProtocol}
[ Info: ESJD = [0.02262425874003432, 0.07044827968512407, 0.4570100000003163]
[ Info: 1.333019873088517

MISC

  • added non-Barker implementation as new issue. I have removed these algorithms as they do not perform well in unit testing and require possibly significant improvements to the tuning mechanisms to work
  • Mac is slow! In unit tests, for sampling with (AbstractGP+ForwardDiff) Mac is over 10x slower, while ReverseDiff Mac is ~5x slower

  • I have read and checked the items on the review checklist.

@odunbar odunbar changed the title MCMC updates (gradient-based variants and ESJD) [WIP] MCMC updates (gradient-based variants and ESJD calculation with the chain) Jan 16, 2025
@odunbar odunbar changed the title [WIP] MCMC updates (gradient-based variants and ESJD calculation with the chain) [WIP] MCMC updates (gradient-based variants and ESJD) Jan 16, 2025
@codecov
Copy link

codecov bot commented Jan 16, 2025

Codecov Report

Attention: Patch coverage is 93.39623% with 7 lines in your changes missing coverage. Please review.

Project coverage is 88.96%. Comparing base (ba5ba3f) to head (b20309c).
Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
src/MarkovChainMonteCarlo.jl 87.03% 7 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #339      +/-   ##
==========================================
+ Coverage   88.82%   88.96%   +0.13%     
==========================================
  Files           7        7              
  Lines        1271     1359      +88     
==========================================
+ Hits         1129     1209      +80     
- Misses        142      150       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@odunbar odunbar changed the title [WIP] MCMC updates (gradient-based variants and ESJD) Add Forward-diff emulator and Barker gradient-based sampler Jan 22, 2025
@odunbar odunbar changed the title Add Forward-diff emulator and Barker gradient-based sampler Add Autodiff compatable emulator and Barker gradient-based sampler Jan 22, 2025
@odunbar odunbar force-pushed the orad/autodiff-mcmc branch from dd0c26a to 18196d1 Compare January 24, 2025 23:09
@odunbar odunbar requested a review from KotaOK-00 January 28, 2025 18:34
@odunbar odunbar merged commit 2ad4cc7 into main Jan 31, 2025
9 checks passed
@odunbar
Copy link
Member Author

odunbar commented Jan 31, 2025

Merging as all tests pass and coverage is high

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment