svi_fused.fut

-- Stochastic VI using Messaging with Algebraic Effect Handlers.
-- See van de Meent et al. (2021).
import "distributions"
import "nn"

module rng_engine = pcg32
module normal_dist = normal_distribution f32 rng_engine
module normal_diag = normal_diagonal f32 rng_engine
module bernoulli = bernoulli f32 rng_engine
module nn = nn f32

def mean [n] (xs: [n]f32) = f32.sum xs / f32.i64 n

-- Expression used in computing the unified ELBO gradient estimator from
-- "Deep Amortized Inference for Probabilistic Programs", Ritchie et al. 2016.
--   Handles both reparameterizable and non-reparameterizable variables. (All
-- discrete random variables are non-reparameterizable).
def elbo log_q_base log_w: f32 =
  -- Refer to Equation 5 in Ritchie et al. 2016 and read importance_sampling
  -- and the handlers. If all random variables are reparameterized, then
  -- log_q_base does not depend on phi meaning grad(log_q_base) is zero and so
  -- the likelihood_ratio term is also zero. If, on the other hand, all random
  -- variables are non-reparameterizable, then log_w is zero in a less obvious
  -- way using identity E_f[grad(log f(x))] = 0; see Appendix B.2 in the
  -- reference.
  --
  -- TODO variance reduction techniques for discrete RVs, see Ritchie 2016.
  let likelihood_ratio = log_q_base * (stop_gradient log_w)
  let pathwise = log_w
  in likelihood_ratio + pathwise


-- The gradient of a function whose codomain is the reals.
def grad 'a (f: a -> f32) (primal: a) = vjp2 f primal (f32.i32 1)

-- Stochastic variational inference.
-- `model` and `guide` are black-box stochastic functions.
def svi [num_batch][batch_sz][n2]
        num_epochs rng model_and_guide f g
        (yss: [num_batch][batch_sz][n2]f32) =
  let lr = 0.001
  let (f_init, f_upd, f_apply) = f
  let (g_init, g_upd, g_apply) = g
  let rngs = rng_engine.split_rng 3 rng
  let (theta, phi) = (f_init nn.normal rngs[0], g_init nn.normal rngs[1])

  let model_and_guide = model_and_guide f_apply g_apply

  let rngsss = rng_engine.split_rng (num_epochs*num_batch*batch_sz) rngs[2]
               |> unflatten_3d num_epochs num_batch batch_sz

  let elbo_loss rngs ys (theta, phi) =
      let elbos = map2 (\rng y ->
        let (log_w, log_q_base) = model_and_guide phi theta rng y
        in elbo log_q_base log_w
      ) rngs ys
      in - mean elbos

  let sgd (w: f32) (grad_w: f32) = w - lr * grad_w

  let epoch theta phi rngss yss =
    foldl (\(_, theta, phi) (rngs, ys) ->
      -- let theta = #[trace(theta)] theta
      -- let phi = #[trace(phi)] phi
      -- Sample the guide and model programs, computing varational objective
      -- w.r.t. both theta and phi.
      let (loss, (grad_theta, grad_phi)) = grad (elbo_loss rngs ys) (theta, phi)
      -- let grad_theta = #[trace(grad_theta)] grad_theta
      -- let grad_phi = #[trace(grad_phi)] grad_phi
      let loss = #[trace(loss)] loss
      in (loss, f_upd sgd theta grad_theta, g_upd sgd phi grad_phi)
    ) (0f32, theta, phi) (zip rngss yss)
  in loop (loss, theta, phi) = (0f32, theta, phi) for rngss in rngsss do
       epoch theta phi rngss yss


def normal (means, stddevs) =
  normal_diag.mk_dist means stddevs

def LATENT_DIM = 100i64
def HIDDEN_DIM = 400i64
def INPUT_DIM = 784i64
def encoder =
  nn.(dense INPUT_DIM HIDDEN_DIM >- softplus
      `split` (dense HIDDEN_DIM LATENT_DIM,
               dense HIDDEN_DIM LATENT_DIM >- elementwise f32.exp))
def decoder =
  nn.(dense LATENT_DIM HIDDEN_DIM >- softplus >- dense HIDDEN_DIM INPUT_DIM >- sigmoid)

--
-- Fuse importance_sampling, model and guide.
--

type^ dist [n] 'x = (distribution [n] rng_engine.rng x f32)

-- Sample equivalent.
def propose [n] rng base_log_prob log_w
                (d_guide: dist [n] f32)
                (d_model: dist [n] f32) =
  -- 1. sample according to guide
  let (rng, x) = d_guide.sample rng
  -- 2. save log_prob_base for guide
  -- let base_log_probs = base_log_probs with [a] = d_guide.log_prob_base c
  let base_log_prob = base_log_prob + d_guide.log_prob_base x
  -- 3. transform and save log_prob for guide and model
  let x = d_guide.transform x
  let log_w = log_w + d_model.log_prob x - d_guide.log_prob x
  in (rng, x, base_log_prob, log_w)

-- Observe equivalent.
def score [n] log_like log_w (d_model: dist [n] f32) obs =
  let log_like = log_like + d_model.log_prob obs
  let log_w = log_w + d_model.log_prob obs
  in (log_like, log_w)

-- Manually fused importance sampling of model and guide programs.
-- Uses the fact that guide does not sample RVs not present in the model
-- (the other way is always the case; any RV in model is sampled in guide).
def p_and_q decoder encoder phi theta rng y =
  let log_like = 0
  let log_q_base = 0
  let log_w = 0

  let params_model = (replicate LATENT_DIM 0f32, replicate LATENT_DIM 1f32)
  let params_guide = encoder phi y
  let (_rng, z, log_q_base, log_w) =
    propose rng log_q_base log_w (normal params_guide) (normal params_model)
  let y_probs = decoder theta z
  let (_log_like, log_w) =
    score log_like log_w (bernoulli.mk_dist y_probs) y
  in (log_w, log_q_base)


-- TODO compiler error when also mnist main present
-- def main num_epochs num_batch =
--   let rngs = rng_engine.rng_from_seed [0] |> rng_engine.split_rng 2
--   -- Unknown data-generating distribution that we wish to learn:
--   let batch_sz = 1
--   let probs =
--     tabulate (num_batch*batch_sz*INPUT_DIM) (\i -> if i%2 == 0 then 0.01 else 0.99)
--   let y = ((bernoulli.mk_dist probs).sample rngs[0]).1
--   let y = #[trace(y______)] (unflatten_3d num_batch batch_sz INPUT_DIM y)
--   let (loss, theta, _phi) = svi num_epochs rngs[1] p_and_q decoder encoder y
--   let (_, z) = (normal (replicate LATENT_DIM 0f32, replicate LATENT_DIM 1f32)).sample rngs[1]
--   in (loss, z, decoder.2 theta z, y[0,0])

def to_img (x: []f32) = intrinsics.unflatten 28 28 (map i8.f32 x) :> [28][28]i8

def binarize rng probs =
  let valid_probabilities = all (\x -> 0 <= x && x <= 1) probs
  let imgs = assert valid_probabilities probs
  in ((bernoulli.mk_dist probs).sample rng).1

def main [n][w][h] epochs batch_sz (imgs: [n][w][h]f32) =
  let num_batch = assert (n % batch_sz == 0) (n / batch_sz)
  let rngs = rng_engine.rng_from_seed [0] |> rng_engine.split_rng 2
  -- Treat pixels as iid Bernoulli random variables.
  let y = flatten_3d imgs :> [num_batch * batch_sz * INPUT_DIM]f32
  let y = binarize rngs[0] y
          |> unflatten_3d num_batch batch_sz INPUT_DIM
  let (loss, _theta, _phi) = svi epochs rngs[1] p_and_q decoder encoder y
  let (_, z) = (normal (replicate LATENT_DIM 0f32, replicate LATENT_DIM 1f32)).sample rngs[1]
  in loss

-- ==
-- entry: mnist_epoch1_batchsz125 mnist_epoch1_batchsz160 mnist_epoch1_batchsz15000
-- compiled input @ mnist-f32
entry mnist_epoch1_batchsz125 = main 1 125
entry mnist_epoch1_batchsz160 = main 1 160
entry mnist_epoch1_batchsz15000 = main 1 15000