Evaluating Robustness to Dataset Shift via Parametric Robustness Sets
Abstract
We give a method for proactively identifying small, plausible shifts in distribution which lead to large differences in model performance.
These shifts are defined via parametric changes in the causal mechanisms of observed variables, where constraints on parameters yield a “robustness set” of plausible distributions and a corresponding worst-case loss over the set.
While the loss under an individual parametric shift can be estimated via reweighting techniques such as importance sampling, the resulting worst-case optimization problem is non-convex, and the estimate may suffer from large variance.
For small shifts, however, we can construct a local second-order approximation to the loss under shift and cast the problem of finding a worst-case shift as a particular non-convex quadratic optimization problem, for which efficient algorithms are available.
We demonstrate that this second-order approximation can be estimated directly for shifts in conditional exponential family models, and we bound the approximation error.
We apply our approach to a computer vision task (classifying gender from images), revealing sensitivity to shifts in non-causal attributes.
Introduction
Predictive models may perform poorly outside of the training distribution, a problem broadly known as dataset shift [Quiñonero-Candela et al., 2008].
In high-stakes applications, such as healthcare, it is important to understand the limitations of a model in advance [Finlayson et al., 2021]: given a model trained on data from one hospital, how will it perform under changes in the population of patients, in the incidence of disease or in the treatment policy?
In this paper our goal is to proactively understand the sensitivity of a predictive model to dataset shift, using only data from the training distribution.
This requires domain knowledge to specify what type of distributional changes are plausible.
Formally, for a model f ( X ) trained on data from P ( X, Y ) , with loss function ℓ ( f ( X ) , Y ) , we seek to understand the loss of the model under a set of plausible future distributions P .
We seek to evaluate the worst-case loss over P ,
[Equation]
and provide an interpretable description of a distribution P which maximizes this objective.
If the value of the worst-case loss is low, this can build confidence prior to deployment, and otherwise, examining the worst-case distribution P can help identify weaknesses of the model.
To illustrate, we use the following running example, inspired by Subbaswamy et al. [Table]
[Figure]
L o s s un d e r ( 0 , 1 ) 0.30.40.50.
Here loss fixed over distributions
Simulation details are given in Appendix A.
Example 1 (Changes in laboratory testing) .
We seek to classify disease ( Y ) based on the age ( A ) of a patient, whether a laboratory test has been ordered ( O ) , and test results ( L ) if a test was ordered.
The performance of a predictive model may be sensitive to changes in testing policies, as the fact that a test has been ordered itself is predictive of disease.
Figure 1 (left) gives a plausible causal relationship between variables.
Let P ( O = 1 | A, Y ) =
σ ( η ( A, Y )) , where σ is the sigmoid function and η ( A, Y ) is the log-odds.
In Figure 1 (right), we show the loss under a set of new distributions parameterized by δ = ( δ 0 , δ 1 ) , where we modify P δ
( O = 1 | A, Y ) =
σ ( η ( A, Y )
+ s ( Y ; δ )) for a shift function s ( Y ; δ ) =
δ 1 · Y + δ 0 · (1 − Y ) , which modifies the log-odds of testing for both sick and healthy patients.
If δ 0 , δ 1 are unconstrained, the worst-case occurs when all healthy patients are tested, and no sick patients are tested.
The first challenge is to define a set of possible distributions P such that each distribution P ∈ P satisfies two desiderata: First, they should be causally interpretable and simple to specify , without placing unnecessary restrictions on the data-generating process.
Second, they should be realistic , which often entails bounding the magnitude of the shift.
We construct causally interpretable shifts by defining perturbed distributions P δ using changes in causal mechanisms, parameterized by a finite-dimensional parameter δ .
Our main requirement is that the shifting mechanisms follow a conditional exponential family distribution.
For discrete variables, this places no restriction on P :
In
Example 1, O is binary and the log-odds η ( A, Y ) can be any function of A, Y .
We also demonstrate that constraining δ can ensure that shifts are realistic:
The unconstrained worst-case shift in Exampleis implausible, where all healthy patients (and no sick patients) are tested.
Equation (1) becomes
[Equation]
where E δ is the expectation in the shifted distribution P δ and ∆ is a bounded set of shifts.
The second challenge is evaluation of the expected loss under shift, as well as finding the worst-case shift.
Under our definition of shifts, we show that the test distribution can always be seen as a reweighting of the training distribution, allowing for reweighting approaches, such as importance sampling, to estimate the expected loss under shifts.
While this is practical for some distribution shifts, for others, importance sampling can lead to extreme variance in estimation.
Further, finding the worst-case shift using a reweighted objective involves maximization over a non-concave objective (see Figure 1), a problem that is generally NP-hard.
We derive a second-order approximation to the expected loss under shift, and show how it can be estimated without the use of reweighting.
When ∆ is a single quadratic constraint, we can approximate the general non-convex optimization problem in Equation (2) with a particular non-convex, quadratically constrained quadratic program (QCQP) for which efficient solvers exist [Conn et al., 2000, Section 7].
We bound the approximation error of this surrogate objective, and show in experiments that it tends to find impactful adversarial shifts.
Our contributions are as follows:
We provide a novel formulation of robustness sets which are defined using parametric shifts.
This formulation only require that the shifting mechanisms (i.e., conditional distributions) can be modelled as a conditional exponential family (see Section 2).
We derive a second-order approximation to the expected loss and provide a bound on the approximation error.
We show that this translates the general non-convex problem into a particular non-convex quadratic program, for which efficient solvers exist (see Section 3).
In a computer vision task, we find that this approach finds more impactful shifts than a reweighting approach, while taking far less time to compute, and that the resulting estimates of accuracy are substantially more reliable (see Section 4).
Distributionally robust optimization/evaluation : Distributionally robust optimization (DRO) seeks to learn models that minimize objectives like Equation (1) with respect to the model [Duchi and
Namkoong, 2021, Duchi et al., 2020, Sagawa et al., 2020]. We focus on proactive worst-case evaluation of a fixed model, not optimization, similar to Subbaswamy et al. [2021], Li et al.
but we also differ in our definition of the set of plausible future distributions P , often called an “uncertainty set” in the optimization literature. Prior work often defines these sets using distributional distances (such as f -divergences): For instance, Joint DRO [Duchi and Namkoong, 2021] allows for shifts in the entire joint distribution (i.e., all distributions in an f -divergence ball around P ( X, Y ) ), which may be overly conservative. Marginal DRO [Duchi et al., 2020] considers shifts in a marginal distribution (e.g., P ( X ) ), while assuming that the remaining conditionals (e.g., P ( Y | X ) ) are fixed.
However, this assumption is not applicable in all scenarios: In Example 1, for instance, this assumption does not hold for a shift in testing policy.
Conditional shifts are considered in recent work that focuses on evaluation [Subbaswamy et al., 2021], using worst-case conditional subpopulations. However, choosing a plausible size of conditional subpopulation is often non-obvious.
In Appendix D we give a simple lab-testing example where taking worst-case 20% conditional subpopulations corresponds to an implausible shift: Healthy patients are always tested, and sick patients never tested. In contrast, our approach uses explicit parametric perturbations to define shifts, as opposed to distributional distances or subpopulations. In addition, our approach allows for shifts in multiple marginal or conditional distributions simultaneously: In Example 1, for instance, we could model a simultaneous change in both the marginal distribution of age P ( A ) , as well as the conditional distribution of lab testing P ( O | A, Y ) , leaving other factors unchanged.
Causality-motivated methods for learning robust models : Several approaches proactively specify shifting causal mechanisms/conditional distributions, and then seek to learn predictors that have good performance under arbitrarily large changes in these mechanisms [Subbaswamy et al., 2019, Veitch et al., 2021, Makar et al., 2022, Puli et al., 2022]. Other approaches use environments [Magliacane et al., 2018, Rojas-Carulla et al., 2018, Arjovsky et al., 2019] or identity indicators [Heinze-Deml and Meinshausen, 2021] to learn models that rely on invariant conditional distributions. However, when shifts are not arbitrarily strong, causality-motivated predictors can be overly conservative.
In Example 1, a model that ignores all test-related features (and only uses age as a predictor) is a particularly simple example of a causality-motivated predictor, with invariant risk over changes in testing policy.
Closer to our setting is a line of work that considers bounded mechanism changes in linear causal models [Rothenhäusler et al., 2021, Oberst et al., 2021], where estimation of the worst-case loss enables learning of worst-case optimal models.
Our work can be seen as extending this idea to more general non-linear causal models, where we focus on evaluation rather than optimization.
Evaluating out-of-distribution performance with unlabelled samples : A recent line of work has focused on predicting model performance in out-of-distribution settings, where unlabelled data is available from the target distribution [Garg et al., 2022, Jiang et al., 2022, Chen et al., 2021]. In contrast, our method operates using only samples from the original source distribution, and seeks to estimate the worst-case loss over a set of possible target distributions. In Appendix F we give a more detailed discussion of these approaches and others.
Defining parametric robustness sets
Notation : Let V denote all observed variables, where ( X, Y ) ⊆
V for features X and labels Y , and use P ( V ) to denote the probability density/mass function in the training distribution.
We also refer to P as simply “the training distribution”.
E [ · ] and cov( · , · ) refer to the mean and covariance in P , and for a shifted distribution
P δ
(Definition 1) we use E δ [ · ] , cov δ ( · , · ) .
For a random variable Z , we use
Z to denote the space of realizations, and d Z for dimension e.g., Z ∈ Z ⊆
R d Z .
For a set of random variables V = { V 1 , . .
. , V d } , we use V i to denote an individual element, and use PA G ( V i ) to denote the set of parents in a directed acyclic graph (DAG) G , omitting the subscript when otherwise clear.
We begin with a general definition of a parameterized robustness set of distributions P .
Definition 1.
A parameterized robustness set around P ( V ) is a family of distributions P with elements P δ ( V ) indexed by δ ∈ ∆ ⊆
R d δ , with 0 ∈ ∆ , where P 0 ( V ) =
P ( V ) .
We give examples shortly that satisfy this general definition.
To construct such a robustness set, we consider distributions P δ that differ from P in one or more conditional distributions (Assumption 1).
We require that the relevant conditional distributions can be described by an exponential family.
Definition 2 (Conditional exponential family (CEF) distribution) .
P ( W | Z ) is a conditional exponential family distribution if there exists a function η ( Z ) :
R d Z → R d T such that the conditional probability density (for continuous W ) or probability mass function (for discrete W ) is given by
[Equation]
where T ( W ) is a vector of sufficient statistics,
T ( W ) ∈ R d T , g ( · ) specifies the density of a base measure and h ( η ( Z )) is the log-partition function.
Definition 2 does not restrict P ( W | Z ) for binary/categorical W , and captures a wide range of distributions, including the conditional Gaussian (see Appendix B.1 for other examples).
Definitionextends to marginal distributions where Z = ∅ and η ( Z ) is a constant function.
Example 1 (Continued) .
Suppose the probability of ordering a test ( O ) depends on age ( A ) and disease ( Y ) , such that P ( O = 1 |
A, Y )
=
σ ( η ( A, Y )) , where σ is the sigmoid, and η is an arbitrary function.
Here, Definition 2 is satisfied with W
= O , Z =
( A, Y ) , and sufficient statistic T ( O )
= O .
We now state our main assumption, where we distinguish between the terms in the joint distribution of P that shift, which we will need to model, and those that remain fixed, which we do not.
Assumption 1 (Factorization into CEF distributions) .
Let W =
{ W 1 , . . .
, W m } ⊆
V be a “intervention set” of variables and let
[Equation]
be a factorization, where Z i , U j , V j ⊆ V are possibly overlapping (or empty) sets of variables, where P ( V j | ∅ ) :
= P ( V j ) .
For each W i we assume Z i is known and P ( W i
| Z i ) satisfies Definition 2.
If P ( V ) factorizes according to a DAG G , the factorization in Assumption 1 is always satisfied by
Z i = PA G ( W i ) .
While we assume data is generated according to Equation (4), we do not require knowledge of the full distribution, but only the conditionals that shift.
In Appendix B.2 we show that we can also consider shifts that extend Z i to include additional variables, subject to an acyclicity constraint.
We now define parametric perturbations and give the general form of the robustness sets that we consider in this work, involving simultaneous perturbations to multiple W i .
Definition 3 (Parameterized shift functions and δ -perturbations) .
Let s ( Z ; δ ) : R d Z → R d T be a parameterized shift function with parameters
δ ∈ ∆ ⊆
R d δ which is twice-differentiable with respect to δ and which satisfies s ( Z ; 0) = 0 for all Z .
For P ( W | Z ) satisfying Equation (3), we refer to P [Equation] |
[Equation]
[Equation]
g [Equation] [Equation] ) ⊤
[Equation]
[Equation]
) − h ( [Equation] [Equation] )) (cid:1) as a δ -perturbation of P ( W | Z ) with shift function s ( Z ; δ ) , where η δ ( Z ) :
= η ( Z ) + s ( Z ; δ ) .
Note that this differs from Equation (3) in that η ( Z ) is replaced by η δ ( Z ) .
Example 1 (Continued) .
A model developer may be concerned about a uniform change in testing rates across all types of patients.
This can be modelled by choosing s ( Z ; δ )
=
δ , for δ ∈ R , an additive intervention on the log-odds scale.
A separate change in testing rates for sick and healthy patients could instead be modeled using s ( Z ; δ ) =
δ 0 (1 − Y )
+
δ 1 Y , using δ ∈ R 2 .
This reasoning extends readily to more complex shifts (e.g., allowing for age-specific changes in testing rates, with a non-linear dependence on age), as long as s ( Z ; δ ) remains a parametric function.
While the shift function s ( Z ; δ ) is parametric, η ( Z ) is unconstrained in Definitions 2 and 3.
Note that this formulation includes multiplicative shifts η δ ( Z ) =
(1 + δ ) η ( Z ) by letting s ( Z ; δ ) =
δ · η ( Z ) .
Definition 4 (CEF parameterized robustness set) .
For a distribution P and intervention set W = { W 1 , . . .
, W m } ⊆
V satisfying Assumption 1, let each P
δ i ( W i | Z i ) be a δ
i -perturbation
(Definition 3) of P ( W i | Z i ) .
Then
[Equation]  is called a δ -perturbation of P ( V ) , and the robustness set P consists of all P δ for δ ∈ ∆ 1 × · · · ∆ m .
To estimate the expected loss under P δ , we will typically 2 need to estimate η ( Z i ) for each W i ∈ W .
However, we make no distributional assumptions on the remaining variables V \ W .
This is useful in applications such as computer vision, where we do not need to restrict the generative model of images given attributes (e.g., background, camera type, etc), but can still model the expected loss under changes in the joint distribution of those attributes.
Remark 1 (Causal Interpretation of Shifts) .
If available, causal knowledge helps identify which factors in the joint distribution are subject to shifts (e.g., P ( O | Y, A ) in Example 1), and which remain stable. It is worth noting, however, that our methodology can be used to model any change in distribution that satisfies Assumption 1, including choices of “non-causal” factorizations and shifting factors. For example, in the context of Example 1, we could choose the factorization P ( Y ) P ( O | Y ) P ( L, A | O, Y ) , and model a change only in the conditional P ( O | Y ) while keeping other factors unchanged.
This shift is not interpretable as a change in causal mechanisms: The shifted distribution would imply a change in the marginal distribution of age, which should be unaffected by a real-world change in laboratory testing. Nonetheless, we can still estimate a worst-case loss over such non-causal shifts in distribution. In short, our machinery can model shifts in non-causal conditionals (for example because the causal structure is unknown), though the resulting shifted distribution is not interpretable as a plausible shift in the ground-truth data generating mechanism.
Evaluation of the worst-case loss
For a fixed predictor and loss function, we can use data from P ( V ) to estimate the expected loss E δ [ ℓ ] :=
E δ [ ℓ ( f ( X ) , Y )] for a fixed δ , and estimate the worst-case loss over all δ of bounded magnitude.
In Section 3.1, we show that P δ shares support with P , suggesting the use of reweighting estimators.
However, these estimators can exhibit high variance for shifts that produce large density ratios (see Appendix C.5 for an example), and maximizing a reweighted objective over δ is generally a non-convex problem.
In Section 3.2 we derive an approximation to the expected loss under P δ , yielding a tractable surrogate optimization problem under quadratic constraints such as ∥ δ ∥ 2 ≤ λ .
Remark 2 .
The methods here can be used with an arbitrary predictor f and loss function ℓ := ℓ ( f ( X ) , Y ) .
We do not even require access to the original predictor f .
Both methods here simply treat ℓ as a random variable in P , for which we have samples from the training distribution.
3.1 Modelling shifted losses using reweighting
The shifts defined in Section 2 share common support, with the following density ratio.
Proposition 1.
For any P δ ( V ) , P ( V ) that satisfy Definition 4, supp( P ) =
supp( P δ ) and the density ratio w δ :=
P δ / P is given by
[Equation]
The proof can be found in Appendix G, along with all proofs for all other claims.
Example 1 (Continued) .
Suppose we perturb the probability of ordering a test O given age A and disease Y with shift function s ( Y ; δ ) =
δ 0 (1 − Y )
+ δ 1 Y , independently changing the conditional probability of testing for healthy and sick patients.
Here, the density ratio is given by
[Equation]
To model the loss E δ [ ℓ ] using data from P , we can consider an importance sampling (IS) estimator [Horvitz and Thompson, 1952, Shimodaira, 2000], observing that E δ [ ℓ ] =
E [ w δ ( V ) · ℓ ] .
This requires estimation of the density ratio w δ ( V ) , and (given a sample { V j } nj =1 from P ) yields the estimator
[Equation]
Equation (6) can have high variance when density ratios are large, and maximizing this equation with respect to δ is a general non-convex optimization problem, which is generally NP-hard to solve.
3.2 Approximating the shifted loss for exponential family models
We now propose an alternative approach for approximating the loss E δ [ ℓ ] .
Recalling that P δ =
= P , we use a second-order Taylor expansion around the training distribution
[Equation]
where E [ ℓ ] denotes the loss in the training distribution and SG 1 , SG 2 are defined as follows.
Definition 5 (Shift gradient and Hessian) .
For a parametric shift satisfying Definition 1 where δ (cid:55)→ E δ [ ℓ ] is twice-differentiable, we denote the shift gradient SG 1 and shift Hessian SG 2 as
[-SG 1 := ∇ δ E δ [ ℓ ] (cid:12)(cid:12) δ =
-][Equation]
and [Equation]
Equation (7) is a local approximation of the loss, whose approximation error we bound in Theorem 2, with smaller approximation error for smaller shifts.
3 For P δ satisfying Definition 4, SG 1 and SG 2 can be computed as expectations in the training distribution, without estimation of density ratios.
Recall that the conditional covariance is given by cov
( A, B | C ) :
=
E [( A − E
[ A | C ])( B − E
[ B | C ]) | C ] .
Theorem 1 (Shift gradients and Hessians as covariances) .
Assume that P δ , P satisfy Definition 4, with intervened variables W = { W 1 , . . .
, W m } and shift functions s i ( Z i ; δ i ) , where δ =
( δ 1 , . .
δ m ) .
Then the shift gradient is given by SG
= (SG 11 , . .
, SG 1 m ) ∈ R
d δ where
[Equation]
and the shift Hessian is a matrix of size ( d δ × d δ ) , where the ( i, j ) th block of size d δ i ×
d δ j equals
[Equation]
where D i,k := ∇ kδ
i s i
( Z i ; δ i )
| δ =0 , is the gradient of the shift function for k = 1 , and the Hessian for k = 2 .
Here, T i ( W i ) is the sufficient statistic of P ( W i | Z i ) and ϵ T i |
Z i :
=
T i ( W i ) − E
[ T ( W i ) | Z i ] .
Theorem 1 handles arbitrary parametric shift functions in multiple variables, but for simple shift functions in a single variable, the notation simplifies substantially, as we show in Corollary 1.
Corollary 1 (Simple shift in a single variable) .
Assume the setup of Theorem 1, restricted to a shift in a single variable W , and that s ( Z ; δ )
= δ .
Then D 1 = 1 , D 2 = 0 , and
SG 1 = E (cid:20) cov (cid:18) ℓ, T
[Equation]
W
[Equation]
and
[Equation]
where T ( W ) is the sufficient statistic of W and ϵ T | Z :
= T ( W ) − E
[ T ( W ) | Z ] .
Example 1 (Continued) .
Suppose that age ( A ) follows a normal distribution with mean µ and variance σ 2 , and consider a shift in the mean (without changing lab testing).
We can parameterize P ( A ) as an exponential family with parameter η = µ/σ and sufficient statistic T ( A ) =
A/σ .
Here, s ( δ ) =
δ implies a shift in the mean of δ standard deviations η δ = η + s ( δ ) =
( µ + σδ ) /σ , and we can write that SG 1 = cov ( ℓ, A ) /σ
and SG 2 = cov (cid:0) ℓ, ( A − E
[ A ]) 2 (cid:1) /σ 2 .
To estimate the shift gradient and Hessian from a sample from P , for each i = 1 , . .
, m we fit models ˆ µ ℓ ( Z i )
≈ E [ ℓ | Z i ] and ˆ µ W
i
( Z i )
≈ E
[ T i ( W i ) | Z i ] and compute residuals on these predictions, which permits estimation of the gradient/Hessian as a sample average of residuals.
A detailed treatment is given in Appendix C.1.
Using estimates of the gradient and Hessian, we estimate the expected loss as
[Equation]
Here, there are two sources of error: Finite-sample error, due to the estimates of SG 1 , SG 2 , as well as approximation error.
The latter is bounded by the norm of δ and a term that depends on the covariance between the loss and the deviations of the sufficient statistic from its shifted mean.
Theorem 2.
Assume that P δ , P satisfy the conditions of Theorem 1, with a shift in a single variable
W , where s ( Z ; δ ) =
δ .
Let E δ, Taylor be the population Taylor estimate (Equation (7) ) and let σ ( M ) denote the largest absolute value of the eigenvalues of a matrix M .
Then
[Equation]
where T ( W ) is the sufficient statistic of W | Z and ϵ t · δ,T | Z = T ( W | Z ) − E t · δ [ T ( W | Z )] .
To build intuition, in Appendix C.8 we give a scenario where this bound can be simplified. In particular, we consider a “covariate shift” setting [Quiñonero-Candela et al., 2008] where X is standard Gaussian, Y = f 0 ( X ) + ϵ with a noise term independent of X and we consider a shift δ in the mean of X . When evaluating a predictor f ( X ) with the loss ℓ being the squared error, the bound in Theorem 2 depends on how the modelling error g ( X ) = f
0 ( X ) − f ( X ) behaves over the domain.
In particular, the bound scales as the supremum (over t ∈ [0 , 1] ) of (cid:112) var( g ( X + t · δ ) 2 − g
( X ) 2 ) .
As a simple corollary, if our predictor is off by an additive constant factor, f = f 0 + C , then the bound is zero, and the approximation is exact for any δ .
On the other hand, if the squared modelling error g ( X ) 2 at one point X tends to be a poor predictor of the squared modelling error at another point X + t · δ , then this variance will be large, and the approximation will be loose.
In exchange for considering a second-order approximation of the loss, we gain two benefits: Variance reduction and tractable optimization.
First, the variance of ˆ E δ, Taylor is O ( ∥ δ ∥ 4 ) for large ∥ δ ∥ , while the variance of ˆ E δ, IS can be much larger: We give a simple case in Appendix C.6 where var( ˆ E δ, Taylor ) =
O ( δ 4 ) while var( ˆ E δ, IS )
=
O ( δ 2 exp( δ 2 )) .
Second, maximizing ˆ E δ, Taylor over the set ∥ δ ∥ ≤ λ can be solved in polynomial time by exploiting the quadratic structure, while maximizing ˆ E δ, IS over the constraints is generally hard, and may be infeasible in high dimensions.
3.3 Identifying worst-case parametric shifts
For λ > 0 , we can locally approximate the worst-case loss over all distributions P δ where ∥ δ ∥ 2 ≤ λ by finding the worst-case loss in the Taylor approximation [Equation]
[Equation]
Since SG 2 is generally not negative definite, the maximization objective is non-concave.
However, this particular problem is an instance of the ‘trust region problem’ 4 which is well-studied in the optimization literature [Conn et al., 2000], and can be solved in polynomial time by specialized algorithms (see Pólik and Terlaky [2007, Section 8.1] for an example).
This follows from the fact that strong duality holds, so that the optimal solution δ ∗ can be characterized in terms of the
Karush-Kuhn-Tucker conditions [Boyd and Vandenberghe, 2004, Section 5.2].
For this problem, we use the trsapp routine from NEWUOA [Powell, 2006], as implemented in the python package trustregion .
Depending on the application and prior knowledge, one may choose constraint sets that differ from ∥ δ ∥ ≤ λ .
In particular, the strong duality of Equation (9) also holds when ∥ δ ∥ 2 ≤ λ is replaced by any single quadratic constraint δ ⊤ Aδ + δ ⊤ b ≤ λ , allowing for e.g., larger shifts in some directions than in others. [Figure]
Experiments 4.1 Illustrative example: Laboratory testing
Y
O
L
Figure
To build intuition, we illustrate our method in a simple generative model, similar to Example 1, where lab tests are more likely to be ordered ( O ) for sick patients ( Y ) , and lab values ( L ) are predictive of Y .
[Equation]
where µ 1 = 0 .
5 , µ 0 = − 0 .
5 , and we initialize with α = − 1 , β
= 2 , so that P ( O = 1 | Y = 0) ≈ 0 .
27 and P ( O = 1 | Y = 1)
≈ 0 .
73 , and the marginal probability of test ordering is P ( O = 1) = 0 .
When O = 0 , we set L to a dummy value of L = 0 .
The underlying causal graph is given in Figure 2.
The predictive model f
( O, L ) is trained on data from P to predict Y using all available features.
If lab tests are not available ( O = 0 ), this model predicts Y based on the observed likelihood of Y given O = 0 , and otherwise uses a logistic regression model trained on cases where O = 1 in the training data.
Defining a shift function: P ( O | Y ) is a conditional exponential family with η ( Y ) =
α + βY .
We consider the shift function s ( Y ; δ ) =
δ
+ δ 1 Y , where δ 0 models an overall change in testing rate, and δ 1 models an additional change in the likelihood of testing sick ( Y = 1) patients.
Estimating the impact of shift using quadratic approximation: To start, we keep δ 1 = 0 fixed and vary only δ 0 , which uniformly increases or decreases testing.
In Figure 3, we show the ground-truth cross-entropy loss of f ( O, L ) under perturbed distributions P δ 0 .
We observe that the direction of the shift matters: In Figure 3, the model performance slightly increases under a small increase in testing rates, but degrades if testing increases too much; moreover, the loss under shift is generally asymmetric, as a decrease hurts more than an increase in testing.
In Figure 3 (left), we demonstrate the use of the quadratic approximation described in Section 3.2.
For illustration, we consider a robustness set of δ 0 ∈ [ − 2 , 2] , and see that the predicted worst-case shift coincides with the actual worst-case shift, and that the quadratic approximation is accurate for smaller values of δ .
In Appendix D, we allow both δ 0 and δ 1 to vary, and compare our approach to that of worst-case (1 − α ) conditional subpopulation shifts [Subbaswamy et al., 2021].
In the context of this example, we demonstrate that for any 1 − α < 0 .
27 , the worst-case conditional subpopulation loss is achieved by having all healthy patients get tested, and no sick patients get tested.
We contrast this with an iterative approach to designing constraints that is made possible by considering parametric shifts, where end-users can restrict the degree to which the shift differs across sick and healthy populations.
4.2 Detecting sensitivity to non-causal correlations
A predictive model may pick up on various problematic dependencies in the data that may not remain stable under dataset shift.
To understand the impact of these dependencies, a model user may wish to understand which changes in distribution pose the greatest threats to model performance, and to measure the impact of these changes.
To illustrate this use-case, we make use of the CelebA dataset [Figure] [Liu et al., 2015], which contains images of faces and binary attributes (e.g., glasses, beard, etc.) encoding several features whose correlations may be unstable (e.g., the relation between gender and being bald).
We consider the task of predicting gender ( Y ) from images of faces ( X ), and assess sensitivity to a shift in the distributions of attributes ( W ).
Setup : To obtain ground-truth shifts in distribution, we generate synthetic datasets of faces using
CausalGAN [Kocaoglu et al., 2018], trained on the CelebA data.
We simulate attributes following the causal graph in Figure 4, and then simulate images from the GAN conditioned on those attributes.
We draw a training sample from this distribution P , and fit a gender classifier f
( X ) using the image data alone, by finetuning a pretrained ResNet50 classifier [Hu et al., 2018].
Each attribute W i is binary, so we consider shifts in the log-odds η i ( Z i ) of each attribute W i given parents Z i .
Here, we use a maximally flexible shift function s i ( Z i ; δ i ) =
(cid:80) z ∈Z i
δ i,z 1 { Z i =
z } , such that for Z i ∈ { 0 , 1 } k there are 2 k parameters.
Across all intervened variables, δ ∈ R 31 .
Due to the synthetic nature of our setup, we can simulate from P δ
( X, W , Y ) to evaluate the ground-truth impact of this shift, simulating first from the shifted attribute distribution, and then simulating images from the GAN conditional on those attributes.
We use the 0/1 loss ℓ = 1 { f ( X ) ̸ = Y } , and constrain δ by ∥ δ ∥ 2 ≤ λ = 2 .
Comparing importance sampling and Taylor across multiple simulations : We simulate K =validation sets from P , in each estimating the worst-case shifts δ Taylor (via the approach in Section 3.3) and δ IS , where the latter corresponds to minimizing ˆ E δ, IS using a standard non-convex solver from the scipy library [Virtanen et al., 2020].
We simulate ground truth data from P δ IS and P δ Taylor , to compare the two shifts.
First, we demonstrate that the Taylor approach finds more impactful shifts, when searching over the space of small, bounded shifts considered here.
In Table 1 (right), we compare the average drop in accuracy using the Taylor shifts (3.8%) and the IS shifts (2.2%).
In Figure(right) we plot the differences in test accuracy E δ Taylor [ 1 { f ( X )
=
Y } ] − E δ IS [ 1 { f ( X ) = Y } ] , where the Taylor approach finds a more impactful shift in 96% of cases.
Second, the Taylor approach has an average run-time of 0 .
01 s , versus 2 .
14 s for the IS approach.
Third, when only used to evaluate the shift δ Taylor , the IS estimator is comparable to the Taylor estimator, with a near-identical average bias (shown in Table 1 (right)) and RMSE (0.0191 and 0.0192 respectively).
Finally, however, in Table(right) we observe that ˆ E δ IS , IS is strongly biased in predicting E δ IS , yielding a mean absolute prediction error (MAPE) of 0 .
069 (not shown in the table).
This can be contrasted with a MAPE of 0 .when using ˆ E δ Taylor , Taylor to predict E δ Taylor .
This may suggest that optimizing the IS objective is prone to “overfitting”, choosing a sub-optimal δ from a region of the search space that has high variance.
Here, where λ = 2 , the drop in accuracy is relatively mild for the shifts found by both approaches.
In Appendix E.4 we show that larger values of λ correspond to more substantial drops in accuracy (e.g., an average drop of 23% for λ
= 8 using the Taylor approach).
Examining a single shift : To illustrate the type of shift found by our approach, we consider the δ Taylor (over the K runs) which yields the P δ with median test accuracy.
We display the largest components of that δ in Table 1 (left).
Among others, this shift entails a 5% increase in the probability of an older woman being bald, and a 5% decrease in the probability of a young woman wearing lipstick.
This suggests that the learned classifier f relies on these associations in the images for prediction.
We validate that this shift leads to a decrease in accuracy of around 3 . 8% , using simulated data from P δ .
To validate that this drop in accuracy is a non-trivial occurrence, we simulate K = 400 random shifts [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
δ k where ∥ δ k ∥ =
λ and evaluate the model accuracy in P δ k (Figure 5, left).
As expected, the chosen δ yields a lower accuracy (red line) than all of the random shifts.
Conclusion
We argue for considering parametric shifts in distribution, to evaluate model performance under a set of changes that are interpretable and controllable.
For parametric shifts in conditional exponential family distributions, we derive a local second-order approximation to the loss under shift.
This approximation enables the use of efficient optimization algorithms (to find the worst-case shift), and empirically provides realistic estimates of the resulting loss.
In a computer vision task, this approach finds more impactful shifts (in far less time) than optimizing a reweighted objective, and the estimates of shifted accuracy under the chosen shift are substantially more reliable.
Of course, our method is not without limitations. Our definition of parametric shifts and resulting approximation relies on the relevant mechanisms P
( W | Z ) being a conditional exponential family, and that the relevant variables are observed.
As illustrated in our experiments, this can be used to model changes in the causal relationships between attributes of an image, but does not immediately extend to modelling changes in the distribution of images given a fixed set of attributes.
As with any method that provides worst-case evaluation, there is potential for misuse and false confidence: If the specified shifts fail to capture important real-world changes, the resulting worst-case loss may be overly optimistic and misleading.
Even if used correctly, our approach examines a narrow measure of model performance, and a small worst-case error should not be used to claim that a model is free of problematic behavior.
For example, implicit dependence on certain attributes (e.g., race in medical imaging [Banerjee et al., 2021]) may be problematic based on ethical grounds, even if it does not lead to major issues with predictive performance under small shifts in distribution.
Acknowledgements
We thank Jonas Peters, Chandler Squires, and Stefan Hegselmann for helpful feedback and discussion, and Irene Chen and Christina X Ji for providing comments on an earlier draft.
MO and DS were supported in part by Office of Naval Research Award No. N00014-21-1-2807.
NT was supported by a
Appendix
This appendix is structured as follows:
• In Appendix A, we provide details on the synthetic lab testing example, including how we generate the loss landscape in Figure 1 (right).
• In Appendix B, we provide a “user’s guide” to defining and interpreting parametric shifts, including worked examples for many common conditional distributions, as well as guidance on how to define and interpret the shift functions s ( Z ; δ ) .
• In Appendix C, we provide additional details on the worst-case optimization problem, as well as comparisons of the reweighting-based approach to the Taylor approximation approach.
We also demonstrate that the quadratic approximation is exact, for particularly simple structural causal models.
• In Appendix D, we compare our approach to that of worst-case conditional subpopulation shifts, in the context of a simpler laboratory testing example where we can explicitly compute the worst-case conditional subpopulations.
Here, we demonstrate that our approach can capture more realistic intuition regarding which shifts are plausible in practice.
• In Appendix E, we give additional experimental details, as well as illustrative samples from the generative model, for the CelebA experiment described in Section 4.
• In Appendix F, we give an extended discussion of related work.
• In Appendix G, we give proofs for all the results in the main paper.
A Details of Figure
In Figure 1 (right), we consider the following, artificial, generative model, which resembles the setup in Section 4.1, but with the addition of age as a continuous variable.
[Equation]
where if Order = 0 , the test result is a placeholder value of zero.
In Figure 1 (right), we consider a simple predictive model: If lab tests are not available (Order = 0 ), this model predicts disease based on an unregularized logistic regression model, which uses age to predict disease.
If a lab test is available, then it uses both age and the lab test for prediction.
This model is trained on 100 ,samples from the training distribution.
To construct the loss landscape shown in Figure 1 (right), we first observe that
[Equation]
where
[Equation]
We construct shifts using the shift function s ( Disease , Age ; δ ) =
δ 0 · (1 − Disease )
+ δ 1 · Disease, and for a grid of values for ( δ 0 , δ 1 ) ∈ [ − 5 , 5] {+
 we+} consider perturbed distributions with a different conditional distribution of testing,
[Equation]
but where all other parts of the generative model are fixed.
For each value of ( δ 0 , δ 1 ) ∈ [ − 5 , 5] 2 , we draw 10 , 000 samples from the corresponding distribution, and compute the negative log-likelihood of the original predictive model under this new distribution.
The resulting surface is plotted in Figure(right). [Table]
B A user’s guide to defining parametric shifts
In this section
[Table]
we discuss practical considerations in designing parametric shift functions for different distributions.
• In Appendix B.1, we give examples of conditional exponential families, illustrative shift functions, and how to interpret them. • In Appendix B.
[Table]
we formalize the idea that one can choose shift functions which depend on additional variables, other than the causal parents of a variable W i .
• In Appendix B.3 we give guidance on how to define shift functions when the parameters η ( Z ) are constrained to lie in a particular domain, which is relevant for considering shifts such as changing the variance of a conditional Gaussian.
B.1 Conditional exponential family models and interpretations of shifts
In this section
[Table]
we give examples of exponential families and their sufficient statistics, and discuss design considerations in specifying the shift function s ( Z ; δ ) .
Here, we restrict attention to shifts in a single variable, for ease of notation.
In Table 2 we give examples of conditional exponential families, along with their typical parameterizations.
In the examples below, we review how shift functions s ( Z ; δ ) impact these parameters, and how they can also be interpreted on the scale of more commonly considered parameters (e.g., conditional means and variances).
Example B.1 (Log-odds shift in a binary variable) .
Consider the distribution of a binary variable W conditioned on variables Z .
Without loss of generality, we can write that
[Equation]
where σ is the sigmoid function, and η ( Z ) is an arbitrary measurable function of Z , taking on values in the extended real line η ( Z ) ∈ R ∪ {−∞ , + ∞} .
This can be written in canonical form as
[Equation]
where η ( Z ) is the canonical parameter (the log-odds ratio), T ( W )
=
W is the sufficient statistic, and h ( θ )
= log(1 + exp η ( Z ) ) is the normalizing constant.
We can consider shifts η δ ( Z ) :
= η ( Z ) + δ , yielding the new conditional distribution
[Equation]
which is well-defined for any δ ∈ R .
Here, we note that these shifts occur on the “natural” parameter scale η ( Z ) (e.g., the log-odds), which at first glance may seem difficult to interpret: Why should we care about changes on the log-odds scale, instead of on the original probability scale?
In addition to mathematical convenience, we argue that in some settings, working with natural parameters is advantageous for retaining a common scale across across multiple variables.
For instance, consider shifts in the two independent variables W 1 and W 2 , where V i ∼ Bernoulli ( p i ) , with p 1 = 10 − 4 and p 2 = 0 .
Suppose we wished to consider an additive shift on the probability scale, e.g., p ′ [-
-]1 = p 1 + 0 .
1 , p ′
= p 2 + 0 . 1 .
Setting aside the inconvenience that we need to ensure p ′ 1 , p ′ 2 ∈ [0 , 1] , we argue that these shifts are not truly of a comparable scale.
In particular, this shift in p 1 may seem implausible in magnitude, while the same shift in p 2 seems more reasonable.
On the other hand, an additive shift in the log-odds captures some aspect of this idea.
Of course, there is some flexibility to incorporate prior expectations of shifts in absolute probabilities.
For instance, in binary variable with no causal parents, we can always construct a one-to-one map of δ to a change in the marginal probability.
For conditional shifts, we can similarly construct a one-to-one map between the value of δ in a shift s ( Z ; δ ) =
δ and the resulting marginal probability of W i , as formalized below.
Proposition B.1.
Consider a binary random variable W with conditional distribution
[Equation]
for an arbitrary measurable function η ( Z ) whose range is the extended real numbers η ( Z ) ∈ R ∪{ + ∞ , −∞} .
Let p + :=
P ( η ( Z ) = + ∞ ) , p − := P ( η ( Z ) = −∞ ) , and assume that p + + p − < 1 .
Then, the marginal probability
[Equation]
is a strictly monotonically increasing function of δ ∈ R whose range is ( p + , 1 − p − ) ,
Proposition B.1 states that, for any achievable marginal probability p δ
=
P δ ( W = 1) , there exists a unique value of δ that achieves this probability.
Because this relationship is strictly monotonic, we can hope to efficiently find such a value by e.g., binary search.
In the laboratory testing example of Example 1, this would allow us to specify a plausible strength for the conditional shift δ in terms of an impact on the overall testing rate, e.g., modelling a scenario where the testing rate decreases from 20% to 15%.
Similar to the binary case, we can (if desired) directly parameterize shifts in terms of the conditional mean of a Gaussian distribution, as illustrated in Example B.2, which operates on the scale of µ ( Z ) alone.
Example B.2 (Mean shift in a conditional Gaussian) .
Consider the distribution of a multi-variate
Gaussian variable W conditioned on a binary variable Z , where we write
[Equation]
where N ( w ; µ ( z ) , Σ( z )) denotes the Gaussian density with mean µ ( z ) and covariance Σ( z ) .
This can be written as an exponential family model with natural parameters η ( Z ) =
[Σ( Z ) − 1 µ ( Z ) , − 12 Σ( Z ) − 1 ] and sufficient statistic T ( W ) =
[ W, WW ⊤ ] .
Here, a shift in the mean can be parameterized by s ( Z ; δ ) =
[Σ( Z ) − 1 δ, 0] , such that
[Equation]
However, shifts of the same magnitude in the conditional mean may not be comparable.
Suppose that
[Equation]
and
[Equation]
such that δ = 1 in Example B.2 corresponds to
[Equation]
While it may seem plausible that the mean of W | Z = 0 can increase by 1 , it may seem unrealistic for W | Z = 1 .
Here, it may be more reasonable to consider a different parameterization of s ( Z ; δ ) , where the impact of the shift in a direction is proportional to the variance in that direction; we discuss this in the next example.
Example B.3 (Variance-scaled mean shift in a conditional Gaussian) .
Consider the distribution of a multi-variate Gaussian variable W conditioned on variables Z , where we write
[Equation]
where N ( w ; µ ( z ) , Σ( z )) denotes the Gaussian density with mean µ ( z ) and covariance Σ( z ) .
This can be written as an exponential family model with natural parameters η ( Z ) =
[Σ( Z ) − 1 µ ( Z ) , − 12 Σ( Z ) − 1 ] and sufficient statistic T ( W ) =
[ W, WW ⊤ ] .
Here, a shift in the mean can be parameterized by s ( Z ; δ ) =
[ δ, 0] , such that
[Equation] [Equation]
In Example B.3, the parameter δ has a different interpretation, as a variance-scaled mean-shift.
If W is one-dimensional, we can see that this becomes
[Equation]
As we demonstrate in Appendix C.2, this particular example of a parameterization has other benefits:
For instance, for estimation of shift gradients and Hessians at δ = 0 can be done without knowledge of Σ( Z ) .
B.2 Adding causal edges to the graph
In Section 2, we consider the case where the shift function s ( Z ; δ ) alters a conditional P ( W | Z ) by a shift function s ( Z ; δ ) .
We now discuss shift functions that use a larger set Z ′ .
In particular, we consider the setting where Z represents the parents in a graph G (that is, Z := PA G ( W ) ), and consider shift functions that correspond to adding additional parents in that causal graph.
Our definitions and results immediately extend to measuring the impact of shifts that add edges to the graph, in the form of shift functions that depend on non-descendants of W .
Building intuition with a simple example: To build intuition, consider the causal graph given in Figure 6.
We consider a shift in X 2 , with a shift function which depends not only on the causal parent Y , but also on X 1 .
Suppose that the distribution P ( X 2 | Y ) is a conditional exponential family, given by
[Equation]
Using that X 2 ⊥⊥ X 1 | Y , we have P ( X 2 | Y ) =
P ( X 2 | Y, X 2 ) , and the joint probability factorizes as
[Equation]
This enables us to consider Z =
( Y, X 1 ) as the conditioning set in the context of Assumption 1.
This is useful, because it allows us to consider shift functions that depend on Z , which includes X 1 in addition to Y .
The δ -perturbation of this conditional distribution under the shift function s ( Y, X 1 ; δ ) is given by
[Equation]
and we can observe that under both graphs, the distribution factorizes in the same fashion, where
[Equation]
keeping the same convention that s ( Y, X 1 ; δ = 0) = 0 , such that P 0 = P .
This is one example of how our results can be applied with shift functions that effectively add edges to the causal graph.
Of course, not all edges are permitted, so we give a more general treatment below.
General guidelines for adding edges : Allowing for the use of non-causal parents in the shift functions is straightforward, and can be done safely as follows, without violating Assumption 1:
Given knowledge of the directed acyclic graph G which generates the observed distribution P , we can add edges to the graph, as long as they do not create cycles.
Formally, let G =
( V , E ) denote the causal DAG which generates the distribution P , where V denotes variables and E denotes the set of edges, where we denote a directed edge by e
=
( V i , V j ) , going from V i to V j .
Let G ′
=
( V ′ , E ′ ) denote another DAG (of our creation) with the constraint that we can only add edges, and that the graph must remain acyclic, such that E ′ ⊇
E , and V ′ =
V .
For any variable W i ∈ V , this implies that PA G ′ ( W i ) ⊇ PA G ( W i ) .
Moreover, any new causal parent
V i of W i in G ′ must have been a non-descendant of W i in the original graph, as otherwise the graph G ′ would have a cycle from W i
→ V i → W i .
For ease of notation, let N ( W i ) :
= PA G ′
( W i ) \ PA G ( W i ) denote the set of new causal parents of W i in G ′ .
For any variable W i such that N ( W i ) ̸
= ∅ , we can write that
[Equation]
by the rules of d-separation [Pearl, 2009].
As in Assumption 1, we use W =
{ W 1 , . . .
, W m } to denote the set of variables to be intervened upon, and accordingly will assume that in the causal graph
G ′ , we have not added new parents to any other variables, i.e., N ( V i ) = ∅ for any V i ⊊ W .
By Equation (10), we can write that the distribution P factorizes as
[Equation]
because P ( W i | PA G ′
( W i ))
=
P ( W i | PA G ( W i ) , and if P ( W i | PA G ( W i )) is a conditional exponential family satisfying Definition 2, then P ( W i | PA G ( W i )) also satisfies this definition, where the function η (PA G ( W i ) , N ( W i )) is constant with respect to fluctuation in the variables
N ( W i ) .
Thus, taking
Z i :=
PA G ′
( W i ) as the conditioning set satisfies Assumption 1, and the rest of our results hold, where the corresponding δ -perturbations in Definition 4 are given by
[Equation]
with shift function s i (PA G ′ ( W i ); δ i ) that are parametric functions of causal parents in the modified graph G ′ .
B.3 Domain-preserving parameterizations of shift
For both of the examples considered above, we did not need to restrict the magnitude of the additive change to η ( Z ) .
However, in some cases, such as changing the variance of a conditional Gaussian, we have the restriction that η δ ( Z ) =
η ( Z ) + s ( Z ; δ ) must lie in the proper domain, e.g., we cannot consider a shift which causes the conditional variance to become negative.
For a conditional
Gaussian, we can consider unrestricted shifts in η ( Z ) 1 , which controls the mean, because the mean has unrestricted domain.
On the other hand, η ( Z ) 2 = ( − 2 σ 2 ( Z ))
− 1 controls the variance, and must remain negative, such that η ( Z ) {+
+} + s ( Z ; δ ) 2 < 0 for the shifts we consider.
This can be resolved in one of two ways.
First, one can consider parameterizations of s ( Z ; δ ) which are guaranteed to preserve the correct domain with an additional constraint on the values of δ , such as the multiplicative shift below, which is sign-preserving for δ > −
[Equation]
To handle the general case, at the expense of some additional complexity in the gradients of s ( Z ; δ ) , one can define the shifts as follows for parameters η ( Z ) that have a lower bound L , with an equivalent formulation for shifts where the parameters have an upper bound, for any desired shift function s ′ ( Z ; δ )
[Equation]
where sigmoid ( γ · ( x − ( L + ϵ ))) is a smooth relaxation of the indicator function 1 { x > L + ϵ } , for a sufficiently large temperature parameter γ > 0 and a small ϵ > 0 .
This transformation preserves the twice-differentiable nature of s ( Z ; δ ) .
In practice, however, we typically evaluate the gradient of s ( Z ; δ ) at δ = 0 , where η ( Z ) does not lie at the boundary of allowable parameter space, such that we can consider simpler parameterizations like
[Equation]
as long as ϵ is taken sufficient small such that η ( Z ) >
L + ϵ almost everywhere in P .
C Considerations and additional results for evaluation of the worst-case loss
In this section, we present additional results on the Taylor approximation and compare how the Taylor approximation compares to the reweighting approach in evaluation and worst-case optimization of the shifted loss.
• In Appendix C.
we give a full treatment of how shift gradients and Hessians are estimated from samples, following Theorem 1.
• In Appendix C.2, we demonstrate in some cases, one does not need to estimate all of η ( Z ) , but only the parts of η ( Z ) that is shifting.
• In Appendix C.3, we demonstrate that the second-order Taylor expansion is exact in a linear-Gaussian setting, which gives a conceptual connection between this work and that of Anchor Regression [Rothenhäusler et al., 2021], which considered a restricted type of additive shift intervention in a globally linear structural causal model.
• In Appendix C.4, we work out the expression for the shift gradient and Hessian when we condition on binary variables.
• In Appendices C.5 to C.7, we provide experiments that compare the variance of the importance sampling estimate ˆ E δ, IS (see Equation (6)) to the variance of the Taylor estimate ˆ E δ, Taylor (see Equation (7)) of the loss in a shifted distribution.
• In Appendix C.8, we consider the bound in Theorem 2 in a covariate shift setting, and give an explicit expression for this under additional assumptions.
C.1 Algorithm for Estimation of Shift Gradients and Hessians
Here, we recall the form of the shift gradients and Hessians in Theorem 1, and demonstrate how to compute them in practice using a set of auxiliary regression functions fit to the validation data.
Theorem 1 (Shift gradients and Hessians as covariances) .
Assume that P δ , P satisfy Definition 4, with intervened variables W = { W 1 , . . .
, W m } and shift functions s i ( Z i ; δ i ) , where δ =
( δ 1 , . .
δ m ) .
Then the shift gradient is given by SG
= (SG 11 , . .
, SG 1 m ) ∈ R
d δ where
[Equation]
and the shift Hessian is a matrix of size ( d δ × d δ ) , where the ( i, j ) th block of size d δ i ×
d δ j equals { SG
[Equation]
i,j
[Equation]
where D i,k := ∇ kδ
i s i
( Z i ; δ i )
| δ =0 , is the gradient of the shift function for k = 1 , and the Hessian for k = 2 .
Here, T i ( W i ) is the sufficient statistic of P ( W i | Z i ) and ϵ T i |
Z i :
=
T i ( W i ) − E
[ T ( W i ) | Z i ] .
Notation and Dimensions : Let W =
W 1 , . .
, W m } denote the set of m intervened variables, and let Z = { Z 1 , . . .
, Z m } denote the conditioning sets.
Note that for a single W i ∈
R d Wi , we will generally have it that Z i ∈ R d Z , where d W is the dimension of W (typically 1) and d Z is the number of conditioning variables, and when considering n samples, W i will be a matrix in R n × d W , and Z i will be a matrix R n × d Z .
The sufficient statistic T i ( W i ) maps from R d W to R d T , where d T is the dimension of the sufficient statistic.
For many common distributions, T i ( W i )
=
W i , the identity function.
For others, like the conditional multi-variate Gaussian, T i ( W i ) =
[ W i ,
W i W ⊤ i ] , where W ∈ R d W and W i W ⊤ i ∈ R d W × d W .
In these cases, we squeeze T i ( W i ) to be a single vector, so in this case d T =
d W + d 2 W .
Auxiliary models : To estimate the shift gradients and Hessians, we first learn auxiliary predictive models, which are required for computing the relevant conditional covariances.
For simplicity, we do not consider sample-splitting in the algorithm given below, but one could employ sample-splitting to learn these predictive models on an independent validation sample.
• For each W i , we learn ˆ µ W
i
( Z i ) as a regression model for E
[ T i ( W i )
| Z i ] .
Because T i ( W i ) may have multiple dimensions, this is a function from R d Z to R d T .
• For each conditioning set Z i , we learn ˆ µ ℓ ( Z i ) as a regression model for E [ ℓ | Z i ] .
Because the loss is one-dimensional, this is a function from R d Z to R .
We then construct the following, which are defined for each data point in the sample
• For each W i , we construct ˆ ϵ T i
| Z i :=
T i ( W i ) − ˆ µ W i ( Z i ) , which is a vector of length d T i .
• For each conditioning set Z i , for the loss ℓ , we construct ˆ ϵ ℓ | Z i :=
ℓ − ˆ µ ℓ ( Z i ) , which is a real number. • For each conditioning set Z i , we compute D i, 1 ( Z i ) as ∇ δ
i s i
( Z i ; δ i ) (cid:12)(cid:12) δ =0 , which is a matrix of size d T × d δ i , and a function of Z i that we can evaluate on each sample.
• For each conditioning set Z i , we compute D i, 2 ( Z i ) as ∇ 2 δ
i s i
( Z i ; δ i ) (cid:12)(cid:12) δ =0 , which is a tensor of size
d T × d δ i ×
d δ i , and a function of Z i that we can evaluate on each sample.
Estimating shift gradients
The shift gradient and Hessian in Theorem 1 are expressed as conditional covariance.
Since E [cov( A, B | C )] =
E [ ϵ A | C ϵ B | C ]
where ϵ A | C :
=
A − E [ A | C ] and ϵ B | C :=
B − E [ B | C ] , we can use the estimated conditional means above, to compute the shift gradient and Hessian.
Suppose that we observe N samples, n ∈ { 1 , . . .
, N } .
For each index i ∈ [ m ] := { 1 , . . .
, m } ,
[Equation]
N
[Equation]
which yields a vector of length d δ i , and these are concatenated together for each i to yield the entire shift gradient.
The shift Hessian is constructed block-wise, for each index i, j ∈ [ m ] × [ m ] as follows:
If i = j , then we construct the corresponding d δ i ×
d δ i block as
[Equation]
where v ⊗ 2 denotes the outer product so that v ⊗ 2 = vv ⊤ , and the transpose of D i, 2 refers to a transpose which has dimension d δ i ×
d δ i × d T .
On the other hand, if i ̸ =
j we have
[Equation]
where ¯ ℓ is the average value of ℓ in the validation sample.
C.2 Shifts where estimating all of η ( Z ) is not necessary for estimating shift gradient and Hessian
The following example shows that when a shift occurs in an exponential conditional distribution with parameter η ( Z ) , we do not necessarily need to model all of η ( Z ) in order to compute the shift gradient and Hessian.
In particular, we only need to model the parts of η ( Z ) that shift.
This is different from estimating the shifted loss using importance sampling, where η ( Z ) needs to be evaluated to evaluate Equation (5).
Example C.1.
Consider the distribution of W conditioned on variables Z that is a multi-variate
Gaussian variable,
[Equation]
for unknown functions µ, Σ .
The sufficient statistic for the multivariate Gaussian distribution is
T ( W ) =
( W, WW ⊤ ) and the canonical parameter is η ( Z ) =
(Σ( Z ) − 1 µ ( Z ) , − 12 Σ( Z ) − 1 ) .
6 The first component of η ( Z ) is a signal-to-variance ratio and the second is the inverse covariance matrix.
For a shift ( δ, 0) that only affects the first component, we show that we do not need to model Σ( Z ) , but only µ ( Z ) .
This is beneficial, since estimating a conditional covariance from data can be challenging, especially if W is high-dimensional.
A
X
H Y
[Figure]
Loss δδδ
For δ ∈ R d W , let s ( Z ; δ ) =
( δ, 0) ⊤ , and suppose that we wish to estimate E δ [ ℓ ] using Equation (7).
The derivative of s is given by
[Equation]
   
[Equation]
where the first block is a d W × d W diagonal matrix, and the second is a d W × d 2 W matrix of zeros.
The second derivative of s is D
=
Hence, using Theorem 1, the shift gradient is
[Equation]
and
[Equation]
Conditional covariances can be computed by only residualizing one of the variables: E [cov( A, B | C )]
=
E [ A ( B − E
[ B | C ])] .
Thus, if we only residualize ℓ , we get
SG
[Equation]
=
E [
[Equation]
E [
[Equation]
[Equation] ] and SG [Equation]
Therefore, given data from P , we can estimate the shift gradients by plugging in estimators ˆ µ ( Z ) of
E [ W | Z ] and ˆ L ( Z ) of E [ ℓ | Z ] .
It follows that we do not need to model Σ( Z ) in order to estimate the shift gradients and Hessian at δ = 0 .
The story is different for a reweighting based estimator that seeks to estimate E δ [ ℓ ] using importance sampling (see Section 3.1), where the weights are given by
[Equation]
and hence estimating w η,δ ( Z ) requires estimation of Σ( Z ) .
C.3 The quadratic approximation is exact, for mean shifts in linear models
We now consider data generated by a linear model, and show that the shifted loss is a quadratic function of δ , meaning that the Taylor approximation E δ, Taylor is globally exact.
Suppose that data is sampled from a linear structural causal model, and a shift in mean occurs in an variable A that does not have any causal parents.
In particular, let A have a normal distribution with mean µ and finite variance and let
[Equation]
This is the model assumed by Rothenhäusler et al.
[2021], and the corresponding graphical model is shown in Figure 7 (left).
We consider the linear predictor f γ ( X ) =
γ ⊤ X and the mean squared loss ℓ ( f γ ( X ) , Y ) =
( Y − f ( X )) 2 .
Due to the linearity of the model, the loss under a mean shift in A is quadratic [Rothenhäusler et al., 2021].
Lemma C.1.
Suppose A ∼ N ( µ, Σ) and that ( X, Y, H ) are generated according to Equation (11) .
For γ ∈ R d X define ℓ :=
( Y − γ ⊤ X ) 2 .
Then there exist v γ , u
µ,γ ∈ R d A such that for all shifts δ ∈ R d A :
[Equation]
where E δ corresponds to taking the mean in the distribution where A ∼ N ( µ + δ, Σ) .
Further u µ,γ =
if µ = 0 .
Proposition C.1 elicits two properties of this linear model:
First the loss is described by a quadratic function globally, i.e. also for very large δ .
In Figure 7 (middle), we plot E δ [ ℓ ] as a function of δ .
We observe a ‘valley’ in the loss, in which the expected loss does not at all change with δ .
This is a consequence of Lemma C.1, and particularly that if δ is orthogonal to both u µ,γ and v γ then E δ [ ℓ ] = E [ ℓ ] .
In higher dimensions d A > 2 , since v γ v ⊤ γ has rank 1 , the ‘valley’ persists in that the loss does not grow at all in d A − 2 dimensions (or d A − 1 if A has mean µ = 0 ), see Figure 7 (right).
We now show that coefficients in the quadratic form in Lemma C.1 is equal to the shift gradient and
Hessian.
We use that the Gaussian distribution with known variance Σ can be parameterized as an exponential family with sufficient statistic T ( A ) =
Σ − 1 A and parameter η = µ .
Proposition C.1.
Suppose A ∼ N ( µ, Σ) and that ( X, Y, H ) are generated according to Equation (11) .
Then the shift gradient and Hessian are given by
SG 1 = cov( ℓ, Σ − 1 A ) and
SG 2 = cov( ℓ, Σ − 1 ( A − µ )( A − µ ) ⊤ Σ −⊤ ) and the loss under a mean shift of δ in A is given by
[Equation]
where ℓ :=
( Y − γ ⊤ X ) 2 and E δ corresponds to taking the mean in the distribution where A ∼
N ( µ + δ, Σ) .
This elicits a connection to anchor regression [Rothenhäusler et al., 2021]: Under the generative model Equation (11) and using the quadratic loss ℓ =
( Y − γ ⊤ X ) 2 for γ ∈ R d X , they show that for any λ ≥ 0 , the worst-case loss E δ [ ℓ ] over a set ∆ =
{ δ | δδ ⊤ ⪯ λ E [ AA ⊤ ] } equals the objective ℓ AR =
E [ ℓ ] + λ E [ E [ Y − γ ⊤ X | A ] 2 ] , which is computable from the observed distribution.
Because of Proposition C.1, ℓ AR also equals the solution of the optimization problem Equation (9) over the constraint set ∆ .
Therefore minimizing the anchor regression objective over γ or minimizing
Equation (9) over γ will lead to the same estimator.
Since our proposed Taylor approximation in
Equation (9) does not assume linearity, one could use the approximation to extend the rationale of anchor regression of minimizing the worst-case loss to non-linear models.
This however comes at the cost of not optimizing the exact worst-case loss, but rather an approximation, whose quality is given by Theorem 2.
Further, this would involving a minimax problem, minimizing Equation (9) over models f , and there are questions, such as convexity and tractability, which would need to be solved.
C.4 Estimating the shift gradient and Hessian for conditional on binary variables
To build intuition for the shift gradient and Hessian, we here give an example where we condition on variables Z that take a finite number of values and write out explicit expressions for the shift gradient and Hessian.
However, we emphasize, that in most practical scenarios, one will not have to work out the shift gradient and Hessian explicitly, but can simply estimate them as covariances from the data (Theorem 1).
Example C.2 (Shift Function of Discrete Parents) .
Consider a conditional distribution W | Z where
Z takes values in a finite set Z .
This is for instance the case if Z = ( Z 1 , . . .
, Z d ) where each Z i is binary, so |Z| = 2 d .
Instead of a shift η ( Z ) + δ , where the parameter increases by the same [Figure] amount for all values of Z , we may consider a shift η ( Z ) + s ( Z ; δ )
where s ( Z ; δ ) =
(cid:80) z ∈Z δ z 1 Z
= z , meaning that the shift is different in each category Z .
Since η ( Z ) only takes a finite number of variables, this shift corresponds to an arbitrary change in η ( Z ) .
s ( Z ; δ ) is a differentiable function in δ , and if d T = 1 the shift gradient is a (1 × 2 d ) -row vector, ∇ δ s ( Z ; δ )
= (1 Z = z ) z ∈Z , and the shift Hessian vanishes, ∇ 2 δ s ( Z ; δ ) = 0 .
Enumerating Z = { 1 , . .
. , 2 d } , the i ’th entry in the shift gradient becomes
[Equation]
and the i, j ’th entry of the shift Hessian becomes 0 if j ̸ = i and else
[Equation]
Consider for example the case where both W and Z are binary.
Then T ( W ) =
W and s ( Z ; δ ) = 1 Z =
δ 0 + 1 Z =1 δ 1 and s (1) =
(1 Z =0 , 1 Z =1 ) and s (2) = 0 .
The conditional covariance can be evaluated by residualizing only one of the variables, E [cov( A, B | C )]
=
E [ A ( B − E
[ B | C ])] , so we can chose to residualize only W (for SG 1 ) or
( W − E
[ W | Z = i ]) 2 (for SG 2 ).
Finally, if we let p i
= P ( W = 1 | Z
= i ) and use that E
[ W | Z = i ]
= p i and E
[( W − p i ) 2 | Z
= i ] =
var( W | Z = i ) =
p i (1 − p i ) , we get that
[Equation]
and
[Equation]
C.5 Comparison of variance of reweighting and Taylor estimates in the lab ordering example
To compare the bias and variance of the Taylor and the importance sampling estimates of the shifted loss, we simulate data from the following, artificial, generative model (which is the same generative model that was used to construct the loss landscape in Figure 1 (right)).
[Equation]
where if Order = 0 , the test result is a placeholder value of zero.
We consider either a shift in the logits of ordering lab tests η δ ( Z ) =
η ( Z ) + δ (Figure 8 left) or a mean shift in the Gaussian distribution of age η δ =
δ (Figure 8 right).
For each δ in a grid, we compute estimates ˆ E δ, IS and ˆ E δ, Taylor of the loss under a shift of size δ , We repeat this n = 1 ,times, and plot the mean and point-wise prediction intervals (the pointwise 0 . 05 and 0 .
95 quantiles) for ˆ E δ, IS and ˆ E δ, Taylor .
We also simulate ground truth data from P δ , to compute the actual loss under shift.
For shifts in the binary variable (Figure 8, left), both estimates capture the loss well for small shifts, but as δ gets larger, the quadratic approximation increasingly deviates from the true mean; the importance sampling estimate remains very close to the ground truth shifted loss.
On the contrary, for the Gaussian mean shift (Figure 8, right), the importance sampling weights are ill-behaved, and the variance dramatically increases as δ becomes larger.
This supports the intuition, that while importance sampling tends to work well for binary variables, the variance can be large in continuous distributions, such as the Gaussian distribution.
C.6 Comparison of theoretical variance of reweighting and Taylor estimates
Example C.3.
To demonstrate the reduction in variance obtained from using the Taylor approximation of the importance weights, we consider a simple example where P ( X ) ∼ N (0 , 1) and P δ ( X ) ∼
N ( δ, 1) and we wish to estimate E δ [ ℓ ( X )] for some loss function ℓ ( X ) .
The importance sampling weights are given by w δ ( X ) =
exp( − 12 δ 2 + X · δ ) , and the shift gradient and Hessians are SG 1 =
E [ ℓ ( X ) X ] and SG 2 = E [ ℓ ( X ) X 2 ] .
Therefore samples X 1 , . .
, X n from P consider the estimators, for any loss function ℓ ( X ) , two estimators of E δ [ ℓ ] are
[Equation]
and the variances of the estimators are
[Equation]
The variance of ˆ µ Taylor grows like δ 4 and the variance of ˆ µ IS grows exponentially fast (unless E [ { ℓ ( X + 2 δ ) } 2 ] also diminishes exponentially fast, which is generally not the case), and so except for small δ , the variance of the importance sampling estimator will be orders of magnitude larger than the variance of the estimator using the Taylor approximation.
While, ˆ µ IS is an unbiased estimator of E δ [ ℓ ( X )] and ˆ µ Taylor is a biased, the overall mean squared error will be smaller for the Taylor approximation, unless the bias of the Taylor approximation also grows exponentially.
For the sake of analysis, consider the simple example ℓ ( X ) =
X .
In this case, the Taylor estimate is unbiased because E δ [ X ]
= δ is a linear function of δ , so the quadratic approximation is adequate.
Further, the variances are given by [Equation] and [Equation]
n . In particular, the variance of the importance sampling estimate grows like exp( δ 2 ) while that of the Taylor estimate grows like δ 4 .
Shift Strength
P r e d i c ti on
E rr o r
Method
IS
IS (clipped)
Taylor
C.7 Comparison of variance of reweighting and Taylor estimates in a simple synthetic example
In this experiment, we compare the variance of importance sampling and Taylor estimates in a simple synthetic example.
We simulate data from P where X ∈ R 3 and Y ∈ R 1 depend either linearly or quadratically on W ∈ R 3 ,
W ∼ N (0 , Id 3 ) and (cid:18)
[Equation]
where ⊙ refers to entrywise multiplication, ϵ ∼ N (0 , Id 4 ) , α is either
(linear) or 12 (nonlinear) and
B :=   2 1 02 2 03 3 04 2 4  and
M :=   2 12 12 24 1  .
On the simulated data from P , we then fit a linear predictor f ( X ) of Y , and consider a shift in the mean of W from P ( W ) ∼ N (0 , Id 3 ) to P δ ( W ) ∼ N ( δ, Id 3 ) , where δ =
[ s, s, s ] ⊤ for some shift strength s > 0 .
We then compute the shift gradient SG 1 = cov( ℓ, W ) and Hessian SG 2 = cov( ℓ, WW ⊤ ) , and approximate E δ [ ℓ ] by ˆ E δ, Taylor (see Equation (7)).
In the linear data, the Taylor approximation is exact (see Appendix C.3), such that any prediction error can be attributed to finite-sample fluctuation, whereas both model misspecification and finite-sample fluctuation contribute to the error in the nonlinear setting.
Similarly, we estimate E δ [ ℓ ] by importance sampling, E δ [ ℓ ] =
E [ w δ ( W ) ℓ ] ≈ 1 n (cid:80) w δ ( W ) ℓ , where w δ ( W ) =
P δ ( W ) P ( W ) =
δ ⊤ W − [-1 
-]12 δ ⊤ δ , and compare this to ground truth data sampled from P δ ; we do the same for an importance sampling estimator with weights ‘clipped’ at the 99% quantile.
We compare the predicted loss E δ [ ℓ ] by actually simulating data from P δ and evaluating E δ [ ℓ ] (where ℓ is still the model trained on data from P ).
We then compute the prediction error, as the difference E δ [ ℓ ] − ˆ E δ, Taylor or E δ [ ℓ ] − ˆ E δ, IS .
For a number of different shift strengths s , we repeat this procedure M = 1 , 000 times, and in
Figure 9 we plot the median and a confidence interval defined by the 2 . 5 and the 97 .
5% quantiles of the prediction error.
In the linear case, bot the importance sampling and the Taylor approximation retains a median error close to 0 , with the variance of ˆ E δ, IS being larger than ˆ E δ, Taylor .
The clipped importance sampling estimate has a smaller variance than that of ordinary importance sampling, though the median deviates further from 0 , and the variance is not smaller than that of the Taylor estimate.
In the non-linear cases, all three models underestimate the shifted loss.
For ˆ E δ, Taylor , this happens because as the mean of W shift, the mean shift is amplified by the non-linearity, such that the quadratic approximation of the loss is an underestimate.
While the variance of the clipped importance sampling is smaller than the variance of the ordinary importance sampling estimate and comparable to the variance of the Taylor estimate, this prediction is further from 0 than the Taylor estimate.
Since importance sampling methods are known to produce very large outliers, the use of the median and quantiles, as opposed to the mean an confidence intervals based on the standard deviation, is favouring importance sampling; the Taylor method looks even more favourable if we instead plot the mean and standard deviations.
C.8 The bound in Theorem 2 under covariate shift
The bound in Theorem 2 is in a general form that applies to any shift in the CEF framework.
In concrete cases, the bound can be made simpler, as we now demonstrate.
Suppose that X is a covariate that is Gaussian distributed N (0 , 1) .
Also consider a prediction target
Y :
= f 0 ( X ) + ϵ for some function f 0 and noise variable ϵ that is independent of X .
Suppose we consider a predictor ˆ Y =
f
( X ) and apply our proposed methodology to estimate the mean squared prediction error when predicting Y ≈ f ( X ) under a mean shift of size δ ∈ R to X .
When we only consider shifts in the mean (and not the variance), the sufficient statistic is T ( X ) =
X .
We can use Theorem 2 to bound the prediction error.
In this setting,
[Equation]
such that the bound in Theorem 2 becomes
[Equation]
The subscript cov t · δ indicates that the covariance is taken in the distribution N ( t · δ, 1) ; instead we can write this in the observed distribution, and add t · δ to X .
Further, the terms relating to ϵ disappear, as they are independent of X .
Thus, if we define the modelling error g ( x ) = f
0 ( x ) − f ( x ) , we can write
[Equation]
We can bound the covariance using the inequality cov( A, B ) ≤ (cid:112) var( A ) var( B ) ,
[Equation]
The first term on the right hand side is the variance of the difference of approximation error in X and in X + tδ .
If we are willing to make assumptions on the quality of the approximation f , we can simplify this further.
For example, we can assume that | g ( x ) 2 − g ( y ) 2 | ≤ C · | x − y | 2 , meaning that the squared error of f 0 ( x ) − f ( x ) does not change faster than quadratically in x .
In that case, we get (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
In some cases, one can sharpen this bound by using prior knowledge about the data generating mechanism (for example, the data generating function f 0 may be bounded).
D Limitations of worst-case conditional subpopulation shift for defining plausible robustness sets
For the example in Section 4.1, we can contrast the type of shift we consider with the worst-case (1 − α ) -conditional subpopulation shift considered by Subbaswamy et al.
In this section, we will make the following points: First, worst-case conditional (1 − α ) -subpopulation shifts can be too pessimistic, with even moderate values of α leading to implausible conditional distributions.
Second, we will argue that parametric robustness sets enable more fine-grained control over the set of plausible shifts, leading to more informative estimates of worst-case risk.
Overall, we argue that the two approaches are complementary, with different strengths.
Before we proceed, we define a conditional (1 − α ) subpopulation shift.
A (1 − α ) subpopulation shift in the conditional distribution P ( O | Y ) is defined by a weighting function h : O × Y (cid:55)→ [0 , 1] , which has the property that E
[ h ( O, Y ) | Y ] = 1 − α for all values of Y .
This can be used to construct a worst-case objective, which measures the worst-case loss under such a shift:
[Equation]
where µ ( O, Y ) :
=
E [ ℓ ( Y, f ) | O, Y ] , for a predictor f and loss ℓ .
This has the effect of leaving the distribution P ( Y ) untouched, while changing the conditional distribution P ( O | Y ) .
Throughout this section, we will use the same predictor f
( O, L ) described in Section 4.1.
The rest of this section is structured as follows:
In Appendix D.1, we derive the feasible set of conditional distributions P ( O | Y ) implicitly considered by this objective in the simple generative model of Section 4.1, which only involves variables
O, L and Y .
We do so by showing that (for discrete O, Y ), maximizing Equation (12) over h is equivalent to solving a linear program, where we can characterize the constraints on h exactly, and translate them into constraints on P ( O = 1 | Y = 1) , P ( O = 1 | Y = 0) .
Here, we show that the resulting feasible set is quite large, even for moderately large subpopulations.
In particular, whenever (1 − α ) < min { P ( O = 1 | Y = 0) , P ( O = 0 | Y = 1) } , all conditional distributions are possible.
In Appendix D.2, we derive the value of h that maximizes Equation (12), and show that, as we vary α , the worst-case shift is always in the same “direction” probability space:
Healthy patients ( Y = 0) are tested more, and sick patients ( Y = 1) are tested less, and for α < 0 .
27 , the worst-case subpopulation shift is the (unrealistic) scenario where healthy patients are always tested, and sick patients are never tested.
In Appendix D.3, we illustrate how this type of behavior can be avoided with our approach.
We first give a parameterized shift function s ( Z ; δ 0 , δ 1 ) such that we can reach any conditional distribution of P ( O | Y ) , for sufficiently large values of δ 0 , δ 1 .
We then demonstrate how an iterative process might play out with domain experts, where we consider different constraint sets until we find a constraint set that contains plausible shifts.
D.1 Feasible conditional subpopulations in Section 4.
For the simple example in Section 4.1, we give a self-contained derivation of the feasible region for 1 − α conditional subpopulations in the distribution P ( O | Y ) .
The advantage of working with this simple generative model is that the conditional distribution can be described by only two numbers, P ( O = 1 | Y = 1) and P ( O = 1 | Y = 0) , and so we can visualize the resulting conditional distribution.
Because O, Y are discrete, the worst-case subpopulation in this simple example can be solved via a linear program, for a fixed α .
We have an optimization problem in two variables, since h 11 P ( O = 1 | Y = 1) + h 01 P ( O = 0 | Y = 1) = 1 − α , and likewise for h 10 , h 00 , where h ij =
h
( O = i, Y = j ) .
We also have the constraint that each variable must live in [0 , 1] .
Meanwhile, the loss to maximize is a linear function, as an expectation of E [ h ( O, Y ) µ ( O, Y )] , where µ ( O, Y ) takes on four possible values, where we write p ij =
P O = i
| Y = j ) , and µ ij similarly.
[Equation]
This linear program is simple enough to solve by hand, and we will do here to build intuition.
In this section, we begin by characterizing the feasible region of h , and then translating that into a feasible region for P h ( O | Y ) , which we can plot in two dimensions.
Characterizing feasible values of h : Here, we focus on characterizing the feasible set that h can lie in, as a way of characterizing the feasible set for P ( O | Y ) .
From the constraints, we can write that
[Equation]
There are only two constraints on h 11 :
Those directly imposed by 0 ≤ h 11 ≤ 1 , and those which are imposed by the equality constraint with h 01 and the fact that 0 ≤ h 01 ≤ 1 .
For the latter, with some algebra we can write that
[Equation]
So that the constraints on h 11 become
[Equation]
which recovers our intuition that if α = 0 , it must be that h 11 = 1 and h
= 1 .
Bounding feasible values of P h
( O | Y )
The parameters h can be understood as importance weights whose expectation is 1 − α instead of 1 , that reweight P to a new distribution P h when appropriately normalized.
To compute conditional probabilities P h
( O = i | Y =
j ) under the new distribution, we can compute the expectation of 1 { O = i, Y = j } , and normalize by P ( Y = j ) .
[Equation]
where the implication follows from the fact that P h ( Y )
=
P ( Y ) .
This allows us to translate bounds on h ij directly into bounds on P h
( O = i
| Y = j ) .
Making use of Equation (14), we can write that
[Equation]
which yields
[Equation]
We can apply a similar logic to h 10 , which is identical except for p 11 being replaced by p 10 , yielding
[Equation]
Visualizing the constraint set : Figure 10 gives feasible conditional distributions under different values of α .
We can observe that when α = 0 .
8 , all conditional distributions are feasible, including the distribution where P ( O = 1 | Y = 0)
= 1 and P ( O = 1 | Y = 1)
= 0 , representing the case where every healthy patient gets tested [Equation] and no sick patients receive a test.
This is generally possible in this example whenever 1 − α < min { P ( O = 1 | Y = 0) , P ( O = 0 | Y = 1) } , as it permits the following subpopulation function, which yields this result.
[Equation]
D.2 Worst-case conditional subpopulation shifts
Given the constraint set which describes the feasible set of conditional distributions under the (1 − α ) - conditional subpopulation objective, we can derive the worst-case conditional distribution.
Here, since Y, O are both binary, the expected loss under a new distribution P h is given by
[Equation]
1.0 P ( O =1 Y =0) 0.
P ( O = 1 Y = 1 )
Robustness Set Original Distribution Worst-Case Direction Worst-Case Distribution (a) (1 − α ) = 0 .0.
1.0 P ( O =1 Y =0) 0.
P ( O = 1 Y = 1 ) (b) (1 − α ) = 0 .0.
1.0 P ( O =1 Y =0) 0.
P ( O = 1 Y = 1 ) (c) (1 − α ) = 0 .
which we can write in terms of the constrained probabilities P h as follows, where q 11 :
=
P h
( O = 1 | Y = 1) and q 10 :=
P h
( O = 1 | Y = 0)
[Equation]
)] which also gives us a direction in which the loss is maximized, since the loss is given by
[Equation]
where C = P ( Y = 1) µ (0 , 1) + P ( Y = 0) µ (0 , 0)
Since q 11 , q 10 can be optimized independently, the worst-case solution is given by taking the maximum value of q 11 if µ (1 , 1) > µ (0 , 1) and the minimum value if µ (1 , 1) < µ (0 , 1) , and likewise taking the maximum value of q 10 if µ (1 , 0) > µ (0 , 0) , and the minimum value otherwise.
If µ (1 , 1) = µ (0 , 1) or µ (1 , 0) = µ (0 , 0) , then the objective is unaffected by the choice of q 11 or q 10 respectively.
Visualizing the worst-case conditional distributions The worst-case directions on the probability scale, and the resulting worst-case conditional distribution obtained by solving Equation (13), are given in Figure 10.
The red line arrow visualizes the direction from Equation (15), and the worst-case distribution is the point which is furthest in this direction in the constraint set.
Here, we are finding the worst-case accuracy of the same predictive model f
( O, L ) described in Section 4.1.
We can observe that the worst-case loss is obtained by seeking to reverse the correlation between Y and O , decreasing the probability that a sick patient ( Y = 1) gets a test ordered, and increasing the probability that a healthy patient ( Y = 0) gets a test ordered.
D.3 Iterating with domain experts to define realistic parametric robustness sets
In the previous sections, we saw that (1 − α ) -conditional subpopulation shift does not always produce realistic worst-case conditional distributions.
Moreover, given only the parameter α , there is limited ability to control the nature of the resulting worst-case conditional distribution P ( O | Y ) .
In this section, we contrast this limitation with the finer-grained control enabled by considering parametric robustness sets.
In particular, we argue that parametric shifts allow for end-users to customize robustness sets, ruling out shifts that represent unrealistic changes.
In practice, we imagine that the following iterative process could be a useful tool in model development: (i) Define a class of shifts with an appropriate s ( Z ; δ ) and constraint set ∆ , and search for a worst-case shift δ .
(ii) Present to domain experts both the worst-case shift δ (in terms of summary statistics of the resulting distribution P δ ) alongside the associated estimate of the worst-case loss.
For instance, report both the worst-case loss, as well as corresponding rate of testing among sick and healthy patients.
(iii)
If the shift itself is unrealistic, further the constrain parameter set or shift function, and repeat the process.
In Figure 11, we give a concrete example.
Each sub-figure shows the set of conditional probability distributions P ( O | Y ) that can be represented by a shift of ( δ 0 , δ 1 ) ∈ ∆ {+
+} × ∆ 1 , along with the 1.0 P ( O =1 Y =0) 0.
P ( O = 1 Y = 1 )
Robustness Set Original Distribution Worst-Case Distribution Worst-Case Direction (a) 0.
1.0 P ( O =1 Y =0) 0. P ( O = 1 Y = 1 ) (b) 0. 1.0 P ( O =1 Y =0) 0. P ( O = 1 Y = 1 ) (c) worst distribution ( s ( δ ) =
δ
+ δ 1 Y , δ controls a general increase or decrease in testing, while δ 1 controls a shift in the testing rate for only sick patients, and allows for a different change in the testing rate of sick vs healthy patients.
Iteration 1: We might imagine starting with a relatively unconstrained robustness set, where δ 0 and δ 1 are unconstrained.
Figure 11a shows the resulting robustness set of conditional distributions, and finds a shift with with a worst-case accuracy of 16%, compared to accuracy of 75% on the original distribution.
However, the corresponding δ -perturbation P δ is unrealistic, where all healthy patients (and no sick patients) are tested.
Luckily, because we have parameterized the shift, we can constrain the robustness set to exclude these types of results.
Iteration 2: A benefit of our approach is that we can refine the robustness set, with this type of feedback in mind.
In Figure 11b, we restrict the support of δ 1 to [ − 1 , 1] , to avoid large changes in the relative probability of testing sick vs healthy patients.
Here, the resulting worst-case accuracy is much higher (50%), but the corresponding worst-case conditional probability distribution is perhaps still unrealistic:
No patients undergo laboratory testing at all!
Iteration 3: Finally, we consider only shifts that affect all patients in a similar way, generally raising or lowering the conditional probability of a lab test, represented by shifts in δ 0 alone.
This may correspond to a more realistic scenario where (in a new hospital) laboratory testing use is more or less constrained.
Additionally, we can specify that this shift should decrease testing rates by at most 20%, which translates directly into a lower-bound on δ 0 .
9 Figure 11c shows the resulting robustness set of distributions, where the worst-case shift may seem more plausible: A reduction in testing rates for both populations.
The worst-case accuracy in this case is 69%, comparable to the accuracy of 75% on the original distribution.
E CelebA: Experiment details and additional results
In this section, we give details of the computer vision experiment in Section 4.2.
E.1 Details for the experiment
Creating the training distribution To construct the training distribution P , we use the conditional
GAN in Kocaoglu et al.
In particular, we use their CausalBEGAN, which is an extends the [Figure] boundary equillibrium GAN [Berthelot et al., 2017] to also take attributes as inputs.
We train the
CausalBEGAN using the default hyper parameters in the implementation provided by Kocaoglu et al. [2018], available under the MIT license.
The model is trained for 250 , 000 iterations on a single GPU, taking around approximately 16 hours.
Similar to Kocaoglu et al.
[2018], we use the CelebA dataset [Liu et al., 2015], which contains approximately 200 , 000 images of faces, along 40 binary attributes.
Of those, we use the following 9 attributes {Male, Young, Wearing Lipstick, Bald, Mustache, Eyeglasses, Narrow Eyes, Smiling,
Mouth Slightly Open}.
The CelebA dataset is licensed for non-commercial research purposes only, and consists of publicly available images of celebrities, which were collected from the internet.
Although the data set has been widely used, Liu et al.
[2015] do not make any mention of consent by the individuals to have the images included in the data set, and it is therefore likely that those celebrities did not provide consent.
Training distribution over attributes For the training distribution, we simulate binary attributes according to the structural causal model in Figure 4 (for convenience also copied to Figure 12), where the model parameters are
P ( Young = 1) =
σ (0 . 0)
P ( Male = 1) =
σ (0 . 0)
P ( Eyeglasses = 1 | Young )
=
σ (0 . 0 − 0 .
4 · Young )
P ( Bald = 1 | Young, Male ) =
σ ( − 3 . 0 + 3 .
· Male − Young )
P ( Mustache = 1 | Young, Male ) =
σ ( − 2 .
5 + 2 . 5 · Male − Young )
P ( Smiling = 1 |
Young, Male )
=
σ (0 . 25 − 0 .
5 · Male + 0 .
5 · Young )
P ( Wearing Lipstick = 1 | Young, Male ) =
σ (3 . 0 − 5 .
· Male − 0 .
5 · Young )
P ( Mouth Slightly Open = 1 | Young, Smiling ) =
σ ( − 1 .
+ 0 .
· Young + Smiling )
P ( Narrow Eyes = 1 | Male, Young, Smiling ) =
σ ( − 0 .
5 + 0 .
3 · Male + 0 .
· Young + Smiling ) , where each variable either takes the value 0 or 1 and σ indicates the sigmoid.
To generate data, we first simulate attributes from this binary Bayesian network, which we then pass as inputs to the GAN to simulate images (in addition to the random noise used by the GANs to simulate different images).
In Figures 13 and 14, we plot examples of the training images that were generated.
Predictive model We simulate a training set of 12 , 000 attribute-image pairs, and a validation set of 2 , 000 pairs.
The training set is used to fit a classifier f , and the validation set is used for model selection.
To build a classifier f , we use the ResNet-
[He et al., 2016] model implemented in the python package torch .
We add a final fully connected layer to adapt the ResNet model to a binary classification task, and fine-tune the model on the training data by (only) learning the weights and bias of the final layer.
The model is trained using the negative log-likelihood criterion and an ADAM optimizer.
The model is trained for 25 epochs and we select the model which after a full epoch had the best validation set performance.
Given the learned model f , we simulate a separate validation dataset of n = 1 , 000 samples, and make model predictions f ( X ) .
We then compute the model accuracy as ℓ = 1 { f ( X )
= Y } , which is the input to computing the shift gradient and Hessian.
Estimation of shifted loss We apply the methods in Section 3.2 to estimate the worst-case shift to the distribution P (given by the binary probabilities above).
For each conditional P ( W i | PA( W i )) , we consider a shift η δ i (PA( W i ))
=
η (PA( W i ))
+ (cid:80) z ∈Z 1 { PA( W i )
=
z } δ i , which corresponds to arbitrarily shifting the conditional distribution (see Appendix C.4).
For example, for W i =
Bald, where η ( Young, Male ) = − 3 .
+ 3 .
· Male − 1 . 0 · Young, the shift would be
[Equation]
For each W i , this means that δ i is R 2 | PA( Wi ) | , and in total δ = ( δ 1 , . . .
, δ 8 ) ∈ R 31 (we do not consider shifts in the distribution of gender, since this is the label we are predicting).
We compute the shift gradient and Hessian using Theorem 1.
In particular, since W i is binary, the sufficient statistic is T ( W i ) =
W i , so the shift gradients and Hessians given by Appendix C.4.
See
Appendix C.1 for a detailed walk through of computing the shift gradient and Hessian from a sample.
For any given δ , the shifted distribution of W i is given by P δ
( W i =
| PA( W i ))
=
σ ( η δ i ) , where η δ i is computed similar to Equation (16), and σ is the sigmoid function.
Then the importance sampling weights are given by
[Equation]
η (PA( W i )) . Using these weights, for any δ , we can estimate E δ [ ℓ ] by ˆ E δ, IS and ˆ E δ, Taylor using Equations (6) and (8), respectively.
E.2 Full table of worst-case shift in Section 4.
In Section 4.2, we find the worst-case shift δ , and display the 5 largest components.
In Table 3, we display the full vector δ ∈ R 31 , sorted by absolute value of the size of the component.
E.3 Sample images from training distribution in Section 4.
In Figure 13, for the 4 attributes { Bald, Smiling, Wearing Lipstick, Male } , we display images generated from the training distribution P (i.e. by the GAN) with that particular attribute.
In Figure 14 we show 10 randomly drawn images from the training distribution P as well as the test distribution P δ corresponding to the worst-case δ found in Section 4.2. [Table]
E.4 Impact of changing λ
The shift considered in the main text yields a relatively small drop in accuracy.
To demonstrate that larger drops in accuracy are possible, we repeated our experimental setup over the same 100 initial validation datasets, while varying the size of the constraint ∥ δ ∥ 2 ≤ λ .
We report results in Table 4 for λ ∈ [2 , 4 , 6 , 8 , 10] , where λ = 2 corresponds to the setting of Table 1 (right). [Table]
Recall that we have two complementary goals: First [Table]
we would like to find a shift that results in a large drop in accuracy. Second, we would like to reliably evaluate the impact of the shift that we find, using only the training data. These two goals can be tackled with different approaches, such as using the Taylor approximation to find a shift
[Table] but using importance sampling (IS) to estimate the loss under that shift. Table 4 allows us to compare three different strategies: (i) using the Taylor approximation for both finding and evaluating the shift, (ii) using IS for both finding and evaluating, and (iii) using Taylor to find, but IS to evaluate the shift. From Table 4, we can observe that using Taylor to find, but IS to evaluate, consistently performs best in terms of reliable evaluation (i.e., predicting the shifted accuracy), across all values of λ .
For λ = 2 , the bias in evaluation is 1% (predicting 86% vs ground truth of 87% on average), and for λ = 10 , the bias of this approach is still only 2% (predicting 63% vs ground truth of 65% on average). In contrast, for λ = 10 , the first strategy (using Taylor to find and evaluate) over-predicts the impact by 18%, and the second strategy (using IS to find and evaluate) over-predicts the impact by 47%. This strategy also tends to find the most impactful shifts, for moderate values of λ .
For λ ≤ 6 , the shifts found by the Taylor approach are more impactful than those found by the IS approach. Moreover, the drop in accuracy remains substantial (e.g., a drop of around 17% at λ = 6 ). For λ > 6 , the story is more subtle: The third approach (using IS to find and evaluate shifts) finds more impactful shifts, but (as noted previously) dramatically over-estimates their impact. F Relationship to other approaches
In this section, we give a more detailed discussion of how our work relates to other approaches for evaluation of distributional robustness and learning of robust models. Much of the content from Section 1.1 is duplicated here, but expanded upon to include other relevant work and detailed discussion. Distributionally Robust Optimization/Evaluation with divergence measures : Distributionally robust optimization (DRO) seeks to learn models that minimize objectives of the form of Equation (1) [Duchi and Namkoong, 2021, Duchi et al., 2020, Sagawa et al., 2020].
We focus on proactive worst-case evaluation of a fixed model, not optimization, similar to Subbaswamy et al. [2021], Li et al. [2021], but major differences between our work and prior work lie in the definition of the set of plausible future distributions P , often called an “uncertainty set” in the optimization literature, where the goal is to specify a set that captures expected shifts, without being overly conservative. Shifts in P ( X, Y ) :
A conservative approach is to include all joint distributions P ( X, Y ) within a certain neighborhood of the training distribution.
Many coherent risk measures can be written as a worst-case loss of this form.
For instance, the Entropic Value-at-Risk (EVaR), with confidence level 1 − α , corresponds to the worst-case loss over a set of distributions P = { P ≪ P 0 : D KL ( P ∥ P 0 ) ≤ − ln α } , where P 0 is the original distribution [Ahmadi-Javid, 2012]. Similarly, the Conditional
Value-at-Risk (CVaR) with parameter α can be seen as the worst-case loss over an uncertainty set obtained from a limiting f -divergence (see Example 3 of Duchi and Namkoong [2021]), including all α -fractions of the original distribution.
These measures are appealing, in that they are straightforward to compute, but can be very conservative.
Indeed, such measures often reduce to only considering the distribution of the loss itself.
CVaR, for instance is equivalent to sorting the training examples by their loss, and taking the average loss of the top α -fraction.
To illustrate these limitations, it is straightforward to see that, using the 0-1 loss and a classifier with 80% accuracy, the worst-case loss under both of these measures is 1 . 0 for any α ≤ 0 .
This is intuitive for CVaR (since over 20% of samples are misclassified in the original distribution), and follows for EVaR from the fact that the binary distribution with probability q =has a KL-divergence to the original distribution p = 0 .
2 of − ln 0 .
Lam [2016] consider a more general problem of estimating the worst-case performance of stochastic systems over infinitesimal changes in distribution, measured by Kullback-Leibler divergence.
Their approach is applicable beyond machine-learning settings, and generalizes to e.g., worst-case waiting times in a queueing system.
They demonstrate that for a sufficiently small neighborhood of distributions, this worst-case performance can be well-approximated by a Taylor expansion whose coefficients can be estimated from the original distribution.
Shifts in P ( X ) alone : Partially due to this overly-conservative behavior, there has been a line of work incorporating additional restrictions on the allowable shift (i.e., adding more assumptions). For instance, Duchi et al. [2020] considers learning predictive models that optimize a worst-case loss similar to CVaR (a “worst-case subpopulation shift”), but where only P ( X ) is allowed to change, and P ( Y | X ) is assumed to be constant. For similar shifts, Li et al.
[2021] considers only the task of evaluation, but provides a novel estimation procedure with dimension-free finite-sample guarantees.
However, many real-world shifts do not fit this framework: In Example 1, both P ( X ) and P ( Y | X ) are changing, where X = ( A, O, L ) , as a result of a shift in P ( O | Y, A ) .
Shifts in a conditional distribution : Closer to our work is Subbaswamy et al.
[2021] who consider evaluating the loss under worst-case changes in a conditional distribution, but while we consider parametric shifts, they estimates the loss under worst-case (1 − α ) conditional subpopulation shifts.
However, it is not obvious how to choose an appropriate level of α : in some settings, seemingly plausible values of α (e.g., a 20% subpopulation) correspond to entirely implausible shifts.
We give a simple lab-testing example in Appendix D, where the worst-case subpopulation is one where healthy patients are always tested, and sick patients never tested.
In contrast to these methods, our approach uses explicit parametric perturbations to define shifts, as opposed to distributional distances or subpopulations.
In addition, our approach allows for shifts in multiple marginal or conditional distributions simultaneously:
In Example 1, for instance, we can model a simultaneous change in both the marginal distribution of age, as well as the conditional distribution of lab testing, while other conditionals are unchanged.
Our main requirement is that each shifting distribution is exponential family, and that the shift can be represented via the natural parameters: For continuous variables this is a non-trivial restriction, but for discrete variables it is true by definition.
Causality-motivated methods for learning robust models
: Several approaches seek to learn models that perform well under arbitrarily large causal interventions (which result in arbitrary changes in selected conditional distributions).
Several approaches proactively specify shifting mechanisms/conditional distributions, and then seek to learn predictors that have good performance under arbitrarily large changes in these mechanisms [Subbaswamy et al., 2019, Veitch et al., 2021,
Makar et al., 2022, Puli et al., 2022].
Less relevant to our setting are approaches that use auxiliary information, such as environments [Magliacane et al., 2018, Rojas-Carulla et al., 2018, Arjovsky et al., 2019] or identity indicators [Heinze-Deml and Meinshausen, 2021] to learn models that rely on invariant conditional distributions.
The worst-case optimality of these approaches is often restricted to cases where the shifts are arbitrarily large: In Example 1, this would correspond to minimizing worst-case loss over all possible lab testing policies.
However, when the causal interventions (i.e., changes in causal mechanisms) are bounded (i.e., not arbitrary), then these approaches are not necessarily optimal.
Closest to our work in motivation is prior work on robustness to bounded shift interventions in linear causal models [Rothenhäusler et al.,
2021, Oberst et al., 2021, Kook et al., 2022].
Our work can be seen as extending those ideas to general non-linear causal models, where our focus is on evaluation rather than learning robust models.
Our work can serve as an aid to deploying these causality-motivated methods in a few ways, by comparing their worst-case performance under bounded shifts:
First, our work can inform whether such methods should be deployed at all, as for sufficiently small shifts, it may be the case that standard training yields better performance.
Second, our work can inform hyperparameter selection for several of these approaches, which include regularization terms that implicitly trade off between robustness and in-distribution performance.
More broadly, our approach is useful for probing (and comparing) the reliability of specific learned models under shift, regardless of the algorithm that produced them.
Evaluating out-of-distribution performance with unlabelled samples : A recent line of work has focused on predicting model performance in out-of-distribution settings, where unlabelled data is available from the target distribution
[Garg et al., 2022, Jiang et al., 2022, Chen et al., 2021].
In contrast, our method operates using only samples from the original source distribution, and seeks to estimate the worst-case loss over a set of possible target distributions.
G Proofs G.1 Proof of Proposition
Proposition 1.
For any P δ ( V ) , P ( V ) that satisfy Definition 4, supp( P ) =
supp( P δ ) and the density ratio w δ :=
P δ / P is given by
[Equation]
Proof.
By Definition 4 and Assumption 1, we have that
[Equation]
It follows that the supports of P δ and P are the same: Since the exponential family density is given by the base measure g i ( W i ) times a exponential term (which is always strictly positive), and since the terms (cid:81) V j ∈ V \ W P
( V j | U j ) are shared between P δ and P , their supports agree.
To get the density ratio, we take the ratio of P δ ( V ) and P ( V ) , and the terms V j ∈ V \ W cancel:
[Equation]
By Definition 4 and Assumption 1, each P δ i ( W i | Z i ) is a δ i -perturbation around the CEF distribution P ( W i | Z i ) , so plugging in the exponential family densities, we get w δ [Equation] i = [Equation]
s i [Equation] exp [Equation]
[Equation]
G.2 Proof of Theorem
Theorem 1 (Shift gradients and Hessians as covariances) .
Assume that P δ , P satisfy Definition 4, with intervened variables W =
{ W 1 , . . .
, W m } and shift functions s i ( Z i ; δ i ) , where δ =
( δ 1 , . .
δ m ) .
Then the shift gradient is given by SG
= (SG 11 , . .
, SG 1 m ) ∈ R
d δ where
[Equation]
and the shift Hessian is a matrix of size ( d δ × d δ ) , where the ( i, j ) th block of size d δ i ×
d δ j equals
[Equation]
̸
[Equation]
where D i,k := ∇ kδ
i s i
( Z i ; δ i )
| δ =0 , is the gradient of the shift function for k = 1 , and the Hessian for k = 2 .
Here, T i ( W i ) is the sufficient statistic of P ( W i | Z i ) and ϵ T i |
Z i :
=
T i ( W i ) − E
[ T ( W i ) | Z i ] .
Proof.
For simplicity throughout, we use h (1) i to denote the gradient of the log-partition function ∇ h
i ( · ) with respect to the arguments, which is a column vector of length d T i , and we use h (2) i to denote the Hessian ∇ 2 h i ( · ) , which is a matrix of size d T i × d T i .
We also use η δ i ( z i ) as short-hand for η
i ( z i ) + s i ( z i ; δ i ) .
Shift Gradient : By Definition 4, the probability density / mass function P δ factorizes as follows, where δ = ( δ 1 , . .
, δ m
[Equation]
and the gradient with respect to shift parameters δ i is given by
[Equation]
where the last equality follows from additivity of the log-likelihood in the conditionals, the factorization above, and the fact that δ i only enters into the given conditional distribution.
Given the assumed form of log p δ i ( w i | z i ) given in Definition 3, we can observe that
[Equation]
18) where ∇ δ i s i
( z i ; δ i ) ∈
R d Ti × d δi , and ∇ h
i ( η ( z i ) + s i ( z i ; δ i )) is the gradient of the function
h
i : R d
T i → R , which is a column vector of length d T i .
It follows from known properties of the logpartition function [Wainwright et al., 2008, Proposition 3.1], that h (1) i ( η δ i ( z i )) =
E δ
[ T i ( W i )
| z i ] .
This gives us that
[Equation]
where the second equality follows from the tower property and Z i -measurability of ∇
δ i s i
( Z i ; δ i ) , and the final equality follows from the definition of the conditional covariance.
This expression, evaluated at δ = 0 , gives us the desired result, that
[Equation]
where D i, 1 = ∇ δ
i s i
( Z i , δ i ) | δ =0 .
The result follows from the definition that gradients are taken entry-wise, giving SG {+
+} = (SG 11 , . .
, SG 1 m ) ∈ R
d δ 1 + ··· d δm .
Shift Hessian (Diagonal) :
For the shift Hessian, we first compute the diagonal entries of ∇ 2 δ E δ [ ℓ ] | δ =0 , which are blocks of size R d δi × d δi . We begin by computing the Hessian of the likelihood.
[Equation]
where we use the notation v ⊗ 2 :
= vv ⊤
[Equation]
and we note that ∇ 2 δ i s
( z i ; δ i ) is a tensor of size d T i ×
d δ i ×
d δ i , and {∇ 2 δ
i s i
( z i ; δ i ) }
⊤ h (1) i ( · ) is a matrix of size d δ i ×
d δ i , where the ( m, n ) ’
th entry is { ∂ ∂δ im ∂∂δ in s ( z i ; δ i ) }
⊤ h (1) ( · ) .
Now, using the fact that h (1) ( η δ i ( z )) =
E δ
[ T i ( W i )
| z i ] and h (2) (
η δ i ( z i ))
=
var δ
[ T i ( W i )
| z i ] [Wainwright et al., 2008, Proposition 3.1], and the definition ϵ T i | Z i
= T i ( W i ) − E δ
[ T i ( W i )
| Z i ] , we obtain
[Equation]
which gives the desired result when we evaluate at δ = 0 .
Shift Hessian (Off-Diagonal)
For i ̸
= j , we have that
[Equation]
where the third line follows from the fact that ∇ δ i ( ∇ δ j log p δ j ( w j | z j ))
= 0 , and the last line follows from the derivation of the gradient of the log-likelihood in Equation (18).
We can again use the fact that h (1) i ( η δ i ( Z i )) =
E δ
[ T i ( W i ) | Z i ] and the shorthand ϵ T i
| Z i :
=
T i ( W i ) − E δ
[ T i ( W i )
|
Z i ] to write that
[Equation]
and when we evaluate this expression at δ = 0 , we obtain
[Equation]
Where the last equality follows because E [ D ⊤ i, 1 ϵ T i | Z
i ϵ ⊤ T
j | Z j D j,i ]
= 0 .
To see this, note that one of
W i , W j must be a non-descendant of the other, and we will assume without loss of generality that W j is a non-descendant of W i in the causal graph consistent with the factorization given in Equation (17), which implies that Z j (the parents of W j in the underlying graph) are also non-descendants of
W i .Thus, W i ⊥⊥ ( W j , Z j ) | Z i , because ( W j , Z j ) are both non-descendants of W i .
Then, observe that D i, 1 is a function of Z i , and
ϵ T i | Z i is a variable with zero-mean conditioned on Z i .
Thus,
E [ D ⊤ i, 1 ϵ T i | Z i | Z i ] = 0 , for all Z i .
Moreover, given Z i , we have that D ⊤ i, 1 ϵ T i | Z i is independent of
D ⊤ j, 1 ϵ T j | Z j .
As a result, we can write that
[Equation]
G.3 Proof of Corollary
Corollary 1 (Simple shift in a single variable) .
Assume the setup of Theorem 1, restricted to a shift in a single variable W , and that s ( Z ; δ )
= δ .
Then D 1 = 1 , D 2 = 0 , and
[Equation]
and
[Equation]
where T ( W ) is the sufficient statistic of W and ϵ T | Z :
= T ( W ) − E
[ T ( W ) | Z ] .
Proof.
We have ∇ δ s ( Z ; δ ) = ∇ δ δ = 1 and ∇ 2 δ s ( Z ; δ )
= ∇ 2 δ δ = 0 .
The result now follows from
Theorem 1.
G.4 Proof of Theorem
Theorem 2.
Assume that P δ , P satisfy the conditions of Theorem 1, with a shift in a single variable
W , where s ( Z ; δ ) =
δ .
Let E δ, Taylor be the population Taylor estimate (Equation (7) ) and let σ ( M ) denote the largest absolute value of the eigenvalues of a matrix M .
Then
[Equation]
where T ( W ) is the sufficient statistic of W | Z and ϵ t · δ,T | Z = T ( W | Z ) − E t · δ [ T ( W | Z )] .
Proof.
The expectation is continuous and twice-differentiable with respect to δ , because of the smoothness of the exponential family in the parameter, the fact that the shift function s is twicedifferentiable, and because the support does not change.
Thus, applying Taylors remainder theorem to the function t (cid:55)→ E t · δ [ ℓ ] , it follows that there exist a t 0 ∈ [0 , 1] such that
[Equation]
We have (cid:18) dd
[Equation]
E
[Equation]
=
= SG 1 and by the same arguments (see the proof of Theorem 1), it follows that
[Equation]
δ .
Plugging this in, and subtracting 12 δ ⊤ SG 2 δ on both sides of Equation (19) yields
[Equation]
Let K := (cid:18)
[Equation]
Since K is symmetric and real valued, it is diagonalizeable, K = U ⊤ Λ U for an orthonormal matrix U and diagonal matrix Λ = diag
( α 1 , . . . , α d ) .
We then have
[Equation]
where Λ 1 / 2 = diag( √ α 1 , . .
. , √ α d ) , ∥·∥ 2 denotes the supremum-norm when applied to matrices and the 2 -norm when applied to vectors and ∥ Uδ ∥ 2 = ∥
δ ∥ 2 because ∥
Uδ ∥ 22 = δ ⊤ U ⊤ Uδ =
δ ⊤ δ = ∥
δ ∥ 22 , using orthonormality of U .
Plugging in this inequality, we get that
[Equation]
which concludes the proof.
G.5 Proof of Proposition B.
[-of Proposition B.
-]
Proposition B.1.
Consider a binary random variable W with conditional distribution
[Equation]
for an arbitrary measurable function η ( Z ) whose range is the extended real numbers η ( Z ) ∈[-
-] R ∪{ + ∞ , −∞} .
Let p + :=
P ( η ( Z ) = + ∞ ) , p − := P ( η ( Z ) = −∞ ) , and assume that p + + p − < 1 .
Then, the marginal probability p
[Equation] = P δ ( W = 1) is a strictly monotonically increasing function of δ ∈ R whose range is ( p + , 1 − p − ) ,
Proof.
Let F denote the event that η ( Z ) is finite (i.e., η ( Z ) ̸∈ {−∞ , + ∞} ).
Under F , the conditional probability function σ ( η ( Z ) + δ ) is a strictly monotonically increasing function of δ , and if η ( Z ) ∈ {−∞ , + ∞} , then the conditional probability is a constant function of δ (zero or one, respectively).
Hence, we can write that
[Equation]
and by assumption, 1 − p + − p − > 0 .
The marginal probability P δ ( W = 1 | F ) is a strictly monotonically increasing function of δ , with a limit of 1 as δ → ∞ , and a limit of 0 as δ → −∞ .
As a result, it is bounded in ( p + , 1 − p − ) .
G.6 Proof of Lemma C.
Lemma C.1.
Suppose A ∼ N ( µ, Σ) and that ( X, Y, H ) are generated according to Equation (11) .
For γ ∈ R d X define ℓ :=
( Y − γ ⊤ X ) 2 .
Then there exist v γ , u
µ,γ ∈ R d A such that for all shifts δ ∈ R d A :
[Equation]
where E δ corresponds to taking the mean in the distribution where A ∼ N ( µ + δ, Σ) .
Further u µ,γ =
if µ = 0 .
Proof.
It follows from Equation (11) that one can write ( X ⊤ , Y ⊤ , H ⊤ ) =
(1 − B ) − 1 ( MA + ϵ ) , and for a given γ , there exist b γ , κ γ
such that Y − γ ⊤ X =
b ⊤ γ A + κ ⊤ γ ϵ
[Rothenhäusler et al., 2021].
In P δ , we can write A = µ + δ + ϵ A , where ϵ A ∼ N (0 , Σ) , for all values of µ and δ
Plugging this in yields
[Equation]
where we do not put a subscript on the expectation in the third line because it is taking expectations over ϵ A and ϵ , both which do not depend on the choice of µ and δ .
The statement of the lemma follows by letting u µ,γ = 2 b ⊤ γ µ and v γ = √ 2 b γ .
G.7 Proof of Proposition C.
[-of Proposition C.
-]
Proposition C.1.
Suppose A ∼ N ( µ, Σ) and that ( X, Y, H ) are generated according to Equation (11) .
Then the shift gradient and Hessian are given by
[Equation]
and
[Equation]
and the loss under a mean shift of δ in A is given by
[Equation]
where ℓ :=
( Y − γ ⊤ X ) 2 and E δ corresponds to taking the mean in the distribution where A ∼
N ( µ + δ, Σ) .
Proof.
Similar to Lemma C.1, we rewrite Y − γ ⊤ X =
b ⊤ γ
A
+ κ ⊤ ϵ , and by rewriting A = µ + δ + ϵ A ,
where ϵ A ∼ N (0 , Σ) , we obtain
[Equation]
We recognize that Equation (20) equals E ( Y − γ ⊤ X ) 2 .
Similarly, we now show that Equations (21) and (22) match the shift gradients (multiplied appropriately with δ ).
First, we assume that Σ
= Id .
Since A is a Gaussian with (known) mean Id , the sufficient statistic is
T ( A ) =
A .
Hence, according to Theorem 1, we can compute the shift gradient as
[Equation]
We can calculate the i ’th entrance of this vector as:
[Equation]
where in the first equality we use that subtracting a constant doesn’t change the covariance, and we use independence of A i from A j A j ′ when i / ∈ { j, j ′ } .
Using the assumption that A i has unit variance, we now get that
[Equation]
By plugging in, we obtain
[Equation]
Since this was element-wise, we obtain that the full vector is SG
= 2 b γ b ⊤ γ µ , which, when multiplied with δ yields Equation (21).
We compute SG 2 similarly.
The diagonal entries are given by
[Equation]
Because Σ = Id , the second through fourth moments of A i are given by E
[ A 2 i ]
= µ 2 i + 1 , E
[ A 3 i ]
= µ 3 i + 3 µ i and E
[ A 4 i ]
= µ 4 i + 6 µ 2 i
+ 3 .
Using this, we get
[Equation]
and for j ̸ = i :
[Equation]
using linearity of the covariance, that A i ⊥⊥ A j and that the first and third moments are zero for a centered Gaussian A i − µ i .
Plugging this in, we get that the diagonal entries are given by
[Equation]
We can compute the off-diagonal entries similarly.
For i ̸
= j , we have:
[Equation]
Using the independence of A i and A j , we have
[Equation]
= 0 , and similarly cov(( A i − µ i )
( A j − µ j ) , A 2 j ) = 0 .
Using the same reasoning, for v / ∈ { i, j }
[Equation]
E [ A v ] = 0 , and the same for cov(( A i − µ i )(
A j − µ j ) , A j A v ) .
Finally, we have
[Equation]
Plugging into Equation (23), we get that
[Equation]
and hence for both diagonal and off-diagonal entries, SG 2 i,j = 2 b γ,i b γ,j , implying that
[Equation]
In particular 12 δ ⊤ SG 2 δ matches Equation (22).
Finally, we consider the case Σ ̸
= Id . Let Σ − 1 / 2 be the ‘square-root’ of Σ − 1 , such that Σ − 1 / 2 Σ −⊤ /(where the latter denotes (Σ − 1 / 2 ) ⊤ .
The sufficient statistics for the mean in a multivariate Gaussian distribution with known variance is given by T ( A ) =
Σ − 1 A .
We then have
[Equation]
where
˜ A = Σ − 1 / 2 A = ∼ N (˜ µ, Id) , ˜ µ = Σ − 1 / 2 µ and ˜ b γ = Σ 1 / 2 b γ .
In particular, since ˜ A has unit variance, we can use the above derivations to obtain
[Equation]
In particular, the first shift gradient is the when Σ ̸
= Id as when Σ = Id .
Similarly,
[Equation]
Hence, also when Σ ̸
= Id , the terms of Equations (21) and (22) matches the expression given by SGand SG 2 .
This concludes the proof.