Adam scale adapter implementation

[jennajiali]

Summary

Adds a third scale-adaptation algorithm to rmcmc, alongside the existing Robbins–Monro (stochastic_approximation) and Nesterov dual-averaging (dual_averaging) schemes. The new adapter applies the Adam optimizer of Kingma & Ba (2014) to the log of the proposal scale, using the acceptance-rate residual target_accept_prob accept_prob as the (stochastic) gradient — the same gradient signal already used by the two existing scale adapters, plugged into the Adam update rule instead of a power-law schedule.

The design follows the C++ implementation in the walnuts library, including its learning_rate / t^learn_rate_decay schedule needed to match the stability of dual averaging (per Bob Carpenter's note on NUTS improvements).

Public API

Two-level exposure as requested in #130:

• High-level wrapper — scale_adapter(algorithm = "adam", ...) accepts the common initial_scale / target_accept_prob arguments plus the single Adam tuning knob users actually reach for:

  scale_adapter(algorithm = "adam", learning_rate = 0.05)

• Low-level function — adam_scale_adapter() exposes all Adam hyperparameters for advanced users:

  adam_scale_adapter(
    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
  )

Defaults and rationale

Parameter	Default	Rationale
learning_rate	0.05	Kingma & Ba's 1e-3 is too slow to hit the shared 2000-iteration / tol = 1e-2 convergence test used by all scale adapters. 0.05 gives fast, stable coercion in typical MCMC warm-up lengths. The docstring notes that 1e-3 matches the original paper.
beta_1	0.9	Standard Adam default.
beta_2	0.999	Standard Adam default.
epsilon	1e-8	Standard Adam default.
learn_rate_decay	0.5	Value recommended in the walnuts implementation; vanilla Adam (0) oscillates.

Algorithm

On each update, with t = sample_index:

grad     ← target_accept_prob − accept_prob
m        ← β₁ · m + (1 − β₁) · grad
v        ← β₂ · v + (1 − β₂) · grad²
m̂        ← m / (1 − β₁ᵗ)
v̂        ← v / (1 − β₂ᵗ)
η_t      ← learning_rate / t^learn_rate_decay
log_scale ← log_scale − η_t · m̂ / (√v̂ + ε)

β₁ᵗ and β₂ᵗ are maintained incrementally as beta_1_pow, beta_2_pow rather than recomputed each step, matching walnuts and avoiding .Machine$double.xmin underflow for long chains.

Sign convention

grad = target − observed combined with log_scale −= η · m̂ / (√v̂ + ε) reproduces the "accept_prob > target ⇒ scale increases" behaviour of the other two scale adapters. This is verified by both the directional and convergence tests in the sweep below. An inline comment in adam_scale_adapter flags the sign for future maintainers.

Changes

R/adaptation.R

• Updated scale_adapter() roxygen:
  • Added "adam" bullet under @param algorithm.
  • Extended @param ... to document the two-level exposure design.
  • Added [adam_scale_adapter()] to @seealso.
  • Added Kingma & Ba (2014) to @references.
• Added adam = adam_scale_adapter to the switch() in scale_adapter().
• Added a new exported adam_scale_adapter() function (roxygen block + implementation) placed immediately after dual_averaging_scale_adapter.
  • State exposed via state(): log_scale, m, v (both beta_*_pow kept internal — they're implementation detail and would only add noise to warm_up_statistics when trace_warm_up = TRUE).
  • initialize() resets all Adam state so the same adapter object can be re-initialised mid-run (used in the tests).
  • finalize = NULL, matching walnuts and the SA adapter — no smoothed-scale tail.

tests/testthat/test-adaptation.R

• Added check_adam_scale_adapter_state() helper (mirrors the SA / dual averaging state checkers).
• Added a parameter sweep over target_accept_prob ∈ {0.2, 0.4, 0.6} × initial_scale ∈ {0.5, 1, 2} × learning_rate ∈ {0.05, 0.1} — 18 test cases — each of which:
  • validates the adapter interface and state schema,
  • checks directional response at t = 1 for both accept_prob > target
    and accept_prob < target,
  • runs the shared check_scale_adapter_coerces_to_target_accept_prob()
    helper (2000 iterations, tol = 1e-2).
• Added a dimension sweep ({1L, 2L, 5L}) using the shared check_scale_adapter_with_default_args_works() helper to verify the default-construction path.

All unit tests passed when I ran devtools::test()

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (65b86ab) to head (2684113).

Additional details and impacted files

@@            Coverage Diff            @@
##              main      #132   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           11        11           
  Lines          620       661   +41     
=========================================
+ Hits           620       661   +41     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[matt-graham]

Thanks @jennajiali - this looks all good to me. I've added one point asking for clarification about a comment in the documentation for scale_adapter about passing through algorithmic parameters but otherwise I think it's ready to merge.

[jennajiali]

Thanks @matt-graham for the review – I have updated the ... documentation to describe the actual pass-through behaviour. There's also a merge commit in the history from resolving a local/remote sync issue on my fork; no substantive changes in it.

[matt-graham]

Thanks @jennajiali this looks all good to me now, will merge this in!