explicit log parameterization for Poisson, NegativeBinomial, etc (Stan-style *_log variants)

[lucidfrontier45]

When we use liner model or neural network with Possison/NegativeBinomial/Bernouli/Binomial as output distribution, we generally apply exp to the predictor. $\mu = \exp(z)$

when z is relatively large number, like 700, it explodes and MCMC sampling becomes unstable.

Stan has explicit log parameterization to address this numerical instability.

• https://mc-stan.org/docs/functions-reference/binary_distributions.html#bernoulli-logit-distribution
• https://mc-stan.org/docs/functions-reference/bounded_discrete_distributions.html#binomial-distribution-logit-parameterization
• https://mc-stan.org/docs/functions-reference/unbounded_discrete_distributions.html#poisson-distribution-log-parameterization
• https://mc-stan.org/docs/functions-reference/unbounded_discrete_distributions.html#neg-binom-2-log-glm

Currently I have encountered this problem and my workaround is to apply normalization to variables whose scale is relatively large.

However if I can do like this it's better.

with pm.Model() as m:
    alpha = pm.Normal("alpha", 0, 1)
    pm.Poisson("y", alpha=alpha, observed=y, log_scale=yes)

[ricardoV94]

@lucidfrontier45 there's a subtlety in that pymc relies on pytensor which does stabilization rewrites automatically.

For Bernoulli we know pytensor can stabilize the expression so that "naive" Bernoulli(p=sigmoid(a)) is stable and we offer Bernoulli(logit_p=a) just as syntactic sugar.

In the Poisson case we write the Poisson logp as -mu + logpow(mu, y) - factln(y), where logpow is just y * log(mu) with a guard for 0 * log(0).

pymc/pymc/distributions/discrete.py

Lines 581 to 597 in bf8abd5

	def logp(value, mu):
	res = pt.switch(
	pt.lt(value, 0),
	-np.inf,
	logpow(mu, value) - factln(value) - mu,
	)
	# Return zero when mu and value are both zero
	res = pt.switch(
	pt.eq(mu, 0) * pt.eq(value, 0),
	0,
	res,
	)
	return check_parameters(
	res,
	mu >= 0,
	msg="mu >= 0",
	)

PyTensor will simplify that log(exp(alpha)) into alpha (with a negative alpha nan guard perhaps). The exp term survives for the guard (not a problem) and the -mu term (unavoidable).

Do you have a case where the log density or gradient come back at -inf / nan, when you know it could be represented in float64?

import pytensor
import pymc as pm

with pm.Model() as m:
    alpha = pm.Flat("alpha")
    pm.Poisson("y", mu=pm.math.exp(alpha))

logp_expr = m.logp()

# it starts with the expression shown above
logp_expr.dprint()
# └─ Mul [id BX]
#    ├─ y [id BG]
#    └─ Log [id BY]
#       └─ Exp [id BD]
#          └─ ···

# excluding fusion just to keep the variable names readable
logp_fn = m.compile_logp(mode=pytensor.get_mode().excluding("fusion"))
# It gets common stabilization, in this case undoes the log(exp(alpha))
logp_fn.dprint()
# └─ Mul [id X]
#    ├─ y [id O]
#    └─ alpha [id F]

print(logp_fn({"y": 3, "alpha": 600}))  # -3.7730203009299397e+260
print(logp_fn({"y": 3, "alpha": 709}))  # -8.218407461554972e+307
print(logp_fn({"y": 3, "alpha": 710}))  # -inf,

The last value is -inf, but the true log density is not expressible in float64. You'd need y around 1e305 for the y * alpha term to offset it, which isn't representable as an int64 to begin with.

---

However we may be missing something for Binomial and NegativeBinomial, those don't seem to be stabilized as of today.

import pymc as pm

def probe(build, zs, label):
    with pm.Model() as m:
        z = pm.Flat("z")
        build(z)
    lp, dlp = m.compile_logp(), m.compile_dlogp()
    for zv in zs:
        print(f"{label}  z={zv:6}  logp={lp({'z': zv}):>22}  dlogp={dlp({'z': zv})[0]:>8}")

# mu = exp(z)
probe(lambda z: pm.NegativeBinomial("y", mu=pm.math.exp(z), alpha=2.0, observed=3),
      (600, 709, 710, 5000), "NegativeBinomial")

# p = sigmoid(z)
probe(lambda z: pm.Binomial("y", n=10, logit_p=z, observed=1),
      (-700, 30, 36, 37, 800), "Binomial        ")

NegativeBinomial  z=   600  logp=   -1197.2274112777604  dlogp=    -2.0
NegativeBinomial  z=   709  logp=   -1415.2274112777604  dlogp=    -inf
NegativeBinomial  z=   710  logp=                   nan  dlogp=     nan
NegativeBinomial  z=  5000  logp=                   nan  dlogp=     nan
Binomial          z=  -700  logp=     -697.697414907006  dlogp=     1.0
Binomial          z=    30  logp=    -267.6974149070069  dlogp=-8.999999999999064
Binomial          z=    36  logp=   -321.69741490700596  dlogp=-8.999999999999998
Binomial          z=    37  logp=                  -inf  dlogp=8.533047625744066e-17
Binomial          z=   800  logp=                  -inf  dlogp=     0.0

---

For your specific case, besides normalizing predictors you can also try pt.softplus, it's positive and becomes linear as x grows

[lucidfrontier45]

Thank you so much for the explanation.

Actually my usecase is exactly NegativeBinomial which isn't stabilized yet.
In the last sentence of your reply did you mean you will also address it by rewriting on PyTensor side?

I'll also try to replace last exp to softplus to see if it is stabilized.

[lucidfrontier45]

I tried softplus and it did make process stabilized. However my current system is tuned for exp case and need large reconfiguration elsewhere for good prediction accuraccy. I'll stick to normalization approach until stabilzed NegativeBinomial is implemented. Looking forward to seeing it.

[ricardoV94]

In the last sentence of your reply did you mean you will also address it by rewriting on PyTensor side?

Yeah I'll give it a try and see if it's good enough. Can you share a bit more details on what regime you're operating that's failing?

[lucidfrontier45]

Thank you so much. Both my model and data are proprietary. Please give me time to remove private factors and reduce them into a simple NegativeBinomial regression.

[ricardoV94]

@lucidfrontier45 trying something out in #8394

[ricardoV94]

@lucidfrontier45 can you check if the last release makes the NB stable in your use case? https://github.com/pymc-devs/pymc/releases/tag/v6.3.0

[lucidfrontier45]

@ricardoV94

Thank you for update. I have also made a minimal project to reproduce the problem.
Since tgz file is not allowed to attach in this github discussion, I created a Google Drive link.
https://drive.google.com/file/d/1wKca8ReGh3FrHcoBEl0kvMHFxLRr9yjU/view?usp=sharing

The anonymized data are in features directory. Essentially my model is Negative Binomial linear regression.
One problem is that feature set c has large scale difference, especially c0 is quite large compared to c1 and c2.
I couldn't stably draw samples even with v6.3.0. The only way to stabilize it is to normalize c0.
Please check fit_nb_model.py script.

If you think it is my modeling problem and normalization is the correct approach, please feel free to close this discussion.

[ricardoV94]

I had claude analyze it, so pasting it's analysis, I tried to verify and keep it brief.

---

Short version: the NegativeBinomial logp is fine in 6.3.0, The density is stable. On your data, logp and dlogp are finite everywhere I probed, including points where mu = exp(z) is literally inf (z ≈ 21600). C and JAX backends agree to ~1e-7. So nothing is overflowing in the distribution anymore.

What breaks is warmup, and only with the JAX samplers. PyMC jitters the initial point, which for your model lands at z ∈ [275, 1592]. logp is finite there, but the gradients are ~1e9, and blackjax/numpyro's dual averaging responds by shrinking the step size monotonically for the whole warmup — down to 1e-226 — so the chain never leaves its starting point. Every draw is identical, and that's the mu explosion your check reports. It's a coin flip on the seed (3 of 4 froze for me). Two things sidestep it:

pm.sample(..., nuts_sampler="blackjax", nuts={"jitter": False})  # works

and nutpie samples your model as-is, no normalization, 0 divergences in ~30s. PyMC's own NUTS also works, just slowly (treedepth pinned at 10).

The real issue is the parameterization. c_0 ranges over [376, 2162] with a Laplace(0, 0.1) prior on its weight, so your prior predictive z spans roughly [-1300, +2050] — 41% of prior draws have some mu > 1e13 — while the data implies z ∈ [7.9, 10.0]. The posterior needs c_w0 = 5.3e-4 ± 3.1e-5, about 3000× narrower than its own prior, next to c_1/c_2 at O(0.3). No sampler should be expected to be robust there.

Suggestion: standardize c_0 (or put its scale in the prior, Laplace(0, 0.1 / c0_scale)).

One unrelated thing worth fixing: b_0..b_6 sum to 1, making them exactly collinear with your intercept (design rank 56 of 57). Drop one column or the intercept.

---

End of reply. I do suggest using nutpie, it's the default if you do pip install pymc[nutpie], and you can use it with jax backend if you want with backend="jax" just fine. You'll still need to be careful about priors and scales, but it will be more robust.

https://www.pymc.io/blog/pymc_v6_ecosystem_updates.html

[lucidfrontier45]

I do know the scale of variables are very different. My product allow any explanatory variables and this situation is inevitable. Even after the new stabilization I'll still apply normalization.

On sampler and backend, when I first evaluated this model with PyMC v5, I found numba+nutpie was so slow especially in model compiling phase and jax+blackjax wasn't. In my application there is also automatic hyper-parameter tuning and I wanted to avoid doing long compiling multiple times. As you posted if in PyMC v6 the numba+nutpie was significantly improved I shall evaluate again.

Thank you so much. Close this thread.

[ricardoV94]

nutpie works with both jax and numba backends so you can have best of both worlds