Add thin wrapper around advi functionality

[PabloRoque]

Description

Introduces a thin wrapper around pm.fit so that advi can be used in CLV models
Some tests related to advi functionality

Related Issue

• Closes #
• Related to Allow MAP fitting for CLV models #130

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Modules affected

• MMM
• CLV
• Customer Choice

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc-marketing--1365.org.readthedocs.build/en/1365/

[codecov]

Codecov Report

Attention: Patch coverage is 20.00000% with 16 lines in your changes missing coverage. Please review.

Project coverage is 55.97%. Comparing base (32f44c7) to head (9b70003).

Files with missing lines	Patch %	Lines
pymc_marketing/clv/models/basic.py	20.00%	16 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (32f44c7) and HEAD (9b70003). Click for more details.

HEAD has 7 uploads less than BASE

Flag	BASE (32f44c7)	HEAD (9b70003)
	11	4

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1365       +/-   ##
===========================================
- Coverage   93.94%   55.97%   -37.98%     
===========================================
  Files          48       48               
  Lines        5137     5156       +19     
===========================================
- Hits         4826     2886     -1940     
- Misses        311     2270     +1959     

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

[ColtAllen]

@wd60622 do you think it's more appropriate for advi support to be added to ModelBuilder instead?

@PabloRoque if your primary motivation is to speed up model fits, have you tried using nutpie? There's example code in this notebook.

[wd60622]

@wd60622 do you think it's more appropriate for advi support to be added to ModelBuilder instead?

Maybe long term but implementing now would require changing the fit api for all

[PabloRoque]

Thanks for the suggestion @ColtAllen.

The issue is mainly that our customer base has 1E8 order of magnitude. I am afraid I'll hit the wall using mcmc when we put things on PROD, regardless of the sampler (normally I use numpyro), so I'm starting to explore advi.

[wd60622]

Why this syntax? Will you add more kwargs here?
Just use method=advi

[PabloRoque]

I was planning to add more kwargs.
I added now functionality to parse the params and kwargs for both fit and sample

[ColtAllen]

Thanks for the suggestion @ColtAllen.

The issue is mainly that our customer base has 1E8 order of magnitude. I am afraid I'll hit the wall using mcmc when we put things on PROD, regardless of the sampler (normally I use numpyro), so I'm starting to explore advi.

I would encourage you to try nutpie before getting any further into this. I was not impressed with the results I got from ADVI fits when I experimented with it several years ago, and I don't think development on the ADVI module in PyMC is highly prioritized either (@ricardoV94)? ADVI is even slower than MCMC when fitting ParetoNBDModel, and is unusable for BetaGeoBetaBinomModel because it's a discrete model instead of a continuous one.

For datasets with millions of customers I usually recommend just using MAP. MAP is generally discouraged in the Bayesian community, but these CLV models are an exception due to the strong population assumptions underlying them. Beyond 30-50k customers, mean values for MAP and MCMC fits are identical. The only downside is that you'll lose the credibility intervals with MAP.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[PabloRoque]

Thanks for the advices @ColtAllen.

I added a gist comparing DEMZ with ADVI and FULLRANK:

• Fullrank does not match mcmc at all
• Advi fit does not match for r and alpha, which is quite worrisome

I the light of your advices, and the little evidence I gathered myself, I'll be closing this PR

Thank you all!

[ColtAllen]

@PabloRoque glad to help! Looks like the gist you did for ParetoNBDModel also revealed a new multiprocessing warning to investigate, which is helpful.

[PabloRoque]

I am sorry to be the bearer of "bad news", but I may be reopening this PR.

I've been able to improve the fit adjusting obj_n_mc following these advices.

The updated gist shows:

• Poor fit using fullrank still. Happy to remove this functionality in the light of the results
• Improved ADVI fit.
  • The mean of the param estimates differ less than 2% in the worst case
  • The sd is much worse, with advi providing narrower distributions

Some considerations, though.

• Fitting time is considerably longer.
• Still, ADVI allows for minibatch training

On a side note, I also tried to use pymc_extras.fit(method="pathfinder"). However jax does not support hyp2f1. [I see now why @ColtAllen is suggesting to use nutpie instead of numpyro]

[wd60622]

I'm not opposed to getting this through even if it is not "recommended" way of sampling the current models.

Support for minibatching is another beast. What work would be required there?

[ColtAllen]

On a side note, I also tried to use pymc_extras.fit(method="pathfinder"). However jax does not support hyp2f1. [I see now why @ColtAllen is suggesting to use nutpie instead of numpyro]

Yes, I briefly looked into a jax implementation, but it would be quite challenging and with uncertain benefits. pytensor is already calling scipy.special.hyp2f1 for this function, which is written in C++. The problem is that Hyp2F1 gradients are very complex to calculate, which is why gradient-free samples are so much faster for it. On that note, Hyp2F1 is also only relevant for model fitting ParetoNBDModel.

My preference would be to make this more of a long-term PR for ModelBuilder so that MMMs are also supported. I could see minibatch ADVI being useful for hierarchical models. @wd60622 how "translatable" do you see this PR being for the ModelBuilder fit API? If this is a good starting point we can proceed, but I don't want it to become too proprietary to CLV?

@PabloRoque I've got a lot on my plate at the moment, but will make time to view your other PRs tomorrow & over the weekend. We're pushing to release pymc-marketing v0.11.0 ASAP.

[PabloRoque]

@ColtAllen No rush!

I had a bit of bandwidth and attempted a few contributions, but I am not expecting a quick review in any case.