# Variational inference

Whittaker's default fitting pipeline uses P-IRLS to find the posterior mode \hat\beta and then treats the **Laplace approximation** q(\beta) = \mathcal{N}(\hat\beta,\\ V\_\beta) as the posterior, where V\_\beta = \phi\\(X^\top W X + \sum_j \lambda_j S_j)^{-1}. For Gaussian response this is exact. For non-Gaussian families (Poisson, Binomial, Gamma) it is an approximation that can underestimate posterior spread when the likelihood surface is skewed.

**Variational inference (VI)** replaces the Laplace approximation with a better-calibrated posterior, chosen by maximizing the Evidence Lower BOund (ELBO):

\text{ELBO}(\phi) = \mathbb{E}\_q\[\log p(y \mid \beta)\] - \text{KL}\bigl(q(\beta) \\\\ p(\beta \mid \boldsymbol\lambda)\bigr)

The key practical benefit is more accurate uncertainty quantification for non-Gaussian families, especially at small-to-moderate sample sizes, without the cost or complexity of MCMC.


# When to use VI

Use `method="VI"` when:

- your response is non-Gaussian (e.g., Poisson, Binomial, Gamma, etc.) **and** you want well-calibrated and credible intervals rather than Wald-style confidence intervals
- sample sizes are moderate (n \lesssim 5000) and the posterior may be skewed
- you want a principled probabilistic posterior (not just a point estimate with a heuristic covariance) but MCMC is too slow for your use case

For large n, the Laplace approximation is already very accurate and `method="REML"` (the default) is both faster and nearly as good. For Gaussian response, VI and Laplace give identical results. Whittaker detects this and skips the optimization entirely.


# Basic usage

Pass `method="VI"` to `fit()` (everything else stays the same).


``` python
import numpy as np
import whittaker as wk
from whittaker.families.poisson import Poisson

rng = np.random.default_rng(23)
n = 300
x = np.linspace(0, 2 * np.pi, n)
lam = np.exp(1.5 * np.sin(x))
y = rng.poisson(lam).astype(float)

data = {"x": x, "y": y}

model = wk.GAM("y ~ s(x)", family=Poisson())
model.fit(data, method="VI")
model.summary()
```


    GAM fit summary
    ============================================================
    Formula:    y ~ s(x)
    Family:     Poisson(link='log')
    Inference:  Variational Bayes
    Observations: 300
    Coefficients: 10

    Total EDF:  6.79
    Scale est:  1.000000
    ELBO:       -413.4276
    VI iters:   88
    Converged:  True
    Deviance:   313.1025
    Null dev:   742.3931
    Dev. expl:  57.8%
    AIC:        809.66
    BIC:        834.80


The summary shows `Inference: Variational Bayes` and the final ELBO instead of GCV/AIC/BIC, which are not defined for VI fits.


# The variational family

Whittaker uses a **full-rank Gaussian** variational family:

q(\beta) = \mathcal{N}(m,\\ C), \qquad C = LL^\top

where m \in \mathbb{R}^p is the variational mean and L is a lower-triangular Cholesky factor with positive diagonal. This parameterization guarantees positive-definiteness throughout optimization and it yields numerically stable gradients.

Full-rank covariance is a good default for GAMs: smooth coefficients within a term are heavily correlated by construction, and a diagonal (mean-field) approximation would produce incorrect uncertainty estimates for the smooth curves.

The parameters (m, L) are initialized from the P-IRLS solution (VI is a refinement of the Laplace approximation and not a replacement) and then optimized with the Adam optimizer.


# Predictions and uncertainty

`predict(se=True)` and confidence intervals work exactly as with any other fitting method. Standard errors are derived from the variational posterior covariance C:

\text{SE}(\hat\eta_i) = \sqrt{x_i^\top C\\ x_i}


``` python
import altair as alt

x_new = np.linspace(0, 2 * np.pi, 200)
preds = model.predict({"x": x_new}, se=True, interval="confidence")

obs_data = [{"x": float(x[i]), "y": float(y[i])} for i in range(n)]
points = alt.Chart({"values": obs_data}).mark_circle(
    size=15, opacity=0.3, color="steelblue"
).encode(
    x=alt.X("x:Q", title="x"),
    y=alt.Y("y:Q", title="count"),
)

fit_data = [
    {
        "x": float(x_new[i]),
        "fit": float(preds.values[i]),
        "lower": float(preds.lower[i]),
        "upper": float(preds.upper[i]),
    }
    for i in range(len(x_new))
]

line = alt.Chart({"values": fit_data}).mark_line(
    color="firebrick", strokeWidth=2
).encode(x="x:Q", y="fit:Q")

band = alt.Chart({"values": fit_data}).mark_area(
    opacity=0.2, color="firebrick"
).encode(x="x:Q", y="lower:Q", y2="upper:Q")

(band + line + points).properties(
    title="Poisson GAM with VI: fitted curve and 95% credible interval",
    width=500,
    height=300,
)
```


<style>
  #altair-viz-0a635fae1e30486c804f609831bcd5bb.vega-embed {
    width: 100%;
    display: flex;
  }

  #altair-viz-0a635fae1e30486c804f609831bcd5bb.vega-embed details,
  #altair-viz-0a635fae1e30486c804f609831bcd5bb.vega-embed details summary {
    position: relative;
  }
</style>


# Inspecting the variational result

The [vi_result](../reference/GAM.md#whittaker.GAM.vi_result) property returns the underlying `VIResult` object with detailed information about the optimization:


``` python
vr = model.vi_result

print(f"ELBO at convergence: {vr.elbo:.4f}")
print(f"Iterations: {vr.n_iter}")
print(f"Converged: {vr.converged}")
print(f"Posterior covariance shape: {vr.posterior_cov.shape}")
```


    ELBO at convergence: -413.4276
    Iterations: 88
    Converged: True
    Posterior covariance shape: (10, 10)


Note that [vi_result](../reference/GAM.md#whittaker.GAM.vi_result) will be `None` if the model was fitted with any other method.


## ELBO trace

The ELBO should increase monotonically (or nearly so) as optimization proceeds. A non-monotone trace suggests that the learning rate is too high. Plot `elbo_history` to diagnose:


``` python
elbo_data = [
    {"iteration": i, "ELBO": float(vr.elbo_history[i])}
    for i in range(len(vr.elbo_history))
]

alt.Chart({"values": elbo_data}).mark_line(color="steelblue").encode(
    x=alt.X("iteration:Q", title="Iteration"),
    y=alt.Y("ELBO:Q", title="ELBO"),
).properties(
    title="ELBO convergence trace",
    width=500,
    height=250,
)
```


<style>
  #altair-viz-b02b5dbd48984cd89acc44d5c15a71c7.vega-embed {
    width: 100%;
    display: flex;
  }

  #altair-viz-b02b5dbd48984cd89acc44d5c15a71c7.vega-embed details,
  #altair-viz-b02b5dbd48984cd89acc44d5c15a71c7.vega-embed details summary {
    position: relative;
  }
</style>


A trace that rises quickly and then flattens is ideal. Oscillations indicate `lr` should be reduced. A trace that rises very slowly may benefit from a larger `lr` or more iterations.


# Posterior samples

`posterior_samples(n)` draws coefficient vectors directly from q(\beta) = \mathcal{N}(m, C):


``` python
beta_samples = model.posterior_samples(n=500, seed=23)
print(f"Shape: {beta_samples.shape}")  # (p, 500)
```


    Shape: (10, 500)


This also works after `method="REML"` or `"GCV"` fits (in that case samples are drawn from the Laplace posterior), so code that uses [posterior_samples()](../reference/GAM.md#whittaker.GAM.posterior_samples) is inference-method agnostic.


# Simulating from the posterior

`simulate()` uses the posterior coefficient samples to propagate uncertainty through the response distribution:


``` python
sims = model.simulate(n_sim=200, seed=0)
print(f"Simulations shape: {sims.shape}")  # (n, 200): integer count draws
```


    Simulations shape: (300, 200)


# Controlling VI

Pass a `vi_options=` dict to `fit()` to override any of the optimizer settings:

``` python
model = wk.GAM("y ~ s(x)", family=Poisson()).fit(
    data,
    method="VI",
    vi_options={
        "lr": 0.005,        # Adam learning rate (default 0.01)
        "max_iter": 2000,   # maximum optimizer steps (default 1000)
        "tol": 1e-5,        # relative ELBO change threshold (default of 1e-4)
        "patience": 10,     # consecutive steps below tol before stopping (the default is 5)
        "n_quad": 30,       # Gauss-Hermite quadrature points (default: 20)
        "seed": 23,         # seed for reproducibility
    },
)
```


## Block-diagonal covariance

For large models with many smooth terms, the full p \times p Cholesky factor can be expensive to run. `cov_structure="block"` assigns one Cholesky block per smooth term, dropping cross-term correlations:

``` python
model = wk.GAM("y ~ s(x1) + s(x2) + s(x3)", family=Poisson()).fit(
    data,
    method="VI",
    vi_options={"cov_structure": "block"},
)
```

The cost drops from O(p^2) to O(\sum_j k_j^2), where k_j is the basis dimension of the j-th smooth. For a model with 5 terms each with k = 10, this is a 5x reduction in parameters.


## Variational scale parameter

For families with a separate scale parameter \phi (Gamma, InverseGaussian), the default is to fix \phi at its P-IRLS estimate. Setting `phi_inference="variational"` includes \log\phi in the variational family as a log-normal marginal q(\phi) = \text{LogNormal}(\mu\_\phi, \sigma\_\phi^2), which gives better-calibrated intervals when n is small:

``` python
from whittaker.families.gamma import Gamma

model = wk.GAM("y ~ s(x)", family=Gamma()).fit(
    data,
    method="VI",
    vi_options={"phi_inference": "variational"},
)

vr = model.vi_result
print(f"log φ mean: {vr.log_phi_mean:.4f}")
print(f"log φ variance: {vr.log_phi_var:.4f}")
```


# Gaussian fast path

For Gaussian response with the identity link, the Laplace approximation is the exact posterior (no optimization is needed). Whittaker detects this and returns immediately with n\\\text{iter} = 0:


``` python
from whittaker.families.gaussian import Gaussian

gauss_model = wk.GAM("y ~ s(x)", family=Gaussian()).fit(
    {"x": np.linspace(0, 1, 200), "y": rng.normal(size=200)},
    method="VI",
)
print(f"Iterations: {gauss_model.vi_result.n_iter}")   # 0
print(f"Converged: {gauss_model.vi_result.converged}") # True
```


    Iterations: 0
    Converged: True


The `VIResult` from a Gaussian fit is identical to what the Laplace approximation would give, so switching between `method="REML"` and `method="VI"` is seamless for Gaussian models.


# Properties not available for VI fits

The `deviance`, [null_deviance](../reference/GAM.md#whittaker.GAM.null_deviance), [deviance_explained](../reference/GAM.md#whittaker.GAM.deviance_explained), `aic`, `bic`, and [gcv_score](../reference/GAM.md#whittaker.GAM.gcv_score) properties are not defined for VI fits and raise `NotImplementedError`. Use the ELBO as the convergence diagnostic and use posterior predictive checks (via `simulate()`) for model comparison.


# Where to go next

- **[Model comparison with LOO](loo.md)**: use PSIS-LOO to compare the predictive accuracy of competing Bayesian GAM fits.
- **[Posterior predictive distributions](posterior-predict.md)**: the full predictive distribution at new data points, including observation noise.
- **[Posterior predictive checks](ppc.md)**: assess whether a model generates data consistent with the observations.
- **[MCMC sampling](mcmc.md)**: exact posterior inference via the No-U-Turn Sampler, for when VI's Gaussian approximation is insufficient.
- **[Prediction and inference](prediction.md)**: confidence intervals, simultaneous bands, and term-level predictions from any fitted model.
- **[Model diagnostics](diagnostics.md)**: residual plots and `model.check()`.
- **[Response families](families.md)**: how the family affects VI convergence and calibration.