This article opens up pigauto’s engine room. It is
intended for readers who want to know exactly what the package
does between impute(traits, tree, covariates) and the
returned pigauto_result. We cover the data flow, the gated
ensemble formula, the phylogenetic baseline (joint MVN / threshold-joint
/ OVR), the GNN’s transformer blocks, the optional within-row
cross-trait attention added in v0.9.3, and the calibration / uncertainty
quantification machinery on top.
The GNN described here runs only when gnn = TRUE. Since
version 0.11.0.9001 the default is gnn = FALSE, which fits
the phylogenetic baseline and the calibration and uncertainty machinery
without a network; at the current baseline defaults the GNN did not
lower imputation error on simulated or real benchmark data. User
covariates are used only by the GNN.
We assume familiarity with maximum-likelihood Brownian motion (BM) on trees and with standard transformer attention. We do not assume familiarity with the package’s internals.
impute() chains six internal stages. Tensor shapes use
\(n\) for observations (= species in
single-obs mode, > species when each species has multiple
measurements), \(p\) for the
latent trait dimension (which expands categorical traits to
\(K\) one-hot columns), \(k\) for spectral Laplacian features, \(c\) for covariate dimensions, and \(T\) for a posterior-tree
prediction-sensitivity sample.
INPUT:
traits data.frame n × p_traits (mixed types, NAs allowed)
tree phylo m tips
covariates data.frame n × c (optional)
STEP 1 — preprocess_traits()
X_scaled n × p per-type encoding + z-score
trait_map list descriptors per trait (type, levels, mean, sd)
obs_to_species n only if n > m (multi-obs mode)
STEP 2 — build_phylo_graph()
coords m × k Laplacian eigenvectors
adj m × m Gaussian kernel on cophenetic distance
D_sq m × m squared cophenetic distance (B2)
STEP 3 — fit_baseline() [FIXED, NOT TRAINED]
mu m × p phylogenetic baseline (BM / LP / joint)
se m × p conditional-MVN standard errors
STEP 4 — fit_pigauto() trains ResidualPhyloDAE on a DAE objective
produces a torch nn_module + per-column gate
STEP 5 — calibrate_gates() picks per-trait r_cal on held-out val cells
STEP 6 — predict.pigauto_fit() applies the blend:
pred = (1 - r_cal) · mu + r_cal · delta_GNNThe two essential ideas are:
The baseline is a per-trait-type dispatcher inside
fit_baseline(). For each trait you get \(\mu_i\) (prior posterior mean for missing
cell \(i\)) and \(\sigma_i\) (its posterior SD). All later
inference is conditional on this.
For one continuous trait \(y\) with \(y_O\) observed and \(y_M\) missing, \(y \sim \mathcal{N}(\beta \mathbf{1},\, \sigma^2 \mathbf{R})\) under BM on the species tree, where \(\mathbf{R} = \mathrm{cov2cor}(\mathrm{vcv}(\text{tree}))\) is the phylogenetic correlation matrix. The conditional posterior is the closed-form GLS solution:
\[ \hat{\mu}_M = \beta \mathbf{1} + \mathbf{R}_{MO} \mathbf{R}_{OO}^{-1} (y_O - \beta \mathbf{1}_O), \quad \widehat{\sigma}^2_{M|O} = \sigma^2\bigl(1 - \mathrm{diag}(\mathbf{R}_{MO} \mathbf{R}_{OO}^{-1} \mathbf{R}_{OM})\bigr). \]
R/bm_internal.R::bm_impute_col() implements this
directly. Count traits are log1p-transformed, ordinal traits are coerced
to integers and z-scored, proportions are logit-z-scored. The same
conditional formula is used in each transformed space.
Discrete traits use a softer phylogenetic prior: each species’s class probability is a kernel-weighted average over the rest of the tree, \(\Pr(y_i = k) \propto \sum_{j} \mathbf{A}_{ij} \, \mathbb{1}(y_j = k)\), with \(\mathbf{A}\) the same Gaussian kernel used in the GNN’s adjacency. For categorical traits with \(K\) classes this returns \(K\) log-probability columns; binary returns one logit column.
When ≥ 2 continuous-like latent columns exist, pigauto upgrades to a
joint multivariate-BM baseline. Stack the \(p\) BM-eligible columns into \(\mathbf{Y}\). Under joint BM, \(\mathrm{vec}(\mathbf{Y}) \sim
\mathcal{N}(\mathrm{vec}(\boldsymbol{\beta}),\,
\mathbf{\Sigma} \otimes \mathbf{R})\). pigauto’s in-house
solver (R/joint_mvn_solver.R) fits \(\hat{\mathbf{\Sigma}}\) jointly across
traits and the conditional posterior is the GLS solution with the
Kronecker covariance. This captures cross-trait phylogenetic correlation
that the per-column path loses.
Two points of precision, since both were previously misstated here:
The solver is pigauto’s own, not Rphylopars.
Earlier versions delegated to Rphylopars::phylopars() and
this section still credited it; joint_mvn_available() now
returns TRUE unconditionally and the joint path has no
Rphylopars dependency. (Rphylopars remains in
Suggests for comparison benchmarks only.)
What the shipped default computes. The solver’s
max_iter argument defaults to 0, which returns
the closed-form \(\hat{\mathbf{\Sigma}} =
\mathbf{L}^\top \mathbf{R}^{-1} \mathbf{L} / n\) estimate —
consistent under matrix-normal BM. What max_iter = 0
switches off is the cross-trait EM cell-refinement loop,
disabled in 2026-05 because it diverged on strong-signal data (the
conditional prior ignores \(\mathbf{R}\)-mediated cross-row
correlation, and \(\mathbf{R}^{-1}\)
amplifies the resulting discrepancy into a multiplicative \(\mathbf{\Sigma}\) blow-up). The non-refined
path was empirically much better on synthetic \(K=5\) BM data (0.93 vs 0.53 argmax accuracy
at \(n=100\), 30% missing). So \(\hat{\mathbf{\Sigma}}\) is
estimated jointly; it is the iterative cell refinement on top that is
off by default, and deliberately so.
The bench in script/bench_joint_baseline.R reported a
33.7% RMSE lift on simulated correlated BM data. That figure predates
the in-house solver and has not been re-measured against it — treat it
as indicative of the joint-vs-per-column direction, not as a current
performance number for the shipped path.
To bring binary traits inside the joint MVN, pigauto’s
fit_joint_threshold_baseline() uses the Wright–Falconer
liability model: each observed binary cell \(y
\in \{0, 1\}\) is replaced by the posterior mean of an underlying
continuous liability \(L\)
truncated by \(y\). With prior \(L \sim \mathcal{N}(0, 1)\) this is the
standard truncated-normal mean. The resulting matrix is then handed to
the same in-house joint solver as the continuous case.
Binary posteriors are decoded back to probabilities via \(p = \Phi\bigl(\mu_L / \sqrt{1 +
\sigma_L^2}\bigr)\), clipped to \([0.01, 0.99]\). Ordinal liability uses K-1
interval cuts (B3 ordinal).
The single-fit approach to categorical liability (K columns into one joint call) is rank-deficient and unstable. pigauto instead runs \(K\) independent threshold-joint fits — class \(k\) vs the rest — and renormalises the resulting \(K\) probabilities into a row-stochastic distribution. That is a numerical-stability choice, not a published accuracy lift for this package.
When each species has multiple observations, baselines run at species level. Phase 10 aggregates obs→species with type-aware rules: mean for continuous, modal class (or argmax one-hot) for discrete. B1 (v0.9.0) adds an opt-in soft path that preserves evidence strength: a species observed as class 1 in 6/10 rows uses the convex combination \(p \cdot E[L \mid L > 0] + (1 - p) \cdot E[L \mid L < 0]\) instead of collapsing to hard class 1.
The optional GNN learns a full per-cell prediction that is blended with the phylogenetic baseline. It is a denoising autoencoder (DAE) trained with masked-cell reconstruction loss.
Name note. The torch class is
ResidualPhyloDAEbecause its internal blocks use ResNet-style residual skip connections. The network outputdeltais not a statistical residual \(y - \mu\) — it is a full per-cell prediction, blended externally with \(\mu\) via the per-trait gate.
Input tensors:
x \(\in \mathbb{R}^{n \times
p}\) — current trait latent matrix (missing cells replaced by a
learnable mask token).coords \(\in \mathbb{R}^{m
\times k}\) — species-level spectral Laplacian features.covs \(\in \mathbb{R}^{n
\times c}\) — [baseline_mu | NA-mask | user_covs]
(the user covariates plus a per-cell mask indicator and the baseline
prediction).If use_trait_attention = TRUE (new in v0.9.3, see §3.5),
a pooled trait-context feature of dimension trait_embed_dim
is also concatenated. The encoder is a two-layer MLP:
\[ h = \mathrm{ReLU}\bigl(\mathbf{W}_2 \, \mathrm{Dropout}(\mathrm{ReLU}(\mathbf{W}_1 \mathrm{concat}(x, \text{coords}, \text{covs}, \ldots)))\bigr), \]
producing \(h \in \mathbb{R}^{n \times h_d}\) with hidden dim \(h_d\) (default 64).
In multi-obs mode, \(h\) is averaged
across observations of the same species (scatter_mean) to
produce a species-level hidden state \(h_\text{species} \in \mathbb{R}^{m \times
h_d}\) for the graph message passing, then broadcast back to
observation level afterwards.
The default path stacks \(L\)
pre-norm transformer encoder blocks (n_gnn_layers, default
2). Each block has:
n_heads default 4) over the \(m\) species, with a per-head
learnable phylogenetic bias added to the attention scores. With
\(\mathbf{D}^2\) the squared
cophenetic-distance matrix and \(\beta_h =
\mathrm{softplus}(\mathrm{log\_bw}_h)\) a learned bandwidth per
head, the bias is \(\mathbf{B}_h =
-\mathbf{D}^2 / (2\beta_h^2)\). One head can attend tightly
(fast-evolving traits), another broadly (conserved traits).cov_h = cov_encoder(user_covs) is added inside
the block via a learnable projection, so covariate features are visible
at every depth, not just the encoder.Legacy single-head attention is retained behind
use_transformer_blocks = FALSE for reconstructing
pre-v0.9.0 fits.
A symmetric two-layer MLP maps the species-broadcast hidden state
back to the latent space: delta = dec2(ReLU(dec1(h))) of
shape \(n \times p\).
The per-column blend gate is
\[ r = \sigma(\rho) \cdot \mathtt{gate\_cap}, \qquad r \in (0, \mathtt{gate\_cap}], \]
with \(\rho \in \mathbb{R}^p\) a learnable per-column parameter. The model output is
\[ \hat{x} = (1 - r) \cdot \mu + r \cdot \mathrm{delta} + \mathrm{cov\_linear}(u), \]
where cov_linear is a small direct linear regression on
user covariates (added outside the blend; gives the GNN a “linear
shortcut” on covariates so it doesn’t have to learn \(\beta\) through nonlinear layers).
The gate is initialised so the GNN contribution starts negligible: \(\rho_\text{init} = -1\) for continuous columns (effective gate ≈ 0.135 × gate_cap), \(\rho_\text{init} \approx 0\) for discrete (fully closed).
For each training batch:
\[ \mathcal{L} = \underbrace{\mathcal{L}_\text{type}(\hat{x}, y_\text{true})}_\text{type-aware reconstruction} + \lambda_\text{shrink} \cdot \mathrm{MSE}(\mathrm{delta} - \mu) + \lambda_\text{gate} \cdot \mathrm{MSE}(r). \]
delta drifting away from the baseline.Together with the architectural cap (gate_cap ≤ 0.8 by
default), these three give a strong inductive bias toward the
phylogenetic baseline — useful when the GNN has nothing extra to
learn on a given dataset.
The encoder above mixes all trait columns into a single hidden vector
via a linear projection, which loses per-trait identity. When
use_trait_attention = TRUE, the model additionally builds a
per-trait token sequence:
\[ \text{tokens}_{i,j} = \mathbf{W}_v \, x_{i,j} + \mathbf{e}_j, \qquad j = 1, \ldots, p, \]
with \(\mathbf{e}_j \in
\mathbb{R}^{e_d}\) a learnable positional embedding per trait
column (dim trait_embed_dim, default 32). One multi-head
self-attention block (n_trait_heads, default 2) mixes these
\(p\) tokens within each row \(i\), followed by mean-pool to a single
\(e_d\)-dim feature. That feature is
concatenated alongside \((x, \text{coords},
\text{covs})\) at the encoder input.
This opt-in mechanism changes the model architecture. Assess it with a held-out comparison on the intended data rather than transferring a result from another trait set or simulation regime.
Backward-compatible: saved fits without the field default to
use_trait_attention = FALSE and reconstruct
identically.
After training, pigauto does not ship the gate \(\sigma(\rho)\) that the optimiser converged to. Instead it overrides \(\rho\) with a per-trait calibrated gate \(r_\text{cal}\) chosen on held-out validation cells.
calibrate_gates() runs a per-trait grid search of \(r \in [0, \mathtt{gate\_cap}]\) minimising
val-set reconstruction loss (MSE for continuous; 0/1 loss for discrete
with an absolute cell floor and a split-validation cross-check that
prevents the GNN from harming baseline accuracy). The output is a single
scalar \(r_\text{cal}\) per latent
column, stored on the fit and used at prediction time.
This is the second layer of safety: even if training somehow pushes the learnable gate wide open, calibration on held-out cells can close it back down. In practice, on datasets with strong phylogenetic signal where the GNN cannot improve on BM, \(r_\text{cal} \approx 0\) and the prediction is exactly the baseline.
pigauto exposes three distinct uncertainty mechanisms; they answer different questions and must not be conflated.
Baseline pred$se (continuous / count / ordinal /
proportion). On the single-column BM path
(R/bm_internal.R) this is the conditional-MVN standard
deviation, delta-method back-transformed — exact under BM,
model-dependent. When ≥ 2 BM-eligible columns use the joint /
threshold-joint path, pred$se is the per-tip conditional SD
\(\sigma\sqrt{v_i}\) from the sparse
Henderson precision (same factorisation as the mean; 2026-08 fix). That
is model-dependent Henderson variance, not the single-trait closed-form
“exact under BM” guarantee, and it does not use the un-restored
cross-trait EM (max_iter = 0).
pred$se (binary / categorical). \(\min(p, 1-p)\) or \(1 - \max_k p_k\) — an uncertainty score,
not a Gaussian standard error. Use it for ranking or
reporting. Do not plug it into Rubin’s rules.
Conformal interval (pred$conformal_lower,
pred$conformal_upper). Split-conformal residual
quantile on the validation set: \(s =
q_{(1-\alpha)}(|y - \hat{y}|_{\mathrm{val}})\). This is a
marginal prediction-coverage diagnostic under exchangeability and a
fixed calibration procedure. It is not an MI draw distribution.
For prediction diagnostics, multi_impute() exposes two
stochastic draw methods. "conformal" (chosen by the default
draws_method = "auto" when the posterior route does not
apply) uses the heuristic \(\mathcal{N}(\mu,
s/1.96)\) on the transformed scale, where \(s\) is the conformal residual quantile. A
conformal prediction interval does not by itself establish that these
Normal draws are proper multiple imputations. "mc_dropout"
runs \(M\) stochastic GNN passes in
training mode (dropout active) on top of BM-draw inputs. These paths are
unsupported for inference.
The separate experimental multi_impute_analysis()
backend declares the substantive model before imputation. It supports
one incomplete continuous covariate under MAR, dispatching to proper
Bayesian normal-regression MI for Gaussian lm,
smcfcs for binomial-logit glm, and
jomo.smc for a Gaussian lmer with one random
intercept. pool_mi() then applies Rubin’s rules to fixed
effects only. The interface remains experimental and deliberately
narrow.
multi_impute_trees() can compare point predictions
across a posterior tree sample. Tree uncertainty was outside the
analysis-aware validation campaign, so its stochastic datasets are
unsupported for downstream inference and cannot currently be combined
with multi_impute_analysis().
A clean end-to-end summary of what predict.pigauto_fit()
returns for one continuous trait on one missing cell \(i\) in a single-tree, single- imputation
call:
\[ \hat{y}_i \;=\; \underbrace{(1 - r_\text{cal}) \cdot \mu_i^{\text{BM/joint-MVN}}}_\text{baseline contribution} + \underbrace{r_\text{cal} \cdot \mathrm{delta}_i^{\text{GNN}}}_\text{GNN contribution} + \underbrace{\mathrm{cov\_linear}(u_i)}_\text{linear cov shortcut}, \]
with \(\mu_i\) from the joint MVN (or per-column BM fallback), \(\mathrm{delta}_i\) from the GNN (transformer blocks + optional within- row cross-trait attention), and \(r_\text{cal}\) from validation calibration. The conformal interval is
\[ \hat{y}_i \pm s_t, \]
where \(s_t\) is the split-conformal residual quantile for trait \(t\). Treat this as a nominal held-out diagnostic under the stated assumptions, not a package-wide interval certification.
| Concept | File |
|---|---|
| BM kernel + conditional MVN | R/bm_internal.R |
| Joint MVN baseline | R/joint_mvn_baseline.R |
| Threshold-joint (binary + ordinal) | R/joint_threshold_baseline.R,
R/liability.R |
| OVR categorical | R/ovr_categorical.R |
| Graph Transformer Block (B2) | R/graph_transformer_block.R |
| ResidualPhyloDAE + B3 trait attention | R/model_residual_dae.R |
| Training loop, gate calibration | R/fit_pigauto.R, R/fit_helpers.R |
| Prediction, conformal intervals | R/predict_pigauto.R |
| Analysis-aware MI and fixed-effect pooling | R/multi_impute_analysis.R,
R/pool_mi.R |
| Prediction-diagnostic draws | R/multi_impute.R, R/predict_pigauto.R |
| Experimental tree-sensitivity workflow | R/multi_impute_trees.R |