Priors Guide Ambiguous Line Fits¶

# Generate scarce, ambiguous training data
num_points = 6
x_train = torch.linspace(-0.015, 0.02, num_points).unsqueeze(-1)
true_intercept = 0.12
true_slope = 0.0
noise = 0.003 * torch.randn_like(x_train)
y_train = true_intercept + true_slope * x_train + noise

print(f"x range: {x_train.min().item():+.3f} to {x_train.max().item():+.3f}")
print(f"mean y: {y_train.mean().item():.3f}")

# Plot the scarce observations
_, ax = plt.subplots()
ax.scatter(x_train, y_train, color="black")
ax.set_xlabel("x")
ax.set_ylim(0, 0.2)
ax.set_ylabel("y")
ax.set_title("Only six points, all near the origin")
plt.tight_layout()

# Build and train stochastic line models under different priors
def build_model(name: str, slope_prior: float) -> StochasticModel:
    surrogate = PolynomialChaosExpansion(num_inputs=1, num_outputs=1, degree=1)
    prior_mean = torch.tensor([0.12, slope_prior])
    prior_std = torch.tensor([0.04, 0.35])
    covariance = torch.diag(prior_std ** 2)
    prior = GaussianPrior(mean=prior_mean, covariance_matrix=covariance)
    group = GaussianGroup(name, surrogate.num_params(), prior=prior)
    return StochasticModel(surrogate, group)

training_options = dict(
    num_epochs=4000,
    lr=3e-3,
    weight_decay=1e-4,
    weight_factor=5e-5,
    num_parameter_samples=128,
)

priors = {
    "positive slope prior": 1,
    "negative slope prior": -1,
}

models = {}
for label, slope_mean in priors.items():
    model = build_model(label, slope_mean)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=training_options["lr"],
        weight_decay=training_options["weight_decay"],
    )

    print(f"Training model with {label} ...")
    for epoch in range(training_options["num_epochs"]):
        optimizer.zero_grad()
        predictions = model(x_train, num_samples=training_options["num_parameter_samples"])
        target = y_train.unsqueeze(0).expand_as(predictions)

        data_loss = F.mse_loss(predictions, target)
        distribution_loss = model.distribution_loss()
        total_loss = data_loss + training_options["weight_factor"] * distribution_loss

        total_loss.backward()
        optimizer.step()

        if (epoch + 1) % 1000 == 0:
            print(
                f"Epoch {epoch + 1:4d}: data_loss={data_loss.item():.6f} | "
                f"distribution_loss={distribution_loss.item():.6f}"
            )

    models[label] = model

# Compare predictive distributions across the domain
x_test = torch.linspace(-0.5, 0.5, 400).unsqueeze(-1)
quantile_levels = torch.tensor([0.05, 0.5, 0.95])

with torch.no_grad():
    fig, axes = plt.subplots(1, 2, figsize=(12, 4), sharey=True)
    for ax, (label, model) in zip(axes, models.items()):
        samples = model(x_test, num_samples=512).squeeze(-1)
        quantiles = torch.quantile(samples, quantile_levels, dim=0)

        ax.fill_between(
            x_test.squeeze(-1).numpy(),
            quantiles[0].numpy(),
            quantiles[2].numpy(),
            color="lightsteelblue",
            alpha=0.6,
            label="90% credible band",
        )
        ax.plot(x_test, quantiles[1], color="royalblue", label="posterior mean")
        ax.scatter(x_train, y_train, color="black", zorder=3, label="observations")
        ax.axvspan(-0.02, 0.02, color="lightgray", alpha=0.3, label="observed region")
        ax.set_title(label)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.legend(loc="best")

    plt.tight_layout()

# Inspect posterior parameter summaries
for label, model in models.items():
    gaussian_group = model.parameter_groups[0]
    posterior = gaussian_group.variational_distribution()
    mean = posterior.mean
    std = posterior.stddev
    print(
        f"{label:>24s} -> intercept={mean[0]:+.3f} ± {std[0]:.3f}, "
        f"slope={mean[1]:+.3f} ± {std[1]:.3f}"
    )