Optimizer Latency

[AnantChandra]

Hey all!
What is the best way to speed this up?

model_granularity = "daily"
allocation_strategy, optimization_result = mmm.optimize_budget(
    budget=time_unit_budget,
    num_periods=num_periods,
    budget_bounds=constrained_bounds,
    utility_function=ut.mean_tightness_score(alpha=0.95),
    minimize_kwargs={
        "method": "SLSQP",
        "options": {"ftol": 1e-9, "maxiter": 200, "disp": True}
    }
)
allocation_strategy

It currently takes me 5 minutes for the number of channels & periods I have.

My understanding from reading the code base is that the optimizer uses SciPy's solvers, which are not JAX-compatible and do not use the GPU. Hence, deploying on a GPU won't help either?

I'd rather not compromise on the maxiter. Is there any other way to speed it up? My goal is to make it run in less than a minute.

[williambdean]

When the objective and its gradient are compiled in pytensor, you should be able to leverage a different pytensor backend via environment variables

Reference:

pymc-marketing/pymc_marketing/mmm/budget_optimizer.py

Line 372 in ae215f0

objective_and_grad_func = function([budgets_flat], [objective, objective_grad])

https://pytensor.readthedocs.io/en/latest/tutorial/modes.html

[AnantChandra]

Is this a battle-tested approach?
I was facing precision issues with JAX