pypomp.Pomp.train

Pomp.train(J: int, M: int, eta: LearningRate, key: Array | None = None, theta: PompParameters | None = None, optimizer: Optimizer = Adam(clip_norm=None, scale=False, ls=False, c=0.1, max_ls_itn=10, beta1=0.9, beta2=0.999, epsilon=1e-08), alpha: float = 0.97, thresh: float = 0.0, alpha_cooling: float = 1.0, n_monitors: int = 1, track_time: bool = True) → None

Optimize parameters via a differentiable particle filter (MOP).

Performs Maximum Likelihood Estimation using the Measurement Off-Parameter (MOP) particle filter (Tan et al. 2024 [1]), treating the particle filter as a differentiable computation graph and applies gradient-based optimizers (e.g. Adam, SGD, Newton) via JAX reverse-mode automatic differentiation.

Warning

MOP gradients are only well-defined for continuous-state models. For discrete-state models, use mif() or dpop_train() instead.

JAX vectorises the computation across all starting parameter sets in theta simultaneously. Results are appended to results_history.

Parameters:

  • J (int) – Number of particles used to estimate the MOP objective and its gradient.

  • M (int) – Number of gradient steps to perform.

  • eta (LearningRate) – Per-parameter learning rate schedules. See LearningRate.

  • key (jax.Array or None, optional) – JAX random key. Defaults to fresh_key.

  • theta (PompParameters or None, optional) – Starting parameter set(s). Defaults to theta.

  • optimizer (Optimizer, optional) – Optimizer configuration object (e.g. Adam, SGD, Newton). Defaults to Adam.

  • alpha (float, optional) – MOP discount factor controlling the bias-variance trade-off. Defaults to 0.97.

  • thresh (float, optional) – ESS-based resampling threshold. Defaults to 0.0.

  • alpha_cooling (float, optional) – Cosine cooling multiplier for alpha. At the end of training, alpha is moved alpha_cooling of the way from its initial value toward 1.0. Defaults to 1.0 (no cooling).

  • n_monitors (int, optional) – Number of unperturbed particle filter runs to average for the log-likelihood monitor. Defaults to 1. Using more than 1 monitor increases computation time but can lead to more stable estimates.

  • track_time (bool, optional) – Whether to record wall-clock execution time. Defaults to True.

Returns:

A Result is appended to results_history, containing log-likelihood and parameter traces over iterations.

Return type:

None

See also

pypomp.functional.train

Pure-functional JAX gradient training.

References

Examples

>>> eta = pp.LearningRate({"beta": 0.01, "gamma": 0.005})
>>> model.fresh_key = jax.random.key(0)
>>> model.train(J=100, M=200, eta=eta)
>>> model.results()