BUG: ADVI calculation is ~4x slower after 5.9.1

[erik-werner]

Describe the issue:

When I run the notebook pymc_examples/case_studies/factor_analysis.ipynb the ADVI calculation trace_vi = pm.fit(n=50000, method="fullrank_advi", obj_n_mc=1).sample() (just before the "Results" section) is very slow on recent versions of PyMC. For versions 5.9.1 and later this step takes ~4 min, compared to ~1 min for 5.9.0.

Reproduceable code example:

Let me know if you want me to try to make a self-contained MRE, or if the notebook reference is enough.

PyMC version information:

conda create -y -n pymc_5.9.1 "pymc==5.9.1" using Miniforge.

[welcome]

🎉 Welcome to PyMC! 🎉 We're really excited to have your input into the project! 💖

If you haven't done so already, please make sure you check out our Contributing Guidelines and Code of Conduct.

[ricardoV94]

Do you get a BLAS warning when you import pymc at the start?

[erik-werner]

@ricardoV94 No.

[ricardoV94]

If you rerun it after (without restarting the kernel) does it take the same time? How is CPU/RAM usage during fitting?

[erik-werner]

No improvement from running a second time.

All 8 cores at 100%.

[ricardoV94]

Did it also max all cores in the older version?

[erik-werner]

Yes.

[erik-werner]

Also, I have now reproduced the behaviour both on my work computer (WSL) and home computer (Linux).

[erik-werner]

I've now run a git bisect, and ce543da is the first bad commit.

[ricardoV94]

We know about that slowdown but should have been fixed after PyTensor 2.17.4
https://github.com/pymc-devs/pytensor/releases/tag/rel-2.17.4

Which version are you running? You can also try PyMC 5.10.0

[erik-werner]

PyMC 5.10.0 is also slow for me. I first tried installing fresh environments with conda create and specifying specific releases. 5.9.1, 5.9.2 and 5.10.0 were slow, but 5.9.0 fast. My latest test was run with pytensor==2.18.1, and doing a git bisect between main and 5.9.0.

[erik-werner]

This is the script I've used for testing. It takes ~1 min to run in the good case and ~4 in the bad one.

import arviz as az
import numpy as np
import pymc as pm
import pytensor.tensor as pt

from numpy.random import default_rng

print(f"Running on PyMC v{pm.__version__}")
az.style.use("arviz-darkgrid")

np.set_printoptions(precision=3, suppress=True)
RANDOM_SEED = 31415
rng = default_rng(RANDOM_SEED)

n = 250
k_true = 4
d = 10

err_sd = 2
M = rng.binomial(1, 0.25, size=(k_true, n))
Q = np.hstack([rng.exponential(2 * k_true - k, size=(d, 1)) for k in range(k_true)]) * rng.binomial(
    1, 0.75, size=(d, k_true)
)
Y = np.round(1000 * np.dot(Q, M) + rng.standard_normal(size=(d, n)) * err_sd) / 1000

k = 2

coords = {"latent_columns": np.arange(k), "rows": np.arange(n), "observed_columns": np.arange(d)}

def expand_packed_block_triangular(d, k, packed, diag=None, mtype="pytensor"):
    # like expand_packed_triangular, but with d > k.
    assert mtype in {"pytensor", "numpy"}
    assert d >= k

    def set_(M, i_, v_):
        if mtype == "pytensor":
            return pt.set_subtensor(M[i_], v_)
        M[i_] = v_
        return M

    out = pt.zeros((d, k), dtype=float) if mtype == "pytensor" else np.zeros((d, k), dtype=float)
    if diag is None:
        idxs = np.tril_indices(d, m=k)
        out = set_(out, idxs, packed)
    else:
        idxs = np.tril_indices(d, k=-1, m=k)
        out = set_(out, idxs, packed)
        idxs = (np.arange(k), np.arange(k))
        out = set_(out, idxs, diag)
    return out

def makeW(d, k, dim_names):
    # make a W matrix adapted to the data shape
    n_od = int(k * d - k * (k - 1) / 2 - k)

    # trick: the cumulative sum of z will be positive increasing
    z = pm.HalfNormal("W_z", 1.0, dims="latent_columns")
    b = pm.HalfNormal("W_b", 1.0, shape=(n_od,), dims="packed_dim")
    L = expand_packed_block_triangular(d, k, b, pt.ones(k))
    W = pm.Deterministic("W", pt.dot(L, pt.diag(pt.extra_ops.cumsum(z))), dims=dim_names)
    return W

coords["observed_columns2"] = coords["observed_columns"]
with pm.Model(coords=coords) as PPCA_scaling:
    W = makeW(d, k, ("observed_columns", "latent_columns"))
    Y_mb = pm.Minibatch(
        Y.T, batch_size=50
    )  # MvNormal parametrizes covariance of columns, so transpose Y
    psi = pm.HalfNormal("psi", 1.0)
    E = pm.Deterministic(
        "cov",
        pt.dot(W, pt.transpose(W)) + psi * pt.diag(pt.ones(d)),
        dims=("observed_columns", "observed_columns2"),
    )
    X = pm.MvNormal("X", 0.0, cov=E, observed=Y_mb)
    trace_vi = pm.fit(n=50000, method="fullrank_advi", obj_n_mc=1).sample()

[ricardoV94]

Thanks for taking a look. I'll check on Monday and see if I can reproduce

[ricardoV94]

Ahh, probably ADVI is using pytensor directly to get the gradient expressions, and not calling this utility:

pymc/pymc/pytensorf.py

Line 1027 in 02cbac6

def rewrite_pregrad(graph):