One-hot encoding or indexed coefficients

Modeling
2

re: the second representation, the explicit if / else branches in transformed parameters collapse if we redefine beta_levels and pad it with an additional artificial beta levels “parameter” that is constrained to always take the constant value zero.

Pinning a parameter to a constant value does not seem to be allowed by Stan, but adding an extra layer of indirection gives something that compiles:

data {
    int n;
    int<lower=2> num_levels;
    array[n] int<lower=1,upper=num_levels> level;

    matrix[num_levels, 3] X;

    vector[n] actual_value;
}
parameters {
    vector[3] beta;
    vector[num_levels-1] beta_levels;
    real intercept;
    real<lower=0> sigma;
}
transformed parameters {
    vector[num_levels] beta_levels_;
    beta_levels_[1:num_levels-1] = beta_levels;
    beta_levels_[num_levels] = 0;

    vector[n] expected_value = intercept + (X * beta) + beta_levels_[level];
}
model {
    beta ~ normal(0, 1);
    beta_levels ~ normal(0, 1);
    sigma ~ student_t(3, 0, 1);
    actual_value ~ normal(expected_value, sigma);
}

I don’t have intuition about relative performance. If you’ve got the time, implement them all and look at what is happening through a CPU profiler!

1 Like