Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

composing lmh and smc #17

Open
nilsbecker opened this issue Mar 7, 2024 · 3 comments
Open

composing lmh and smc #17

nilsbecker opened this issue Mar 7, 2024 · 3 comments
Labels
enhancement New feature or request

Comments

@nilsbecker
Copy link

nilsbecker commented Mar 7, 2024

disclaimer: this may be just me not seeing a solution

i'm currently trying to build an inference scheme combining lmh within smc. i believe this is similar to what's called rm-smc in this paper, cited in the docs.

roughly it goes like this:

  • in an outer (Smc_inference) iteration, a population of particles is (custom-)resampled in a number of steps; each time the target likelihood is modified.
  • between any two reweighting steps, particles move according to a metropolis-hastings mcmc, which is meant to achieve decorrelation between particles ('mixing', 'rejuvenation'). as each particle represents a stochastic model, i would like to use Lmh_inference to do for the mixing phases.

i'm facing the following difficulty: i don't know how to carry the mcmc state of a particle through the resampling step, and then continue Lmh sampling from that same state -- but with scoring functions that have changed after the Smc resampling. what i've tried is to pass an Lmh_inference.t computation as particle state across the resampling. but as far as i could see, when i define an Lmh_inference.t computation, calls to Lmh_inference.score within it are already fixed and i have to way to modify them after a resampling. this would give no way to update the likelihood on the fly. is that correct?

as i'm writing this i could imagine parametrizing the log-likelihood to score with using mutable state, e.g. references, which i would have to manage manually between smc iterations. that seems clunky and potentially thread-unsafe.

in the paper mentioned above there appears to be an elegant solution to this kind of composability -- does dagger also enjoy this? is there a way to do it that i don't see? do i have to switch the nesting around? ie. call Smc_inference.yield inside the Lmh model?

@igarnier
Copy link
Owner

igarnier commented Mar 8, 2024

dagger is not structured in the same way as MonadBayes, so I don't think we can just plug the LMH inference directly inside SMC though. A bit more thinking is required but having rm-smc would indeed be nice!

@igarnier igarnier added the enhancement New feature or request label Mar 8, 2024
@nilsbecker
Copy link
Author

I tried around a bit but have not found a way so far. I keep getting conflicting requirement of returning both a Smc.t and a Lmh.t It almost seems like one would need to fuse the monads together -- unclear to me how.

@nilsbecker
Copy link
Author

nilsbecker commented Mar 21, 2024

I have been trying an approach to achieve the combination of lmh and smc, but ultimately this failed. It may be interesting for reference. I try to sketch the setup in what follows:


open Dagger
open Containers

module L = struct include Lmh_inference include Infix end
module T = struct
  type particle_output = float L.t
  type resampling_state = {std:float} end
module Sm = Smc_inference.Make (T)
module S = struct include Sm include Sm.Infix end

let rng = RNG.make [|10;10;|]

let model n init =
  let open L in
  let* start = init in
  Format.printf "start %g@." start;
  let b = Stats_dist.brownian ~start ~std:0.1 in
  let rec loop m =
    let* v = sample b in
    if m <= 0
    then (Format.printf "%g@."v; return v)
    else loop (pred m) in
  loop n

let init_flat = L.sample (Stats_dist.flat 0. 10.)

let lmh_score std m =
  let open L in
  let* v = m in
  let s = Stats_dist.Pdfs.gaussian_ln ~mean:1. ~std v in
  Format.printf "score %g@." (Log_space.to_float s);
  let* () = log_score s in (* add a final scoring call to m *)
  return v

let rec outer_loop inner_iterations m rng_state outer_iterations =
  let open S in
  let* {T.std} = yield m in
  let final_output =
    let m_scored = lmh_score std m in (* decorate m with scoring *)
    L.stream_samples m_scored rng_state
    |> Seq.drop (pred inner_iterations) |> Seq.head_exn in
  let final_score = Stats_dist.Pdfs.gaussian_ln ~mean:1. ~std final_output in
  let* _score_outer = S.log_score final_score in 
  (* also score for smc (not actually checking for statistical correctness
   * here; this is just for demo purposes) *)
  if outer_iterations < 1
  then return (final_score, final_output)
  else outer_loop inner_iterations m rng_state (pred outer_iterations)


let strategy =
  (* systematic resampling with stepwise sharpening of the likelihood parameter
   * std *)
  let resample = Smc_inference.systematic_resampling ~ess_threshold:1. in
  let step_down {T.std} = 
    let std = Float.max 1. (0.5 *. std) in
    Format.printf "reducing likelihood width to %g@." std; 
    {T.std} in
  fun ~target_size m rss rng -> resample ~target_size m rss rng |> step_down

let run () =
  let model_run = outer_loop 2 (model 3 init_flat) rng 6 in
  S.run ~nthreads:1 strategy {T.std=100.} ~npart:10 model_run rng

let final_pop = (run () |> List.of_seq_rev |> List.hd).S.terminated

the idea was to set up an smc loop where the particle state is an lmh.t, which is propagated between smc iterations. i then tried to add a final score to the lmh model that depends on the smc iteration; this was possible in my use case as the scoring only depended on the final output of the lmh model; in this way i tried to implement a stepwise sharpening of the likelihood over smc runs.

however, this turned out not the way i expected. i expected that the model m:float L.t would retain its internal lmh sampling state from the previous smc iteration, so that it would stream further lmh samples from that state but with a new internal lmh weight. instead what happens is that on each new call to L.stream_samples the model m gets initialized fresh from the initial condition. this invalidates my stepwise sharpening strategy which only would make sense if the population remains pre-adapted to an intermediate sharpness from the previous iteration.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

2 participants