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.
To ensure that these shifts are plausible, we parameterize them in terms of interpretable changes in causal mechanisms of observed variables.
This defines a parametric robustness set of plausible distributions and a corresponding worst-case loss.
We construct a local approximation to the loss under shift, and show that problem of finding worst-case shifts can be efficiently solved.
1. Introduction
Predictive models may perform poorly outside of the training distribution, a problem broadly known as dataset shift (Qui˜nonero-Candela et al., 2008).
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.
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.
To illustrate, we use a running example inspired by Subbaswamy et al.
Example 1.1 (Shifts in laboratory testing) .
We seek to classify disease ( Y ) based on the age ( A ) of a patient, whether a lab test has been ordered ( O ) , and test results ( L ) if 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 and Stability .
Baltimore, Maryland, USA.
Copyright 2022 by the author(s).
[Figure]
{+Test result δ age δ order 4 2 0 20 4 2 
+}
L o ss un d e r ( 0 , 1 ) 0.30.40.50.60.
variables. Let P ( O | 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 | A, Y ) =
σ ( η ( A, Y )
+ s ( Y ; δ )) for a shift function s ( Y ; δ ) =
δ 1 · Y + δ 0 · (1 − Y ) , which modifies the logodds of testing for both sick and healthy patients.
If δ 0 , δare unconstrained, the worst-case occurs when all healthy patients are tested, and no sick patients are tested.
We address two challenges in this paper 1 : The first challenge (Section 2) is to define a set of possible distributions
P where each P ∈ P is (i) causally interpretable and simple to specify , without unnecessary restrictions on the data-generating process, and (ii) realistic , which often entails bounding the magnitude of shift.
We construct causally interpretable shifts by defining perturbed distributions P δ via changes in causal mechanisms parameterized by a finitedimensional parameter δ .
Our main requirement is that the shifting mechanisms follow a conditional exponential family distribution.
Constraining δ can then ensure that shifts are realistic:
The unconstrained worst-case shift in Example 1.1 is 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 (Section 3) 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 nonconcave 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.
For quadratic constraints ∆ , we can approximate the general non-convex optimization problem in Equation 2 with a 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 experimentally that it is well-behaved.
2. Defining parametric robustness sets
Notation : Let V denote all observed variables, where ( X, Y ) ⊆
V for features X and labels Y , and let P ( V ) denote the training distribution.
E [ · ] and cov( · , · ) refer to the mean and covariance in P , and for a shifted distribution
P δ
(Definition 2.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 2.1.
A parameterized robustness set around P ( V ) is a family of distributions P with elements
P δ ( V ) indexed by δ ∈ ∆ ⊆
R d δ , with 0 ∈ ∆ , and 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 2.3).
We require that the relevant conditional distributions can be described by an exponential family.
Definition 2.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 a normalizing constant.
Definition 2.2 extends to marginal distributions where Z = ∅ and η ( Z ) is a constant function.
Example 1.1 (Continued) .
Let ordering a test ( O ) depend 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.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 2.3 (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 sets of variables.
We assume for each W i that Z i is known and
that P ( W i | Z i ) satisfies Definition 2.2.
If P ( V ) factorizes according to a DAG G , the factorization in Assumption 2.3 is always satisfied by Z i = PA G ( W i ) .
Here, we require limited knowledge of the underlying graph, and only need to know the parents PA( W i ) for the variables
W i that shift.
We now define parametric perturbations and give the general form of the robustness sets that we consider, involving simultaneous perturbations to multiple W i .
Definition 2.4 (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 satisfy s ( Z ; 0) = 0 for all Z .
For P ( W | Z ) satisfying Equation (3), we refer to
[Equation]
as a δ -perturbation of P ( W | Z ) with shift function s ( Z ; δ ) , where η δ ( Z ) :
= η ( Z ) + s ( Z ; δ ) .
Example 1.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 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.2 and 2.4.
Note that this formulation includes multiplicative shifts η δ ( Z ) =
(1 + δ ) η ( Z ) by letting s ( Z ; δ ) =
δ · η ( Z ) .
Definition 2.5 (CEF parameterized robustness set) .
For a distribution P and intervention set W = { W 1 , . . .
, W m } ⊆
V satisfying Assumption 2.3, let each P
δ i ( W i | Z i ) be a δ iperturbation (Definition 2.4) 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 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 2.6 (Causal Interpretation of Shifts) .
If the DAG G represents a causal graph (Pearl, 2009), then P δ can interpreted as a change in causal mechanisms.
We see this as an important perspective for interpreting and specifying shifts, but our methods do not require a causal interpretation.
3. 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.
Remark 3.1 .
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 3.2.
For any P δ ( V ) , P ( V ) that satisfy Definition 2.5, supp( P ) =
supp( P δ ) and the density ratio w δ :=
P δ / P is given by
[Equation]
Example 1.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 use an importance sampling (IS) estimator (Horvitz & 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]
In practice, Equation (5) 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 secondorder 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 3.3 (Shift gradient and Hessian) .
For a parametric shift satisfying Definition 2.1 where δ (cid:55)→ E
δ
[ ℓ ] is twicedifferentiable, we denote the shift gradient SG 1 and Hessian
SG 2 as SG 1 := ∇ δ E δ [ ℓ ] (cid:12)(cid:12) δ =0 and SG 2 := ∇ 2 δ E δ [ ℓ ] (cid:12)(cid:12) δ =0 .
Equation (6) is a local approximation of the loss, whose approximation error we bound in Theorem B.2, with smaller approximation error for smaller shifts.
For P δ satisfying Definition 2.5, 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 3.4 (Simple shift in a single variable) .
Assume the setup of Theorem B.1, restricted to a shift in a single variable W , and that s ( Z ; δ )
= δ .
[Equation]
and
[Equation]
where T ( W ) is the sufficient statistic of W and ϵ T | Z :=
T ( W ) − E
[ T ( W ) | Z ] .
In Theorem B.1, we state a general form of Theorem 3.4, which allows for shifts in multiple variables and for arbitrary shift functions s ( Z ; δ ) .
Example 1.1 (Continued) .
Suppose that age ( A ), which has no causal parents, follows a normal distribution with mean µ and variance σ 2 , and that we wish to consider a shift in the mean.
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 , we fit models ˆ µ ℓ ( Z )
≈ E [ ℓ | Z ] and ˆ µ W ( Z )
≈ E [ T ( W ) | Z ] and compute residuals on these predictions, which permits estimation of the gradient and Hessian as a sample average of residuals.
Using these, we can 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; in Theorem B.2 we give a bound on the latter.
In exchange for considering a second-order approximation of the loss, we gain two benefits: Variance reduction and tractable optimization.
First, as SG 1 , SG 2 are not functions of δ , the variance of ˆ E δ, Taylor is O ( ∥ δ ∥ 4 ) , while the variance of ˆ E δ, IS can be much larger (see Section 4).
Second, maximizing ˆ
E δ, Taylor over the set ∥ δ ∥ ≤ λ can be solved in polynomial time by exploiting the quadratic structure (see Section 3.3), 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]
Since SG 2 is generally not negative definite, the maximization objective is non-concave.
However, this particular problem is an instance of the well-studied ‘trust region problem’ (Conn et al., 2000), which can be solved in polynomial time (P´olik & Terlaky, 2007, Section 8.1).
4. Experiments
Synthetic Example: Lab Tests : To compare the bias and variance of the Taylor and the importance sampling estimates of the shifted loss, we simulate synthetic data from Test [Figure] Age
[Figure]
Example 1.1.
We consider either a shift in the logits of ordering lab tests (Figure 2 left) or a mean shift in the Gaussian distribution of age (Figure 2 right).
We compute estimates ˆ
E δ, IS and ˆ E δ, Taylor of the loss under a shift of size δ , and compare this to ground truth data simulated from P δ .
Additional details are given in Appendix C.
For shifts in binary test ordering, both estimates capture the loss well for small shifts, but as δ gets larger, the quadratic approximation increasingly deviates from the true mean.
For the Gaussian mean shift, the importance sampling weights quickly become ill-behaved, and the variance dramatically increases as δ grows.
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.
Finding worst-case shifts in computer vision, given image attributes : In Appendix A we demonstrate the application of the approach outlined in Section 3.3, to find bounded worst-case shifts in a gender classification task using a synthetic variant of the CelebA dataset.
Here, we uncover (and validate) sensitivity to interpretable changes in distribution, such as reduced rates of wearing lipstick among young women.
5. 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.
0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods , 17:261–272, 2020. Age Gender
Eyeglasses
Bald
Mustache
Smiling Wearing Lipstick
[Figure]
Wainwright, M. J., Jordan, M. I., et al.
Graphical models, exponential families, and variational inference.
Foundations and Trends® in Machine Learning , 1(1–2):1–305, 2008.
A. Experiment: Detecting sensitivity to spurious 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 (Liu et al., 2015), containing images of faces along and binary attributes (e.g., glasses, beard, etc.)
encoding several non-causal features whose correlations may be unstable (e.g., the relation between gender and wearing lipstick).
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 3, 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), with a test accuracy of 0.919.
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 .
Per Section 3.3, we identify the shift δ causing the largest drop in accuracy, and compute ˆ E δ, Taylor and ˆ E δ, IS to estimate accuracy under that shift.
Due to the synthetic nature of our setup, we can simulate from P δ
( X, W , Y ) to estimate the ground-truth impact of this shift, by simulating from the shifted attribute distribution, and then simulating images from the GAN conditional
[Table]
on those attributes.
Finding and validating a single, high impact shift : Using a validation sample from P [Table] we estimate the shift gradient and Hessian (Theorem B.1).
Solving the quadratic optimization problem in Section 3.
[Table]
we find the worst-case shift δ such that ∥ δ ∥ ≤ λ = 2 .
We display the largest components of δ in Table 1 (top).
Among others, this shift entails a 10% increase in the probability of a young man wearing lipstick, and a similar decrease for young women.
This suggests that the learned classifier f relies on this non-causal association in the images for prediction.
We validate that this shift leads to a measurable decrease in accuracy, from 92% to 89%, using simulated data from P δ .
To validate that this shift is indeed a worst-case shift, we simulate K = 400 random shifts δ k and evaluate the model accuracy in P
δ k (Figure 4, top).
As expected, our chosen shift δ (red line) is in the left tail of the distribution.
We compare the ground-truth accuracy under this shift (89%) to the original estimates ˆ
E δ, Taylor = 0 .
878 and ˆ E δ, IS = 0 .
792 (Table 1), and observe that both correctly predict a drop in accuracy, although ˆ E δ, IS overestimates the size of the drop.
Comparing importance sampling and Taylor across multiple simulations : We simulate K = 400 validation sets from P , in each estimating the worst-case shifts δ Taylor and δ IS , where the latter corresponds to maximizing ˆ 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 , and in Figure 4 (bottom) we plot the differences E δ Taylor [ ℓ ] − E δ IS [ ℓ ] , showing that in 73% of cases, the
Taylor method finds a more impactful shift.
Moreover, the average run-time for the Taylor approach is 0 .
02 s while that of the importance sampling approach is 2 .
52 s .
Finally, the optimal value of the ˆ E δ, Taylor objective tends to a reasonably accurate estimate of the shifted accuracy, while the optimal
[Figure]
value of the ˆ E δ, IS objective (both on the validation set) is a poor predictor, with a mean absolute prediction error (in predicting E δ [ ℓ ] ) of 0 .
176 for ˆ E δ, IS and 0 .
for ˆ E δ, Taylor .
B. Additional theorems
Theorem B.1 (Shift gradients and Hessians as covariances) .
Assume that P δ , P satisfy Definition 2.5, 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, i ) th diagonal block of size d δ i ×
d δ i equals
E
[Equation]
Z [Equation] and the ( i, j ) th off-diagonal block of size d δ i × d δ j equals cov( ℓ, D ⊤ i, 1 ϵ T i | Z
i ϵ ⊤ T
j | Z j D j [Equation] 1 ) 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 B.2.
Assume that P δ , P satisfy the conditions of Theorem B.1, with a shift in a single variable W ,
where s ( Z ; δ ) =
δ .
Let E δ, Taylor be the population Taylor estimate (Equation (6) ) 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 ϵ
[Equation] T ( W [Equation]
C. Simulation details for Section
The data in Section 4 (and also the loss land-scape in Figure 1 (right)), was simulated using the following generative model
[Equation]
Result | Order = 1 , Disease ∼ N ( − 0 .
5 + Disease , 1) , where σ denotes the sigmoid function and if Order = 0 , the test result is a placeholder value of zero.
For the lab test ordering, we consider the shift η δ ( Z ) =
η ( Z )+ δ (i.e. a linear shift in the logits of test ordering)
and for age, we consider a shift in the marginal mean η δ = δ .
To construct the plot, we simulate data n = 1 , 000 times, and for each dataset and each δ in a grid, we compute estimates ˆ
E δ, IS and ˆ E δ, Taylor .
We then plot the mean and point-wise 90% prediction intervals for ˆ E δ, IS and ˆ E δ, Taylor .
D. Proofs D.1. Proof of Proposition 3.
Proposition 3.2.
For any P δ ( V ) , P ( V ) that satisfy Definition 2.5, supp( P ) =
supp( P δ ) and the density ratio w δ :=
P δ / P is given by
[Equation]
Proof.
By Definition 2.5 and Assumption 2.3, 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 2.5 and Assumption 2.3, each P δ i ( W i | Z i ) is a δ i -perturbation around the CEF distribuition P ( W i | Z i ) , so plugging in the exponential family densities, we get w δ ( V )
[Equation]
D.2. Proof of Theorem B.
Theorem B.1 (Shift gradients and Hessians as covariances) .
Assume that P δ , P satisfy Definition 2.5, 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, i ) th diagonal block of size d δ i ×
d δ i equals
E
[Equation]
and the ( i, j ) th off-diagonal 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 ] .
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 2.5, the probability density / mass function P δ factorizes as follows, where δ = (
[Equation]
[Equation]
m )
[Equation]
and the gradient with respect to shift parameters
δ i is given by
[Equation]
where the last equality follows from additivity of the loglikelihood 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 2.4, we can observe that
[Equation]
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 log-partition 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 ∇ δ i E δ [ ℓ ] = E δ (cid:104) ℓ · ( ∇ δ i s i ( Z i ; δ i )) ⊤ ( T i ( W i ) − E δ [ T i ( W i )
| Z i ]) (cid:105) =
E δ (cid:20) ( ∇ δ
i s i ( Z i ; δ i )) ⊤ E δ [ ℓ · ( T i ( W i ) − E δ [ T i ( W i ) | Z i ]) | Z i ] (cid:21) =
E δ (cid:104) ( ∇ δ
i s i ( Z i ; δ i )) ⊤ cov δ ( ℓ, T i ( W i ) | Z i ) (cid:105) , 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
[Equation]
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 ⊤ , and we note that ∇
δ
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]
E
[Equation]
[ ℓ ]
[Equation]
which gives the desired result when we evaluate at δ = 0 .
Shift Hessian (Off-Diagonal)
For i ̸
= j , we have that
[Equation]
p
[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 (10).
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 ∇ δ i ∇ δ j E δ [ ℓ ]
[Equation]
δ
[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 (9), 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]
D.3. Proof of Theorem 3.
Theorem 3.4 (Simple shift in a single variable) .
Assume the setup of Theorem B.1, restricted to a shift in a single variable W , and that s ( Z ; δ )
= δ .
[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 B.1.
D.4. Proof of Theorem B.
Theorem B.2.
Assume that P δ , P satisfy the conditions of Theorem B.1, with a shift in a single variable W ,
where s ( Z ; δ ) =
δ .
Let E δ, Taylor be the population Taylor estimate (Equation (6) ) 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
| W = T ( W | Z ) − E t · δ [ T ( W | Z )] .
Proof.
The expectation is continuous and twicedifferentiable with respect to δ , because of the smoothness of the exponential family in the parameter, the fact that the shift function s is twice-differentiable, 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
[Equation]
t [-=
-]=
[Equation]
= SG 1 and by the same arguments (see the proof of Theorem B.1), it follows that
[Equation]
this in, and subtracting 12 δ ⊤ SG 2 δ on both sides of Equation (11) yields [Equation] (cid:12)(cid:12)(cid:12)(cid:12) [Equation]
(cid:19)
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
Let K := (cid:18)
[Equation]
(cid:19) .
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.