Add MRP model

[ricardoV94]

No description provided.

[ricardoV94]

@velochy do you want to tweak this model just a bit. Right now it's maybe too simplified in that the mu is broadcast across the softmax axis and is therefore useless, so all priors are disconnected from the likelihood. For perf benchmarking it didn't matter but it would for benchmarking samplers.

[velochy]

Yeah... it went a bit too far with the simplification. We use stereotype regression for likert-scale responses, but it streamlined it out to just a linear model that cancels out.

Instead of

        # 2. likelihood: normalise over the K answer categories -> (N, G, K)
        alpha = pm.Normal("alpha", 0.0, 2.0, shape=(G, K))  # per-question scale-point intercepts
        pm.Categorical("y", logit_p=mu[:, :, None] + alpha[None, :, :], observed=y)

it should have

        # 2. stereotype head: a monotone loading phi turns the single latent into K logits.
        # phi_g = (0, ..., 1) from a Dirichlet's cumulative sums, so mu orders the scale points
        # rather than shifting all of them equally (a constant shift cancels under the softmax).
        phi_diffs = pm.Dirichlet("phi_diffs", np.ones(K - 1), shape=(G, K - 1))
        phi = pt.concatenate([pt.zeros((G, 1)), pt.cumsum(phi_diffs, axis=-1)], axis=-1)  # (G, K)
        # Per-question base rates, zero-sum over the scale so they don't fight the softmax.
        alpha = pm.ZeroSumNormal("alpha", sigma=2.0, shape=(G, K), n_zerosum_axes=1)
        pm.Categorical("y", logit_p=alpha[None, :, :] + phi[None, :, :] * mu[:, :, None], observed=y)

Or, if you want to be more standard, use either OrdinalProbit or OrdinalLogistic which is likely what most other people would put there (but they are both much less expressive).

[ricardoV94]

@velochy does it look good now? The updated version still had the same performance delta from jax (and benefit from the patches)

[velochy]

To nitpick, for stereotype regression, intercept creates a redundant DoF, so it's best to remove it.
In practice, having tried fitting this model with and without, it makes almost zero difference. But still, if it's in your model library, its probably better if it is "clean"

[ricardoV94]

Good now?