Adam scale adapter implementation (#132)

* Add Adam-based scale adapter (#130)

* Add tests for adam_scale_adapter

* Update namespace and manual files associated with adam_scale_adapter

* Clarify scale_adapter() '...' documentation

Fix inaccurate docstring that the wrapper only exposes learning_rate for the Adam adapter. In fact, '...' passes through to the underlying adapter constructor, so any of its arguments can be set via the wrapper. Addresses review comment on #132.

* Update adaptation.R to correct the '...' documentation for scale_adapter, and update the manual file

Author: jennajiali

NAMESPACE

@@ -1,5 +1,6 @@
 # Generated by roxygen2: do not edit by hand
 
+export(adam_scale_adapter)
 export(barker_proposal)
 export(bimodal_barker_proposal)
 export(chain_state)

R/adaptation.R

@@ -3,12 +3,20 @@
 #' @param algorithm String specifying algorithm to use. One of:
 #'   * "stochastic_approximation" to use a Robbins-Monro (1951) based scheme,
 #'   * "dual_averaging" to use dual-averaging scheme of Nesterov (2009).
+#'   * "adam" to use the Adam optimizer of Kingma and Ba (2014) applied to the
+#'     acceptance-rate residual, following the implementation in the `walnuts`
+#'     library.
 #' @param initial_scale Initial value to use for scale parameter. If not set
 #'   explicitly a proposal and dimension dependent default will be used.
 #' @param target_accept_prob Target value for average accept probability for
 #'   chain. If not set a proposal dependent default will be used.
-#' @param ... Any additional algorithmic parameters to pass through to
-#'   [dual_averaging_scale_adapter()] or [stochastic_approximation_scale_adapter()].
+#' @param ... Any additional algorithmic parameters to pass through to the
+#'   selected adapter constructor: see [dual_averaging_scale_adapter()],
+#'   [stochastic_approximation_scale_adapter()] or [adam_scale_adapter()] for
+#'   the full list of parameters accepted by each. In practice, most users
+#'   tuning the Adam adapter only need to adjust `learning_rate`; the
+#'   moment-decay parameters `beta_1`, `beta_2`, `epsilon` and
+#'   `learn_rate_decay` have sensible defaults that rarely need adjustment.
 #'
 #' @return List of functions with entries
 #' * `initialize`, a function for initializing adapter state and proposal
@@ -21,12 +29,15 @@
 #' * `state` a zero-argument function for accessing current values of adapter
 #'   state variables.
 #'
-#' @seealso [dual_averaging_scale_adapter()], [stochastic_approximation_scale_adapter()]
+#' @seealso [dual_averaging_scale_adapter()],
+#'   [stochastic_approximation_scale_adapter()], [adam_scale_adapter()]
 #'
 #' @references Nesterov, Y. (2009). Primal-dual subgradient methods for convex
 #'   problems. _Mathematical Programming_, 120(1), 221-259.
 #' @references Robbins, H., & Monro, S. (1951). A stochastic approximation
 #'   method. _The Annals of Mathematical Statistics_, 400-407.
+#' @references Kingma, D. P., & Ba, J. (2014). Adam: A method for stochastic
+#'   optimization. _arXiv preprint_ arXiv:1412.6980.
 #'
 #' @export
 #'
@@ -43,6 +54,7 @@ scale_adapter <- function(
   adapter_function <- switch(algorithm,
     dual_averaging = dual_averaging_scale_adapter,
     stochastic_approximation = stochastic_approximation_scale_adapter,
+    adam = adam_scale_adapter,
     stop(sprintf("Unrecognized algorithm choice %s", algorithm))
   )
   adapter_function(initial_scale, target_accept_prob, ...)
@@ -187,6 +199,116 @@ dual_averaging_scale_adapter <- function(
   )
 }
 
+#' Create object to adapt proposal scale to coerce average acceptance rate
+#' using the Adam optimizer of Kingma and Ba (2014).
+#'
+#' Applies the Adam stochastic-gradient optimizer to the log of the proposal
+#' scale, using the acceptance-rate residual `target_accept_prob - accept_prob`
+#' as the (stochastic) gradient. This corresponds to treating
+#' `-0.5 * (accept_prob - target_accept_prob)^2` as the objective and is the
+#' same gradient signal used by [dual_averaging_scale_adapter()] and
+#' [stochastic_approximation_scale_adapter()], but plugged into the Adam
+#' update rule instead of a dual-averaging or Robbins-Monro schedule. Follows
+#' the reference implementation in the `walnuts` library (see references
+#' below).
+#'
+#' To match the stability provided by dual averaging, a learning-rate decay
+#' `learning_rate / t^learn_rate_decay` is applied to the per-iteration Adam
+#' step; a value of `learn_rate_decay = 0` recovers standard Adam, while
+#' `learn_rate_decay = 0.5` (the default here) is recommended by the `walnuts`
+#' authors for stable convergence.
+#'
+#' @inherit scale_adapter params return
+#'
+#' @param learning_rate Learning rate for Adam optimizer (the `alpha`
+#'   hyperparameter in the original Adam paper). Controls the magnitude of the
+#'   update applied to the log-scale on each iteration. Should be positive.
+#'   Defaults to `0.05`, which is a practical setting that gives fast
+#'   convergence of the acceptance-rate coercion within typical MCMC warm-up
+#'   lengths; the original Adam paper uses `1e-3`.
+#' @param beta_1 Exponential decay rate for the first-moment estimate of the
+#'   gradient (the `beta_1` hyperparameter in the original Adam paper). Should
+#'   be in `[0, 1)`. Defaults to `0.9`.
+#' @param beta_2 Exponential decay rate for the second-moment (squared
+#'   gradient) estimate (the `beta_2` hyperparameter in the original Adam
+#'   paper). Should be in `[0, 1)`. Defaults to `0.999`.
+#' @param epsilon Small positive constant added to the square root of the
+#'   second-moment estimate in the denominator of the Adam update for
+#'   numerical stability. Defaults to `1e-8`.
+#' @param learn_rate_decay Exponent controlling the decay of the effective
+#'   learning rate across iterations: on iteration `t` the effective learning
+#'   rate is `learning_rate / t^learn_rate_decay`. Should be in `[0, 1]`. A
+#'   value of `0` recovers standard Adam; the default of `0.5` matches the
+#'   value recommended in the `walnuts` implementation.
+#'
+#' @references Kingma, D. P., & Ba, J. (2014). Adam: A method for stochastic
+#'   optimization. _arXiv preprint_ arXiv:1412.6980.
+#' @references Reference implementation in `walnuts`:
+#'   <https://github.com/flatironinstitute/walnuts/blob/main/include/walnuts/adam.hpp>
+#'
+#' @export
+#'
+#' @examples
+#' proposal <- barker_proposal()
+#' adapter <- adam_scale_adapter(
+#'   initial_scale = 1., target_accept_prob = 0.4
+#' )
+#' adapter$initialize(proposal, chain_state(c(0, 0)))
+adam_scale_adapter <- function(
+  initial_scale = NULL,
+  target_accept_prob = NULL,
+  learning_rate = 0.05,
+  beta_1 = 0.9,
+  beta_2 = 0.999,
+  epsilon = 1e-8,
+  learn_rate_decay = 0.5
+) {
+  log_scale <- NULL
+  m <- 0
+  v <- 0
+  beta_1_pow <- 1
+  beta_2_pow <- 1
+  initialize <- function(proposal, initial_state) {
+    if (is.null(initial_scale)) {
+      initial_scale <- proposal$default_initial_scale(initial_state$dimension())
+    }
+    log_scale <<- log(initial_scale)
+    m <<- 0
+    v <<- 0
+    beta_1_pow <<- 1
+    beta_2_pow <<- 1
+    proposal$update(scale = initial_scale)
+  }
+  update <- function(proposal, sample_index, state_and_statistics) {
+    if (is.null(target_accept_prob)) {
+      target_accept_prob <- proposal$default_target_accept_prob()
+    }
+    accept_prob <- state_and_statistics$statistics$accept_prob
+    # Gradient of the squared-error objective in log-scale space. Note the
+    # sign: `grad = target - observed` combined with the `log_scale -= ...`
+    # update below reproduces the "observed > target => scale increases"
+    # behaviour of the other scale adapters.
+    grad <- target_accept_prob - accept_prob
+    beta_1_pow <<- beta_1_pow * beta_1
+    beta_2_pow <<- beta_2_pow * beta_2
+    m <<- beta_1 * m + (1 - beta_1) * grad
+    v <<- beta_2 * v + (1 - beta_2) * grad^2
+    m_hat <- m / (1 - beta_1_pow)
+    v_hat <- v / (1 - beta_2_pow)
+    effective_learning_rate <- learning_rate / sample_index^learn_rate_decay
+    log_scale <<- log_scale - (
+      effective_learning_rate * m_hat / (sqrt(v_hat) + epsilon)
+    )
+    proposal$update(scale = exp(log_scale))
+  }
+  list(
+    initialize = initialize,
+    update = update,
+    finalize = NULL,
+    state = function() list(log_scale = log_scale, m = m, v = v)
+  )
+}
+
 #' Create object to adapt proposal shape.
 #'
 #' @param type Type of shape adapter to use. One of:

man/adam_scale_adapter.Rd

@@ -210,6 +210,83 @@ for (dimension in c(1L, 2L, 5L)) {
   )
 }
 
+check_adam_scale_adapter_state <- function(adapter_state) {
+  expect_named(adapter_state, c("log_scale", "m", "v"), ignore.order = TRUE)
+  expect_length(adapter_state$log_scale, 1)
+  expect_length(adapter_state$m, 1)
+  expect_length(adapter_state$v, 1)
+}
+
+for (target_accept_prob in c(0.2, 0.4, 0.6)) {
+  for (initial_scale in c(0.5, 1., 2.)) {
+    for (learning_rate in c(0.05, 0.1)) {
+      test_that(
+        sprintf(
+          paste0(
+            "Adam scale adapter works with target_accept_prob %.1f ",
+            "initial_scale %.1f learning_rate %.2f"
+          ),
+          target_accept_prob, initial_scale, learning_rate
+        ),
+        {
+          adapter <- scale_adapter(
+            algorithm = "adam",
+            initial_scale = initial_scale,
+            target_accept_prob = target_accept_prob,
+            learning_rate = learning_rate
+          )
+          check_adapter(adapter)
+          proposal <- dummy_proposal_with_scale_parameter()
+          adapter$initialize(proposal, chain_state(rep(0, dimension)))
+          adapter_state <- adapter$state()
+          check_adam_scale_adapter_state(adapter_state)
+          # After a fresh initialisation, one update with accept_prob higher
+          # than target should push the scale up (m_hat = grad, v_hat = grad^2
+          # at t = 1, so the update is deterministic in sign).
+          adapter$update(
+            proposal,
+            1,
+            list(statistics = list(accept_prob = target_accept_prob + 0.1))
+          )
+          expect_type(adapter$state(), "list")
+          expect_gt(proposal$parameters()$scale, initial_scale)
+          # After re-initialisation, one update with accept_prob lower than
+          # target should push the scale down.
+          adapter$initialize(proposal, chain_state(rep(0, dimension)))
+          adapter$update(
+            proposal,
+            1,
+            list(statistics = list(accept_prob = target_accept_prob - 0.1))
+          )
+          expect_lt(proposal$parameters()$scale, initial_scale)
+          # Long-run convergence: reuse shared helper. Re-initialise so the
+          # convergence run starts from a clean adapter state.
+          adapter$initialize(proposal, chain_state(rep(0, dimension)))
+          check_scale_adapter_coerces_to_target_accept_prob(
+            adapter, proposal, target_accept_prob, initial_scale
+          )
+        }
+      )
+    }
+  }
+}
+
+for (dimension in c(1L, 2L, 5L)) {
+  test_that(
+    sprintf(
+      "Adam scale adapter with default args works in dimension %i",
+      dimension
+    ),
+    {
+      check_scale_adapter_with_default_args_works(
+        scale_adapter(algorithm = "adam"),
+        dimension,
+        check_adam_scale_adapter_state
+      )
+    }
+  )
+}
+
 for (dimension in c(1L, 2L, 5L)) {
   for (kappa in c(0.5, 0.6, 0.8)) {
     for (correlation in c(0.0, 0.5, 0.8)) {