multiply_log(0, 0) has infinite gradient

[jsocolar]

Description

For more details, see https://discourse.mc-stan.org/t/non-linear-modelling-michaelis-menten/22736/8

Example

The gradient for poisson_lpmf(0 | lambda) is infinite at lambda = 0.

Suggested fix

@nhuurre writes on discourse:

We should just special-case a=0.0 like we do with a=1.0

Current Version:

v4.0.1

[SteveBronder]

Ah interesting, it looks like for reverse mode we don't check this, but we do when working with doubles so we def should be doing it for reverse mode!

https://github.com/stan-dev/math/blob/develop/stan/math/rev/fun/multiply_log.hpp

[mhollanders]

Any chance this could be addressed? I think accounting for this is making formulating complex HMMs on the log scale a bit more verbose than it needs to be.

[mhollanders]

@bob-carpenter, is this an edge case which makes it not easy to fix? It's quite pain for fitting hidden Markov models with log-scale matrix multiplication!

[SteveBronder]

Sorry this has not been addressed yet. So for the gradient we have

class multiply_log_vv_vari : public op_vv_vari {
 public:
  multiply_log_vv_vari(vari* avi, vari* bvi)
      : op_vv_vari(multiply_log(avi->val_, bvi->val_), avi, bvi) {}
  void chain() {
    avi_->adj_ += adj_ * log(bvi_->val_);
    bvi_->adj_ += adj_ * avi_->val_ / bvi_->val_;
  }
};

Where the gradient accumulation becomes infinity / nan when b is zero because we are doing log(b) and a / b.

    avi_->adj_ += adj_ * log(bvi_->val_);
    bvi_->adj_ += adj_ * avi_->val_ / bvi_->val_;

Here if b is 0, aka when lambda in @jsocolar's example is 0, the gradient will want to do both a log and division where the input is 0.0. So the gradient is undefined here. We can fix this so that if and b is zero the gradient will be zero, but this will still have issues when b is zero. Though I think that is fine with how we've defined how multiply_log should work. If that's okay with you two then I can put up a PR to fix this for just the case a and b are zero

[jsocolar]

If b is zero and a is nonzero, I think it makes sense that the gradient would be undefined. The important thing here is that when a is zero the limit of the gradient as b approaches zero is clearly defined and is zero. So we want to return zero gradient whenever a is zero, including when b is zero.

[bob-carpenter]

Just so it's clear to others following along, the computation is

${\large \text{multiplyLog}(a, b) = a \cdot \log(b).}$

I believe the intent was to set $\text{multiplyLog}(0, 0) = 0$.

The gradient is

${\large \nabla \ \text{multiplyLog}(a, b) = \left[\log(b) \quad \dfrac{a}{b}\right]}$

There are two elements of the gradient. I'm not sure which one is being discussed or both. If $a = b = 0$, then the naive numerical answer is

${\large \nabla \ \text{multiplyLog}(0, 0) = \left[\log(0) \quad \dfrac{0}{0}\right]} = \left[ -\infty \quad \text{NaN} \right].$

If we just fix $\text{multiplyLog}(0, 0) = 0$, as if it were a constant, then the gradients would both be zero.

[jsocolar]

Thanks for the clarity @bob-carpenter . This is only relevant when a is a literal zero. I'm out of my depth here, but I think what we care about therefore is setting the partial with respect to b to zero when a is a literal zero. I don't know enough to say whether it's important to do anything particular with the partial wrt a. It depends on whether there will be situations where a is constant (and zero) but autodiff still thinks it needs to differentiate through a. My thinking is that if a is continuous and not constant, it will never be a literal zero. But I could be committing major sins of fragility and brittleness here--I just don't know enough to say.

[mhollanders]

I too am out of my depth. I came to this issue because @jsocolar linked it in a discourse thread about running into infinite gradient errors with a hidden Markov model. I can replicate the issue like this, where I get infinite gradient errors:

parameters {
  real<lower=0, upper=1> theta;
}

model {
  vector[2] x = rep_vector(theta, 2) + rep_vector(negative_infinity(), 2);
  real out = log_sum_exp(x);
}

Is this the same issue as the multiply_log() issue described herein?

[bob-carpenter]

My thinking is that if a is continuous and not constant, it will never be a literal zero.

That's true for real numbers, but floating point can underflow. That's where all the "scale is zero, but must be nonzero" errors at run time come from.

Is this the same issue as the multiply_log()

Note quite---it's a different function and there's no multiply-log involved in log-sum-exp. This is a case of not dealing with the negative infinity. I think for x you meant [theta, negative_infinity()]'---what you have will work out to be two negative infinities, which should be an error.

[mhollanders]

Hey Bob, I might have to ask in a different thread, but matrix multiplication on the probability scale works fine, but not on the log scale. I suspect it could be that the gradient of 0 is handled well but not of negative_infinity()?

[jsocolar]

Thanks again @bob-carpenter . As I mentioned above, I'm out of my depth, but I'll leave one more comment in case it's helpful:

I don't understand Stan at the C++ level well enough to say whether we can trust a to be recognized as a constant type when a is in fact a constant. If that were the case, then I think we can get what we want out of multiply_log by either:

• when a = 0, special-casing the partial only with respect to b (and not wrt a). This would work if Stan never attempts to take a partial wrt a when a is a constant.
• only special casing scenarios where the input for a is constant and is zero. This would work if we can trust a to come in as a constant type when a is in fact a constant.

These sorts of ideas would not fix the problems that raised this issue unless we can trust that a will be recognized as constant by multiply_log. Having limited experience with strongly typed languages, I'll see myself out now.

Here's the in-the-wild example that originally raised the issue:
https://discourse.mc-stan.org/t/non-linear-modelling-michaelis-menten/22736

[bob-carpenter]

trust a to be recognized as a constant type when a is in fact a constant.

Yes, you can trust that. We have template metaprograms to distinguish primitives (int, double) from autodiff variables (class stan::math::var).

We have tried very hard to avoid differing return or exception behavior of functions when passed autodiff variables (parameters) versus primitives (data).

strongly typed languages,

C++ is declaratively strongly typed, but you can always cast away const modifiers, convert to raw memory, and start bit fiddling. So while it's strongly typed, it's not safe in the way it is in other languages that don't let you access the raw memory.