fit_dadvi#

pymc_extras.inference.fit_dadvi(model: Model | None = None, n_fixed_draws: int = 30, n_draws: int = 1000, include_transformed: bool = False, optimizer_method: Literal['nelder-mead', 'powell', 'CG', 'BFGS', 'Newton-CG', 'L-BFGS-B', 'TNC', 'COBYLA', 'SLSQP', 'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact', 'trust-krylov'] = 'trust-ncg', use_grad: bool | None = None, use_hessp: bool | None = None, use_hess: bool | None = None, gradient_backend: str = 'pytensor', compile_kwargs: dict | None = None, random_seed: None | int | Sequence[int] | ndarray = None, progressbar: bool = True, **optimizer_kwargs) → DataTree[source]#

Does inference using Deterministic ADVI (Automatic Differentiation Variational Inference), DADVI for short.

For full details see the paper cited in the references: https://www.jmlr.org/papers/v25/23-1015.html

Parameters:

  • model (pm.Model) – The PyMC model to be fit. If None, the current model context is used.

  • n_fixed_draws (int) – The number of fixed draws to use for the optimisation. More draws will result in more accurate estimates, but also increase inference time. Usually, the default of 30 is a good tradeoff between speed and accuracy.

  • random_seed (int) – The random seed to use for the fixed draws. Running the optimisation twice with the same seed should arrive at the same result.

  • n_draws (int) – The number of draws to return from the variational approximation.

  • include_transformed (bool) – Whether or not to keep the unconstrained variables (such as logs of positive-constrained parameters) in the output.

  • optimizer_method (str) – Which optimization method to use. The function calls scipy.optimize.minimize, so any of the methods there can be used. The default is trust-ncg, which uses second-order information and is generally very reliable. Other methods such as L-BFGS-B might be faster but potentially more brittle and may not converge exactly to the optimum.

  • gradient_backend (str) – Which backend to use to compute gradients. Must be one of “jax” or “pytensor”. Default is “pytensor”.

  • compile_kwargs (dict, optional) – Additional keyword arguments to pass to pytensor.function

  • use_grad (bool, optional) – If True, pass the gradient function to scipy.optimize.minimize (where it is referred to as jac).

  • use_hessp (bool, optional) – If True, pass the hessian vector product to scipy.optimize.minimize.

  • use_hess (bool, optional) – If True, pass the hessian to scipy.optimize.minimize. Note that this is generally not recommended since its computation can be slow and memory-intensive if there are many parameters.

  • progressbar (bool) – Whether or not to show a progress bar during optimization. Default is True.

  • optimizer_kwargs – Additional keyword arguments to pass to the scipy.optimize.minimize function. See the documentation of that function for details.

Returns:

The inference data containing the results of the DADVI algorithm.

Return type:

DataTree

References

Giordano, R., Ingram, M., & Broderick, T. (2024). Black Box Variational Inference with a Deterministic Objective: Faster, More Accurate, and Even More Black Box. Journal of Machine Learning Research, 25(18), 1–39.