composing lmh and smc

[nilsbecker]

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]

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!

[nilsbecker]

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]

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.