Offset multiplier initialization

Thanks for the super careful follow-up! I’d really like to get to the bottom of this.

Values that are simulated non-zero can be consistent with zero in the posterior. This won’t typically happen in a small model, but in larger models, things get a lot more complicated, typically by allowing alternative ways to get non-zero variance.

What’s happening with the epsilon-style lower bound is that it will “save” a trajectory that underflows a scale to zero (-infinity) on the unconstrained scale by boosting epsilon above zero. You lose the gradients due to the underflow, which is why it’s not ideal.

As far as I know, the only other difference is that the offset/multiplier version will reject when the multiplier is zero. So it makes sense that’s the place where behavior differs.

Arguably we should’ve just let this case go and let the pathological behavior randomize its way out. It kind of blows out the theory, but probably not enough to make a difference in practice—we probably wouldn’t even be able to measure the amount of bias introduced if there isn’t a lot of probability mass near zero in the posterior.

I was expecting to see two different sets of inits. One to use with the affine transform and one to use without. I suggest using the exact same generated values to test. That is, create an init for the unmarked case of

init ~ uniform(-2, 2)

when you implement manually without the affine transform and then match it with

init_affine = offset + multiplier * init

for the case with the affine transform. Then the behavior should be similar up to minor arithmetic ordering differences. The place where it will vary is if multiplier ever underflows to zero—in that case the affine transform will automatically reject, whereas the manually reparameterized will underflow and lose connection to gradients.

I was expecting to see two different sets of inits. One to use with the affine transform and one to use without. I suggest using the exact same generated values to test. That is, create an init for the unmarked case of

init ~ uniform(-2, 2)

when you implement manually without the affine transform and then match it with

init_affine = offset + multiplier * init

for the case with the affine transform. Then the behavior should be similar up to minor arithmetic ordering differences. The place where it will vary is if multiplier ever underflows to zero—in that case the affine transform will automatically reject, whereas the manually reparameterized will underflow and lose connection to gradients.

I think that is what I’m doing. That function generates initial values for the logit_p_raw parameters as uniform(-2, 2), but then logit_p is supplied through transformation.

In any case, I think I figured out the difference between the two models I’ve written compared to the models mentioned earlier in this thread and perhaps a clue to why I’m seeing such different behavior. binom_heir_raw.stan features a transformation, while binom_heir_affine.stan features a change of variables. (Relevant User’s Guide link). A different section of the Users guide suggests that these two approaches are equivalent when doing a non-centered parameterization. But the affine transformation is doing a Jacobian adjustment under the hood, while in my version of manual non-centering the raw parameter is sampled, so no Jacobian is necessary.

If I fit this version of the manual non-centered model, I see similar poor behavior (I think I did this right):

data {
  int<lower = 0>  I;
  array[I] int    N, y;
}
parameters {
  real                     mu;
  real<lower = 0>          sigma;
  vector[I]                logit_p_raw;
}
transformed parameters{
  vector[I] logit_p = fma(sigma, logit_p_raw, mu);
  jacobian += log(sigma) * I;
}
model{
  mu      ~ student_t(3, 0, 2.5);
  sigma   ~ loglogistic(1, 4);
  logit_p ~ normal(mu, sigma);
  
  y ~ binomial_logit(N, logit_p);
}

Going back to what @martinmodrak said above:

In addition the offset-multiplier logic (both manual and controled by Stan) adds unnecessary operations to the model: log(sigma) is added in the Jacobian correction, only to be subtracted in x ~ normal(0, sigma) , with a penalty for both performance and numerical precision.

I wonder if maybe the issue is something funky happening in adding log(sigma) to the Jacobian when sigma is close to zero?

Deleted my previous message because the reprex didn’t actually work (I was misinterpreting the plot). I’ll produce a proper reprex and re-post at some point.

I’m really glad this is getting picked back up–I never really understood the resolution to it when it came up years ago.

Just want to chime in now to make one point that is conceptually tangential but practically important.

In models with big random effect vectors, the early phases of exploration routinely drop down to tiny values for the random effect variances for every single chain. This only happens when using the non-centered parameterization, but happens regardless of whether we non-center manually or with offset/multiplier. When using offset/multiplier, there’s a chance that we get stuck and don’t recover, but this odd early warmup behavior is a general feature of the non-centered paramterization irrespective of whether or not we use offset/multiplier. I’ve never had a confident understanding of why this happens; my guess is that the sampler falls down into the global mode that sits on zero variance for the random effect before clawing its way back out. I’ve always wondered whether, even in cases where the sampler recovers, this poses any challenges for inference (either with full MCMC or with optimization/pathfinder) due to the fact that effectively we are not getting a reasonable range of initial values for the random effect sd–we are always directing our exploration through tiny values and then expanding back out from there.

Here are some examples using the following data:

n <- 1000
mu <- 0
sigma <- .5
n_trials <- rep(5, n)

re <- rnorm(n, mu, sigma)
y <- rbinom(n, n_trials, boot::inv.logit(re))

re_data <- list(
  N = n,
  y = y,
  n_trials = n_trials
)

Below, we’ll always produce plots for various models with

fit <- mod$sample(
  data = re_data, chains = 8, parallel_chains = 4, 
  save_warmup = T, iter_warmup = 50, iter_sampling = 1
  )
samps <- fit$draws(variables = "sigma", inc_warmup = TRUE)
bayesplot::mcmc_trace(samps)

The centered parameterization:

data {
  int<lower=0> N;
  array[N] int<lower=0> n_trials;
  array[N] int y;
}

parameters {
  real mu;
  real<lower=0> sigma;
  vector[N] random_effects;
}

model {
  random_effects ~ normal(0, sigma);
  y ~ binomial(n_trials, inv_logit(mu + random_effects));
}

The manually non-centered parameterization:

data {
  int<lower=0> N;
  array[N] int<lower=0> n_trials;
  array[N] int y;
}

parameters {
  real mu;
  real<lower=0> sigma;
  vector[N] random_effects;
}

model {
  random_effects ~ std_normal();
  y ~ binomial(n_trials, inv_logit(mu + sigma * random_effects));
}

The offset-multiplier parameterization:

data {
  int<lower=0> N;
  array[N] int<lower=0> n_trials;
  array[N] int y;
}

parameters {
  real mu;
  real<lower=0> sigma;
  vector<offset=mu, multiplier=sigma>[N] random_effects;
}

model {
  random_effects ~ normal(mu, sigma);
  y ~ binomial(n_trials, inv_logit(random_effects));
}

Note that this isn’t unique to binomial models. We get the same behavior with a normal model as well.

data {
  int<lower=0> N;
  int<lower=0> n_grp;
  vector[N] y;
  array[N] int<lower=1, upper=n_grp> grp_id;
}

parameters {
  real mu;
  real<lower=0> sigma;
  real<lower=0> sigma_residual;
  vector[n_grp] random_effects_raw;
}

model {
  random_effects_raw ~ std_normal();
  y ~ normal(mu + sigma * random_effects_raw[grp_id], sigma_residual);
}
n_grp <- 1000
n_per_grp <- 5
mu <- 0
sigma_re <- .5
sigma_residual <- 1

re <- rnorm(n_grp, mu, sigma_re)
y <- rnorm(n_grp * n_per_grp, re, sigma_residual)
grp_id <- rep(1:n_grp, n_per_grp)

re_data <- list(
  N = n_grp * n_per_grp,
  n_grp = n_grp,
  y = y, 
  grp_id = grp_id
)


fit <- re_mod$sample(data = re_data, chains = 8, parallel_chains = 4, save_warmup = T, 
                     iter_warmup = 50, iter_sampling = 1)
samps <- fit$draws(variables = "sigma", inc_warmup = TRUE)
bayesplot::mcmc_trace(samps)

I’d LOVE it if somebody could explain to me why we see this early pinch through tiny values of sigma in the NCP but not the centered! (edit: and hopefully assuage my fear that this could cause problems for convergence diagnostics by making it effectively impossible to specify overdispersed inits for sigma)

Let’s work out the effect, simplifying away the intercept which doesn’t modify the geometry.

parameters {
  real<lower=0> sigma;
  real<multiplier = sigma> alpha;
model {
  alpha ~ normal(0, sigma);

we will wind up with an unconstrained parameter alpha_unc,

alpha = sigma * alpha_unc;
jacobian += log(sigma);  // Jacobian is d/d.alpha_unc sigma * alpha_unc = sigma
target += normal_lpdf(alpha | 0, sigma);

The Jacobian adjustment is d / d.alpha_unc (sigma * alpha_unc) = sigma, so its log is log(sigma). Conveniently, that’s just what we need to cancel out

log p(alpha_unc)
    = log(normal(sigma * alpha_unc | 0, sigma)) + log(sigma)
    = ((sigma * alpha_unc) / sigma)^2 - log(sigma) + log(sigma)
    = alpha_unc^2

This is why the User’s Guide says it does the same thing. It’s a little less efficient with the cancelling log(sigma), but we find it much easier to code if not so easy to understand at a fundamental level.

That is strange, especially as it looks like you’re randomly initializing the random effects. But you have set iter_warmup=50, which changes how everything works because it’s too few iterations to do our normal default stages. I actually don’t know what it does in terms of either step size or mass matrix adaptation. Do you see this with our usual 75 Phase I warmup iterations using the defaults?

Although it doesn’t explain the behavior in the non-centered case, with the centered parameterization, the geometry makes it almost impossible to get down to near zero values of the hierarchical scale parameters.

I used a short period to make the simulations fast and easy to plot without extra code. However, this happens for normal warmup schedules as well. When not using the offset/multiplier syntax, we don’t notice except by saving and examining the warmup (which we rarely do). With offset/multiplier, sometimes we hit this sticky boundary and get nonsense out. When I first noticed this (see first post on the thread) it was because it came up “in the wild” in a model I was working with.

This crams a 38-iteration window between a 7-iteration init buffer and a 5-iteration term buffer. But note that we crash down towards zero sigma before the window ends, and therefore before the mass matrix updates, and before the step size is boosted at the end of the window. So by the time the adaptation actually does anything different from the default 75-iteration init buffer the behavior has already manifested. Indeed, in many cases above the behavior has already manifested before we even end the foreshortened init buffer.

@Bob_Carpenter just for completeness, here’s the normal model from above, run for the default 1000 warmup iterations, with the first 50 warmup iterations plotted on the traceplot:

Let’s work out the effect, simplifying away the intercept which doesn’t modify the geometry.

parameters {
  real<lower=0> sigma;
  real<multiplier = sigma> alpha;
model {
  alpha ~ normal(0, sigma);

we will wind up with an unconstrained parameter alpha_unc,

alpha = sigma * alpha_unc;
jacobian += log(sigma);  // Jacobian is d/d.alpha_unc sigma * alpha_unc = sigma
target += normal_lpdf(alpha | 0, sigma);

The Jacobian adjustment is d / d.alpha_unc (sigma * alpha_unc) = sigma, so its log is log(sigma). Conveniently, that’s just what we need to cancel out

log p(alpha_unc)
    = log(normal(sigma * alpha_unc | 0, sigma)) + log(sigma)
    = ((sigma * alpha_unc) / sigma)^2 - log(sigma) + log(sigma)
    = alpha_unc^2

I’m probably in over my head, but this is helpful for me in thinking through all of the steps that might be causing this. I’d like to lay out my understanding of all of what happens when you declare an affine transform and I’m hoping somebody can correct me where I’m wrong.

So supposing we have something like this:

parameters {
 real<lower=0> sigma;
 real<multiplier = sigma> alpha;
}
model {
  sigma ~ lognormal(0, 1);
  alpha ~ normal(0, sigma);
}

then what is happening under the hood is that Stan will initialize two unconstrained variables to deal with the two tranformations: alpha_unc and sigma_unc. So that looks something like this, no?

sigma = exp(sigma_unc) + 0;
jacobian += sigma_unc; // Jacobian is d/d.sigma_unc exp(sigma_unc) + 0 = exp(sigma_unc)
target += lognormal_lpdf(sigma|0, 1);

alpha = sigma * alpha_unc;
jacobian += log(sigma);  // Jacobian is d/d.alpha_unc sigma * alpha_unc = sigma
target += normal_lpdf(alpha | 0, sigma);

So far I think I understand that, except for what happens in the math with the normal_lpdf(alpha| 0, sigma) statement.
I’m getting a little lost in this bit:

log p(alpha_unc)
    = log(normal(sigma * alpha_unc | 0, sigma)) + log(sigma)
    = ((sigma * alpha_unc) / sigma)^2 - log(sigma) + log(sigma)
    = alpha_unc^2

Does that mean that normal_lpdf(alpha | 0, sigma) = std_normal_lpdf(alpha_unc) - log(sigma) and then we add the log(sigma) from the Jacobian?

If so then what is the joint log probability of sigma_unc and alpha_unc after everything cancels out? Would it be something like:

sigma = exp(sigma_unc) + 0;
alpha = sigma * alpha_unc;

log p(alpha_unc)
     = log(normal(sigma * alpha_unc | 0, sigma)) + log(sigma)
     = ((sigma * alpha_unc) / sigma)^2 - log(sigma) + log(sigma)
     = alpha_unc^2

log p(sigma_unc) 
  = log(lognormal(exp(sigma_unc)|0, 1)) + sigma_unc
  = normal(sigma_unc|0, 1) + sigma_unc
  = exp(-sigma_unc^2/2) + sigma_unc


log p(alpha_unc, sigma_unc) 
   = log(normal(exp(sigma_unc) * alpha_unc | 0, sigma)) + sigma_unc + 
       log(lognormal(exp(sigma_unc)|0, 1)) + sigma_unc
   = alpha_unc^2 + exp(-sigma_unc^2/2) + sigma_unc

Maybe I’m thinking about this wrong, but I’m similarly curious as @jsocolar as to why these models tend to be drawn towards zero values for sigma and sometimes get stuck there. This doesn’t seem to me like the model responding reasonably to the data; it seems like a bug. Given that an affine transform requires sigma to positive and that we see problems when sigma approaches zero (and log(sigma) approaches negative_infinity()), my hunch is that it’s got something to do with Jacobians of the transforms from unconstrained to constrained variables, all of which involve log(sigma) in someway. But that’s just a hunch and, like I said, I might be in way over my head.

This message has been edited to fix a bug pointed out to me by @Dalton. Fixing the bug didn’t substantially change my conclusions.

There are two separate behaviors here. One is the tendency for the chains to explore very tiny values of sigma early during warmup (“the pinch”). The other is their tendency to get stuck there (“the sticky boundary”). The pinch is weird, and I don’t understand why it happens, but I’m fairly certain that it’s not a bug. It’s just what happens in non-centered parameterizations of high-dimensional random effects when Stan’s non-buggy HMC implementation and adaptation strategy operate on Stan’s non-buggy specification of the target density, in the context of Stan’s non-buggy choices of initial values. I’d still love to understand why it happens and whether it poses any dangers (other than the possibility of encountering the sticky boundary). However, the pinch is the only way that I’ve seen/heard of anyone actually hitting the sticky-boundary in practice. So if there’s a way to initialize or parameterize that reliably avoids the pinch, that could be very useful.

The sticky boundary is clearly undesirable and we’d prefer that it not happen, and in this sense one might think to call it a “bug”, but it might better be thought of as an unfortunate limitation of the implementation. Unlike the pinch, the sticky boundary can be avoided by using an NCP that doesn’t involve a change of variables. @Dalton’s strategy of using a tiny positive lower bound for sigma also seems reasonable to me in this particular case. Perhaps an even better/safer option would be to run an initialization using a tiny positive bound, and then switch to the zero-bound model while passing through the last iteration from the previous model as inits. This should avoid the sticky boundary unless tiny values for sigma are actually an important region of the posterior to explore, in which case you’d want the model to complain if it is unable to explore there effectively.

I think that @martinmodrak and @Dalton are likely on the right track in attributing the sticky boundary the Jacobian adjustments, but I still don’t get why.

Below are two of Stan programs and the associated traceplots for sigma (showing the first 50 warmup iterations, fit under the default 1000-iteration warmup schedule).

  • The first is a manual NCP that doesn’t use any change of variables and doesn’t require a Jacobian (except the one applied under-the-hood due to the positive constraint on sigma). This one displays the pinch through tiny values of sigma, but never gets stuck there.
  • The second is a manual NCP that includes the change of variables and the log(sigma) Jacobian term. This one often gets stuck, just like actually using the offset, multiplier syntax.

Data, fitting, plotting

Data sim, fitting, and plotting in both cases performed as

set.seed(3)
n_grp <- 1000
n_per_grp <- 5
mu <- 0
sigma_re <- .5
sigma_residual <- 1

re <- rnorm(n_grp, mu, sigma_re)
y <- rnorm(n_grp * n_per_grp, re, sigma_residual)
grp_id <- rep(1:n_grp, n_per_grp)

re_data <- list(
  N = n_grp * n_per_grp,
  n_grp = n_grp,
  y = y, 
  grp_id = grp_id
)

fit <- re_mod$sample(data = re_data, chains = 8, parallel_chains = 4, save_warmup = T, 
                     iter_warmup = 1000, iter_sampling = 1)
samps <- fit$draws(variables = "sigma", inc_warmup = TRUE)
bayesplot::mcmc_trace(samps[1:50,,])

First one (no change of variables)

data {
  int<lower=0> N;
  int<lower=0> n_grp;
  vector[N] y;
  array[N] int<lower=1, upper=n_grp> grp_id;
}

parameters {
  real mu;
  real<lower=0> sigma;
  real<lower=0> sigma_residual;
  vector[n_grp] random_effects_raw;
}

model {
  random_effects_raw ~ std_normal();
  vector[n_grp] random_effects_transformed = mu + sigma * random_effects_raw;
  y ~ normal(random_effects_transformed[grp_id], sigma_residual);
}

Second one (change of variables)

data {
  int<lower=0> N;
  int<lower=0> n_grp;
  vector[N] y;
  array[N] int<lower=1, upper=n_grp> grp_id;
}

parameters {
  real mu;
  real<lower=0> sigma;
  real<lower=0> sigma_residual;
  vector[n_grp] random_effects_raw;
}

model {
  vector[n_grp] random_effects_transformed = mu + sigma * random_effects_raw;
  random_effects_transformed ~ normal(mu, sigma);
  target += log(sigma) * n_grp;
  y ~ normal(random_effects_transformed[grp_id], sigma_residual);
}


(note that one chain is stuck here)

The values of sigma where the stuck chain ended up after 2000 iterations was on the order of 10^{-22}

Check out a pairs plot of log(sigma) and lp__. Are you seeing a slope on the order lp__ = const + n_grp/2 * log(sigma)?

Here is the pairs plot from two versions of the binomial model that dropped code for earlier. These are model fits that actually converged. Affine transform and manual NCP without jacobian looked about the same. Strong correlation between lp__ and log(sigma) even after convergence.

Last thing from me tonight, just playing a bit more. Here are some insights from 60-chain runs (otherwise identical to my previous post except that the seed in the R script was changed from 3 to 5.

no change of variables

  • all chains experience the pinch
  • no chains experience the sticky boundary
  • during the pinch, 17 chains reached values lower than 1e-17, with the lowest reaching 1.39e-37.

change of variables

  • all chains experience the pinch
  • 9 chains experience the sticky boundary
  • Among the 51 chains that did not experience the sticky boundary, the lowest value for sigma reached by any chain was 1.00311e-16
  • Among the 9 chains that did experience the sticky boundary, all reached minimum values at least as low as 3.49616e-17, reached these values within the first 50 warmup iterations, and ended after 2000 iterations on values that differed from the minimum value reached during the first 50 warmup iterations by no more than 7.9032e-18.

I don’t have the answer for the behavior, but I did test that initializing MCMC with Pathfinder improves the behavior (I love that init can be a CmdStanFit object)

fitp <- re_mod$pathfinder(data = re_data, num_paths=60, refresh=0)

fit <- re_mod$sample(data = re_data, 
                     chains = 60, parallel_chains = 8, save_warmup = T, 
                     iter_warmup = 1000, iter_sampling = 1, refresh=0,
                     init=fitp)

The lowest sigma among all 60 chains is 6.6e-4

Furthermore, it seems the default init from [-2,2] is just too wild. Switching to init from [-0.1, 0.1] makes the results better, too

fit <- re_mod$sample(data = re_data,
                     chains = 60, parallel_chains = 8, save_warmup = T, 
                     iter_warmup = 1000, iter_sampling = 1, refresh=0,
                     init=0.1)

with minimum sigma in any chain being 0.38

I’d hazard a guess that this explains the problem.

When we do random_effects_transformed ~ normal(mu, sigma), the autodiff stack for the derivative w.r.t. sigma will add ((random_effects_transformed - mu) * (1/sigma))^2 * (1/sigma) - (1/sigma) (code link).

The derivative of log(x) is 1/x, so when log(sigma) is added to target, the autodiff stack for the derivative wrt. sigma will add 1/sigma (code link).

So on the reverse pass, you’d be adding and then subtracting 1/sigma (maybe immediately? or after some operations?? not 100% sure on the order in autodiff). For small sigma this is a large number, so whenever the absolute value of “derivative so-far” is small, we have a catastrophic cancellation. For sigma = 3e-17 this starts posing problems when the abs of “derivative so-far” is smaller then 10:

sigma <- 3e-17
1 + (1/sigma) - (1/sigma)
# 0
10 + (1/sigma) - (1/sigma)
# 12
20 + (1/sigma) - (1/sigma)
# 20 

So I’d guess that if you get the gradient of both versions of the model, when sigma is small, the two gradients will differ substantially, even though they should agree. I don’t have time to dig more deeply to see if the result is that gradient becomes too large (and thus all steps are almost always rejected) or that the gradient becomes 0 (and thus you cannot move away from this point).

I’d LOVE it if somebody could explain to me why we see this early pinch through tiny values of sigma in the NCP but not the centered! (edit: and hopefully assuage my fear that this could cause problems for convergence diagnostics by making it effectively impossible to specify overdispersed inits for sigma )

I don’t have an explanation, but I do have a pretty good clue. When I fit a binomial hierarchical logit-normal model (haven’t tested for the normal model) and examine the pairs plot for post warm-up draws there is a strong correlation between log(sigma) and lp__ for both centered and non-centered models. But the correlation is negative for the centered model and positive for the non-centered model. Maybe somebody can tell me how this correlation plays out in the gradient calculation.

Non-centered:

Centered (affine transform):

The negative correlation in the centered case is the classic funnel geometry. The global posterior mode in this parameterization sits in the neck of the funnel, but the neck is narrow enough that little posterior mass sits there.

In the NCP, my instinct is to search for the cause of the pinch outside of the typical set, because initializations in or near the typical set don’t manifest the pinch. It seems that the key issue somehow involves the interaction between sigma, the random effects vector, and the data. Above, @avehtari showed that initializing all unconstrained parameters on the interval [-.1, .1] avoids the pinch. A key thing to notice here is that the resulting initial values for sigma are in a range that still manifests the pinch when initializing all parameters on the interval [-2, 2]. I think we can conclude that a key determinant of whether or not the early warmup manifests the pinch is not necessarily the initialization of sigma, but rather the initialization of the random effect vector, which accounts for 1000 of the 1003 parameters in this model.

If I had to make a guess about the pinch, here’s what I’d guess currently:

  • When the random effects vector is initialized to nonsense, it’s easy to see why early exploration might explore small values of sigma. That is, from the point of initialization, the partial with respect to sigma is negative.
  • As we move towards smaller values of sigma (for the sake of argument holding the random effects vector fixed), the gradient gets much smaller but never becomes positive. Therefore, we might expect the Hamiltonian trajectory to fall down towards small sigma and then continue skating along the flat towards negative-infinite log(sigma) until a stopping criterion (either a u-turn or a treedepth limit) is triggered.
  • Of course we aren’t holding the random effects vector fixed–it’s also evolving along the trajectory. But maybe for high-dimensional random effects vectors, this evolution is slow enough that sigma reaches its tiny values before the random effects vector can take on a configuration that would actually favor larger sigma. Note that this would explain why the pinch only shows up when the random effects vector is high-dimensional.
  • Once we reach the tiny-sigma region, all gradients are weak. Changing unconstrained sigma doesn’t do much of anything when sigma is tiny, and changing the raw random effects vector also doesn’t do much when sigma is tiny.
  • This also explains why initializing the random effecrts vector in a much narrower range also avoids the pinch. It’s because in these cases, the initial gradients with respect to sigma aren’t so steeply negative. This also helps to explain why the pinch only shows up for high-dimensional random effects vectors, because only when the vector is high dimensional would the gradient at a nonsense initialization be (very) steeply negative with respect to sigma.

If the above is right, then it should be the case that specifying a much lower treedepth lets chains fall into the pinch, but avoids having them skate out to absurdly small values for sigma, because at low treedepth we will terminate the trajectory and resample the momentum before a trajectory gets a chance to do much “skating” across the flat towards minuscule sigma. I just ran this experiment, and the results are consistent with my hunch.

fit <- re_mod$sample(data = re_data, chains = 8, parallel_chains = 4, save_warmup = T, 
                     max_treedepth = 1,
                     iter_warmup = 20000, 
                     iter_sampling = 1, 
                     init_buffer = 19000
                     )

In this case, no chain ever explores sigma lower than 0.001. The chains all drop down into the region where sigma is between about 0.01 and 0.05, which I suppose is the best fitting region for sigma before the random effect vector has a chance to equilibrate. In the traceplot below, the chains don’t manage to start recovering until the metric adaptation kicks in after iteration 19000, but note that these 19000 iterations only represent as many leapfrog steps as 19 iterations that saturate the default max treedepth of 10.

Thanks—this is super useful. The interesting thing is that the problem goes away for numbers an order of magnitude or two larger than 3e-17.

We just discussed this and worked through your math during the Stan meeting today. The problem is that we don’t know how to fix this with the offset/multiplier design because the canceling +/- 1 / sigma terms come from different sources. If we were doing something like PyMC then we could fix this as we wouldn’t separate the constraint from the distribution—we’d just have a non-centered normal distribution that could do the cancellation algebraically under the hood.

Thanks—this is really interesting and very confusing.

I was about to say that it should depend on the value of sigma. I would have thought that if the current value of sigma is less than the standard deviation of the coefficients distributed normal(0, sigma), the derivative would be positive. But that’s not right for reasons I still don’t quite understand, but following the actual math, @jsocolar is right:

\dfrac{\partial}{\partial \sigma} \left( \dfrac{\alpha_1^2}{\sigma} + \dfrac{\alpha_2^2}{\sigma} \right) = -\dfrac{\alpha_1^2 + \alpha_2^2}{\sigma^2}

Similarly, if you look at the derivatives of the coefficients, they’re always aligned with the sign leading to higher magnitude.

\dfrac{\partial}{\partial \alpha_1} \left( \dfrac{\alpha_1^2}{\sigma} + \dfrac{\alpha_2^2}{\sigma} \right) = \dfrac{2 \alpha_1}{\sigma}.

So the derivatives for scale are always negative and the derivative for a coefficient always has the same sign as the coefficient.

How can this work?

This is what @andrewgelman concluded for actual problems he was working on—lower tree depth worked better in wall time.

What’s going on with the sigma gradient is easier to see in the NCP without change of variables:

model {
  random_effects_raw ~ std_normal();
  vector[n_grp] random_effects_transformed = mu + sigma * random_effects_raw;
  y ~ normal(random_effects_transformed[grp_id], sigma_residual);
}

If random_effects_raw is initialized to nonsense, then the likelihood is improved by removing random_effects_raw from random_effects_transformed by sending sigma to zero.

I don’t think that’s an easily fixable problem - and it is not the only inefficiency/potential loss of precision in the offset-multiplier version model. Fundamentally, whenever you have <multiplier=sigma> x and then normal_XXXX(x, whatever, sigma) you are adding work, because you evaluate the multiplier, just to undo it when computing normal_XXXX as all those functions shift and scale the values to obtain a standard normal (dtto for offset, and also any other location-scale family). So I’d argue non-centered parametrization is just not a good use case for the multiplier transformation.

I think this just signals that what I lazily call “derivative so far” (due to unwillingness to really parse out the precise order of operations in the reverse pass) in this model/dataset tends to be roughly between 0.1 and 10 in absolute value in the relevant parameter neighborhood. I’d expect the scale of the other contributions to the derivative w.r.t sigma and thus also the values of sigma when the numerical instability manifests to differ across models/datasets.

This is a really cool insight. I think that this thread has two things going on though and it might be worth it at this point to acknowledge the forks in the conversation and give people some idea of how to proceed with modelling random effects.

The two issues are:

  1. The general tendency of the NCP to dive towards very small values during warmup, particularly with high-dimensional random effect vectors.
  2. The offset/multiplier version of a NCP getting “stuck” at the boundary, resulting in chains never converging.

The cause of issue 2 seems to have been identified by @martinmodrak and may be related to the catastrophic cancellation of + 1/sigma - 1/sigma in gradient calculation. This is an issue that happens in the change of variables case using either the built in affine-transformation or a manual transformation with a Jacobian adjustment. My takeaway as non-developer user is that I should avoid using the offset/multiplier (change of variables) form for RE models and stick to the transformation (no change of variables) version, at least until there is an update that addresses this. That kind of information was what I was seeking when I resurrected this thread.

Issue 1 is really interesting and it seems like it is well worth diving into that deeper because it seems like there opportunities to figure out better ways to initialize NCP structures.

I don’t want to derail further discussion of issue 1, but as this old thread is getting quite long in it’s new life, I thought it might be worth summarizing what I’ve learned from following the conversation to this point, especially since it seems like issue 1 is not exclusive to the offset/multiplier form (which is the title of this thread).