This guide introduces the main ideas of WitDraw through one running example, Bayesian linear regression. By the end you will have written a model, drawn samples from it, and fitted it to data with the No-U-Turn Sampler (NUTS).
A WitDraw program has four ingredients:
- Sites name the random variables of a model and fix their tensor types.
- A model is an ordinary Scala method that describes how data is
generated. It needs a
ProbContextto run, which it receives as ausingparameter. - Interpreters provide that context and so give the model its meaning.
ForwardSamplerruns it and draws a value for every random variable;ModelScorercomputes its log-density for given values. - Handlers sit between a model and the interpreter and change what the
model's random variables do.
observeis a handler: it fixes some sites to known values.
Everything else - prior predictive checks, posterior inference - is built from these pieces.
We want to model outcomes
We import dimwit, which provides the tensors and distributions, and WitDraw itself:
import dimwit.*
import dimwit.stats.*
import witdraw.*
import witdraw.inference.{Diagnostics, NUTS}Every program that uses dimwit must initialise the dimwit runtime first:
dimwit.initialize()dimwit tensors have labelled axes, so we first define a label for each kind of axis: one entry per feature and one per sample (observation).
Then we declare one site per random variable. A site is an object, so it
has a name that we can refer to later, for example when we condition on
observed data. Its type parameter is the tensor type of the variable.
trait Feature derives Label
trait Sample derives Label
object Alpha extends Site[Tensor0[Float32]] // intercept
object Beta extends Site[Tensor1[Feature, Float32]] // regression weights
object Sigma extends Site[Tensor0[Float32]] // noise scale (> 0)
object Y extends Site[Tensor1[Sample, Float32]] // outcomesA model is a method with a (using ProbContext) parameter. The line
val beta = Beta ~ Normal.isotropic(...)reads "the site Beta is distributed according to this normal distribution".
It gives back a plain tensor of the site's type (Tensor1[Feature, Float32]),
which we can use in ordinary dimwit code. Which tensor depends on the
context the model runs in: a sampler draws one, a scorer uses a value we
supply.
The model takes the feature matrix x as an argument: it uses the covariates
both to compute the mean of y and to find out how many weights there are.
The outcomes y, on the other hand, are a random variable of the model. We
pass in the observed values later, by conditioning.
def model(x: Tensor2[Sample, Feature, Float32])(using ProbContext) =
val alpha = Alpha ~ Normal(0f, 10f)
val beta = Beta ~ Normal.isotropic(
Tensor1(x.extent(Axis[Feature])).fill(0f),
10f
)
val sigma = Sigma ~ HalfCauchy(0f, 5f)
val mu = x.dot(Axis[Feature])(beta) +! alpha
val y = Y ~ Normal.isotropic(mu, sigma)
(alpha = alpha, beta = beta, sigma = sigma, y = y)The model returns a named tuple of all its variables. This is what
ForwardSampler gives back, so naming the fields makes the samples easy to
use.
To check that inference works, we simulate data where we know the answer. dimwit uses explicit random keys (as in JAX). We split one root key into independent keys, one for each use, so that no key is used twice:
val (dataKey, noiseKey, priorKey, fixedKey, nutsKey) =
val ks = Key(42).split(5)
(ks(0), ks(1), ks(2), ks(3), ks(4))
val sampleExtent = Axis[Sample] -> 100
val featureExtent = Axis[Feature] -> 2
// independent standard-normal features
val xData = Normal.standardNormal(Shape(sampleExtent, featureExtent)).sample(dataKey)
// y = alpha + x · beta + sigma * eps
val alphaTrue = Tensor0(0.5f)
val betaTrue = Tensor(Shape(featureExtent)).fromArray(Array(-1.5f, 2.0f))
val sigmaTrue = Tensor0(0.8f)
val eps = Normal.standardNormal(Shape(sampleExtent)).sample(noiseKey)
val yData = (xData.dot(Axis[Feature])(betaTrue) +! alphaTrue) + (eps *! sigmaTrue)The features are drawn independently on purpose. If one feature were a
multiple of the other, only a combination of the weights could be learned
from the data, and no inference algorithm could recover betaTrue.
A model can only run inside a context, so to draw from it we hand
model(xData) to ForwardSampler. (Calling model(xData) on its own, outside
an interpreter, is a compile error.) With no site conditioned, this samples
every variable from its prior, so the y values are draws from the prior
predictive distribution:
val priorPredictive = ForwardSampler.sample(priorKey, 5)(model(xData))ForwardSampler.sample runs the model 5 times, each with its own key split
from priorKey; ForwardSampler.run(key)(model) runs it once.
Each sample is a named tuple, so we can print a short summary of each:
priorPredictive.foreach { s =>
println(f"alpha=${s.alpha.item}%7.2f sigma=${s.sigma.item}%7.2f mean(y)=${s.y.mean.item}%8.2f")
}
// alpha= 3.57 sigma= 2.96 mean(y)= 3.51
// alpha= 0.02 sigma= 3.84 mean(y)= -1.06
// alpha= 23.41 sigma= 1.85 mean(y)= 22.30
// alpha= 0.68 sigma= 53.42 mean(y)= 1.33
// alpha= -10.14 sigma= 13.13 mean(y)= -10.02The spread of these draws shows how vague our priors are.
The handler observe runs a block of code with some sites fixed to given
values. Here we fix the noise scale to 1 and let everything else vary:
val withFixedSigma =
ForwardSampler.sample(fixedKey, 5):
observe(Sigma := Tensor0(1f)):
model(xData)withFixedSigma.foreach { s =>
println(f"sigma=${s.sigma.item}%4.2f mean(y)=${s.y.mean.item}%8.2f")
}
// sigma=1.00 mean(y)= -2.25
// sigma=1.00 mean(y)= 11.51
// sigma=1.00 mean(y)= 8.88
// sigma=1.00 mean(y)= 2.48
// sigma=1.00 mean(y)= -16.95Now we condition on the observed outcomes and ask which parameter values are
plausible. The model observed on the data has three latent sites left,
Alpha, Beta and Sigma: the parameters whose posterior we want.
The inference methods in WitDraw work on a Latent: the latent sites, a
parameter type for them, and the posterior's log-density as a function of
that type. Latent.discover builds it. It runs the observed model forward
once, records every site with the withTrace handler, and collects the sites
that are not observed:
val latent = Latent.discover:
observe(Y := yData):
model(xData)latent.sites.map(_.name)
// res3: Seq[String] = Vector("Alpha", "Beta", "Sigma")latent.Params is a parameter type made for this model: one unconstrained
value per latent site, in one vector, which is where inference moves. We do
not read it directly: latent.result(params) gives what the model returns at
given parameters, the same named tuple as the prior predictive draws above.
latent.logDensity is the log-density as a function of latent.Params, and
latent.init is a starting point for inference: the values of the traced
run.
Handlers nest like any other block, and each one sees every random variable
that is sampled inside its block. To compute logDensity, Latent wraps the
model in one more observe that fixes the latent sites to the given
parameters, so ModelScorer finds every site observed and adds up their
log-densities.
NUTS moves freely through all of
Every distribution has a support, the values it can take, and each support
comes with a Bijector, a change of variables from unconstrained values to
the support. HalfCauchy takes
positive values, so Latent samples Sigma through latent.logDensity includes the
Jacobian correction for the change of variables. We never see latent.result(params).sigma always gives
NUTS is a variant of Hamiltonian Monte Carlo that chooses the length of each
trajectory automatically. It still has two tuning parameters: the step size,
and the mass matrix, which rescales each direction of the parameter space and
works best when it approximates the posterior covariance. NUTS.adapt runs a
warmup that tunes both (Stan's window adaptation, run by BlackJax). The warmup
also moves the chain away from its starting point, so no draws need to be
discarded afterwards: adapted.sample continues where the warmup ended, and
returns a chain of draws; samples gives the model's result at each, the same
named tuple as the prior predictive draws above.
val (warmupKey, sampleKey) = nutsKey.split2()
val adapted = NUTS.adapt(latent, warmupKey, latent.init, numSteps = 1000)
val nSamples = 2000
val nutsChain = adapted.sample(sampleKey, nSamples)
val samples = nutsChain.samplesThe tuned step size is adapted.stepSize. By default the mass matrix is
diagonal, and adapted.inverseMassMatrix holds the estimated posterior
variance of each parameter; dense = true estimates the full covariance
instead, at the cost of a longer warmup. To choose the parameters yourself,
use NUTS(latent, stepSize = 0.05), optionally with an InverseMassMatrix.
Finally we stack the draws along a new axis and compare the posterior mean and standard deviation of each parameter with the true values:
trait Draw derives Label
def summary[T <: Tuple: Labels](draws: Seq[Tensor[T, Float32]]) =
val stacked = stack(draws, Axis[Draw])
(stacked.mean(Axis[Draw]), stacked.std(Axis[Draw]))
val (alphaMean, alphaSd) = summary(samples.map(_.alpha))
val (betaMean, betaSd) = summary(samples.map(_.beta))
val (sigmaMean, sigmaSd) = summary(samples.map(_.sigma))println(s"alpha: mean $alphaMean, sd $alphaSd (true $alphaTrue)")
// alpha: mean 0.40361282, sd 0.06946172 (true 0.5)
println(s"beta: mean $betaMean, sd $betaSd (true $betaTrue)")
// beta: mean [-1.412222 2.04343 ], sd [0.07319704 0.07117827] (true [-1.5 2. ])
println(s"sigma: mean $sigmaMean, sd $sigmaSd (true $sigmaTrue)")
// sigma: mean 0.6899296, sd 0.05095555 (true 0.8)The true values lie within roughly two posterior standard deviations of the posterior means. With only 100 noisy observations we cannot expect an exact match.
Each draw also comes with diagnostic information from the sampler. A healthy run has no divergent transitions and a reasonably high acceptance rate:
val divergences = nutsChain.numDivergent
// divergences: Int = 0
val meanAcceptance = nutsChain.infos.map(_.acceptanceRate).sum / nutsChain.size
// meanAcceptance: Double = 0.9308549618422985One chain cannot tell us whether it has explored the whole posterior. The
standard check is to run several chains from different starting points and
compare them. latent.initRandom draws a starting point (uniformly between -2
and 2 in unconstrained space, as Stan does). A chain is ordinary code: here a
warmup that tunes NUTS for this chain alone, followed by sampling.
def chain(key: Key, start: latent.Params) =
val (warmupKey, sampleKey) = key.split2()
val adapted = NUTS.adapt(latent, warmupKey, start)
adapted.sample(sampleKey, 1000)sample returns a Chain: its draws, with the sampler's info for each.
chain.samples gives the model's result at each draw, and
chain.numDivergent the number of divergent transitions. To run the chains at the same time, we run each in a
Future:
import scala.concurrent.{Await, Future, blocking}
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration.Duration
val Seq(initKey, chainsKey) = nutsKey.split(2)
val starts = initKey.split(4).map(latent.initRandom(_))
val runs = chainsKey.split(4).zip(starts).map((key, start) => Future(blocking(chain(key, start))))
val chains = Await.result(Future.sequence(runs), Duration.Inf)A chain mostly waits for compiled code, and blocking lets each chain have a
thread of its own. The warmup and the sampling loop are compiled for the first
chain, and the others reuse them. A chain can run any algorithm, use other
settings, or return more than its draws, e.g. its adapted step size.
Diagnostics.rhat compares the chains: values close to 1 mean they agree, and
values above about 1.01 are a warning. Diagnostics.ess estimates how many
independent draws the chains are worth. Both take the draws of one quantity,
per chain, and have its shape, so beta gets one value per predictor:
val chainSamples = chains.map(_.samples)
// chainSamples: Seq[Seq[NamedTuple[Tuple4["alpha", "beta", "sigma", "y"], Tuple4[Tensor[EmptyTuple, Float32], Tensor[Tuple1[Feature], Float32], Tensor[EmptyTuple, Float32], Tensor[Tuple1[Sample], Float32]]]]] = Vector(
// List(
// (
// 0.36428267,
// [-1.4184844 2.0454156],
// 0.72128636,
// [-0.10196236 0.2429781 -0.05705082 -2.4873717 3.7834146 -0.4873079
// -1.9699229 -1.8418556 1.2729169 3.7633083 4.183006 3.387136
// 2.9561925 -0.62688136 -4.299474 -0.25279456 -0.8641423 3.6897504
// -0.2727195 2.7278574 -2.4174492 5.498894 -0.34408772 0.33337837
// 0.7967427 4.4882426 5.1004615 1.6002898 1.3427472 -1.447418
// 1.5399214 -2.2170792 -2.14231 5.986033 4.02204 -1.9723366
// -1.5114689 -1.2178509 -0.759506 -0.05417639 6.4333787 3.2024684
// 0.0425368 3.8634524 -1.4387481 -3.3665395 0.03691113 0.46876243
// 4.0792418 1.9930805 2.0883217 1.9663954 -1.1308525 -2.3779635
// 2.5628986 -2.0016458 -0.8894644 -1.4009279 1.2943025 -1.9712675
// -1.7932141 2.2704911 -1.9712567 0.6256285 -2.7002976 -4.2193174
// -5.6854215 -2.1527383 3.252094 0.551646 0.18111685 3.3525226
// -2.0822215 -3.8152156 -0.26669765 0.40129253 3.6129751 5.231741
// -0.8901636 2.5462198 -2.3932333 -2.3989074 -0.94802153 3.0354598
// 2.3858023 -0.73757464 0.09980711 2.6494377 5.2678933 -1.1561251
// 0.56709236 2.849967 0.48100537 -0.5164595 3.6530075 1.91774
// -1.0295032 -0.8064156 2.4187832 0.6456901 ]
// ),
// (
// 0.37039635,
// [-1.4348248 2.1133378],
// 0.69500756,
// [-0.10196236 0.2429781 -0.05705082 -2.4873717 3.7834146 -0.4873079
// -1.9699229 -1.8418556 1.2729169 3.7633083 4.183006 3.387136
// 2.9561925 -0.62688136 -4.299474 -0.25279456 -0.8641423 3.6897504
// -0.2727195 2.7278574 -2.4174492 5.498894 -0.34408772 0.33337837
// 0.7967427 4.4882426 5.1004615 1.6002898 1.3427472 -1.447418
// 1.5399214 -2.2170792 -2.14231 5.986033 4.02204 -1.9723366
// -1.5114689 -1.2178509 -0.759506 -0.05417639 6.4333787 3.2024684
// 0.0425368 3.8634524 -1.4387481 -3.3665395 0.03691113 0.46876243
// 4.0792418 1.9930805 2.0883217 1.9663954 -1.1308525 -2.3779635
// 2.5628986 -2.0016458 -0.8894644 -1.4009279 1.2943025 -1.9712675
// -1.7932141 2.2704911 -1.9712567 0.6256285 -2.7002976 -4.2193174
// -5.6854215 -2.1527383 3.252094 0.551646 0.18111685 3.3525226
// -2.0822215 -3.8152156 -0.26669765 0.40129253 3.6129751 5.231741
// -0.8901636 2.5462198 -2.3932333 -2.3989074 -0.94802153 3.0354598
// 2.3858023 -0.73757464 0.09980711 2.6494377 5.2678933 -1.1561251
// 0.56709236 2.849967 0.48100537 -0.5164595 3.6530075 1.91774
// -1.0295032 -0.8064156 2.4187832 0.6456901 ]
// ),
// (
// 0.36322737,
// [-1.3989373 2.1478271],
// ...
Diagnostics.rhat(chainSamples.map(_.map(_.beta)))
// res7: Tensor[Tuple1[Feature], Float32] = [0.9999283 0.9999028]
Diagnostics.ess(chainSamples.map(_.map(_.sigma)))
// res8: Tensor[EmptyTuple, Float32] = 4152.4NUTS is not the only inference method: MAP(latent) finds the posterior mode,
and Laplace(latent) approximates the posterior by a Gaussian around its mode
and is run like NUTS. Both are in witdraw.inference and take the same
latent.
model above tells the whole story in one method. As models grow, it helps
to split them into parts, for example to reuse a likelihood with different
priors. Since a model is an ordinary method, this is ordinary Scala: a part of
a model is a method with a (using ProbContext) parameter, and calling it
from another model runs it in the same context.
The prior is a part of its own. It returns the parameters as a named tuple:
type Params = (alpha: Alpha.Value, beta: Beta.Value, sigma: Sigma.Value)
def prior(nFeatures: Int)(using ProbContext): Params =
val alpha = Alpha ~ Normal(0f, 10f)
val beta = Beta ~ Normal.isotropic(
Tensor1(Axis[Feature] -> nFeatures).fill(0f),
10f
)
val sigma = Sigma ~ HalfCauchy(0f, 5f)
(alpha = alpha, beta = beta, sigma = sigma)The likelihood takes the parameters as ordinary tensors. It does not know where they come from, only that they have the right types:
def likelihood(
x: Tensor2[Sample, Feature, Float32],
params: Params
)(using ProbContext) =
val mu = x.dot(Axis[Feature])(params.beta) +! params.alpha
Y ~ Normal.isotropic(mu, params.sigma)The full model takes the prior as a parameter, so that we can swap it. A
model passed around as a value has the type ProbProg[A], short for the
context function type ProbContext ?=> A. Writing prior inside the model
runs it, in the model's own context:
def regression(
x: Tensor2[Sample, Feature, Float32],
prior: ProbProg[Params]
)(using ProbContext) =
val params = prior
val y = likelihood(x, params)
(alpha = params.alpha, beta = params.beta, sigma = params.sigma, y = y)regression(xData, prior(nFeatures)) contains the same sites with the same
distributions as model(xData), so it can be used in exactly the same way,
from ForwardSampler to Latent.discover and NUTS. For example, both give
the same log-density at the true parameter values:
val nFeatures = xData.shape(Axis[Feature])
// nFeatures: Int = 2
def logDensityAtTruth(m: ProbProg[Any]): Float =
ModelScorer.run:
observe(Y := yData, Alpha := alphaTrue, Beta := betaTrue, Sigma := sigmaTrue):
m
.item
(logDensityAtTruth(model(xData)), logDensityAtTruth(regression(xData, prior(nFeatures))))
// res9: Tuple2[Float, Float] = (-117.938705F, -117.938705F)Because the prior is a parameter of regression, trying a different prior
does not touch the likelihood. Any program that yields Params fits,
for example one with tighter priors on the weights:
def tightPrior(nFeatures: Int)(using ProbContext): Params =
val alpha = Alpha ~ Normal(0f, 10f)
val beta = Beta ~ Normal.isotropic(
Tensor1(Axis[Feature] -> nFeatures).fill(0f),
1f
)
val sigma = Sigma ~ HalfCauchy(0f, 5f)
(alpha = alpha, beta = beta, sigma = sigma)
val tightSamples = priorKey.split(5).map { k =>
ForwardSampler.run(k)(regression(xData, tightPrior(nFeatures)))
}Two things to keep in mind when composing:
- A part that samples sites needs a
ProbContext, so it can only be called from a model, a handler or an interpreter. Calling it anywhere else is a compile error. - A site is identified by its object: two different objects are different sites, even if they share a name. Each site may be sampled only once per run when it is observed, traced or inferred; there, sampling it twice is an error.
Methods are enough to split a model into parts, but each part still uses
the global sites Alpha, Beta and Sigma. A part that samples sites can
therefore be used only once per model. Calling likelihood twice, say for
two outcomes, would sample Y twice.
A component is a part of a model that owns its sites. It is written as a
module: a trait that extends Component declares the part, and each object
that extends the trait is one instance of it, with its own sites.
As an example we let the noise scale depend on the features too, a location-scale regression:
Both the mean and the log-scale are linear in the features, so we write the linear predictor once, as a component. Its prior scale is a trait parameter:
trait Linear(priorScale: Float) extends Component:
object Intercept extends Site[Tensor0[Float32]]
object Weights extends Site[Tensor1[Feature, Float32]]
def apply(x: Tensor2[Sample, Feature, Float32])(using ProbContext) =
val intercept = Intercept ~ Normal(0f, priorScale)
val weights = Weights ~ Normal.isotropic(
Tensor1(x.extent(Axis[Feature])).fill(0f),
priorScale
)
(
output = x.dot(Axis[Feature])(weights) +! intercept,
values = Linear.Values(intercept, weights)
)
object Linear:
case class Values(intercept: Tensor0[Float32], weights: Tensor1[Feature, Float32])apply returns the linear predictor and a typed record of the component's
values, Linear.Values; we come back to the record below.
A component can contain other components. The location-scale regression
has two instances of Linear, and its own outcome site:
trait LocationScale extends Component:
object mean extends Linear(priorScale = 10f)
object logScale extends Linear(priorScale = 1f)
object Y extends Site[Tensor1[Sample, Float32]]
def apply(x: Tensor2[Sample, Feature, Float32])(using ProbContext) =
val m = mean(x)
val s = logScale(x)
val y = Y ~ Normal(m.output, s.output.exp)
(y = y, values = LocationScale.Values(m.values, s.values))
object LocationScale:
case class Values(mean: Linear.Values, logScale: Linear.Values)
object outcome extends LocationScaleA site is named by the path of the components it is declared in. This is how sites with the same declaration are told apart, in error messages as well as here:
val lsLatent = Latent.discover:
observe(outcome.Y := yData):
outcome(xData)
// lsLatent: Latent {
type Result >: NamedTuple[Tuple2["y", "values"], Tuple2[Tensor[Tuple1[Sample], Float32], Values]] <: NamedTuple[Tuple2["y", "values"], Tuple2[Tensor[Tuple1[Sample], Float32], Values]]
} = witdraw.Latent$$anon$4@852222c
lsLatent.sites.map(_.name)
// res10: Seq[String] = Vector(
// "outcome.mean.Intercept",
// "outcome.mean.Weights",
// "outcome.logScale.Intercept",
// "outcome.logScale.Weights"
// )All parameters can take any real value: the scale is positive through
exp.
With many sites it helps to read all values of a component at once, as a
typed record. Each component above returns its record, and the records nest
like the components do: LocationScale.Values holds one Linear.Values for
the mean and one for the log-scale.
The record is part of the model's result, so it comes with every forward run,
with lsLatent.result(params) at a point of the posterior, and with every
sample of a chain. Here we find the posterior mode with MAP and read its
record. We start the search at 0,
i.e. mean 0 and scale 1 for every sample, rather than at the traced values,
a draw from the prior whose scales can be extreme:
import witdraw.inference.MAP
val zeros = Tensor1(xData.extent(Axis[Feature])).fill(0f)
val start = Seq(
outcome.mean.Intercept := Tensor0(0f),
outcome.mean.Weights := zeros,
outcome.logScale.Intercept := Tensor0(0f),
outcome.logScale.Weights := zeros
)val mode = lsLatent.result(MAP(lsLatent, start)).values
// mode: Values = Values(
// mean = Values(intercept = 0.40074044, weights = [-1.4111495 2.0541945]),
// logScale = Values(
// intercept = -0.40187716,
// weights = [ 0.0855608 -0.02117608]
// )
// )
mode.mean.weights
// res11: Tensor[Tuple1[Feature], Float32] = [-1.4111495 2.0541945]
mode.logScale.weights
// res12: Tensor[Tuple1[Feature], Float32] = [ 0.0855608 -0.02117608]
mode.logScale.intercept.exp
// res13: Tensor[EmptyTuple, Float32] = 0.669063The weights of the mean are close to betaTrue. The data was simulated
with a constant noise scale, so the weights of the log-scale are close to 0,
and the scale itself,
Records of tensors are DimWit tensor trees, so draws can also be averaged,
mapped over or saved as a whole. The example
LocationScale.scala
fits a model with two outcomes, each a LocationScale, with NUTS, and
averages the posterior draws as records.
A component that is not an object, e.g. one of several created in a loop,
needs an explicit name: new Linear(1f) { override def name = s"region$i" }.