mmm.fit Method Hangs Indefinitely During Model Fitting

[stdp95]

I’m encountering an issue where the mmm.fit method seems to run indefinitely and does not complete. Here’s the call I’m using:

mmm.fit(X_train, y_train, target_accept=0.95, chains=4, draws=2000, tune=1000)

Problem Details

• The process appears to hang during the fitting phase, with no progress or errors reported.
• I’ve ensured that my input data (X_train and y_train) are properly preprocessed and scaled as required.

Question

Has anyone else encountered this behavior? If so, are there any debugging steps or fixes I could try? Any insights into what might be causing the infinite runtime would be greatly appreciated.

PS: My config for reference

mmm_config = {
    'intercept': {'dist': 'Normal', 
                  'kwargs': {'mu': 0, 'sigma': 2}},
   
    'beta_channel': {'dist': 'LogNormal',
                     'kwargs': {'mu': 1, 'sigma': prior_sigma}},  
    
    'alpha': {'dist': 'Beta',
              'kwargs': {'alpha': 1, 'beta': 3}},
    
    'lam': {'dist': 'Gamma',
            'kwargs': {'alpha': 3, 'beta': 1}},
    
    'likelihood': {'dist': 'Normal',
                   'kwargs': {'sigma': {'dist': 'HalfNormal', 'kwargs': {'sigma': 2}}}},
    
    'gamma_control': {'dist': 'Normal',
                      'kwargs': {'mu': 0, 'sigma': 1}},
    
    'gamma_fourier': {'dist': 'Laplace',
                      'kwargs': {'mu': 0, 'b': 1}}
}

[stdp95]

When I previously tried the following configuration with fewer draws and a shorter tuning phase:

mmm.fit(X_train, y_train, target_accept=0.95, chains=4, draws=1000, tune=500)

I received the following warning message:

WARNING:pymc.stats.convergence:Chain 0 reached the maximum tree depth. 
Increase `max_treedepth`, increase `target_accept` or reparameterize.

Do you think this warning might be related to the issue I'm facing with the indefinite runtime? Could this indicate a problem with the priors or model configuration that’s causing sampling difficulties?

[AlfredoJF]

@stdp95 Have you tried Other NUTS samplers? PyMC's default sampler is quite slow IMO. Numpyro and Nutpie have these nice progress bars for each chain and estimations of runtime.

[stdp95]

I hadn’t tried other NUTS samplers before, but after your suggestion, I gave NumPyro a try, and it seems to work much better. Thanks for the recommendation.

[williambdean]

Thanks @AlfredoJF for the suggestion. nutpie is also a good alternative

[stdp95]

Since NumPyro has worked well for my needs, I haven’t prioritized exploring nutpie. While it seems like a good alternative, its reliance on Rust dependencies might complicate the setup in Databricks.

[williambdean]

Can you share the pymc marketing, pymc version as well as how you installed and your machine

[stdp95]

I installed pymc-marketing using pip, which also installed pymc as a dependency. Here are the versions I am using:

• pymc-marketing version: 0.8.0
• pymc version: 5.15.1

I am running this on Databricks.

[williambdean]

Are you able to sample simple models with pymc? Does that hang?

import numpy as np
import pymc as pm

seed = sum(map(ord, "Simple model"))
rng = np.random.default_rng(seed)

mu = np.array([1, 2, 3])
sigma = 0.25

n = 500
gdp = pm.Normal.dist(mu=mu, sigma=sigma, shape=(n, 3))
data = pm.draw(gdp, random_seed=rng)

coords = {"channel": ["A", "B", "C"], "idx": range(n)}
with pm.Model(coords=coords) as model:
    mu = pm.Normal("mu", dims="channel")
    sigma = pm.HalfNormal("sigma")
    pm.Normal("obs", mu=mu, sigma=sigma, observed=data, dims=("idx", "channel"))
   
idata = pm.sample(model=model)

Do you get any warnings or messages at any point?

[stdp95]

Nope, it ran fine without any issues. I didn’t encounter any warnings or messages during the sampling process.

[williambdean]

How does it do with the default priors? Are you changing the MMM configurations in any other way?

[stdp95]

Even with the default priors, the process is taking an unusually long time. I ended up canceling the run after 3 hours because it stalled at 28%, which was the same progress it had made in the first 10 minutes. I’m pretty sure that I haven’t made any changes to the MMM configurations.

[williambdean]

Can you show us the graphviz output?

[stdp95]

Thank you for pointing that out! I hadn’t realized that the keys beta_channel, alpha, and lam are not directly recognized as parameters in the model. Since the model remains silent when unused keys are present, it would be really helpful if you could share what the correct parameter names should be. This will ensure I align my configuration properly. Appreciate your guidance!

[stdp95]

I assume these are the correct parameter names:

{
    'intercept': Prior("Normal", mu=0, sigma=2),
    'likelihood': Prior("Normal", sigma=Prior("HalfNormal", sigma=2)),
    'gamma_control': Prior("Normal", mu=0, sigma=2, dims="control"),
    'gamma_fourier': Prior("Laplace", mu=0, b=1, dims="fourier_mode"),
    'adstock_alpha': Prior("Beta", alpha=1, beta=3, dims="channel"),
    'saturation_lam': Prior("Gamma", alpha=3, beta=1, dims="channel"),
    'saturation_beta': Prior("HalfNormal", sigma=2, dims="channel")
}

[williambdean]

Yes, appears right. I created an issue for this

[stdp95]

I updated the priors, but the progress bar remains stuck at 28% with no further progress.
model graph.pdf

[williambdean]

But no with the numpyro sampler though? We must not use the default C backend.

How does that sampling do? How are the parameter estimates

[williambdean]

Can you provide some information about your data, number of rows, numer of channels, number of controls, etc

[stdp95]

The dataset contains weekly data spanning 105 weeks. It includes 8 marketing channels and 1 control variable.