BRM.jl: simple formulas, room for custom models

Start with a familiar regression

Use formulas where they help

gaussian = (@brm begin
    sigma ~ Exponential(2)
    mu ~ 1 + x
    effect(:, :) ~ Normal(0, 2)
    effect(mu, x) ~ Normal(0, 0.25)
    y ~ Normal(mu, sigma)
end)((; x=[-1.0, 0.5, 2.0], y=[0.2, 1.1, -0.4]))

The formula handles the regression.

mu ~ 1 + x gives an intercept and a slope.

effect(...) assigns priors to those coefficients.

The noise parameter and observation distribution are written alongside them.

The aim: keep this convenience when a model needs more than regression.

Existing executable example from the feature atlas. Julia setup: using BayesianRegressionModels, Distributions.

Add the part that describes your science

Example: drug concentration in each subject

sigma ~ Exponential(1)
log(CL) ~ 1 + (1 | pk | subject)
log(V)  ~ 1 + (1 | pk | subject)

predicted_concentration ~ kernel(
    ragged(time, obs_subject), dose, CL, V,
) do ts, d, cl, volume
    d / volume * exp((-cl / volume) * ts)
end

ragged(concentration, obs_subject) ~
    Normal(predicted_concentration, sigma)

Formula lines: describe how clearance (CL) and volume (V) vary between subjects.

Custom calculation: the do block computes the concentration curve.

Observation line: compares that curve with measured concentrations.

The shared pk label allows clearance and volume effects to be correlated.

Formulas and a custom calculation live in the same model declaration.

StanBlocks example. Turing currently rejects this ragged response. Body of the existing @brm begin … end declaration; complete code and data in appendix B. kernel runs the calculation for each subject.

Subjects and measurements can fit together

The same PK example

# Subject axis: one row per subject.
subject=["alice", "bob"],
dose=[100.0, 80.0],

# Observation axis: one row per sample, interleaved by subject.
obs_subject=["alice", "bob", "alice", "bob", "alice"],
time=[0.5, 0.25, 1.5, 1.0, 3.0],
concentration=[8.1, 7.6, 5.2, 4.9, 2.1],

Two subjects, five measurements.

Each subject has a dose, clearance, and volume. Each measurement has a time and a subject label.

ragged(time, obs_subject) groups the times for the calculation. The response uses the same grouping.

Describe subject differences with formulas; evaluate the curve at each subject’s measurement times.

Data fields passed to the declaration on the previous slide. The interleaved rows are intentional.

The same approach reaches larger models

Existing BRM examples with custom scientific calculations

Warfarin PK/PD

Subject effects feed concentration calculations and an ODE for drug response.

Full equations, code, and audit

StanBlocks

Public two-stage model: compilation and finite density/gradient checks. A joint extension is documented separately.

Wastewater

Regression and group effects combine with infection renewal, shedding, and measurement models.

Formula and calculation definitions

StanBlocks

Structural port: compilation and finite density/gradient checks; not a numerically equivalent CDC reproduction.

Grey-seal population

Covariate and group effects combine with a population process and eight observation streams.

Formula and population-process code

StanBlocks

Compilation checks on an example dataset; no completed posterior fit. Turing does not execute this structural model.

Keep the useful regression formulas as the scientific calculation becomes more involved.

Compare the same modeling task

Subject effects + a custom concentration curve

brms

Formula-based regression and nonlinear parameter models.

Nonlinear expressions and custom Stan functions; custom families have additional prediction and likelihood hooks.

BRM

Formula-based regression and subject effects inside the model block.

The PK example passes their values directly into a custom calculation, then uses its result in the likelihood.

Hand-written Stan or Turing

Write the coefficients, group distributions, and predictor calculations explicitly.

Write the scientific calculation in the modeling language and its supported libraries.

BRM’s attraction is the combination: concise regression code with room for the rest of the model.

brms nonlinear models · custom Stan code · Turing with differential equations. This compares authoring approaches, not measured development time or fitting speed.

Custom code has specific requirements

Flexibility in practice

A Julia function or distribution

Turing can call ordinary Julia code retained by BRM.

It still needs the operations used by the fit: a valid density, a supported gradient path, and sampling support when generating predictions.

A custom Stan calculation

StanBlocks needs code it can translate to Stan.

A function written only for that path needs a Julia implementation before Turing can use it.

Adding a calculation and adding a reusable model term are different amounts of work.

A new parameter structure or fitted transform may need extra methods. See the current extension contracts.

A different parameterization can help sampling

Same probability model, different variables for the sampler

Sample the coefficient directly

$β_{k} \sim N (0, s_{k})$

Often useful when the data strongly identify the coefficient.

Sample a scaled coefficient

$z_{k} \sim N (0, 1), β_{k} = s_{k} z_{k}$

Often useful when the data provide less information.

BRM’s HSGP example can choose between these forms separately for each coefficient.

The example models motorcycle-impact acceleration over time, with a Gaussian process for the mean and another for the noise level. It uses all 133 observations and 20 basis functions per process.

One model, three recorded fits

Corrected HSGP study · StanBlocks sampling

Centeredness for each of the 20 basis frequencies in the mean and log-noise Gaussian processes, comparing online warmup and offline pilot selection

Each coefficient gets its own choice. Zero is fully noncentered; one is fully centered. The two selection procedures need not agree.

Noncentered pilot10,000 draws · 34 divergencesmax split R-hat 1.0012 · min bulk/tail ESS 2,514 / 2,886

Fresh partially centered fit10,000 draws · 16 divergencesmax split R-hat 1.0007 · min bulk/tail ESS 2,533 / 1,797

Separate online extension10,000 draws · 0 divergencesmax split R-hat 1.0018 · min bulk/tail ESS 2,318 / 2,775

One chain per fit. These counts do not establish a general improvement in speed or accuracy. Split R-hat compares parts of the same chain.

No corrected Turing samples. Target-value agreement and warmed gradient runtime still need verification.

Make predictions with the fitted model

Keep the meaning of the regression when the data change

During fitting

Store the category levels, scaling constants, spline knots, and GP basis used to construct the model.

For new measurements

Apply those same transformations to new inputs, so the posterior coefficients still have the same meaning.

For new subjects

Choose explicitly how new group effects are generated. Check that the model and execution path support the requested prediction.

Custom model components should fit into the prediction workflow as well as the model declaration.

Current operations have documented limits; this is not a claim that every custom component already supports every prediction task.

A possible future: Julia models on GPUs

A concrete direction to investigate

The compiler path exists

Reactant compiles Julia computations through MLIR and XLA for CPU, GPU, and TPU execution.

It also documents a probabilistic-programming interface, including custom log-density functions.

BRM integration remains to be shown

A possible next step is to compile suitable Julia model code through this path.

Compatibility, gradients, and useful performance need to be demonstrated on actual BRM models.

Keep the model declaration useful as more of its computation can run on accelerators.

Reactant · probabilistic programming · control-flow requirements. brms already exposes Stan/OpenCL support for some operations; GPU access alone is not the proposed distinction.

Simple formulas, room for custom models

End of main talk

Use concise formulas for the regression parts, and combine them with custom model code when the science needs more.

The PK example shows the idea today. The larger examples show where it can go—and which execution and prediction details still need work.

Current execution support

Appendix A

Population and distributional regression

StanBlocks + Turing

Shared formulas, links, and prior addresses.

Group effects, advanced terms, modified responses

Both, within limits

Individual features and their combinations have specific supported subsets.

Ordinary Julia callables without a Stan translation

Turing

Density, gradient, transform, and prediction requirements still apply.

Custom structural code written for StanBlocks

StanBlocks

A Julia implementation is needed before Turing can execute that code.

Corrected HSGP centering study

StanBlocks sampled

No corrected Turing samples; see the study-specific checks in appendix E.

This table covers the two paths discussed here. BRM also has vectorized Julia execution and a narrower experimental NativePPL path; neither is presented as a verified GPU backend.

The complete PK example

Appendix B

population_pk = (@brm begin
    sigma ~ Exponential(1)
    log(CL) ~ 1 + (1 | pk | subject)
    log(V)  ~ 1 + (1 | pk | subject)

    predicted_concentration ~ kernel(
        ragged(time, obs_subject), dose, CL, V,
    ) do ts, d, cl, volume
        d / volume * exp((-cl / volume) * ts)
    end

    ragged(concentration, obs_subject) ~
        Normal(predicted_concentration, sigma)
end)((;
    # Subject axis: one row per subject.
    subject=["alice", "bob"],
    dose=[100.0, 80.0],

    # Observation axis: one row per sample, interleaved by subject.
    obs_subject=["alice", "bob", "alice", "bob", "alice"],
    time=[0.5, 0.25, 1.5, 1.0, 3.0],
    concentration=[8.1, 7.6, 5.2, 4.9, 2.1],
))

StanBlocks example; Turing does not yet support the ragged-response wrapper. Julia setup: using BayesianRegressionModels, Distributions. Code and data are read from the executable feature atlas.

Inspect the actual model code

Appendix C

Start from the declaration

The feature atlas executes each displayed model.

Its four views show the BRM declaration, emitted StanBlocks code, generated Stan, and the Turing model or its construction error.

Check the operation you need

A model that can evaluate a density may still lack a prediction method or a required gradient.

For new fitted terms, check new-data transformations too. For a second execution path, check the same statistical model and parameter mapping.

Read generated code from the executable example, not from a manually maintained illustration.

How the comparison is generated · Turing operations and limits · Extension contracts

The centering study in detail

Appendix D

Mean and noise both vary with time

$\begin{aligned} μ (t) & = {HSGP}_{μ} (t), \\ \log σ (t) & = {HSGP}_{σ} (t), \\ y (t) & \sim N (μ (t), σ (t)) . \end{aligned}$

Each GP uses 20 basis functions. Its log length scale and log marginal standard deviation have $N (0, 4)$ priors.

Partial centering

$u_{k} \sim N (0, s_{k}^{c_{k}}), β_{k} = s_{k}^{1 - c_{k}} u_{k} .$

Offline pilot selects one centeredness value for each of the 20 frequencies in both Gaussian processes

1 noncentered pilot2 choose each

c_{k}

3 fresh partial fit

The source reproduction has two fits. Online adaptation is a separate BRM extension.

What the centering results establish

Appendix E

Recorded evidence

All 133 rows; 20 basis functions per GP.
Xoshiro(1); 10,000 retained draws per fit; unchanged WarmupHMC defaults.
Original Stan density and gradient checks: 26 for the noncentered target; 51 for the partial target.
All 40 offline choices match the source-literal calculation.
Complete-record and figure verification: 8,221 checks.

Limits of the conclusion

One chain per fit; split R-hat is a within-chain diagnostic.
Divergences: 34, 16, and 0 in these runs.
The partial refit has lower minimum tail ESS than the pilot.
No general speed or accuracy advantage is established.
No corrected Turing samples exist.

Corrected case ea27c18ed2517e3061f441971f561235f48cceb8, integrated as 2db645e5e9bdc59465ecdaa18d039460f682cda2. Committed diagnostics, provenance, and audits.

Sources and reproduction

Appendix F

BRM examples and exports

Primary external sources

Cited authors and prospective guests have not reviewed or endorsed this deck. Immutable scientific-source revisions and further references are in the README.