Repository: normal-computing/posteriors

Issue:
Add support for `Callable` learning rates and temperature
Would be very useful to be able to support `lr = lambda i: 1e-3 * i ** -0.5` to support [Robbins-Monro](https://en.wikipedia.org/wiki/Stochastic_approximation) style scheduling (and others)

Up for debate but could be useful to add this support for the `temperature` parameter in the spirit of [simulated annealing](https://en.wikipedia.org/wiki/Simulated_annealing)

Discussion:

Solution patch:
diff --git a/posteriors/ekf/dense_fisher.py b/posteriors/ekf/dense_fisher.py
index 31eaaf0..02ef77c 100644
--- a/posteriors/ekf/dense_fisher.py
+++ b/posteriors/ekf/dense_fisher.py
@@ -6,7 +6,7 @@ from optree.integration.torch import tree_ravel
 from tensordict import TensorClass
 
 from posteriors.tree_utils import tree_size, tree_insert_
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.utils import (
     per_samplify,
     empirical_fisher,
@@ -17,7 +17,7 @@ from posteriors.utils import (
 
 def build(
     log_likelihood: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     transition_cov: torch.Tensor | float = 0.0,
     per_sample: bool = False,
     init_cov: torch.Tensor | float = 1.0,
@@ -43,6 +43,7 @@ def build(
             returns the log-likelihood value as well as auxiliary information,
             e.g. from the model call.
         lr: Inverse temperature of the update, which behaves like a learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         transition_cov: Covariance of the transition noise, to additively
             inflate the covariance before the update.
         per_sample: If True, then log_likelihood is assumed to return a vector of
@@ -74,11 +75,13 @@ class EKFDenseState(TensorClass["frozen"]):
         cov: Covariance matrix of the
             Normal distribution.
         log_likelihood: Log likelihood of the data given the parameters.
+        step: Current step count.
     """
 
     params: TensorTree
     cov: torch.Tensor
     log_likelihood: torch.Tensor = torch.tensor([])
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -122,6 +125,7 @@ def update(
             returns the log-likelihood value as well as auxiliary information,
             e.g. from the model call.
         lr: Inverse temperature of the update, which behaves like a learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         transition_cov: Covariance of the transition noise, to additively
             inflate the covariance before the update.
         per_sample: If True, then log_likelihood is assumed to return a vector of
@@ -137,6 +141,8 @@ def update(
     if not per_sample:
         log_likelihood = per_samplify(log_likelihood)
 
+    lr = lr(state.step) if callable(lr) else lr
+
     with torch.no_grad(), CatchAuxError():
 
         def log_likelihood_reduced(params, batch):
@@ -163,9 +169,12 @@ def update(
         tree_insert_(state.params, update_mean)
         tree_insert_(state.cov, update_cov)
         tree_insert_(state.log_likelihood, log_liks.mean().detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
 
-    return EKFDenseState(update_mean, update_cov, log_liks.mean().detach()), aux
+    return EKFDenseState(
+        update_mean, update_cov, log_liks.mean().detach(), state.step + 1
+    ), aux
 
 
 def sample(
diff --git a/posteriors/ekf/diag_fisher.py b/posteriors/ekf/diag_fisher.py
index b5beffe..44cdeda 100644
--- a/posteriors/ekf/diag_fisher.py
+++ b/posteriors/ekf/diag_fisher.py
@@ -5,7 +5,7 @@ from torch.func import jacrev
 from optree import tree_map
 from tensordict import TensorClass
 
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import (
     diag_normal_sample,
@@ -17,7 +17,7 @@ from posteriors.utils import (
 
 def build(
     log_likelihood: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     transition_sd: float = 0.0,
     per_sample: bool = False,
     init_sds: TensorTree | float = 1.0,
@@ -44,6 +44,7 @@ def build(
             returns the log-likelihood value as well as auxiliary information,
             e.g. from the model call.
         lr: Inverse temperature of the update, which behaves like a learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         transition_sd: Standard deviation of the transition noise, to additively
             inflate the diagonal covariance before the update.
         per_sample: If True, then log_likelihood is assumed to return a vector of
@@ -76,11 +77,13 @@ class EKFDiagState(TensorClass["frozen"]):
         sd_diag: Square-root diagonal of the covariance matrix of the
             Normal distribution.
         log_likelihood: Log likelihood of the data given the parameters.
+        step: Current step count.
     """
 
     params: TensorTree
     sd_diag: TensorTree
     log_likelihood: torch.Tensor = torch.tensor([])
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -110,7 +113,7 @@ def update(
     state: EKFDiagState,
     batch: Any,
     log_likelihood: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     transition_sd: float = 0.0,
     per_sample: bool = False,
     inplace: bool = False,
@@ -126,6 +129,7 @@ def update(
             returns the log-likelihood value as well as auxiliary information,
             e.g. from the model call.
         lr: Inverse temperature of the update, which behaves like a learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         transition_sd: Standard deviation of the transition noise, to additively
             inflate the diagonal covariance before the update.
         per_sample: If True, then log_likelihood is assumed to return a vector of
@@ -142,6 +146,8 @@ def update(
     if not per_sample:
         log_likelihood = per_samplify(log_likelihood)
 
+    lr = lr(state.step) if callable(lr) else lr
+
     predict_sd_diag = flexi_tree_map(
         lambda x: (x**2 + transition_sd**2) ** 0.5, state.sd_diag, inplace=inplace
     )
@@ -167,9 +173,12 @@ def update(
 
     if inplace:
         tree_insert_(state.log_likelihood, log_liks.mean().detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
 
-    return EKFDiagState(update_mean, update_sd_diag, log_liks.mean().detach()), aux
+    return EKFDiagState(
+        update_mean, update_sd_diag, log_liks.mean().detach(), state.step + 1
+    ), aux
 
 
 def sample(
diff --git a/posteriors/laplace/dense_fisher.py b/posteriors/laplace/dense_fisher.py
index 9e2f350..e2f7220 100644
--- a/posteriors/laplace/dense_fisher.py
+++ b/posteriors/laplace/dense_fisher.py
@@ -5,7 +5,7 @@ from optree import tree_map
 from optree.integration.torch import tree_ravel
 from tensordict import TensorClass
 from posteriors.types import TensorTree, Transform, LogProbFn
-from posteriors.tree_utils import tree_size
+from posteriors.tree_utils import tree_size, tree_insert_
 from posteriors.utils import (
     per_samplify,
     empirical_fisher,
@@ -61,10 +61,12 @@ class DenseLaplaceState(TensorClass["frozen"]):
     Attributes:
         params: Mean of the Normal distribution.
         prec: Precision matrix of the Normal distribution.
+        step: Current step count.
     """
 
     params: TensorTree
     prec: torch.Tensor
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -127,9 +129,10 @@ def update(
 
     if inplace:
         state.prec.data += fisher
+        tree_insert_(state.step, state.step + 1)
         return state, aux
     else:
-        return DenseLaplaceState(state.params, state.prec + fisher), aux
+        return DenseLaplaceState(state.params, state.prec + fisher, state.step + 1), aux
 
 
 def sample(
diff --git a/posteriors/laplace/dense_ggn.py b/posteriors/laplace/dense_ggn.py
index fee5c22..1639345 100644
--- a/posteriors/laplace/dense_ggn.py
+++ b/posteriors/laplace/dense_ggn.py
@@ -11,8 +11,8 @@ from posteriors.types import (
     ForwardFn,
     OuterLogProbFn,
 )
+from posteriors.tree_utils import tree_insert_, tree_size
 from posteriors.utils import (
-    tree_size,
     ggn,
     is_scalar,
     CatchAuxError,
@@ -73,10 +73,12 @@ class DenseLaplaceState(TensorClass["frozen"]):
     Attributes:
         params: Mean of the Normal distribution.
         prec: Precision matrix of the Normal distribution.
+        step: Current step count.
     """
 
     params: TensorTree
     prec: torch.Tensor
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -142,9 +144,12 @@ def update(
 
     if inplace:
         state.prec.data += ggn_batch
+        tree_insert_(state.step, state.step + 1)
         return state, aux
     else:
-        return DenseLaplaceState(state.params, state.prec + ggn_batch), aux
+        return DenseLaplaceState(
+            state.params, state.prec + ggn_batch, state.step + 1
+        ), aux
 
 
 def sample(
diff --git a/posteriors/laplace/dense_hessian.py b/posteriors/laplace/dense_hessian.py
index 6aee84b..6c7d329 100644
--- a/posteriors/laplace/dense_hessian.py
+++ b/posteriors/laplace/dense_hessian.py
@@ -6,7 +6,7 @@ from optree.integration.torch import tree_ravel
 from tensordict import TensorClass
 
 from posteriors.types import TensorTree, Transform, LogProbFn
-from posteriors.tree_utils import tree_size
+from posteriors.tree_utils import tree_size, tree_insert_
 from posteriors.utils import (
     is_scalar,
     CatchAuxError,
@@ -59,10 +59,12 @@ class DenseLaplaceState(TensorClass["frozen"]):
     Attributes:
         params: Mean of the Normal distribution.
         prec: Precision matrix of the Normal distribution.
+        step: Current step count.
     """
 
     params: TensorTree
     prec: torch.Tensor
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -131,9 +133,10 @@ def update(
 
     if inplace:
         state.prec.data += hess
+        tree_insert_(state.step, state.step + 1)
         return state, aux
     else:
-        return DenseLaplaceState(state.params, state.prec + hess), aux
+        return DenseLaplaceState(state.params, state.prec + hess, state.step + 1), aux
 
 
 def sample(
diff --git a/posteriors/laplace/diag_fisher.py b/posteriors/laplace/diag_fisher.py
index 6535400..3f9754a 100644
--- a/posteriors/laplace/diag_fisher.py
+++ b/posteriors/laplace/diag_fisher.py
@@ -6,7 +6,7 @@ from optree import tree_map
 from tensordict import TensorClass
 
 from posteriors.types import TensorTree, Transform, LogProbFn
-from posteriors.tree_utils import flexi_tree_map
+from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import (
     diag_normal_sample,
     per_samplify,
@@ -60,10 +60,12 @@ class DiagLaplaceState(TensorClass["frozen"]):
     Attributes:
         params: Mean of the Normal distribution.
         prec_diag: Diagonal of the precision matrix of the Normal distribution.
+        step: Current step count.
     """
 
     params: TensorTree
     prec_diag: TensorTree
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -131,8 +133,9 @@ def update(
     )
 
     if inplace:
+        tree_insert_(state.step, state.step + 1)
         return state, aux
-    return DiagLaplaceState(state.params, prec_diag), aux
+    return DiagLaplaceState(state.params, prec_diag, state.step + 1), aux
 
 
 def sample(
diff --git a/posteriors/laplace/diag_ggn.py b/posteriors/laplace/diag_ggn.py
index dbf1d8a..10685ec 100644
--- a/posteriors/laplace/diag_ggn.py
+++ b/posteriors/laplace/diag_ggn.py
@@ -10,7 +10,7 @@ from posteriors.types import (
     ForwardFn,
     OuterLogProbFn,
 )
-from posteriors.tree_utils import flexi_tree_map
+from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import (
     diag_normal_sample,
     diag_ggn,
@@ -71,10 +71,12 @@ class DiagLaplaceState(TensorClass["frozen"]):
     Attributes:
         params: Mean of the Normal distribution.
         prec_diag: Diagonal of the precision matrix of the Normal distribution.
+        step: Current step count.
     """
 
     params: TensorTree
     prec_diag: TensorTree
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -145,8 +147,9 @@ def update(
     )
 
     if inplace:
+        tree_insert_(state.step, state.step + 1)
         return state, aux
-    return DiagLaplaceState(state.params, prec_diag), aux
+    return DiagLaplaceState(state.params, prec_diag, state.step + 1), aux
 
 
 def sample(
diff --git a/posteriors/sgmcmc/baoa.py b/posteriors/sgmcmc/baoa.py
index 704a352..bfdebaa 100644
--- a/posteriors/sgmcmc/baoa.py
+++ b/posteriors/sgmcmc/baoa.py
@@ -5,17 +5,17 @@ from torch.func import grad_and_value
 from optree import tree_map
 from tensordict import TensorClass
 
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import is_scalar, CatchAuxError
 
 
 def build(
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     alpha: float = 0.01,
     sigma: float = 1.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     momenta: TensorTree | float | None = None,
 ) -> Transform:
     """Builds BAOA transform.
@@ -49,9 +49,11 @@ def build(
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
         lr: Learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         alpha: Friction coefficient.
         sigma: Standard deviation of momenta target distribution.
         temperature: Temperature of the joint parameter + momenta distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         momenta: Initial momenta. Can be tree like params or scalar.
             Defaults to random iid samples from N(0, 1).
 
@@ -77,11 +79,13 @@ class BAOAState(TensorClass["frozen"]):
         params: Parameters.
         momenta: Momenta for each parameter.
         log_posterior: Log posterior evaluation.
+        step: Current step count.
     """
 
     params: TensorTree
     momenta: TensorTree
     log_posterior: torch.Tensor = torch.tensor(torch.nan)
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(params: TensorTree, momenta: TensorTree | float | None = None) -> BAOAState:
@@ -113,10 +117,10 @@ def update(
     state: BAOAState,
     batch: Any,
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     alpha: float = 0.01,
     sigma: float = 1.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     inplace: bool = False,
 ) -> tuple[BAOAState, TensorTree]:
     """Updates parameters and momenta for BAOA.
@@ -132,9 +136,11 @@ def update(
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
         lr: Learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         alpha: Friction coefficient.
         sigma: Standard deviation of momenta target distribution.
         temperature: Temperature of the joint parameter + momenta distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         inplace: Whether to modify state in place.
 
     Returns:
@@ -147,6 +153,8 @@ def update(
             state.params, batch
         )
 
+    lr = lr(state.step) if callable(lr) else lr
+    temperature = temperature(state.step) if callable(temperature) else temperature
     prec = sigma**-2
     gamma = torch.tensor(alpha * prec)
     zeta2 = (temperature * (1 - torch.exp(-2 * gamma * lr))) ** 0.5
@@ -167,5 +175,6 @@ def update(
 
     if inplace:
         tree_insert_(state.log_posterior, log_post.detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
-    return BAOAState(params, momenta, log_post.detach()), aux
+    return BAOAState(params, momenta, log_post.detach(), state.step + 1), aux
diff --git a/posteriors/sgmcmc/sghmc.py b/posteriors/sgmcmc/sghmc.py
index 4c8f7e1..f8d5d2a 100644
--- a/posteriors/sgmcmc/sghmc.py
+++ b/posteriors/sgmcmc/sghmc.py
@@ -5,18 +5,18 @@ from torch.func import grad_and_value
 from optree import tree_map
 from tensordict import TensorClass
 
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import is_scalar, CatchAuxError
 
 
 def build(
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     alpha: float = 0.01,
     beta: float = 0.0,
     sigma: float = 1.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     momenta: TensorTree | float | None = None,
 ) -> Transform:
     """Builds SGHMC transform.
@@ -41,11 +41,13 @@ def build(
         log_posterior: Function that takes parameters and input batch and
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
-        lr: Learning rate.
+        lr: Learning rate,
+            scalar or schedule (callable taking step index, returning scalar).
         alpha: Friction coefficient.
         beta: Gradient noise coefficient (estimated variance).
         sigma: Standard deviation of momenta target distribution.
         temperature: Temperature of the joint parameter + momenta distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         momenta: Initial momenta. Can be tree like params or scalar.
             Defaults to random iid samples from N(0, 1).
 
@@ -72,11 +74,13 @@ class SGHMCState(TensorClass["frozen"]):
         params: Parameters.
         momenta: Momenta for each parameter.
         log_posterior: Log posterior evaluation.
+        step: Current step count.
     """
 
     params: TensorTree
     momenta: TensorTree
     log_posterior: torch.Tensor = torch.tensor(torch.nan)
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(params: TensorTree, momenta: TensorTree | float | None = None) -> SGHMCState:
@@ -108,11 +112,11 @@ def update(
     state: SGHMCState,
     batch: Any,
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     alpha: float = 0.01,
     beta: float = 0.0,
     sigma: float = 1.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     inplace: bool = False,
 ) -> tuple[SGHMCState, TensorTree]:
     """Updates parameters and momenta for SGHMC.
@@ -126,11 +130,13 @@ def update(
         log_posterior: Function that takes parameters and input batch and
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
-        lr: Learning rate.
+        lr: Learning rate,
+            scalar or schedule (callable taking step index, returning scalar).
         alpha: Friction coefficient.
         beta: Gradient noise coefficient (estimated variance).
         sigma: Standard deviation of momenta target distribution.
         temperature: Temperature of the joint parameter + momenta distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         inplace: Whether to modify state in place.
 
     Returns:
@@ -142,6 +148,8 @@ def update(
             state.params, batch
         )
 
+    lr = lr(state.step) if callable(lr) else lr
+    temperature = temperature(state.step) if callable(temperature) else temperature
     prec = sigma**-2
 
     def transform_params(p, m):
@@ -163,5 +171,6 @@ def update(
 
     if inplace:
         tree_insert_(state.log_posterior, log_post.detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
-    return SGHMCState(params, momenta, log_post.detach()), aux
+    return SGHMCState(params, momenta, log_post.detach(), state.step + 1), aux
diff --git a/posteriors/sgmcmc/sgld.py b/posteriors/sgmcmc/sgld.py
index 237d279..39ed2d8 100644
--- a/posteriors/sgmcmc/sgld.py
+++ b/posteriors/sgmcmc/sgld.py
@@ -1,19 +1,20 @@
 from typing import Any
 from functools import partial
 import torch
+from torch import Tensor
 from torch.func import grad_and_value
 from tensordict import TensorClass
 
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import CatchAuxError
 
 
 def build(
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     beta: float = 0.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
 ) -> Transform:
     """Builds SGLD transform.
 
@@ -32,9 +33,11 @@ def build(
         log_posterior: Function that takes parameters and input batch and
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
-        lr: Learning rate.
+        lr: Learning rate,
+            scalar or schedule (callable taking step index, returning scalar).
         beta: Gradient noise coefficient (estimated variance).
         temperature: Temperature of the sampling distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
 
     Returns:
         SGLD transform (posteriors.types.Transform instance).
@@ -55,10 +58,12 @@ class SGLDState(TensorClass["frozen"]):
     Attributes:
         params: Parameters.
         log_posterior: Log posterior evaluation.
+        step: Current step count.
     """
 
     params: TensorTree
-    log_posterior: torch.Tensor = torch.tensor(torch.nan)
+    log_posterior: Tensor = torch.tensor(torch.nan)
+    step: Tensor = torch.tensor(0)
 
 
 def init(params: TensorTree) -> SGLDState:
@@ -78,9 +83,9 @@ def update(
     state: SGLDState,
     batch: Any,
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     beta: float = 0.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     inplace: bool = False,
 ) -> tuple[SGLDState, TensorTree]:
     """Updates parameters for SGLD.
@@ -94,9 +99,11 @@ def update(
         log_posterior: Function that takes parameters and input batch and
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
-        lr: Learning rate.
+        lr: Learning rate,
+            scalar or schedule (callable taking step index, returning scalar).
         beta: Gradient noise coefficient (estimated variance).
         temperature: Temperature of the sampling distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         inplace: Whether to modify state in place.
 
     Returns:
@@ -108,6 +115,9 @@ def update(
             state.params, batch
         )
 
+    lr = lr(state.step) if callable(lr) else lr
+    temperature = temperature(state.step) if callable(temperature) else temperature
+
     def transform_params(p, g):
         return (
             p
@@ -120,5 +130,6 @@ def update(
 
     if inplace:
         tree_insert_(state.log_posterior, log_post.detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
-    return SGLDState(params, log_post.detach()), aux
+    return SGLDState(params, log_post.detach(), state.step + 1), aux
diff --git a/posteriors/sgmcmc/sgnht.py b/posteriors/sgmcmc/sgnht.py
index 383011c..198c02d 100644
--- a/posteriors/sgmcmc/sgnht.py
+++ b/posteriors/sgmcmc/sgnht.py
@@ -5,18 +5,18 @@ from torch.func import grad_and_value
 from optree import tree_map
 from optree.integration.torch import tree_ravel
 from tensordict import TensorClass
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import is_scalar, CatchAuxError
 
 
 def build(
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     alpha: float = 0.01,
     beta: float = 0.0,
     sigma: float = 1.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     momenta: TensorTree | float | None = None,
     xi: float = None,
 ) -> Transform:
@@ -43,10 +43,12 @@ def build(
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
         lr: Learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         alpha: Friction coefficient.
         beta: Gradient noise coefficient (estimated variance).
         sigma: Standard deviation of momenta target distribution.
         temperature: Temperature of the joint parameter + momenta distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         momenta: Initial momenta. Can be tree like params or scalar.
             Defaults to random iid samples from N(0, 1).
         xi: Initial value for scalar thermostat ξ. Defaults to `alpha`.
@@ -75,12 +77,14 @@ class SGNHTState(TensorClass["frozen"]):
         momenta: Momenta for each parameter.
         xi: Scalar thermostat.
         log_posterior: Log posterior evaluation.
+        step: Current step count.
     """
 
     params: TensorTree
     momenta: TensorTree
     xi: torch.Tensor
     log_posterior: torch.Tensor = torch.tensor(torch.nan)
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -117,11 +121,11 @@ def update(
     state: SGNHTState,
     batch: Any,
     log_posterior: LogProbFn,
-    lr: float,
+    lr: float | Schedule,
     alpha: float = 0.01,
     beta: float = 0.0,
     sigma: float = 1.0,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     inplace: bool = False,
 ) -> tuple[SGNHTState, TensorTree]:
     """Updates parameters, momenta and xi for SGNHT.
@@ -136,10 +140,12 @@ def update(
             returns the log posterior value (which can be unnormalised)
             as well as auxiliary information, e.g. from the model call.
         lr: Learning rate.
+            Scalar or schedule (callable taking step index, returning scalar).
         alpha: Friction coefficient.
         beta: Gradient noise coefficient (estimated variance).
         sigma: Standard deviation of momenta target distribution.
         temperature: Temperature of the joint parameter + momenta distribution.
+            Scalar or schedule (callable taking step index, returning scalar).
         inplace: Whether to modify state in place.
 
     Returns:
@@ -151,6 +157,8 @@ def update(
             state.params, batch
         )
 
+    lr = lr(state.step) if callable(lr) else lr
+    temperature = temperature(state.step) if callable(temperature) else temperature
     prec = sigma**-2
 
     def transform_params(p, m):
@@ -176,5 +184,6 @@ def update(
     if inplace:
         tree_insert_(state.xi, xi_new)
         tree_insert_(state.log_posterior, log_post.detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
-    return SGNHTState(params, momenta, xi_new, log_post.detach()), aux
+    return SGNHTState(params, momenta, xi_new, log_post.detach(), state.step + 1), aux
diff --git a/posteriors/types.py b/posteriors/types.py
index 47eb8b1..6fad1dc 100644
--- a/posteriors/types.py
+++ b/posteriors/types.py
@@ -17,10 +17,15 @@ LogProbFn = Callable[[TensorTree, TensorTree], Tuple[float, TensorTree]]
 ForwardFn = Callable[[TensorTree, TensorTree], Tuple[Tensor, TensorTree]]
 OuterLogProbFn = Callable[[TensorTree, TensorTree], float]
 
+TensorLike = float | int | bool | Tensor
+Schedule = Callable[
+    [TensorLike], TensorLike
+]  # Learning rate schedule (input is step index, output is scalar learning rate)
+
 
 class InitFn(Protocol):
-    @staticmethod
     def __call__(
+        self,
         params: TensorTree,
     ) -> TensorClass:
         """Initiate a posteriors state with unified API:
@@ -48,8 +53,8 @@ class InitFn(Protocol):
 
 
 class UpdateFn(Protocol):
-    @staticmethod
     def __call__(
+        self,
         state: TensorClass,
         batch: Any,
         inplace: bool = False,
diff --git a/posteriors/vi/dense.py b/posteriors/vi/dense.py
index 86e0ee1..72ba070 100644
--- a/posteriors/vi/dense.py
+++ b/posteriors/vi/dense.py
@@ -7,7 +7,7 @@ from optree.integration.torch import tree_ravel
 import torchopt
 from tensordict import TensorClass
 
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import tree_size, tree_insert_
 from posteriors.utils import (
     is_scalar,
@@ -20,7 +20,7 @@ from posteriors.utils import (
 def build(
     log_posterior: Callable[[TensorTree, Any], float],
     optimizer: torchopt.base.GradientTransformation,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     n_samples: int = 1,
     stl: bool = True,
     init_L: torch.Tensor | float = 1.0,
@@ -42,6 +42,7 @@ def build(
         optimizer: TorchOpt functional optimizer for updating the variational
             parameters. Make sure to use lower case like torchopt.adam()
         temperature: Temperature to rescale (divide) log_posterior.
+            Scalar or schedule (callable taking step index, returning scalar).
         n_samples: Number of samples to use for Monte Carlo estimate.
         stl: Whether to use the stick-the-landing estimator
             from [Roeder et al](https://arxiv.org/abs/1703.09194).
@@ -73,12 +74,14 @@ class VIDenseState(TensorClass["frozen"]):
         opt_state: TorchOpt state storing optimizer data for updating the
             variational parameters.
         nelbo: Negative evidence lower bound (lower is better).
+        step: Current step count.
     """
 
     params: TensorTree
     L_factor: torch.Tensor
     opt_state: torchopt.typing.OptState
     nelbo: torch.Tensor = torch.tensor([])
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -140,6 +143,7 @@ def update(
         optimizer: TorchOpt functional optimizer for updating the variational
             parameters. Make sure to use lower case like torchopt.adam()
         temperature: Temperature to rescale (divide) log_posterior.
+            Scalar or schedule (callable taking step index, returning scalar).
         n_samples: Number of samples to use for Monte Carlo estimate.
         stl: Whether to use the stick-the-landing estimator
             from (Roeder et al](https://arxiv.org/abs/1703.09194).
@@ -149,6 +153,8 @@ def update(
         Updated DenseVIState and auxiliary information.
     """
 
+    temperature = temperature(state.step) if callable(temperature) else temperature
+
     def nelbo_L_factor(m, L_flat):
         return nelbo(m, L_flat, batch, log_posterior, temperature, n_samples, stl)
 
@@ -169,9 +175,12 @@ def update(
 
     if inplace:
         tree_insert_(state.nelbo, nelbo_val.detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
 
-    return VIDenseState(mean, L_factor, opt_state, nelbo_val.detach()), aux
+    return VIDenseState(
+        mean, L_factor, opt_state, nelbo_val.detach(), state.step + 1
+    ), aux
 
 
 def nelbo(
diff --git a/posteriors/vi/diag.py b/posteriors/vi/diag.py
index 5247e9d..e49d326 100644
--- a/posteriors/vi/diag.py
+++ b/posteriors/vi/diag.py
@@ -6,7 +6,7 @@ from optree import tree_map
 import torchopt
 from tensordict import TensorClass
 
-from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import tree_insert_
 from posteriors.utils import (
     diag_normal_log_prob,
@@ -19,7 +19,7 @@ from posteriors.utils import (
 def build(
     log_posterior: Callable[[TensorTree, Any], float],
     optimizer: torchopt.base.GradientTransformation,
-    temperature: float = 1.0,
+    temperature: float | Schedule = 1.0,
     n_samples: int = 1,
     stl: bool = True,
     init_log_sds: TensorTree | float = 0.0,
@@ -41,6 +41,7 @@ def build(
         optimizer: TorchOpt functional optimizer for updating the variational
             parameters. Make sure to use lower case like torchopt.adam()
         temperature: Temperature to rescale (divide) log_posterior.
+            Scalar or schedule (callable taking step index, returning scalar).
         n_samples: Number of samples to use for Monte Carlo estimate.
         stl: Whether to use the stick-the-landing estimator
             from [Roeder et al](https://arxiv.org/abs/1703.09194).
@@ -72,12 +73,14 @@ class VIDiagState(TensorClass["frozen"]):
         opt_state: TorchOpt state storing optimizer data for updating the
             variational parameters.
         nelbo: Negative evidence lower bound (lower is better).
+        step: Current step count.
     """
 
     params: TensorTree
     log_sd_diag: TensorTree
     opt_state: torchopt.typing.OptState
     nelbo: torch.Tensor = torch.tensor([])
+    step: torch.Tensor = torch.tensor(0)
 
 
 def init(
@@ -139,6 +142,7 @@ def update(
         optimizer: TorchOpt functional optimizer for updating the variational
             parameters. Make sure to use lower case like torchopt.adam()
         temperature: Temperature to rescale (divide) log_posterior.
+            Scalar or schedule (callable taking step index, returning scalar).
         n_samples: Number of samples to use for Monte Carlo estimate.
         stl: Whether to use the stick-the-landing estimator
             from (Roeder et al](https://arxiv.org/abs/1703.09194).
@@ -147,6 +151,7 @@ def update(
     Returns:
         Updated DiagVIState and auxiliary information.
     """
+    temperature = temperature(state.step) if callable(temperature) else temperature
 
     def nelbo_log_sd(m, lsd):
         sd_diag = tree_map(torch.exp, lsd)
@@ -169,9 +174,12 @@ def update(
 
     if inplace:
         tree_insert_(state.nelbo, nelbo_val.detach())
+        tree_insert_(state.step, state.step + 1)
         return state, aux
 
-    return VIDiagState(mean, log_sd_diag, opt_state, nelbo_val.detach()), aux
+    return VIDiagState(
+        mean, log_sd_diag, opt_state, nelbo_val.detach(), state.step + 1
+    ), aux
 
 
 def nelbo(


Test patch:
diff --git a/tests/ekf/test_dense_fisher.py b/tests/ekf/test_dense_fisher.py
index fe7bb02..6571b8f 100644
--- a/tests/ekf/test_dense_fisher.py
+++ b/tests/ekf/test_dense_fisher.py
@@ -22,7 +22,11 @@ def test_ekf_dense_inplace():
         for i, x in enumerate(target_mean)
     }
     batch = torch.arange(3).reshape(-1, 1)
-    transform = ekf.dense_fisher.build(log_prob, lr=1e-1)
+
+    def lr(step):
+        return 1e-1 * (step + 1) ** -0.33
+
+    transform = ekf.dense_fisher.build(log_prob, lr=lr)
 
     verify_inplace_update(transform, init_mean, batch)
 
diff --git a/tests/ekf/test_diag_fisher.py b/tests/ekf/test_diag_fisher.py
index e239e62..c34b937 100644
--- a/tests/ekf/test_diag_fisher.py
+++ b/tests/ekf/test_diag_fisher.py
@@ -24,7 +24,11 @@ def test_ekf_diag_inplace():
         for i, x in enumerate(target_mean)
     }
     batch = torch.arange(3).reshape(-1, 1)
-    transform = ekf.diag_fisher.build(log_prob, lr=1e-1)
+
+    def lr(step):
+        return 1e-1 * (step + 1) ** -0.33
+
+    transform = ekf.diag_fisher.build(log_prob, lr=lr)
 
     verify_inplace_update(transform, init_mean, batch)
 
diff --git a/tests/sgmcmc/test_baoa.py b/tests/sgmcmc/test_baoa.py
index a899d46..76ba518 100644
--- a/tests/sgmcmc/test_baoa.py
+++ b/tests/sgmcmc/test_baoa.py
@@ -34,7 +34,8 @@ def test_baoa_inplace_step():
     log_prob, _ = scenarios.get_multivariate_normal_log_prob(dim)
 
     # Set inference parameters
-    lr = 1e-2
+    def lr(step):
+        return 1e-2 * (step + 1) ** -0.33
 
     # Build transform
     transform = baoa.build(log_prob, lr)
diff --git a/tests/sgmcmc/test_sghmc.py b/tests/sgmcmc/test_sghmc.py
index 14a56f5..321be53 100644
--- a/tests/sgmcmc/test_sghmc.py
+++ b/tests/sgmcmc/test_sghmc.py
@@ -42,7 +42,8 @@ def test_sghmc_inplace_step():
     log_prob, _ = scenarios.get_multivariate_normal_log_prob(dim)
 
     # Set inference parameters
-    lr = 1e-2
+    def lr(step):
+        return 1e-2 * (step + 1) ** -0.33
 
     # Build transform
     transform = sghmc.build(log_prob, lr)
diff --git a/tests/sgmcmc/test_sgld.py b/tests/sgmcmc/test_sgld.py
index 4f37850..adcc76e 100644
--- a/tests/sgmcmc/test_sgld.py
+++ b/tests/sgmcmc/test_sgld.py
@@ -27,7 +27,8 @@ def test_sgld_inplace_step():
     log_prob, _ = scenarios.get_multivariate_normal_log_prob(dim)
 
     # Set inference parameters
-    lr = 1e-2
+    def lr(step):
+        return 1e-2 * (step + 1) ** -0.33
 
     # Build transform
     transform = sgld.build(log_prob, lr)
diff --git a/tests/sgmcmc/test_sgnht.py b/tests/sgmcmc/test_sgnht.py
index 024ee29..cc8dd78 100644
--- a/tests/sgmcmc/test_sgnht.py
+++ b/tests/sgmcmc/test_sgnht.py
@@ -42,7 +42,8 @@ def test_sgnht_inplace_step():
     log_prob, _ = scenarios.get_multivariate_normal_log_prob(dim)
 
     # Set inference parameters
-    lr = 1e-2
+    def lr(step):
+        return 1e-2 * (step + 1) ** -0.33
 
     # Build transform
     transform = sgnht.build(log_prob, lr)
diff --git a/tests/utils.py b/tests/utils.py
index 241e837..e0ce748 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -18,14 +18,21 @@ def kl_gaussians(
 def verify_inplace_update(transform: Transform, params: TensorTree, batch: Any):
     init_state = transform.init(params)
 
+    assert init_state.step == 0
+
     # One step not inplace
     torch.manual_seed(42)
     state_not_inplace, _ = transform.update(init_state, batch)
 
+    assert state_not_inplace.step == 1
+
     # One step in place = same as not inplace
     torch.manual_seed(42)
     state_inplace, _ = transform.update(init_state, batch, inplace=True)
 
+    assert state_inplace.step == 1
+    assert init_state.step == 1
+
     # Iterate through leaves of params and check they are the same
     for init_leaf, not_inplace_leaf, inplace_leaf in zip(
         tree_leaves(init_state.to_dict()),