ProBot: Back and Forth

Overview

ProBot provides a compact interface for training mixture density networks (MDNs) with torch. This vignette demonstrates the main workflow on a toy problem where a generative function maps input parameters to the coefficients of a polynomial response, and then shows the inverse problem where we recover plausible inputs from an observed output.

What is a Mixture Density Network?

A standard neural network outputs a single point prediction for each input. A Mixture Density Network (MDN) instead outputs the parameters of a mixture of Gaussian distributions, allowing the network to represent an entire probability distribution over the target rather than a single value. Concretely, for a given input x the MDN produces:

The predictive density is then p(y | x) = Σ_k π_k N(y; μ_k, σ_k²). This is especially powerful when the mapping from inputs to outputs is one-to-many or genuinely stochastic, because the mixture can represent multimodal, skewed, or otherwise non-Gaussian posteriors. In the forward direction this lets us capture prediction uncertainty; in the inverse direction it lets the network represent the full posterior over latent parameters given an observation.

Posterior calibration diagnostics

After training we want to verify that the uncertainty reported by the MDN is well-calibrated — that is, that its credible intervals actually contain the truth at the stated rate. ProBot provides two complementary diagnostics for this:

PIT (Probability Integral Transform). For each test observation and each output dimension, the PIT value is the fraction of posterior samples that fall below the true value. If the posterior is perfectly calibrated, the true value is equally likely to land anywhere in the predictive distribution, so the PIT values should be uniformly distributed on [0, 1]. Systematic deviations from uniformity reveal bias (values clustered near 0 or 1) or over/under-dispersion (U-shaped or Λ-shaped histograms).

TARP (Test of Accuracy with Random Points). PIT is computed independently for each output dimension and therefore cannot detect calibration failures that only appear in the joint posterior. TARP addresses this by projecting both the posterior samples and the true parameter vector onto a random unit direction before computing the coverage fraction. Averaging over many random directions provides an omnibus test of joint calibration: under a well-calibrated posterior the TARP values are also uniform on [0, 1]. A tendency of the values toward 0.5 indicates over-dispersion (posteriors too wide), and bimodal concentrations near 0 and 1 indicate under-dispersion (posteriors too narrow).

Example 1: infer the polynomial output from inputs

We define a generative function that turns two input parameters (u, v) into three polynomial coefficients (a0, a1, a2). Those coefficients are then evaluated on a fixed grid of nine equally spaced points in [−1, 1], giving a nine-dimensional output vector. The nonlinear dependence on u (via sin and squaring) means the forward mapping is not trivially invertible, which is precisely the kind of problem where MDNs excel.

make_polynomial_output <- function(theta) {
  a0 <- 0.3 + 0.8 * theta[, 1] - 0.4 * theta[, 2]
  a1 <- -1.2 + 0.5 * theta[, 1]^2 + 0.7 * theta[, 2]
  a2 <- 0.6 * sin(pi * theta[, 1]) - 0.3 * theta[, 2]^2

  x_grid <- seq(-1, 1, length.out = 9)
  y <- vapply(
    x_grid,
    function(x) a0 + a1 * x + a2 * x^2,
    numeric(nrow(theta))
  )

  colnames(y) <- paste0("x", seq_along(x_grid))
  list(output = y, x_grid = x_grid)
}

n_train <- 1e5
n_test <- 1e4

theta_train <- cbind(
  u = runif(n_train, -1, 1),
  v = runif(n_train, -1, 1)
)

theta_test <- cbind(
  u = runif(n_test, -1, 1),
  v = runif(n_test, -1, 1)
)

poly_train <- make_polynomial_output(theta_train)
poly_test <- make_polynomial_output(theta_test)

It is worth seeing how directly (or not) the various input parameters predict the outputs:

magtri(cbind(theta_train, poly_train$output))

So the question then becomes how well can our inference network (which in this context can be thought of like an emulator with errors) predict x1-9 given u and v and vica versa? Well, spoiling the punchline the answer is ‘remarkably well’ but let’s go through how we get there in small steps.

Scaling and data loading

Neural networks train most efficiently when inputs and outputs are on comparable scales. probotScaleForward() standardises each column to zero mean and unit variance and attaches the centering and scaling parameters as attributes so that the same transform can later be applied to new data or reversed on predictions.

input_scaled <- probotScaleForward(theta_train)
output_scaled <- probotScaleForward(poly_train$output)

dataloader <- probotDataLoader(
  input = input_scaled,
  output = output_scaled,
  batch = 1024
)

probotDataLoader() wraps the training matrices in a torch dataset and returns a dataloader that shuffles and batches observations during training. The batch size of 1024 is a sensible default for datasets of this size; smaller batches introduce more gradient noise (which can help escape local optima) while larger batches give smoother but slower updates.

Building and training the forward MDN

We build an MDN with a three-layer encoder (64 → 128 → 64 hidden units) and five Gaussian mixture components. Five components give the network enough flexibility to represent multimodal or skewed predictive distributions without being unnecessarily large.

forward_model <- probotMakeMDN(
  input_dim = ncol(input_scaled),
  output_dim = ncol(output_scaled),
  mdn_components = 5,
  hidden_dims = c(64, 128, 64)
)()

forward_optimizer <- optim_adam(forward_model$parameters, lr = 0.01)

forward_fit <- probotTrainMDN(
  model = forward_model,
  dataloader = dataloader,
  optimizer = forward_optimizer,
  epochs = 100,
  mdn_components = 5,
  verbose = TRUE,
  early_stop = FALSE,
  stop_window = 5,
  stop_delta = 1e-3
)
#> Epoch 1 Loss -9.308 MAE 1.502 RMSE 0.868 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 10 Loss -31.758 MAE 0.575 RMSE 0.342 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 20 Loss -34.808 MAE 0.782 RMSE 0.539 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 30 Loss -31.112 MAE 0.563 RMSE 0.306 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 40 Loss -30.206 MAE 0.628 RMSE 0.274 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 50 Loss -43.539 MAE 0.246 RMSE 0.123 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 60 Loss -47.397 MAE 0.276 RMSE 0.138 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 70 Loss -36.165 MAE 0.370 RMSE 0.197 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 80 Loss -42.006 MAE 0.472 RMSE 0.278 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 90 Loss -48.021 MAE 0.307 RMSE 0.135 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 100 Loss -50.172 MAE 0.165 RMSE 0.077 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]

early_stop = FALSE runs for the full 100 epochs here; in practice you would set early_stop = TRUE and tune stop_window and stop_delta so that training halts automatically once the loss has plateaued.

Inspecting the training history

probotTrainMDN() returns a data frame recording the negative log-likelihood loss at each epoch. Inspecting the first and last few rows lets you check that the loss decreased steadily and did not plateau prematurely or diverge.

head(forward_fit$history)
#>   epoch               loss               mae              rmse sigma mix_sd
#> 1     1  -9.30808421066284 1.501725139465332 0.868480889641040     0      0
#> 2     2 -17.32022755310059 1.111796351623535 0.663803802543879     0      0
#> 3     3 -22.80797286743164 0.615129773864746 0.334823463498281     0      0
#> 4     4 -25.17761605407715 0.673840639343262 0.434298422944984     0      0
#> 5     5 -17.90195058410644 1.238304868927002 0.679408544679092     0      0
#> 6     6 -21.30663478759766 0.848171200866699 0.424939007185956     0      0
tail(forward_fit$history)
#>     epoch              loss               mae               rmse sigma mix_sd
#> 95     95 -46.1191718225098 0.238092448387146 0.1128445429640708     0      0
#> 96     96 -47.9795334826660 0.202158998565674 0.1148516754322134     0      0
#> 97     97 -46.6236337170410 0.177489519348145 0.0841928017227657     0      0
#> 98     98 -46.0588050207520 0.234969417114258 0.1114440771676907     0      0
#> 99     99 -50.0830497937012 0.143120989837646 0.0670678424235951     0      0
#> 100   100 -50.1717548327637 0.165383726043701 0.0773857163948796     0      0

A monotonically decreasing loss that flattens toward the end of training is healthy. If the loss stops decreasing early, try a lower learning rate or more epochs. If it decreases but then increases (overfitting), reduce the network size or enable early stopping. Possibly use a different loss function.

Predicting on the test set

Before predicting we apply the same scaling that was fitted on the training inputs, using the stored col_means and col_sds attributes. This is critical: using different scaling at prediction time is a common source of subtle errors.

theta_test_scaled <- theta_test
theta_test_scaled <- probotScaleForward(
  theta_test_scaled,
  col_means = attr(input_scaled, "col_scale")$col_means,
  col_sds = attr(input_scaled, "col_scale")$col_sds
)

forward_pred <- probotPredictMDN(
  input = theta_test_scaled,
  model = forward_fit$model,
  mdn_components = 5
)

forward_summary <- probotMarginalPostMDN(
  mdn_output = forward_pred,
  col_means = attr(output_scaled, "col_scale")$col_means,
  col_sds = attr(output_scaled, "col_scale")$col_sds,
  col_names = colnames(output_scaled)
)

round(forward_summary$post_mean[1:3,], 3)
#>         x1    x2    x3    x4    x5     x6     x7     x8     x9
#> [1,] 1.268 0.993 0.766 0.584 0.447  0.356  0.314  0.319  0.369
#> [2,] 1.076 0.840 0.612 0.385 0.161 -0.058 -0.275 -0.491 -0.703
#> [3,] 1.276 0.958 0.688 0.456 0.265  0.116  0.010 -0.054 -0.078

probotMarginalPostMDN() collapses the full mixture distribution to per-dimension marginal summaries (posterior mean and standard deviation) back-transformed to the original scale. This is convenient for point predictions and uncertainty bars, but remember that it discards joint information about correlations between output dimensions.

A quick mean absolute error on the test set shows that the model can recover the polynomial output well.

mean(abs(forward_summary$post_mean - poly_test$output))
#> [1] 0.0161663241343009

Posterior samples and corner plot

Rather than relying only on marginal summaries, we can draw samples directly from the joint predictive distribution for any single observation. This is the most faithful representation of what the MDN has learned.

post_samples <- probotSamplePostMDN(
  mdn_output = forward_pred,
  index = 1,
  n_samples = 2000,
  col_means = attr(output_scaled, "col_scale")$col_means,
  col_sds = attr(output_scaled, "col_scale")$col_sds,
  col_names = colnames(output_scaled)
)

head(post_samples)
#>                    x1                x2                x3                x4
#> [1,] 1.22729367367912 1.002157225713962 0.767353303249470 0.579690699598635
#> [2,] 1.26518573410772 0.983527735991897 0.761583354016096 0.583927569547006
#> [3,] 1.24671134451697 0.905813950221932 0.732127517864297 0.581617921781439
#> [4,] 1.27052526749412 0.967943570051682 0.756843376868781 0.567485744246119
#> [5,] 1.26251941684364 0.984640699853467 0.753465350563364 0.581372718126063
#> [6,] 1.23713143828714 0.968137720566707 0.752830340782203 0.566535195796083
#>                     x5                x6                x7                x8
#> [1,] 0.442075243521205 0.337830517885361 0.304072251451985 0.316012738141691
#> [2,] 0.436529392211126 0.329252085901508 0.294908194170247 0.301600029261156
#> [3,] 0.423957997091565 0.321668425870347 0.269497496116087 0.309936760895435
#> [4,] 0.434413399096889 0.336143635080935 0.300674736207764 0.300560654595664
#> [5,] 0.436964819666944 0.339788582686241 0.294479900834708 0.309450166659099
#> [6,] 0.443991208311972 0.345273468595138 0.299243952121502 0.311036205974620
#>                     x9
#> [1,] 0.382818459998418
#> [2,] 0.345501653306979
#> [3,] 0.367713999493093
#> [4,] 0.358562147704912
#> [5,] 0.358287037121194
#> [6,] 0.368076090068712

The triangle (corner) plot below shows the pairwise marginal distributions of the nine polynomial outputs for the first test observation. The blue reference points/lines mark the true values. Well-calibrated samples should scatter symmetrically around the truth.

magtri(post_samples, refvals = poly_test$output[1,])

Truth vs. predicted scatter plots

For a population-level view we plot the predicted posterior mean against the true value for each of the nine output dimensions. The grey error bars show ±1 posterior standard deviation. Points should scatter symmetrically about the diagonal, and the error bars should span the scatter.

for(i in 1:9){
  magplot(poly_test$output[,i], forward_summary$post_mean[,i],
          xlab=paste0('Truth output: ',i), ylab=paste0('Pred output: ',i), pch='.')
  magerr(poly_test$output[,i], forward_summary$post_mean[,i],
         ylo = forward_summary$post_sd[,i], col=hsv(v=0, alpha=0.05))
  
}

PIT diagnostic: checking marginal calibration

The Probability Integral Transform (PIT) provides a rigorous marginal calibration check. For each test observation and each output dimension the PIT value is the fraction of posterior samples that fall below the true value. Under a perfectly calibrated posterior this quantity is uniformly distributed on [0, 1]: the truth is equally likely to appear anywhere within the predictive distribution.

PIT_out = probotPIT(forward_pred,
            params = poly_test$output,
            col_means = attr(output_scaled, "col_scale")$col_means,
            col_sds = attr(output_scaled, "col_scale")$col_sds,
            col_names = colnames(output_scaled)
)
#> 1000 
#> 2000 
#> 3000 
#> 4000 
#> 5000 
#> 6000 
#> 7000 
#> 8000 
#> 9000 
#> 10000

The triangle plot of the PIT matrix shows each output dimension on the diagonal as a histogram and pairwise scatter plots off-diagonal. What to look for:

magtri(PIT_out)

TARP diagnostic: checking joint calibration

PIT is computed independently for each output dimension and therefore cannot detect miscalibration that only manifests in the joint posterior. TARP (Test of Accuracy with Random Points) addresses this by projecting both the posterior samples and the true parameter vector onto a fresh random unit direction for each observation, then computing the coverage fraction along that direction. Averaging over many random directions provides a sensitive omnibus test of the full joint posterior.

TARP_out_forward = probotTARP(
  forward_pred,
  params = poly_test$output,
  col_means = attr(output_scaled, "col_scale")$col_means,
  col_sds = attr(output_scaled, "col_scale")$col_sds,
  col_names = colnames(output_scaled)
)
#> 1000 
#> 2000 
#> 3000 
#> 4000 
#> 5000 
#> 6000 
#> 7000 
#> 8000 
#> 9000 
#> 10000

The TARP values should be uniform on [0, 1] under a well-calibrated joint posterior. We can assess this with a histogram and compute the Kolmogorov–Smirnov statistic against a Uniform(0,1) reference.

maghist(TARP_out_forward, breaks = 20,
        xlab = "TARP coverage", ylab = "Count",
        main = "Forward model: TARP diagnostic")
#> Summary of used sample:
#>       Min.    1st Qu.     Median       Mean    3rd Qu.       Max. 
#> 0.00020000 0.28140000 0.49440000 0.49648056 0.71060000 0.99970000
#> Pop Std Dev: 0.25568
#> MAD: 0.31757
#> Half 16-84 Quan (1s): 0.30531
#> Half 02-98 Quan (2s): 0.42663
#> Using 10000 out of 10000
abline(h = length(TARP_out_forward) / 20, lty = 2, col = "red")


# Kolmogorov-Smirnov test against Uniform(0, 1)
ks.test(TARP_out_forward, "punif")
#> Warning in ks.test.default(TARP_out_forward, "punif"): ties should not be
#> present for the one-sample Kolmogorov-Smirnov test
#> 
#>  Asymptotic one-sample Kolmogorov-Smirnov test
#> 
#> data:  TARP_out_forward
#> D = 0.0549, p-value < 2.220446049e-16
#> alternative hypothesis: two-sided

Interpreting the TARP histogram:

Example 2: invert the mapping with MDN uncertainty

For the inverse problem, the polynomial values are the inputs and the latent parameters are the targets. This is where the MDN is especially useful because multiple parameter settings can explain similar outputs — the nonlinear, non-injective nature of the forward mapping means the inverse problem is genuinely ill-posed. A standard regression network could only return a single best-guess parameter vector and would silently discard the inherent ambiguity; the MDN instead returns a full posterior distribution over plausible parameter configurations for each observed polynomial.

Perhaps the most important aspect of this workflow is that, once the MDN is trained, evaluating the posterior for a new observation is essentially free (a single forward pass through the network), compared to the thousands of likelihood evaluations required by traditional MCMC or nested sampling.

inverse_loader <- probotDataLoader(
  input = output_scaled, #we reuse these, but swap the order
  output = input_scaled, #we reuse these, but swap the order
  batch = 1024
)

inverse_model <- probotMakeMDN(
  input_dim = ncol(output_scaled),
  output_dim = ncol(input_scaled),
  mdn_components = 5,
  hidden_dims = c(64, 128, 64)
)()

inverse_optimizer <- optim_adam(inverse_model$parameters, lr = 0.01)

inverse_fit <- probotTrainMDN(
  model = inverse_model,
  dataloader = inverse_loader,
  optimizer = inverse_optimizer,
  epochs = 100,
  mdn_components = 5,
  verbose = TRUE,
  early_stop = FALSE,
  stop_window = 5,
  stop_delta = 1e-3
)
#> Epoch 1 Loss -1.191 MAE 0.393 RMSE 0.438 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 10 Loss -16.853 MAE 0.175 RMSE 0.196 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 20 Loss -31.145 MAE 0.066 RMSE 0.075 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 30 Loss -29.131 MAE 0.089 RMSE 0.106 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 40 Loss -38.598 MAE 0.050 RMSE 0.058 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 50 Loss -42.222 MAE 0.044 RMSE 0.044 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 60 Loss -37.105 MAE 0.057 RMSE 0.066 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 70 Loss -35.628 MAE 0.057 RMSE 0.065 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 80 Loss -37.108 MAE 0.090 RMSE 0.175 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 90 Loss -41.902 MAE 0.048 RMSE 0.057 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]
#> Epoch 100 Loss -44.850 MAE 0.032 RMSE 0.034 Sigma 0.000 MixSD 0.000 Mix [0.00 0.00 0.00 0.00 0.00]

Note that the only change relative to the forward model is that the roles of input and output are swapped in probotDataLoader() and the corresponding dimensions are swapped in probotMakeMDN(). Everything else — training loop, diagnostics, prediction — is identical.

Predicting the inverse posterior

Given an observed polynomial output, we predict a posterior distribution over the original inputs. As before, we apply the training scaling to the test data using the stored attributes before passing it to the model.

output_test_scaled <- poly_test$output
output_test_scaled <- probotScaleForward(
  output_test_scaled,
  col_means = attr(output_scaled, "col_scale")$col_means,
  col_sds = attr(output_scaled, "col_scale")$col_sds
)

inverse_pred <- probotPredictMDN(
  input = output_test_scaled,
  model = inverse_fit$model,
  mdn_components = 5
)

inverse_summary <- probotMarginalPostMDN(
  mdn_output = inverse_pred,
  col_means = attr(input_scaled, "col_scale")$col_means,
  col_sds = attr(input_scaled, "col_scale")$col_sds,
  col_names = colnames(theta_train)
)

head(round(inverse_summary$post_mean, 3))
#>           u      v
#> [1,]  0.572  0.810
#> [2,]  0.048  0.438
#> [3,]  0.306  0.682
#> [4,]  0.509 -0.373
#> [5,] -0.371 -0.621
#> [6,]  0.530  0.605
head(round(inverse_summary$post_sd, 3))
#>          u     v
#> [1,] 0.017 0.010
#> [2,] 0.012 0.007
#> [3,] 0.014 0.008
#> [4,] 0.030 0.017
#> [5,] 0.012 0.014
#> [6,] 0.017 0.010

post_sd gives the marginal posterior standard deviation for each parameter. Larger values indicate dimensions where the data are less informative or where the inverse problem is more degenerate.

Truth vs. predicted parameter scatter plots

We plot the recovered posterior means against the true generating parameters for all 10,000 test observations. A tight, unbiased scatter around the diagonal indicates accurate recovery. The grey error bars show ±1 posterior standard deviation; roughly 68% of the true values should fall within the bars if the posterior is well calibrated.

for(i in 1:2){
  magplot(theta_test[,i], inverse_summary$post_mean[,i],
          xlab=paste0('Truth param: ',i), ylab=paste0('Pred param: ',i), pch='.')
  magerr(theta_test[,i], inverse_summary$post_mean[,i], ylo=inverse_summary$post_sd[,i],
         col=hsv(v=0, alpha=0.05))
}

Posterior samples and corner plot

The posterior samples provide a direct way to inspect uncertainty and possible multimodality in the recovered inputs. Because the forward mapping is nonlinear and non-injective, some observed polynomials may be consistent with multiple distinct (u, v) configurations; the MDN can represent this as a multimodal mixture.

post_samples <- probotSamplePostMDN(
  mdn_output = inverse_pred,
  index = 1,
  n_samples = 2000,
  col_means = attr(input_scaled, "col_scale")$col_means,
  col_sds = attr(input_scaled, "col_scale")$col_sds,
  col_names = colnames(theta_train)
)

head(post_samples)
#>                      u                 v
#> [1,] 0.559752255727298 0.815323366241411
#> [2,] 0.575515806407412 0.810990001588867
#> [3,] 0.577259703162403 0.802594376144006
#> [4,] 0.571663119126920 0.806058231020345
#> [5,] 0.577798647686285 0.816594674739335
#> [6,] 0.579819828538640 0.804088770266276
magtri(post_samples, refvals = theta_test[1, ])

Summary comparison for one observation

We can also compare the posterior mean and standard deviation directly against the true generating inputs for a single held-out observation, giving a concise numerical summary of the recovery quality.

comparison <- data.frame(
  parameter = colnames(theta_test),
  truth = as.numeric(theta_test[1, ]),
  estimate = as.numeric(inverse_summary$post_mean[1, ]),
  sd = as.numeric(inverse_summary$post_sd[1, ])
)
comparison
#>   parameter             truth          estimate                 sd
#> 1         u 0.571313630789518 0.572076034468495 0.0172140772407963
#> 2         v 0.807143969461322 0.810350936773078 0.0100606152158296

PIT diagnostic for the inverse model

We repeat the PIT calibration check for the inverse posterior. In the inverse problem the output space is the two-dimensional parameter space (u, v), so PIT now has two columns. Uniform histograms on the diagonal confirm that the MDN’s marginal posteriors over each parameter are well calibrated.

PIT_inverse = probotPIT(
  inverse_pred,
  params = theta_test,
  col_means = attr(input_scaled, "col_scale")$col_means,
  col_sds = attr(input_scaled, "col_scale")$col_sds,
  col_names = colnames(theta_train)
)
#> 1000 
#> 2000 
#> 3000 
#> 4000 
#> 5000 
#> 6000 
#> 7000 
#> 8000 
#> 9000 
#> 10000
magtri(PIT_inverse)

TARP diagnostic for the inverse model

Finally, we run the TARP check on the inverse posterior. The inverse problem is genuinely harder to calibrate well because the posterior over (u, v) can be multimodal or banana-shaped, and TARP is sensitive to these joint features in a way that PIT is not.

TARP_inverse = probotTARP(
  inverse_pred,
  params = theta_test,
  col_means = attr(input_scaled, "col_scale")$col_means,
  col_sds = attr(input_scaled, "col_scale")$col_sds,
  col_names = colnames(theta_train)
)
#> 1000 
#> 2000 
#> 3000 
#> 4000 
#> 5000 
#> 6000 
#> 7000 
#> 8000 
#> 9000 
#> 10000
maghist(TARP_inverse, breaks = 20,
        xlab = "TARP coverage", ylab = "Count",
        main = "Inverse model: TARP diagnostic")
#> Summary of used sample:
#>       Min.    1st Qu.     Median       Mean    3rd Qu.       Max. 
#> 0.02070000 0.35470000 0.50075000 0.50107485 0.64652500 0.98510000
#> Pop Std Dev: 0.20189
#> MAD: 0.21639
#> Half 16-84 Quan (1s): 0.22301
#> Half 02-98 Quan (2s): 0.37905
#> Using 10000 out of 10000
abline(h = length(TARP_inverse) / 20, lty = 2, col = "red")


ks.test(TARP_inverse, "punif")
#> Warning in ks.test.default(TARP_inverse, "punif"): ties should not be present
#> for the one-sample Kolmogorov-Smirnov test
#> 
#>  Asymptotic one-sample Kolmogorov-Smirnov test
#> 
#> data:  TARP_inverse
#> D = 0.125, p-value < 2.220446049e-16
#> alternative hypothesis: two-sided

A flat TARP histogram and a non-significant KS p-value provide strong evidence that the inverse MDN posterior is jointly well calibrated — a demanding standard that simultaneously validates the predicted means, uncertainties, and correlations between parameters.

Summary

This forward-and-inverse workflow captures the core ProBot interface:

Function Role
probotScaleForward() / probotScaleBackward() Standardise columns and reverse the transform
probotDataLoader() Assemble batched, shuffled training data for Torch
probotMakeMDN() Define the MDN architecture (hidden layers, components)
probotTrainMDN() Optimise the network via negative log-likelihood
probotPredictMDN() Run the trained MDN on new inputs
probotMarginalPostMDN() Compute per-dimension posterior means and SDs
probotSamplePostMDN() Draw joint samples from the predictive mixture
probotPIT() Check marginal calibration via the Probability Integral Transform
probotTARP() Check joint calibration via the Test of Accuracy with Random Points

The two calibration diagnostics are complementary: PIT catches dimension-wise miscalibration (bias, over/under-dispersion in individual outputs), while TARP catches miscalibration in the joint distribution (e.g. incorrect correlations or missed multimodality) that is invisible to marginal checks alone. Together they provide a comprehensive assessment of whether the MDN posterior can be trusted for downstream inference.