Turing.jl icon indicating copy to clipboard operation
Turing.jl copied to clipboard

Interface for MAP/MLE/VI

Open yebai opened this issue 10 months ago • 12 comments

Proposal for a unified interface

optimise(model, MAP(), ...)
optimise(model, MLE(), ...)
optimise(model, VI(variational_distribution=q, ...), ...)

This resembles the sample(model, sampler, ...) interface for Monte Carlo sampling-based inference algorithms.

yebai avatar Mar 16 '25 16:03 yebai

Relates to #2506; cc @Red-Portal

yebai avatar Mar 16 '25 16:03 yebai

I think one downside would be that referring to the documentation becomes a little tricky since everything requires specialization. For instance, if I only care about VI, doing ?optimise shows the documentation for unrelated functionalities.

Red-Portal avatar Mar 16 '25 17:03 Red-Portal

How would the other arguments to optimise look? I'm not familiar with what sorts of arguments VI needs, but if they are different from what MAP/MLE needs, then that would be an argument in favour of separate functions. Another thing that I wonder about is the return value, if that would be very different (point-estimate vs posterior approximation).

I should also check what other PPLs do for this.

mhauru avatar Mar 17 '25 16:03 mhauru

@mhauru I think that could be accommodated by pushing everything VI or ML/MAP specific to the object representing the algorithm like vi( ... ). So, the general interface for optimise could be left simple. But this also raises the question of what's the point of making everything go through optimise.

Red-Portal avatar Mar 17 '25 18:03 Red-Portal

For what it's worth (i.e., from my basic user perspective), I really hope there is some kind of unified API (and it might emerge through some other packages acting like an interface to Turing, but might as well think of it early) à-la-rstan (brms / rstanarm):

brm(y ~ x,, data = df, algorithm = "sampling")

The algorithm can be swapped for anything ("meanfield", "fullrank", and "pathfinder") through global options. This makes is super convenient to prototype a script, test it and make sure it runs using "fast" algorithms, and then just change the default algo for MCMC and run it without any edits needed to the script.

In the case described above there is some additional considerations if one wants to incorporate approaches that return point-estimates (MLE/MAP optimization), but I don't think it poses that big of a challenge

DominiqueMakowski avatar Mar 17 '25 18:03 DominiqueMakowski

Hi Dominique. I think here, the situation is slightly different since the unified interface is over ML and MAP. And the potential users of ML and MAP are most likely very different from VI. (On the other hand, I think choosing between MCMC, Pathfinder, or VI, for example, is a more likely scenario.) Therefore, I think there is less incentive the unify the interface.

Red-Portal avatar Mar 17 '25 18:03 Red-Portal

A unified interface for all optimisation-based approximate inference algorithms (MAP/MLE/VI) would be nice. There are some implementation details on the optimise(model, InferAlgo(...), ...) function signature, but that could be figured out.

yebai avatar Mar 20 '25 21:03 yebai

A more ambitious goal would be infer(model, data = df, algorithm = ...; kwargs):

  1. unify sampling, optimisation, and VI-based inference algorithms
  2. automatically construct a conditioned model, i.e. perform model() | data under the hood

yebai avatar Mar 28 '25 15:03 yebai

Nice!

Can I also give a shout out for Batch based (e.g. https://turinglang.org/DynamicPPL.jl/stable/api/#DynamicPPL.MiniBatchContext ) and (nearly equivalently) tempered (e.g. https://github.com/TuringLang/Turing.jl/discussions/2524)

SamuelBrand1 avatar Mar 28 '25 16:03 SamuelBrand1

Hi @SamuelBrand1 , yes, minibatching is on our agenda.

Red-Portal avatar Mar 28 '25 17:03 Red-Portal

@yebai

A more ambitious goal would be infer(model, data = df, algorithm = ...; kwargs):

this is how RxInfer does it too

bvdmitri avatar Mar 31 '25 09:03 bvdmitri

Imagine a future where models are written in a universal backend-agnostic way (e.g., as a string) and then infer(model_as_str, ..., engine="Turing|RxInfer")

DominiqueMakowski avatar Mar 31 '25 10:03 DominiqueMakowski