Pupil: numeric scale predictor and automatic totals

Result

BRM can express the pupil model from post 3 of the Stan discussion, integrate its population mean coefficients, and expose its total coefficients to WHMC adaptation automatically.

All six total-coefficient variants completed with zero sampling divergences. In this pilot, online position adaptation gave 318× the total-gradient efficiency of ordinary brms NCP + native Stan, and 1.64× that of brms S2Z auto + WHMC. Online gradient adaptation gave 224× the native baseline. Fixed CP and S2Z auto were also effective; neither marginalization made fixed NCP competitive with those centered alternatives.

The table compares one common set of 46 scientific quantities. These are single-chain experiments; the numerical ordering of the two losses needs replication.

Model and data

The data contain 2,228 observations, 20 subjects with IDs 701–720, and numeric load values 0–5. We preserve observation order and use independent subject intercept and slope deviations:

bf(p_size ~ load + (load || subj), sigma ~ subj)

In these data subj is numeric. Thus the residual model is a regression of log residual SD on subject ID. The second pupil study instead assigns each subject a hierarchical residual SD.

$$y_n\sim N({\mu }_n,{\sigma }_n^2),{\mu }_n={\beta }_0+{\beta }_1(x_n-\overline{x})+a_{j[n]}+b_{j[n]}{x}_n,\log {\sigma }_n={\gamma }_0+{\gamma }_1(s_n-\overline{s}).$$

The population load and subject-ID predictors use observation-weighted centering. The random slope uses raw load. Priors match the generated brms target:

• Population intercept: Student-t(3, location 5651.9, scale 2026.1).

• Population load slope: flat.

• Independent deviations: a[j] ~ Normal(0, tau_a), b[j] ~ Normal(0, tau_b).

• Both group SDs: half-Student-t(3, 0, 2026.1).

• Log-SD intercept: Student-t(3, 0, 2.5); numeric-ID slope: flat.

Flat priors apply to these population slopes because that is the source model's specification. The group effects retain their Gaussian hierarchical priors. Independence of the mean random effects is an explicit simplification of the forum model, shared by every row below.

BRM implementation

The authoring pane reads the actual declarations used for the fits. Its backend panes are generated during the documentation build. The measurements use StanBlocks/BridgeStan and WHMC; the Turing pane is a generated-model comparison only.

Pupil means with a numeric subject-ID scale predictor

using BayesianRegressionModels, Distributions, JSON, DelimitedFiles

function pupil_numeric_brm_model()
    reference = JSON.parsefile(joinpath(pkgdir(BayesianRegressionModels),
        "research", "pupil_scale_totals", "reference", "ordinary_ncp.json"))
    data = (;p_size=Float64.(reference["Y"]), load=Float64.(reference["Z_1_2"]),
             subj=Int.(reference["J_1"]))
    data = merge(data, (;subject_id=700 .+ data.subj))
    
    builder = @brm begin
        mu ~ 1 + center(load) + (1 | mean_intercept | subj) + (0 + load | mean_slope | subj)
        logsigma ~ 1 + center(subject_id)
        effect(mu, Intercept) ~ LocationScale(5651.9,2026.1,TDist(3))
        effect(mu, center_load) ~ Flat()
        effect(logsigma, Intercept) ~ LocationScale(0.,2.5,TDist(3))
        effect(logsigma, center_subject_id) ~ Flat()
        sd(:, mean_intercept) ~ LocationScale(0.,2026.1,TDist(3))
        sd(:, mean_slope) ~ LocationScale(0.,2026.1,TDist(3))
        p_size ~ Normal(mu,exp(logsigma))
    end
    
    builder(data)
end

BRMI:
  load: data (eltype=Float64, n=2228)
  subj: data (eltype=Int64, n=2228)
  mu ~ 1 + center(load) + (1 | mean_intercept | subj) + ((0 + load) | mean_slope | subj)
  subject_id: data (eltype=Int64, n=2228)
  logsigma ~ 1 + center(subject_id)
  effect(mu, Intercept) ~ AffineDistribution(5651.9, 2026.1, TDist(3))
  effect(mu, center_load) ~ Flat()
  effect(logsigma, Intercept) ~ AffineDistribution(0.0, 2.5, TDist(3))
  effect(logsigma, center_subject_id) ~ Flat()
  effect(sd, mean_intercept) ~ AffineDistribution(0.0, 2026.1, TDist(3))
  effect(sd, mean_slope) ~ AffineDistribution(0.0, 2026.1, TDist(3))
  p_size ~ Normal(mu, exp(logsigma))

SBBRMI with data keys = [:center_subject_id, :load, :p_size, :subj, :subject_id, :total_A_mu, :total_group_mu, :total_location_mu, :total_mixture_shape_mu, :total_ng_mu, :total_nk_mu, :total_nm_mu, :total_np_mu, :total_precision_mu]
configured submodels:
_brm_total_scales_configured_1 = Base.merge(BayesianRegressionModels._brm_total_scales, quote
            tau ~ (ValueFamily(brm_vector_prior_eecc99814bef47ac))(3.0, 0.0, 2026.1, 3.0, 0.0, 2026.1; n = 2)
        end)
_popefs_generic_configured_1 = Base.merge(BayesianRegressionModels._popefs_generic, quote
            beta_pop::vector[n_covariates] ~ (ValueFamily(brm_vector_prior_38a4c9e64638c44f))(3.0, 0.0, 2.5; n = 2)
        end)
emitted @slic body:
begin
    total_scale_mu ~ _brm_total_scales_configured_1(; n = total_nk_mu)
    total_mixture_mu::vector[total_nm_mu] ~ gamma(total_mixture_shape_mu, total_mixture_shape_mu; lower = 0.0)
    total_conditional_precision_mu = [total_precision_mu[1] * total_mixture_mu[1], total_precision_mu[2]]
    total_mu::matrix[total_ng_mu, total_nk_mu] ~ brm_total(total_scale_mu, total_A_mu, total_location_mu, total_conditional_precision_mu)
    population_mu = brm_total_recover_rng(total_mu, total_scale_mu, total_A_mu, total_location_mu, total_conditional_precision_mu)
    deviation_mu = brm_total_deviations(total_mu, total_A_mu * population_mu)
    total_Z_mu = hcat(rep_vector(1.0, num_elements(total_group_mu)), load)
    mu = rows_dot_product(total_mu[total_group_mu, :], total_Z_mu)
    X_logsigma = hcat(rep_vector(1.0, num_elements(subject_id)), center_subject_id)
    pop_logsigma ~ _popefs_generic_configured_1(; X = X_logsigma)
    logsigma = pop_logsigma
    p_size ~ normal(mu, (exp)(logsigma))
end

functions {
// value UDF brm_vector_prior_eecc99814bef47ac_lpdf
real brm_vector_prior_eecc99814bef47ac_lpdf(
    vector x,
    real arg_1,
    real arg_2,
    real arg_3,
    real arg_4,
    real arg_5,
    real arg_6
) {
    if((x[1] < 0.0)) {
        return negative_infinity();
    }
    if((x[2] < 0.0)) {
        return negative_infinity();
    }
    return (student_t_lpdf(x[1] | arg_1, arg_2, arg_3) + student_t_lpdf(x[2] | arg_4, arg_5, arg_6));
}
real brm_total_lpdf(
    matrix total,
    vector tau,
    matrix A,
    vector location,
    vector precision
) {
    int j = dims(total)[1];
    int k = dims(total)[2];
    int p = dims(A)[2];
    if (dims(tau)[1] != k) reject("brm_total_lpdf: dim mismatch — `tau` dim 1 (= ", dims(tau)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(A)[1] != k) reject("brm_total_lpdf: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(location)[1] != p) reject("brm_total_lpdf: dim mismatch — `location` dim 1 (= ", dims(location)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
    if (dims(precision)[1] != p) reject("brm_total_lpdf: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
    matrix[dims(precision)[1], dims(precision)[1]] Q = brm_total_precision(tau, A, precision, j);
    vector[dims(total)[2]] average = brm_total_mean(total);
    vector[dims(precision)[1]] beta = brm_total_conditional_mean(total, tau, A, location, precision, Q);
    vector[dims(total)[2]] residual = (average - (A * beta));
    real quadratic = 0.0;
    real lp = ((-0.5 * ((j * k) - p) * 1.8378770664093453) - (j * sum(log(tau))));
    for(c in 1:k) {
        quadratic += ((j * square(residual[c])) / square(tau[c]));
        for(g in 1:j) {
            quadratic += (square((total[g, c] - average[c])) / square(tau[c]));
        }
    }
    for(a in 1:p) {
        if((precision[a] > 0.0)) {
            lp += (0.5 * (log(precision[a]) - 1.8378770664093453));
            quadratic += (precision[a] * square((beta[a] - location[a])));
        }
    }
    return (lp - (0.5 * (log_determinant(Q) + quadratic)));
}
matrix brm_total_precision(
    vector tau,
    matrix A,
    vector precision,
    int n_groups
) {
    int k = dims(tau)[1];
    int p = dims(A)[2];
    if (dims(A)[1] != k) reject("brm_total_precision: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `tau` dim 1. `k` sizes: `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(precision)[1] != p) reject("brm_total_precision: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `precision` dim 1 (= ", dims(precision)[1], ").");
    matrix[dims(precision)[1], dims(precision)[1]] out = diag_matrix(precision);
    for(a in 1:p) {
        for(b in 1:p) {
            for(c in 1:k) {
                out[a, b] += ((n_groups * A[c, a] * A[c, b]) / square(tau[c]));
            }
        }
    }
    return out;
}
vector brm_total_mean(
    matrix total
) {
    int j = dims(total)[1];
    int k = dims(total)[2];
    vector[k] out = rep_vector(0.0, k);
    for(c in 1:k) {
        out[c] = (sum(total[:, c]) / j);
    }
    return out;
}
vector brm_total_conditional_mean(
    matrix total,
    vector tau,
    matrix A,
    vector location,
    vector precision,
    matrix Q
) {
    int j = dims(total)[1];
    int k = dims(total)[2];
    int p = dims(A)[2];
    if (dims(tau)[1] != k) reject("brm_total_conditional_mean: dim mismatch — `tau` dim 1 (= ", dims(tau)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(A)[1] != k) reject("brm_total_conditional_mean: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(location)[1] != p) reject("brm_total_conditional_mean: dim mismatch — `location` dim 1 (= ", dims(location)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
    if (dims(precision)[1] != p) reject("brm_total_conditional_mean: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
    if (dims(Q)[1] != p) reject("brm_total_conditional_mean: dim mismatch — `Q` dim 1 (= ", dims(Q)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
    if (dims(Q)[2] != p) reject("brm_total_conditional_mean: dim mismatch — `Q` dim 2 (= ", dims(Q)[2], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], "), `Q` dim 1 (= ", dims(Q)[1], "), `Q` dim 2 (= ", dims(Q)[2], ").");
    vector[dims(total)[2]] average = brm_total_mean(total);
    vector[dims(location)[1]] natural = (precision .* location);
    for(a in 1:p) {
        for(c in 1:k) {
            natural[a] += ((j * A[c, a] * average[c]) / square(tau[c]));
        }
    }
    return mdivide_left_spd(Q, natural);
}
vector brm_total_recover_rng(
    matrix total,
    vector tau,
    matrix A,
    vector location,
    vector precision
) {
    int j = dims(total)[1];
    int k = dims(total)[2];
    int p = dims(A)[2];
    if (dims(tau)[1] != k) reject("brm_total_recover_rng: dim mismatch — `tau` dim 1 (= ", dims(tau)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(A)[1] != k) reject("brm_total_recover_rng: dim mismatch — `A` dim 1 (= ", dims(A)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `tau` dim 1 (= ", dims(tau)[1], "), `A` dim 1 (= ", dims(A)[1], ").");
    if (dims(location)[1] != p) reject("brm_total_recover_rng: dim mismatch — `location` dim 1 (= ", dims(location)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
    if (dims(precision)[1] != p) reject("brm_total_recover_rng: dim mismatch — `precision` dim 1 (= ", dims(precision)[1], ") does not match `p` (= ", p, "), inferred from `A` dim 2. `p` sizes: `A` dim 2 (= ", dims(A)[2], "), `location` dim 1 (= ", dims(location)[1], "), `precision` dim 1 (= ", dims(precision)[1], ").");
    matrix[dims(precision)[1], dims(precision)[1]] Q = brm_total_precision(tau, A, precision, j);
    vector[dims(precision)[1]] beta = brm_total_conditional_mean(total, tau, A, location, precision, Q);
    return multi_normal_rng(beta, inverse_spd(Q));
}
matrix brm_total_deviations(
    matrix total,
    vector mu
) {
    int j = dims(total)[1];
    int k = dims(total)[2];
    if (dims(mu)[1] != k) reject("brm_total_deviations: dim mismatch — `mu` dim 1 (= ", dims(mu)[1], ") does not match `k` (= ", k, "), inferred from `total` dim 2. `k` sizes: `total` dim 2 (= ", dims(total)[2], "), `mu` dim 1 (= ", dims(mu)[1], ").");
    matrix[dims(total)[1], dims(total)[2]] out = total;
    for(c in 1:k) {
        out[:, c] = (total[:, c] - rep_vector(mu[c], j));
    }
    return out;
}
matrix hcat(
    vector x,
    vector y
) {
    int n = dims(x)[1];
    if (dims(y)[1] != n) reject("hcat: dim mismatch — `y` dim 1 (= ", dims(y)[1], ") does not match `n` (= ", n, "), inferred from `x` dim 1. `n` sizes: `x` dim 1 (= ", dims(x)[1], "), `y` dim 1 (= ", dims(y)[1], ").");
    return append_col(x, y);
}
// value UDF brm_vector_prior_38a4c9e64638c44f_lpdf
real brm_vector_prior_38a4c9e64638c44f_lpdf(
    vector x,
    real arg_1,
    real arg_2,
    real arg_3
) {
    return (student_t_lpdf(x[1] | arg_1, arg_2, arg_3) + 0.0);
}
vector normal_lpdfs(
    vector obs,
    vector loc,
    vector scale
) {
    return jbroadcasted_normal_lpdfs(obs, loc, scale);
}
vector jbroadcasted_normal_lpdfs(
    vector x1,
    vector x2,
    vector x3
) {
    int n = dims(x1)[1];
    vector[n] rv;
    for(i in 1:n) {
        rv[i] = normal_lpdfs(broadcasted_getindex(x1, i), broadcasted_getindex(x2, i), broadcasted_getindex(x3, i));
    }
    return rv;
}
real normal_lpdfs(
    real args1,
    real args2,
    real args3
) {
    return normal_lpdf(args1 | args2, args3);
}
real broadcasted_getindex(vector x, int i) {
    return x[i];
}
vector normal_vector_rng(
    int anontok__1,
    vector a,
    vector b
) {
    int n = anontok__1;
    if((n == 0)) {
        vector[n] rv;
        return rv;
    } else {
        return to_vector(normal_rng(a, b));
    }
}
}
data {
    int total_nm_mu;
    int total_mixture_shape_mu_n;
    vector[total_mixture_shape_mu_n] total_mixture_shape_mu;
    int total_precision_mu_n;
    vector[total_precision_mu_n] total_precision_mu;
    int total_ng_mu;
    int total_nk_mu;
    int total_A_mu_m;
    int total_A_mu_n;
    matrix[total_A_mu_m, total_A_mu_n] total_A_mu;
    int total_location_mu_n;
    vector[total_location_mu_n] total_location_mu;
    int load_n;
    int total_group_mu_n;
    array[total_group_mu_n] int total_group_mu;
    vector[load_n] load;
    int center_subject_id_n;
    int subject_id_n;
    array[subject_id_n] int subject_id;
    vector[center_subject_id_n] center_subject_id;
    int p_size_n;
    vector[p_size_n] p_size;
}
transformed data {
    matrix[load_n, 2] total_Z_mu = hcat(rep_vector(1.0, num_elements(total_group_mu)), load);
    matrix[center_subject_id_n, 2] X_logsigma = hcat(rep_vector(1.0, num_elements(subject_id)), center_subject_id);
    int pop_logsigma_n_covariates = 2;
}
parameters {
    vector<lower=0.0>[2] total_scale_mu_tau;
    vector<lower=0.0>[total_nm_mu] total_mixture_mu;
    matrix[total_ng_mu, total_nk_mu] total_mu;
    vector[pop_logsigma_n_covariates] pop_logsigma_beta_pop;
}
transformed parameters {
    vector<lower=0.0>[2] total_scale_mu = total_scale_mu_tau;
    vector[2] total_conditional_precision_mu = [(total_precision_mu[1] * total_mixture_mu[1]), total_precision_mu[2]]';
    vector[load_n] mu = rows_dot_product(total_mu[total_group_mu, :], total_Z_mu);
    vector[center_subject_id_n] pop_logsigma = (X_logsigma * pop_logsigma_beta_pop);
    vector[center_subject_id_n] logsigma = pop_logsigma;
}
model {
    total_scale_mu_tau ~ brm_vector_prior_eecc99814bef47ac(3.0, 0.0, 2026.1, 3.0, 0.0, 2026.1);
    total_mixture_mu ~ gamma(total_mixture_shape_mu, total_mixture_shape_mu);
    total_mu ~ brm_total(total_scale_mu, total_A_mu, total_location_mu, total_conditional_precision_mu);
    pop_logsigma_beta_pop ~ brm_vector_prior_38a4c9e64638c44f(3.0, 0.0, 2.5);
    p_size ~ normal(mu, exp(logsigma));
}
generated quantities {
    vector[2] population_mu = brm_total_recover_rng(
        total_mu,
        total_scale_mu,
        total_A_mu,
        total_location_mu,
        total_conditional_precision_mu
    );
    matrix[total_ng_mu, total_A_mu_m] deviation_mu = brm_total_deviations(total_mu, (total_A_mu * population_mu));
    vector[p_size_n] p_size_likelihood = normal_lpdfs(p_size, mu, exp(logsigma));
    vector[p_size_n] p_size_gen = normal_vector_rng(p_size_n, mu, exp(logsigma));
}

#= line 0 =# Turing.@model(function brm_model(y, callable_1, X_mu, group_effects_mu_1, callable_2, group_effects_mu_2, callable_3, callable_4, X_logsigma)
        beta_pop ~ Distributions.product_distribution([callable_1(5651.9, 2026.1, Distributions.TDist(3)), BayesianRegressionModels.Flat()])
        eta_mu = X_mu * beta_pop
        group_effect_1 = Base.zeros(Base.length(y))
        group_1_1 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_1, (callable_2(0.0, 2026.1, Distributions.TDist(3)),), nothing))
        group_effect_1 = group_effect_1 + group_1_1.effect
        group_1_2 ~ DynamicPPL.to_submodel(BRM.turing_group_effect(group_effects_mu_2, (callable_3(0.0, 2026.1, Distributions.TDist(3)),), nothing))
        group_effect_1 = group_effect_1 + group_1_2.effect
        eta_mu = eta_mu + group_effect_1
        mu = eta_mu
        beta_pop_logsigma ~ Distributions.product_distribution([callable_4(0.0, 2.5, Distributions.TDist(3)), BayesianRegressionModels.Flat()])
        eta_logsigma = X_logsigma * beta_pop_logsigma
        logsigma = eta_logsigma
        begin
            for i = Base.eachindex(y)
                y[i] ~ Distributions.Normal(mu[i], Base.exp(logsigma[i]))
            end
        end
        (; mu = mu, logsigma = logsigma, response = y)
    end)

BRM recognizes that the population intercept and slope share their design with the subject effects. It replaces them with subject totals

$$A_j={\beta }_0-\overline{x}{\beta }_1+a_j,B_j={\beta }_1+b_j,{\mu }_n=A_{j[n]}+B_{j[n]}{x}_n.$$

The Student-t intercept uses an exact Gaussian scale mixture,

$$\lambda \sim \mathrm{Gamma}(3/2,\mathrm{rate}=3/2),{\beta }_0\mid \lambda \sim N(5651.9,{2026.1}^2/\lambda ).$$

Conditional on lambda, both population mean coefficients are integrated analytically. All subjects are treated symmetrically. The two residual-model coefficients remain explicit parameters because this predictor has no random-effects block. Ordinary brms samples 46 parameters; BRM totals and brms S2Z sample 45, including one mixture variable.

The induced prior depends on the sums and centered sums of squares of the totals, so evaluation is O(J) for fixed coefficient dimension. BRM's generated likelihood evaluates the observations. No numerical quadrature is used during sampling.

Centering and automatic WHMC integration

For each total q[j], BRM supplies the family

$$r_j(c_j)=c_j{m}_j+(q_j-m_j){\tau }_{k(j)}^{c_j-1},0\le c_j\le 1.$$

The fixed reference location is 5651.9 for intercept totals and zero for slopes. CP uses c=1, NCP uses c=0, and partial centering chooses an intermediate value. Scaling the totals does not whiten their integrated joint prior; total NCP is a different construction from ordinary standardized-deviation NCP.

using StanBlocks, BridgeStan, WarmupHMC, Enzyme, Random
using DifferentiationInterface: AutoEnzyme

sb = SBBRMI(pupil_numeric_brm_model(); mod=@__MODULE__)
problem = StanBlocks.stan_instantiate(sb.model)
names = BridgeStan.param_unc_names(problem.model)
adaptive = adaptive_centering_problem(sb, problem, AutoEnzyme(); centeredness=0.0)
fit = adaptive_warmup_mcmc(Xoshiro(1), adaptive;
    n_draws=2000, monitor_ess=true, nonlinear_adapt=true)

draws = permutedims(fit.posterior_position) # compiled-model coordinates
recovered = recover_population_draws(sb, draws, names; rng=Xoshiro(101))

Fixed CP/NCP uses the corresponding centeredness and nonlinear_adapt=false. Post-hoc selection uses a completed NCP pilot:

selection = select_total_centeredness(sb, pilot_draws, names; criterion=:position)
partial = adaptive_centering_problem(sb, problem, AutoEnzyme();
    centeredness=selection.centeredness)
partial_fit = adaptive_warmup_mcmc(Xoshiro(1), partial;
    n_draws=2000, monitor_ess=true, nonlinear_adapt=false)

Both post-hoc arms use the same pilot and the built-in grid 0:0.1:1. The position loss minimizes log coordinate SD plus the mean inverse log-Jacobian contribution. The gradient loss minimizes the correlation between candidate position and its density gradient. It uses the pilot's saved gradients after transport into the compiled model frame. No new target evaluations are needed for selection. Neither loss directly optimizes ESS.

Online adaptation uses warmup windows and freezes the centering controls for retained sampling. Both losses are tested; the harness selects WHMC's existing position-loss weights locally, while its default is the position–gradient criterion. This loss switch is not a new public BRM option.

Matched comparisons and gradient accounting

Every row measures the same 46 scientific quantities: the two population mean coefficients, two group SDs, two residual-model coefficients, and 20 subject total intercepts and slopes. Population coefficients are sampled directly by ordinary brms and recovered conditionally by the marginalized methods. Totals are deterministic functions of each sampled state. Deviations, standardized coordinates and mixture variables do not enter the minimum ESS.

Both efficiency columns are relative to ordinary brms NCP + native Stan. Sampling efficiency is minimum bulk ESS divided by sampling gradients. Total efficiency includes initialization, warmup and all required pilots. Each post-hoc workflow is charged its full 84,335-gradient NCP pilot. S2Z auto includes its 789-gradient precursor.

Method	Total gradients	Sampling efficiency	Total efficiency
brms NCP · Native Stan	1,063,829	1×	1×
brms NCP · WHMC	274,051	1.43×	1.91×
brms CP · WHMC	80,373	2.54×	3.68×
brms S2Z CP · Native Stan	319,693	135×	40×
brms S2Z CP · WHMC	188,667	137×	66×
brms S2Z NCP · Native Stan	771,035	1.17×	1.22×
brms S2Z NCP · WHMC	123,267	3.19×	4.48×
brms S2Z auto · Native Stan	150,012	187×	62.6×
brms S2Z auto · WHMC	64,875	173×	194×
BRM total NCP · WHMC	84,335	4.54×	6.57×
BRM total CP · WHMC	44,130	133×	194×
BRM total post-hoc position · WHMC	118,655	232×	91.5×
BRM total post-hoc gradient · WHMC	118,674	215×	84.5×
BRM total online position · WHMC	34,526	235×	318×
BRM total online gradient · WHMC	52,951	159×	224×

The post-hoc rows improve sampling efficiency but pay for a separate pilot. Online adaptation avoids that cost. The strong CP and S2Z-auto controls show that a comparison against ordinary NCP alone would give an incomplete picture.

Each arm uses one chain, seed 1 and 2,000 retained draws. Native Stan uses 1,000 warmup iterations, target acceptance 0.8 and maximum depth 10. WHMC uses its adaptive warmup and Pathfinder initialization. Ordinary and total arms start at the same physical per-subject OLS coefficients; S2Z starts at pooled coefficients with zero contrasts and the OLS scales. Initialization matches within native/WHMC pairs but is not fully controlled across parameterization families.

The nine brms fits are reused from the verified comparison because their targets and sampler configurations are unchanged. All six BRM-total fits use the built-in implementation and repaired WHMC transport. The comparison measures gradient work, not identical wall cost per gradient. WHMC counts target-gradient requests. Native Stan uses a zero-contribution C++ counter compiled into the model; its sampler is unchanged. Retained increments match leapfrog counts plus one, and final process counts include initialization, warmup and the auto precursor.

Saved-draw geometry

The total-coefficient figure shows the same 2,000 post-hoc position draws in CP, NCP and ACP coordinates. Rows choose distinct coordinates by minimum selected centeredness, closest to 0.5, and maximum. CP is the visualization baseline; the ACP values were checked against the saved source coordinates.

The S2Z figure reuses the unchanged auto-WHMC fit. It shows centered contrasts, standardized contrasts, and the branch's subject-labelled Q*z coordinates. Those J displayed values represent J−1 independent contrast directions. The auto transform includes its weighted centering correction.

brms S2Z auto uses the pinned branch's Pathfinder/Fisher selection, then holds its weights fixed. WHMC receives that same resolved target and weights. These weights differ from BRM's per-total centering controls and loss functions; S2Z auto here is not online WHMC adaptation.

Recovery and validation

BRM recovers population coefficients from their exact Gaussian conditional given the totals, scales and mixture variable. For total block T with population-to-group design matrix A, its precision is

$$Q=P+J{A}^{\mathsf{T}}\mathrm{diag}({\tau }^{-2})A,$$

and its mean is Q \ (P*location + A' * diag(tau^-2) * sum(T)). Adding an independent conditional draw recovers the original population parameters; the subject totals stay fixed.

Conditional recovery adds genuine posterior variation and can raise ESS by diluting autocorrelation. It also adds noise to posterior-mean estimation. We therefore retain per-quantity MCSEs, repeat recovery across ten seeds, and compare the 44 quantities that require no stochastic recovery as a sensitivity. The main table consistently includes all 46 scientific quantities.

All 15 rows have zero sampling divergences. The new total-NCP fit has minimum ESS 88.8 and maximum within-chain split R-hat 1.013; the online position and gradient fits have maxima 1.007 and 1.004. Across all quantities, mean differences from the total-CP fit are below three combined estimated MCSEs except one ordinary-NCP WHMC quantity at 3.08 and one native S2Z-auto quantity at 3.48. These are descriptive consistency checks, not convergence guarantees.

Excluding recovered population coefficients leaves 44 invariant quantities. Online position still gives 124× sampling efficiency and 168× total efficiency against that baseline set. Across ten recovery seeds, its minimum ESS ranges from 1,707 to 1,759; online gradient ranges from 1,874 to 1,903. The other new total rows have unchanged minima. The main conclusions do not depend on one recovery draw.

The automatic target passed 25 density, gradient, centering-coordinate and recovery checks against the independent analytic implementation. Adding two constant half-Student-t normalizers aligns the absolute density; gradients and the posterior already agree. Generated S2Z and ordinary brms targets have their separate retained audits. These checks establish target equivalence, not a replicated sampling-performance ranking.

Inspect and reproduce

• BRM declarations and scientific quantities.

• Independent target and recovery audit.

• Six total fits and centering selection.

• Reused brms fits.

• Complete comparison table.

• Per-quantity ESS and MCSE.

• Recovery sensitivity.

• brms model sources and original fit archives.

• New generated target and fit manifest.

The data revision is d90fc01e6f6fcdced7ee64c9d2ed607d212ec77c; brms is pinned to PR #1919 revision 73cf607889879cb2a55f50b88d8141d76ff43279. New total fits use BRM's built-in planner with the retained-flat-prior correction and WarmupHMC 7aed40b18bd4cdabb75330f285d5c9b355575ab9, Julia 1.10.11. Native fits use CmdStan 2.39.0 and CmdStanR 0.9.0. The archive manifests retain the exact source and raw-draw provenance for reused and new fits.