Distributional anchor regression

Abstract

Prediction models often fail if train and test data do not stem from the same distribution. Out-of-distribution (OOD) generalization to unseen, perturbed test data is a desirable but difficult-to-achieve property for prediction models and in general requires strong assumptions on the data generating process (DGP). In a causally inspired perspective on OOD generalization, the test data arise from a specific class of interventions on exogenous random variables of the DGP, called anchors. Anchor regression models, introduced by Rothenhäusler et al. (J R Stat Soc Ser B 83(2):215-246, 2021. 10.1111/rssb.12398), protect against distributional shifts in the test data by employing causal regularization. However, so far anchor regression has only been used with a squared-error loss which is inapplicable to common responses such as censored continuous or ordinal data. Here, we propose a distributional version of anchor regression which generalizes the method to potentially censored responses with at least an ordered sample space. To this end, we combine a flexible class of parametric transformation models for distributional regression with an appropriate causal regularizer under a more general notion of residuals. In an exemplary application and several simulation scenarios we demonstrate the extent to which OOD generalization is possible.

Keywords: Anchor regression; Covariate shift; Diluted causality; Distributional regression; Out-of-distribution generalization; Transformation models.

Fig. 1

Graphical models for the response variable $Y$, covariates $\mathit{X}$ and hidden confounders $\mathit{H}$: IV regression with instruments $\mathit{A}$ (left) and anchor regression with anchor $\mathit{A}$ (right). In anchor regression, $\mathit{A}$ is only required to be a source node but is allowed to directly influence response, covariates and hidden confounders

Fig. 2

Illustration of an unconditional transformation model $(1-exp(-exp(·)),{\mathit{b}}_{\text{Bs},6},\mathit{\vartheta })$ for the Old Faithful Geyser data (Azzalini and Bowman 1990) using a Bernstein polynomial basis expansion of order 6 for the transformation function, $h(y)={\mathit{b}}_{\text{Bs},6}{(y)}^{\top }\mathit{\vartheta }$. The colored regions indicate the transport of probability mass from ${\mathbb{P}}_Y$ (lower right) to ${\mathbb{P}}_Z$ (upper left) via the transformation function $h(y)$ (upper right). If $h$ is continuously differentiable, the density of $Y$ is given by $f_Y(y)=f_Z(h(y))h^{\prime }(y)$

Fig. 3

Structural equation model for a transformation model. Instead of setting up the SEM on the scale of $Y$, it is defined on the scale of the inverse transformation function $h^{-1}$. The conditional distribution of $Y:=h^{-1}(Z|\mathit{X},\mathit{H},\mathit{A})$ is still fully determined by $h$ and $F_Z$. The circle around $Y$ emphasizes that its distribution is a deterministic function of its parents

Fig. 4

Leave-one-environment-out cross validation under increasing causal regularization for the BostonHousing2 data, with town as anchors. A linear (Lm), continuous probit (c-probit) and continuous logit (c-logit, using the exact and censored response) model is fitted on 91 towns and used to predict the left out town. a Mean out-of-sample NLL for the left-out census tracts. Beacon Hill, Back Bay and North End are consistently hardest to predict. Consequently, for these towns the cross-validated NLL improves with increasing causal regularization up to a certain point. For the majority of the remaining towns prediction performance decreases. We thus indeed improve worst-case prediction, in analogy to Eq. (2). Note that ${log}_{10}\xi =-\infty$ corresponds to the unpenalized model. b Scaled regression coefficients, which are interpretable as difference in means (Lm), difference in transformed means (c-probit) and log odds-ratios (c-logit) per standard deviation increase in a covariate. Solely the c-logit (censored) model accounts for right-censored observations. With increasing causal regularization the estimates shrink towards zero, indicating that town may be a weak instrument (see Appendix E)

Fig. 5

Density estimates for the three census tracts (Loc 1, Loc 2, Loc 3) in Boston Beacon Hill, the hardest to predict town in terms of LOEO cross-validated NLL for $\xi =10$ (cf. Fig. 4). The dashed gray line indicates the observed outcomes for all three locations, which were all censored at $\mathrm{\$}{50}^{\prime }000$. Lm assumes equal variances and conditional normality, whereas c-probit loosens this assumption leading to more accurate, skewed distributions. Only c-logit (censored) takes into account right censoring in the data and puts a smaller probability density on $\mathrm{\$}{50}^{\prime }000$ than the c-logit (exact) model, which ignores the censoring

Fig. 6

Test performance (thin lines) over 100 simulations for scenario la with $n_{\text{train}}=300$ and $n_{\text{test}}=2000$. Median test performance over all simulations is indicated by the thick line. The $\alpha$-quantiles of test absolute prediction error $\text{APE}:=|y-\widehat{y}|$, where $\widehat{y}$ denotes the conditional median, is shown for linear $L_2$ anchor regression (a) using $\gamma =13$ and the negative log-likelihood contributions for distributional (conditionally Gaussian) linear anchor regression (b) with $\xi =(\gamma -1)/2=6$. The two models are equivalent up to estimating the residual variance via maximum likelihood in the distributional anchor TM. The change in perspective from an $L_2$ to a distributional loss requires different evaluation metrics, of which the log-likelihood, being a proper scoring rule, is the most natural choice

Fig. 7

Test performance over 100 simulations for scenario nla with $n_{\text{train}}=300$ and $n_{\text{test}}=2000$. Mean (a) and $\alpha$-quantiles of the negative log-likelihood contributions (b) for the c-probit anchor TM. The test data are generated under strong push-interventions on the distribution of the anchors (cf. Table 2). The strength of causal regularization was chosen as $\xi =6$

Fig. 8

Test performance over 100 simulations for scenario iv1 with $n_{\text{train}}=1000$ and $n_{\text{test}}=2000$. Quantiles of the individual negative log-likelihood contributions (a) and estimates of $\beta$ (b) for increasingly strong causal regularization. The ground truth is indicated by a dashed line. The test data are generated under the intervention $\mathrm{do}(\mathit{A}=3.6)$

Fig. 9

Test performance and coefficient estimates over 200 simulations for scenario iv2. Because the results are comparable for differing sample sizes and numbers of classes, solely the results for $n_{\text{train}}=1000$ and $K=10$ are displayed. a: Test log-likelihood contributions for varyingly strong instruments (columns) and perturbation sizes (rows). b: Parameter estimates $\widehat{\beta }$ for all intervention scenarios together, since they do not influence estimation. The simulated ground truth $\beta =0.5$ is indicated with a dashed line

Fig. 10

Test performance (quantile functions) over 100 simulations for scenario nla with the same parameters as in Fig. 7, but using a mis-specified linear $L_2$ anchor (a) and normal linear anchor transformation (b) model. The quantile function of the absolute prediction errors (a) and negative log-likelihood (b) contributions are shown with $\gamma =13$ and $\xi =6$, respectively. The point-wise median over all simulations is indicated by a thick line. As a reference, we show the point-wise median NLL of the c-probit model from Fig. 7 in b

Fig. 11

Test performance (quantile functions) over 100 simulations for scenario nla with the same parameters as in Fig. 7, using conditional mean predictions from $L_2$ anchor boosting ($\gamma =13$) and conditional median predictions from a c-probit model ($\xi =6$). The point-wise median over all simulations is indicated by a thick line

Fig. 12

Test performance (quantile functions) over 100 simulations for scenario iv1 with the same parameters as in Fig. 8, but using a mis-specified linear $L_2$ anchor (a) and normal linear anchor transformation (b) model. The quantile function of the absolute prediction errors (a) and negative log-likelihood (b) contributions are shown with $\gamma =15$ and $\xi =7$, respectively. The point-wise median over all simulations is indicated by a thick line. The original point-wise median NLL of the correctly specified c-probit model from Fig. 8 is depicted in black

Fig. 13

Simulation results for scenario iv2 for $K\in \{4,6\}$ repeated 100 times. a The average test NLL is displayed for each simulation, varying $M_X$ and constant $\mathrm{do}(A=3)$. b Coefficient estimates for each model, where the simulated ground truth ($\beta =0.5$) is indicated by a dashed line. For details, see Fig. 9