Question: Predictive models that generalize well under distributional shift are often desirable and sometimes crucial to machine learning applications. One example is the estimation of treatment effects from observational data, where a subtask is to predict the effect of a treatment on subjects that are systematically different from those who received the treatment in the data. A related kind of distributional shift appears in unsupervised domain adaptation, where we are tasked with generalizing to a distribution of inputs that is different from the one in which we observe labels. We pose both of these problems as prediction under a shift in design. Popular methods for overcoming distributional shift are often heuristic or rely on assumptions that are rarely true in practice, such as having a well-specified model or knowing the policy that gave rise to the observed data. Other methods are hindered by their need for a pre-specified metric for comparing observations, or by poor asymptotic properties. In this work, we devise a bound on the generalization error under design shift, based on integral probability metrics and sample re-weighting. We combine this idea with representation learning, generalizing and tightening existing in this space. Finally, we propose an algorithmic framework inspired by our bound and verify is effectiveness in causal effect estimation. A long-term goal in artificial intelligence is for agents to learn how to act. This endeavor relies on accurately predicting and optimizing for the outcomes of actions, and fundamentally involves estimating counterfactuals-what would have happened if the agent acted differently? In many applications, such as the treatment of patients in hospitals, experimentation is infeasible or impractical, and we are forced to learn from biased, observational data. Doing so requires adjusting for the distributional shift between groups of patients that received different treatments. A related kind of distributional shift arises in unsupervised domain adaptation, the goal of which is to learn predictive models for a target domain, observing ground truth only in a source domain. In this work, we pose both domain adaptation and treatment effect estimation as special cases of prediction across shifting designs, referring to changes in both action policy and feature domain. We separate policy from domain as we wish to make causal statements about the policy, but not about the domain. Learning from observational data to predict the counterfactual outcome under treatment B for a patient who received treatment A, one must adjust for the fact that treatment A was systematically given to patients of different characteristics from those who received treatment B. We call this predicting under a shift in policy. Furthermore, if all of our observational data comes from hospital P, but we wish to predict counterfactuals for patients in hospital Q, with a population that differs from P, an additional source of distributional shift is at play. We call this a shift in domain. Together, we refer to the combination of domain and policy as the design. The design for which we observe ground truth is called the source, and the design of interest the target. The two most common approaches for addressing distributional shift are to learn shift-invariant representations of the data BID0 or to perform sample re-weighting or matching (; BID13 . Representation learning approaches attempt to extract only information from the input that is invariant to a change in design and predictive of the variable of interest. Such representations are typically learned by fitting deep neural networks in which activations of deeper layers are regularized to be distributionally similar across designs BID0 BID15 . Although representation learning can be shown to reduce the error associated to distributional shift BID15 in some cases, standard approaches are biased, even in the limit of infinite data, as they penalize the use also of predictive information. In contrast, re-weighting methods correct for distributional shift by assigning higher weight to samples from the source design that are representative of the target design, often using importance sampling. This idea has been well studied in, for example, the causal inference BID20, domain adaptation and reinforcement learning BID19 literature. For example, in causal effect estimation, importance sampling is equivalent to re-weighting units by the inverse probability of observed treatments (treatment propensity). Re-weighting with knowledge of importance sampling weights often leads to asymptotically unbiased estimators of the target outcome, but may suffer from high variance in finite samples .A significant hurdle in applying re-weighting methods is that optimal weights are rarely known in practice. There are a variety of methods to learn these weights. Weights can be estimated as the inverse of estimated feature or treatment densities BID20 BID7 but this plug-in approach can lead to highly unstable estimates. More stable methods learn weights by minimizing distributional distance metrics BID8 BID13 BID4 ). Closely related, matching produces weights by finding units in the source design that are close in some metric to units in the target design. Specifying a distributional or unit-wise metric is challenging, especially if the input space is high-dimensional where no metric incorporating all features can ever be made small. This has inspired heuristics such as first performing variable selection and then finding a matching in the selected covariates. Our key algorithmic contribution is to show how to combine the intuition behind shift-invariant representation learning and re-weighting methods by jointly learning a representation Φ of the input space and a weighting function w(Φ) to minimize a) the re-weighted empirical risk and b) a re-weighted measure of distributional shift between designs. This is useful also for the identity representation Φ(x) = x, as it allows for principled control of the variance of estimators through regularization of the re-weighting function w(x), mitigating the issues of exact importance sampling methods. Further, this allows us to evaluate w on hold-out samples to select hyperparameters or do early stopping. Finally, letting w depend on Φ alleviates the problem of choosing a metric by which to optimize sample weights, as Φ is trained to extract information predictive of the outcome. We capture these ideas in an upper bound on the generalization error under a shift in design and specialize it to the case of treatment effect estimation. We bring together two techniques used to overcome distributional shift between designs-re-weighting and representation learning, with complementary robustness properties, generalizing existing methods based on either technique. We give finite-sample generalization bounds for prediction under design shift, without assuming access to importance sampling weights or to a well-specified model, and develop an algorithmic framework to minimize these bounds. We propose a neural network architecture that jointly learns a representation of the input and a weighting function to improve balance across changing settings. Finally, we apply our proposed algorithm to the task of predicting causal effects from observational data, achieving state-of-the art on a widely used benchmark. The goal of this work is to accurately predict outcomes of interventions T ∈ T in contexts X ∈ X drawn from a target design p π (X, T). The outcome of intervening with t ∈ T is the potential outcome Y (t) ∈ Y (, Ch. 1-2), which has a stationary distribution p t (Y | X) given context X. Assuming a stationary outcome is akin to the covariate shift assumption , often used in domain adaptation.1 For example, in the classical binary setting, Y represents the outcome under treatment and Y the outcome under control. The target design consists of two components: the target policy p π (T | X), which describes how one intends to map observations of contexts (such as patient prognostics) to interventions (such as pharmacological treatments) and the target domain p π (X), which describes the population of contexts to which the policy will be applied. The target design is known to us only through m unlabeled sam-ples (x 1, t 1),..., (x m, t m) from p π (X, T). Outcomes are only available to us in labeled samples from a source domain: (x 1, t 1, y 1),..., (x n, t n, y n), where (x i, t i) are draws from a source design p µ (X, T) and y i = y i (t i) is a draw from p T (Y | X), corresponding only to the factual outcome Y (T) of the treatment administered. Like the target design, the source design consists of a domain of contexts for which we have data and a policy, which describes the (unknown) historical administration of treatment in the data. Only the factual outcomes of the treatments administered are observed, while the counterfactual outcomes y i (t) for t = t i are, naturally, unobserved. Our focus is the observational or off-policy setting, in which interventions in the source design are performed non-randomly as a function of X, p µ (T | X) = p µ (T). This encapsulates both the covariate shift often observed between treated and control populations in observational studies and the covariate shift between the domain of the study and the domain of an eventual wider intervention. Examples of this problem are plentiful: in addition to the example given in the introduction, consider predicting the return of an advertising policy based on the historical of a different policy, applied to a different population of customers. We stress that we are interested in the causal effect of an intervention T on Y, conditioned on X. As such, we cannot think of X and T as a single variable. Without additional assumptions, it is impossible to deduce the effect of an intervention based on observational data alone BID18, as it amounts disentangling correlation and causation. Crucially, for any unit i, we can observe the potential outcome y i (t) of at most one intervention t. In our analysis, we make the following standard assumptions. Assumption 1 (Consistency, ignorability and overlap). For any unit i, assigned to intervention t i, we observe Y i = Y (t i). Further, {Y (t)} t∈T and the data-generating process p µ (X, T, Y) satisfy strong ignorability: {Y (t)} t∈T ⊥ ⊥ T | X and overlap: DISPLAYFORM0 Assumption 1 is a sufficient condition for causal identifiability BID20. Ignorability is also known as the no hidden confounders assumption, indicating that all variables that cause both T and Y are assumed to be measured. Under ignorability therefore, any domain shift in p(X) cannot be due to variables that causally influence T and Y, other than through X. Under Assumption 1, potential outcomes equal conditional expectations: DISPLAYFORM1, and we may predict Y (t) by regression. We further assume common domain support, ∀x ∈ X: p π (X = x) > 0 ⇒ p µ (X = x) > 0. Finally, we adopt the notation p(x):= p(X = x). We attempt to learn predictors f: DISPLAYFORM0 Recall that under Assumption 1, this conditional expectation is equal to the (possibly counterfactual) potential outcome Y (t), conditioned on X. Our goal is to ensure that hypotheses f are accurate under a design p π that deviates from the data-generating process, p µ. This is unlike standard supervised learning for which p π = p µ. We measure the (in)ability of f to predict outcomes under π, using the expected risk, DISPLAYFORM1 is an appropriate loss function, such as the squared loss, L(y, y):= (y − y) 2 or the log-loss, depending on application. As outcomes under the target design p π are not observed, even through a finite sample, we cannot directly estimate using the empirical risk under p π. A common way to resolve this is to use importance sampling -the observation that if p µ and p π have common support, with w DISPLAYFORM2 Hence, with access to w *, an unbiased estimator of R π (f) may be obtained by re-weighting the (factual) empirical risk under µ, DISPLAYFORM3 Unfortunately, importance sampling weights can be very large when p π is large and p µ small, ing in large variance inR & ). More importantly, p µ (x, t) is rarely known in practice, and neither is w *. In principle, however, any re-weighting function w with the following property yields a valid risk under the re-weighted distribution p w µ. DISPLAYFORM4 DISPLAYFORM5 We denote the re-weighted density p w µ (x, t):= w(x, t)p µ (x, t).A natural candidate in place of w * is an estimateŵ * based on estimating densities p π (x, t) and p µ (x, t). In this work, we adopt a different strategy, learning parameteric re-weighting functions w from observational data, that minimize an upper bound on the risk under p π. An important special case of our setting is when treatments are binary, T ∈ {0, 1}, often interpreted as treating (T = 1) or not treating (T = 0) a unit, and the domain is fixed across designs, p µ (X) = p π (X). This is the classical setting for estimating treatment effects-the effect of choosing one intervention over another BID17. 2 The effect of an intervention T = 1 in context X, is measured by the conditional average treatment effect (CATE), DISPLAYFORM0 Predicting τ for unobserved units typically involves prediction of both potential outcomes 3. In a clinical setting, knowledge of τ is necessary to assess which medication should be administered to a certain individual. Historically, the (population) average treatment effect, ATE = E x∼p [τ (x)], has received comparatively much more attention BID20, but is inadequate for personalized decision making. Using predictors f (x, t) of potential outcomes Y (t) in contexts X = x, we can estimate the CATE byτ (x) = f (x, 1) − f (x, 0) and measure the quality using the mean squared error (MSE), DISPLAYFORM1 In Section 4, we argue that estimating CATE from observational data requires overcoming distributional shift with respect to the treat-all and treat-none policies, in predicting each respective potential outcome, and show how this can be used to derive generalization bounds for CATE. A large body of work has shown that under assumptions of ignorability and having a wellspecified model, various regression methods for counterfactual estimation are asymptotically consistent BID4 BID1 BID3. However, consistency like these provide little insight into the case of model misspecification. Under model misspecification, regression methods may suffer from additional bias when generalizing across designs due to distributional shift. A common way to alleviate this is importance sampling, see Section 2. This idea is used in propensity-score methods BID2, that use treatment assignment probabilities (propensities) to re-weight samples for causal effect estimation, and more generally in re-weighted regression, see e.g. . A major drawback of these methods is the assumption that the design density is known. To address this, others BID8 BID13, have proposed learning sample weights w to minimize the distributional distance between samples under p π and p w µ, but rely on specifying the data representation a priori, without regard for which aspects of the data actually matter for outcome prediction and policy estimation. On the other hand, BID12 proposed learning representations for counterfactual inference, inspired by work in unsupervised domain adaptation BID16. The drawback of this line of work is that the generalization bounds of Shalit et al. FORMULA4 and BID15 are loose-even if infinite samples are available, they are not guaranteed to converge to the lowest possible error. Moreover, these approaches do not make use of important information that can be estimated from data: the treatment/domain assignment probabilities. We give a bound on the risk in predicting outcomes under a target design p π (T, X) based on unlabeled samples from p π and labeled samples from a source design p µ (T, X). Our combines representation learning, distribution matching and re-weighting, ing in a tighter bound than the closest related work,. The predictors we consider are compositions f (x, t) = h(Φ(x), t) where Φ is a representation of x and h an hypothesis. We first give an upper bound on the risk in the general design shift setting, then show how this can be used to bound the error in prediction of treatment effects. In Section 5 we give a about the asymptotic properties of the minimizers of this upper bound. Risk under distributional shift Our bounds on the risk under a target design capture the intuition that if either a) the target design π and source design µ are close, or b) the true outcome is a simple function of x and t, the gap between the target risk and the re-weighted source risk is small. These notions can be formalized using integral probability metrics (IPM) that measure distance between distributions w.r.t. a normed vector space of functions H. Definition 2. The integral probability metric (IPM) distance, associated with a normed vector space of functions H, between distributions p and q is, DISPLAYFORM0 Important examples of IPMs include the Wasserstein distance, for which H is the family of functions with Lipschitz constant at most 1, and the Maximum Mean Discrepancy for which H are functions in the norm-1 ball in a reproducing kernel Hilbert space. Using definitions 1-2, and the definition of re-weighted risk, see FORMULA4, we can state the following (see Appendix A.2 for a proof). Lemma 1. For hypotheses f with loss f such that f / f H ∈ H, and p µ, p π with common support, there exists a valid re-weighting w of p µ, see Definition 1, such that, DISPLAYFORM1 The first inequality is tight for importance sampling weights, w(DISPLAYFORM2 The bound of Lemma 1 is tighter if p µ and p π are close (the IPM is smaller), and if the loss lives in a small family of functions H (the supremum is taken over a smaller set). Lemma 1 also implies that there exist weighting functions w(x, t) that achieve a tighter bound than the uniform weighting w(x, t) = 1, implicitly used by. While importance sampling weights in a tight bound in expectation, neither the design densities nor their ratio are known in general. Moreover, exact importance weights often in large variance in finite samples BID6. Here, we will search for a weighting function w, that minimizes a finite-sample version of, trading off bias and variance. We examine the empirical value of this idea alone in Section 6.1. We now introduce the notion of representation learning to combat distributional shift. Representation learning The idea of learning representations that reduce distributional shift in the induced space, and thus the source-target error gap, has been applied in domain adaptation BID0, algorithmic fairness and counterfactual prediction . The hope of these approaches is to learn predictors that predominantly exploit information that is common to both source and target distributions. For example, a face detector should be able to recognize the structure of human features even under highly variable environment conditions, by ignoring , lighting etc. We argue that re-weighting (e.g. importance sampling) should also only be done with respect to features that are predictive of the outcome. Hence, in Section 5, we propose using re-weightings that are functions of learned representations. We follow the setup of , and consider learning twice-differentiable, invertible representations Φ: X → Z, where Z is the representation space, and Ψ: Z → X is the inverse representation, such that Ψ(Φ(x)) = x for all x. Let E denote space of such representation functions. For a design π, we let p π,Φ (z, t) be the distribution induced by Φ over Z × T, with p w π,Φ (z, t):= p π,Φ (z, t)w(Ψ(z), t) its re-weighted form andp w π,Φ its re-weighted empirical form, following our previous notation. Finally, we let G ⊆ {h : Z × T → Y} denote a set of hypotheses h(Φ, t) operating on the representation Φ and let F the space of all compositions, F = {f = h(Φ(x), t): h ∈ G, Φ ∈ E}. We can now relate the expected target risk R π (f) to the re-weighted empirical source riskR w µ (f). Theorem 1. Given is a labeled sample (x 1, t 1, y 1),..., (x n, t n, y n) from p µ, and an unlabeled sample (x 1, t 1),..., (x m, t m) from p π, with corresponding empirical measuresp µ andp π. Suppose that Φ is a twice-differentiable, invertible representation, that h(Φ, t) is an hypothesis, and DISPLAYFORM3 where L is the squared loss, L(y, y) = (y − y) 2, and assume that there exists a constant B Φ > 0 such that h,Φ /B Φ ∈ H ⊆ {h : Z × T → Y}, where H is a reproducing kernel Hilbert space of a kernel, k such that k((z, t), (z, t)) < ∞. Finally, let w be a valid re-weighting of p µ,Φ. Then with probability at least 1 − 2δ, DISPLAYFORM4 where C DISPLAYFORM5. A similar bound exists where H is the family of functions Lipschitz constant at most 1, and IPM H the Wasserstein distance, but with worse sample complexity. See Appendix A.2 for a proof of Theorem 1 that involves applying finite-sample generalization bounds to Lemma 1, as well as moving to the space induced by the representation Φ.Theorem 1 has several implications: non-identity feature representations, non-uniform sample weights, and variance control of these weights can all contribute to a lower bound. Using uniform weights w(x, t) = 1 in, in a bound similar to that of Shalit et al. FORMULA4 and BID15. When π = µ, minimizing uniform-weight bounds in biased hypotheses, even in the asymptotical limit, as the IPM term does not vanish when the sample size increases. This is an undesirable property, as even k-nearest-neighbor classifiers are consistent in the limit of infinite samples. We consider minimizing with respect to w, improving the tightness of the bound. Theorem 1 indicates that even though importance sampling weights w * yield estimators with small bias, they can suffer from high variance, as captured by the factor V µ (w, f). The factor B Φ is not known in general as it depends on the true outcome, and is determined by f H as well as the determinant of the Jacobian of Ψ, see Appendix A.2. Qualitatively, B Φ measures the joint complexity of Φ and f and is sensitive to the scale of Φ-as the scale of Φ vanishes, B Φ blows up. To prevent this in practice, we normalize Φ. As B Φ is unknown, substituted a hyperparameter α for B Φ, but discussed the difficulties of selecting its value without access to counterfactual labels. In our experiments, we explore a heuristic for adaptively choosing α, based on measures of complexity of the observed held-out loss as a function of the input. Finally, the term C Theorem 1 is immediately applicable to the case of unsupervised domain adaptation in which there is only a single potential outcome of interest, T = {0}. In this case, DISPLAYFORM6 Conditional average treatment effects A simple argument shows that the error in predicting the conditional average treatment effect, MSE(τ) can be bounded by the sum of risks under the constant treat-all and treat-none policies. As in Section 2.2, we consider the case of a fixed domain p π (X) = p µ (X) and binary treatment T = {0, 1}. Let R πt (f) denote the risk under the constant policy π t such that ∀x ∈ X: p πt (T = t | X = x) = 1. Proposition 1. We have with MSE(τ) as in and R πt (f) the risk under the constant policy π t, DISPLAYFORM7 The proof involves the relaxed triangle inequality and the law of total probability. By Proposition 1, we can apply Theorem 1 to R π1 and R π0 separately, to obtain a bound on MSE(τ). For brevity, we refrain from stating the full , but emphasize that it follows from Theorem 1. In Section 6.2, we evaluate our framework in treatment effect estimation, minimizing this bound. Motivated by the theoretical insights of Section 4, we propose to jointly learn a representation Φ(x), a re-weighting w(x, t) and an hypothesis h(Φ, t) by minimizing a bound on the risk under the target design, see. This approach improves on previous work in two ways: it alleviates the bias of when sample sizes are large, see Section 4, and it increases the flexibility of the balancing method of BID8 by learning the representation to balance. For notational brevity, we let w i = w(x i, t i). Recall thatp w π,Φ is the re-weighted empirical distribution of representations Φ under p π. The training objective of our algorithm is the RHS of, with hyperparameters β = (α, λ h, λ w) substituted for model (and representation) complexity terms, DISPLAYFORM0 where R(h) is a regularizer of h, such as 2 -regularization. We can show the following . Theorem 2. Suppose H is a reproducing kernel Hilbert space given by a bounded kernel. Suppose weak overlap holds in that DISPLAYFORM1 Consequently, under the assumptions of Thm. 1, for sufficiently large α and λ w, DISPLAYFORM2 In words, the minimizers of converge to the representation and hypothesis that minimize the counterfactual risk, in the limit of infinite samples. Implementation Minimization of L π (h, Φ, w; β) over h, Φ and w is, while motivated by Theorem 2, a difficult optimization problem to solve in practice. For example, adjusting w to minimize the empirical risk term may in overemphasizing "easy" training examples, ing in a poor local minimum. Perhaps more importantly, ensuring invertibility of Φ is challenging for many representation learning frameworks, such as deep neural networks. In our implementation, we deviate from theory on these points, by fitting the re-weighting w based only on imbalance and variance terms, and don't explicitly enforce invertibility. As a heuristic, we split the objective, see, in two and use only the IPM term and regularizer to learn w. In short, we adopt the following alternating procedure. DISPLAYFORM3 The re-weighting function w(x, t) could be represented by one free parameter per training point, as it is only used to learn the model, not for prediction. However, we propose to let w be a parametric function of Φ(x). Doing so ensures that information predictive of the outcome is used for balancing, and lets us compute weights on a hold-out set, to perform early stopping or select hyperparameters. This is not possible with existing re-weighting methods such as BID8 BID13. An example architecture for the treatment effect estimation setting is presented in FIG2. By Proposition 1, estimating treatment effects involves predicting under the two constant policiestreat-everyone and treat-no-one. In Section 6, we evaluate our method in this task. As noted by , choosing hyperparameters for counterfactual prediction is fundamentally difficult, as we cannot observe ground truth for counterfactuals. In this work, we explore setting the balance parameter α adaptively. α is used in in place of B Φ, a factor measuring the complexity of the loss and representation function as functions of the input, a quantity that changes during training. As a heuristic, we use an approximation of the Lipschitz constant of f, with f = h(Φ(x), t), based on observed examples: DISPLAYFORM4 We use a moving average over batches to improve stability. We create a synthetic domain adaptation experiment to highlight the benefit of using a learned reweighting function to minimize weighted risk over using importance sampling weights w * (x) = DISPLAYFORM0 for small sample sizes. We observe n labeled source samples, distributed according to p µ (x) = N (x; m µ, I d) and predict for n unlabeled target samples drawn according to DISPLAYFORM1 and c ∼ N and let y = σ(β x + c) where σ(z) = 1/(1 + e −z). Importance sampling weights, w * (x) = p π (x)/p µ (x), are known. In experiments, we vary n from 10 to 600. We fit (misspecified) linear models 4 f (x) = β x + γ to the logistic outcome, and compare minimizing a weighted source risk by a) parameterizing sample weights as a small feed-forward neural network to minimize (ours) b) using importance sampling weights (baseline), both using gradient descent. For our method, we add a small variance penalty, λ w = 10 −3, to the learned weights, use MMD with an RBF-kernel of σ = 1.0 as IPM, and let α = 10. We compare to exact importance sampling weights (IS) as well as clipped IS weights (ISC), w M (x) = min(w(x), M ) for M ∈ {5, 10}, a common way of reducing variance of re-weighting methods .In FIG4, we see that our proposed method behaves well at small sample sizes compared to importance sampling methods. The poor performance of exact IS weights is expected at smaller samples, as single samples are given very large weight, ing in hypotheses that are highly sensitive to the training set. While clipped weights alleviates this issue, they do not preserve relevance ordering of high-weight samples, as many are given the truncation value M, in contrast to the reweighting learned by our method. True domain densities are known only to IS methods. We evaluate our framework in the CATE estimation setting, see Section 2.2. Our task is to predict the expected difference between potential outcomes conditioned on pre-treatment variables, for a held-out sample of the population. We compare our to ordinary least squares (OLS) (with one regressor per outcome), OLS-IPW (re-weighted OLS according to a logistic regression estimate of propensities), Random Forests, Causal Forests , BID5, and CFRW (with Wasserstein penalty). Finally, we use as baseline (IPM-WNN): first weights are found by IPM minimization in the input space BID8 BID13, then used in a re-weighted neural net regression, with the same architecture as our method. Our implementation, dubbed RCFR for Re-weighted CounterFactual Regression, parameterizes representations Φ(x), weighting functions w(Φ, t) and hypotheses h(Φ, t) using neural networks, trained by minimizing. We use the RBF-kernel maximum mean discrepancy as the IPM BID9. For a description of the architecture, training procedure and hyperparameters, see Appendix B. We compare using uniform w = 1 and learned weights, setting the balance parameter α either fixed, by an oracle (test-set error), or adaptively using the heuristic described in Section 5. To pick other hyperparameters, we split training sets into one part used for function fitting and one used for early stopping and hyperparameter selection. Hyperparameters for regularization are chosen based on the empirical loss on a held-out source (factual) sample. CATE Error, RMSE(τ) DISPLAYFORM0 (b) For small imbalance penalties α, re-weighting (low λw) has no effect. For moderate α, less uniform re-weighting (smaller λw) improves the error, c) for large α, weighting helps, but overall error increases. Best viewed in color. FORMULA4, and a synthesized continuous outcome that can be used to compute the ground-truth CATE error. Average over 100 different realizations/settings of the outcome are presented in TAB0. We see that our proposed method achieves state-of-the-art , and that adaptively choosing α does not hurt performance much. Furthermore, we see a substantial improvement from using nonuniform sample weights. In FIG4 we take a closer look at the behavior of our model as we vary its hyperparameters on the IHDP dataset. Between the two plots we can draw the following : a) For moderate to large α ∈, we observe a marginal gain from using the IPM penalty. This is consistent with the observations of. b) For large α ∈, we see a large gain from using a non-uniform re-weighting (small λ w). c) While large α makes the factual error more representative of the counterfactual error, using it without re-weighting in higher absolute error. We believe that the moderate sample size of this dataset is one of the reasons for the usefulness of our method. See Appendix C.2 for a complementary view of these . We have proposed a theory and an algorithmic framework for learning to predict outcomes of interventions under shifts in design-changes in both intervention policy and feature domain. The framework combines representation learning and sample re-weighting to balance source and target designs, emphasizing information from the source sample relevant for the target. Existing reweighting methods either use pre-defined weights or learn weights based on a measure of distributional distance in the input space. These approaches are highly sensitive to the choice of metric used to measure balance, as the input may be high-dimensional and contain information that is not predictive of the outcome. In contrast, by learning weights to achieve balance in representation space, we base our re-weighting only on information that is predictive of the outcome. In this work, we apply this framework to causal effect estimation, but emphasize that joint representation learning and re-weighting is a general idea that could be applied in many applications with design shift. Our work suggests that distributional shift should be measured and adjusted for in a representation space relevant to the task at hand. Joint learning of this space and the associated re-weighting is attractive, but several challenges remain, including optimization of the full objective and relaxing the invertibility constraint on representations. For example, variable selection methods are not covered by our current theory, as they induce a non-ivertible representation, but a similar intuition holds there-only predictive attributes should be used when measuring imbalance. We believe that addressing these limitations is a fruitful path forward for future work. We denote the re-weighted density p w µ (x, t):= w(x, t)p µ (x, t).Expected & empirical risk We let the (expected) risk of f measured by h under p µ be denoted DISPLAYFORM0 where l h is an appropriate loss function, and the empirical risk over a sample DISPLAYFORM1 We use the superscript w to denote the re-weighted risks DISPLAYFORM2 Definition A1 (Importance sampling). For two distributions p, q on Z, of common support, ∀z ∈ Z: p(z) > 0 ⇐⇒ q(z) > 0, we call DISPLAYFORM3 the importance sampling weights of p and q. Definition 2 (Restated). The integral probability metric (IPM) distance, associated with the function family H, between distributions p and q is defined by DISPLAYFORM4 We begin by bounding the expected risk under a distribution p π in terms of the expected risk under p µ and a measure of the discrepancy between p π and p µ. Using definition 2 we can show the following . Lemma 1 (Restated). For hypotheses f with loss f such that f / f H ∈ H, and p µ, p π with common support, there exists a valid re-weighting w of p µ, see Definition 1, such that, DISPLAYFORM5 The first inequality is tight for importance sampling weights, w(x, t) = p π (x, t)/p µ (x, t). The second inequality is not tight for general f, even if f ∈ H, unless p π = p µ.Proof. The follows immediately from the definition of IPM. DISPLAYFORM6 Further, for importance sampling weights w IS (x, t) = π(t;x) µ(t;x), for any h ∈ H, DISPLAYFORM7 and the LHS is tight. We could apply Lemma 1 to bound the loss under a distribution q based on the weighted loss under p. Unfortunately, bounding the expected risk in terms of another expectation is not enough to reason about generalization from an empirical sample. To do that we use Corollary 2 of BID6, restated as a Theorem below. Theorem A1 (Generalization error of re-weighted loss BID6). For a loss function h of any hypothesis h ∈ H ⊆ {h : X → R}, such that d = Pdim({ h : h ∈ H}) where Pdim is the pseudo-dimension, and a weighting function w(x) such that E p [w] = 1, with probability 1 − δ over a sample (x 1, ..., x n), with empirical distributionp, DISPLAYFORM8 we get the simpler form DISPLAYFORM9 We will also need the following about estimating IPMs from finite samples from.Theorem A2 (Estimation of IPMs from empirical samples ). Let M be a measurable space. Suppose k is measurable kernel such that sup x∈M k(x, x) ≤ C ≤ ∞ and H the reproducing kernel Hilbert space induced by k, with ν:= sup x∈M,f ∈H f (x) < ∞. Then, witĥ p,q the empirical distributions of p, q from m and n samples respectively, and with probability at least 1 − δ, DISPLAYFORM10 We consider learning twice-differentiable, invertible representations Φ: X → Z, where Z is the representation space, and Ψ: Z → X is the inverse representation, such that Ψ(Φ(x)) = x for all x. Let E denote space of such representation functions. For a design π, we let p π,Φ (z, t) be the distribution induced by Φ over Z × T, with p w π,Φ (z, t):= p π,Φ (z, t)w(Ψ(z), t) its re-weighted form andp w π,Φ its re-weighted empirical form, following our previous notation. Note that we do not include t in the representation itself, although this could be done in principle. Let G ⊆ {h : Z × T → Y} denote a set of hypotheses h(Φ, t) operating on the representation Φ and let F denote the space of all compositions, F = {f = h(Φ(x), t): h ∈ G, Φ ∈ E}. We now restate and prove Theorem 1. Given is a labeled sample D µ = {(x 1, t 1, y 1),..., (x n, t n, y n)} from p µ, and an unlabeled sample D π = {(x 1, t 1),..., (x m, t m)} from p π, with corresponding empirical measuresp µ andp π. Suppose that Φ is a twice-differentiable, invertible representation, that h(Φ, t) is an hypothesis, and DISPLAYFORM0 where L is the squared loss, L(y, y) = (y − y) 2, and assume that there exists a constant B Φ > 0 such that h,Φ /B Φ ∈ H ⊆ {h : Z × T → Y}, where H is a reproducing kernel Hilbert space of a kernel, k such that k((z, t), (z, t)) < ∞. Finally, let w be a valid re-weighting of p µ,Φ. Then with probability at least 1 − 2δ, DISPLAYFORM1 where C F n,δ measures the capacity of F and has only logarithmic dependence on n, D H m,n,δ measures the capacity of H, σ 2 Y is the expected variance in potential outcomes, and DISPLAYFORM2 A similar bound exists where H is the family of functions Lipschitz constant at most 1, but with worse sample complexity. Proof. We have by definition DISPLAYFORM3 f (x, t, y)p(y | t, x)(p π (x, t) − p Proof. Let f * = Φ * • h * ∈ arg min f ∈F R π (f) and let w * (x, t) = p π,Φ (Φ * (x), t)/p µ,Φ (Φ * (x), t). Since min h,Φ,w L π (h, Φ, w; β) ≤ L π (h *, Φ *, w * ; β), it suffices to show that L π (h *, Φ *, w * ; β) = R π (f *) + O(1/ √ n + 1/ √ m). We will work term by term: DISPLAYFORM4 For term D, letting w * i = w * (x i, t i), we have that by weak overlap DISPLAYFORM5 For term A, under ignorability, each term in the sum in the first term has expectation equal to R π (f *) and so, so by weak overlap and bounded second moments of loss, we have A = R π (f *) + O p (1/ √ n). For term B, since h * is fixed we have deterministically that DISPLAYFORM6 Finally, we address term C, which when expanded can be written as DISPLAYFORM7 w * i h(Φ * (x i), t i )).10 −4 10 −2 10 −0 10 2 ∞ Re-weighting regularization λw (uniformity) Figure 3: Error in CATE estimation on IHDP as a function of re-weighting regularization strength λ w (left) and source prediction error (right). We see in the left-hand plot that a) for small imbalance penalties α, re-weighting (low λ w) has no effect, b) for moderate α, less uniform re-weighting (smaller λ w) improves the error, c) for large α, weighting helps, but overall error increases. In the right-hand plot, we compare the ratio of CATE error to source error. Color represents α (see left) and size λ w. We see that for large α, the source error is more representative of CATE error, but does not improve in absolute value without weighting. Here, α was fixed. Best viewed in color. In Figure 3, we see two different views of the IHDP .

Represent this text to find a much much longer article about it with say hundreds-thousands of words

A theory and algorithmic framework for prediction under distributional shift, including causal effect estimation and domain adaptation

Represent the TLDR for finding the corresponding scientific article

This paper identifies classes of problems for which adversarial examples are inescapable, and derives fundamental bounds on the susceptibility of any classifier to adversarial examples.


Question: Represent the scientific article to find a short TLDR

Question: Deep ensembles have been empirically shown to be a promising approach for improving accuracy, uncertainty and out-of-distribution robustness of deep learning models. While deep ensembles were theoretically motivated by the bootstrap, non-bootstrap ensembles trained with just random initialization also perform well in practice, which suggests that there could be other explanations for why deep ensembles work well. Bayesian neural networks, which learn distributions over the parameters of the network, are theoretically well-motivated by Bayesian principles, but do not perform as well as deep ensembles in practice, particularly under dataset shift. One possible explanation for this gap between theory and practice is that popular scalable approximate Bayesian methods tend to focus on a single mode, whereas deep ensembles tend to explore diverse modes in function space. We investigate this hypothesis by building on recent work on understanding the loss landscape of neural networks and adding our own exploration to measure the similarity of functions in the space of predictions. Our show that random initializations explore entirely different modes, while functions along an optimization trajectory or sampled from the subspace thereof cluster within a single mode predictions-wise, while often deviating significantly in the weight space. We demonstrate that while low-loss connectors between modes exist, they are not connected in the space of predictions. Developing the concept of the diversity--accuracy plane, we show that the decorrelation power of random initializations is unmatched by popular subspace sampling methods. Consider a typical classification problem, where x n ∈ R D denotes the D-dimensional features and y n ∈ [1, . . ., K] denotes the class label. Assume we have a parametric model p(y|x, θ) for the conditional distribution where θ denotes weights and biases of a neural network, and p(θ) is a prior distribution over parameters. The Bayesian posterior over parameters is given by p(y n |x n, θ). (Computing the exact posterior distribution over θ is computationally expensive (if not impossible) when p(y n |x n, θ) is a deep neural network. A variety of approximations have been developed for Bayesian neural networks, including Laplace approximation , Markov chain Monte Carlo methods (; ;), variational Bayesian methods (; ; ;) and Monte-Carlo dropout . While computing the posterior is challenging, it is usually easy to perform maximum-a-posteriori (MAP) estimation, which corresponds to a mode of the posterior. The MAP solution can be written as the minimizer of the following loss (negative log likelihood + negative log prior): log p(y n |x n, θ). The MAP solution is computationally efficient, but only gives a point estimate and not a distribution over parameters. Deep ensembles, proposed by , train an ensemble of neural networks by initializing at M different values and repeating the minimization multiple times which could lead to M different solutions, if the loss is non-convex. (found adversarial training provides additional benefits in some of their experiments, but we will ignore adversarial training and focus only on ensembles with random initialization in this paper.) Given finite training data, many parameter values could equally well explain the observations, and capturing these diverse solutions is crucial for quantifying epistemic uncertainty . Bayesian neural networks learn a distribution over weights, and a good posterior approximation should be able to learn multi-modal posterior distributions in theory. Deep ensembles were inspired by the bootstrap , which has nice theoretical properties. However, it has been empirically observed by; that training individual networks with just random initialization is sufficient in practice and using the bootstrap even hurts performance in some cases (e.g. for small ensemble sizes). and independently benchmarked existing methods for uncertainty quantification on a variety of datasets and architectures, and observed that ensembles tend to outperform approximate Bayesian neural networks in terms of both accuracy and uncertainty, particularly under dataset shift. ) on train and validation data. These empirical observations raise an important question: Why do ensembles trained with just random initialization work so well in practice? One possible hypothesis is that ensembles tend to sample from different modes 1 in function space, whereas variational Bayesian methods (which minimize)) might fail to explore multiple modes even though they are effective at capturing uncertainty within a single mode. See Figure 1 for a cartoon illustration. Note that while the MAP solution is a local minima for the training loss by definition, it may not necessarily be a local minima for the validation loss. Recent work on understanding loss landscapes (; ; allows us to investigate this hypothesis. Note that prior work on loss landscapes has focused on mode-connectivity and low-loss tunnels, but has not explicitly focused on how diverse the functions from different modes are, beyond an initial exploration in . Our findings show that: • The functions sampled along a single training trajectory or subspace thereof (e.g. diagonal Gaussian, low-rank Gaussian and Dropout subspaces) tend to be very similar in predictions (while potential far away in the weight space), whereas functions sampled from different randomly initialized trajectories tend to be very diverse. • Solution modes are connected in the loss landscape but they are distinct in the space of predictions. Low-loss tunnels create functions with near-identical low values of loss along the path, however these functions tend to be very different in function space, changing significantly in the middle of the tunnel. The loss landscape of neural networks (also called the objective landscape) -the space of weights and biases that the network navigates -is typically a very high dimensional function and therefore could potentially be very complicated. However, many empirical show interesting properties of the loss surface. observed that the loss along a linear path from an initialization to the corresponding optimum is monotonically decreasing, encountering no significant obstacles along the way. demonstrated that constraining optimization to a random, low-dimensional hyperplane in the weight space leads to comparable to full-space optimization, provided that the dimension exceeds a modest threshold. This was geometrically understood and extended in .; demonstrate that while a linear path between two independent optima hits a high loss area in the middle, there in fact exist continuous, low-loss paths connecting any pair of optima. These observations are unified into a single phenomenological model in . While independent, low-loss optima in the loss landscape are connected, provide an early indication that in fact they represent very different functions in terms of their predictions. Therefore the connectivity cannot be due to trivial symmetries of the network which would keep the input-output mapping intact. We train convolutional neural networks on the CIFAR-10 dataset: • SmallCNN: channels for 10 epochs which achieves 64% test accuracy. • MediumCNN: channels for 20 epochs which achieves 70% test accuracy. • ResNet20v1: for 200 epochs which achieves 90% test accuracy. We use the Adam optimizer for training and to make sure the effects we observe are general, we validate that our hold for vanilla stochastic gradient descent (SGD) as well, which we do not show in this paper. We use batch size 128 and dropout 0.03 for training SmallCNN and MediumCNN. To generate weight space and prediction space similarity , we use a constant learning rate of 1.6 × 10 −3, unless specified otherwise. We do not use any data augmentation with those two architectures. For ResNet20v1, we use the data augmentation and learning rate schedule used in Keras examples 2. The overall trends are consistent across all architectures, datasets, and other hyperparameter and non-linearity choices we explored. First, we compute the similarity between different checkpoints along a single trajectory. We plot the cosine similarity in weight space in Figure 2 (a) and the disagreement in function space, defined as the fraction of points the checkpoints disagree on, in Figure 2 (b). We observe that the checkpoints along a trajectory are largely similar both in the weight space and the function space. Next, we evaluate how diverse the final solutions from different random initializations are. The functions from different initialization are different, as demonstrated by the similarity plots in Figure 3. Comparing this with Figures 2(a) and 2(b), we see that functions within a single trajectory exhibit higher similarity and functions across different trajectories exhibit much lower similarity. Next, we take the predictions from different checkpoints along the individual training trajectories from multiple initializations and compute a t-SNE plot to visualize their similarity in function space. More precisely, we take the softmax output for a set of points, flatten the vector and use it as the input to the t-SNE plot. Figure 2 (c) shows that the functions explored by different trajectories (denoted by circles with different colors) are far away, while functions explored within a single trajectory (circles with the same color) tend to be much more similar. In addition to the checkpoints along a trajectory, we also construct subspaces based on each individual trajectory. Scalable Bayesian methods typically compute statistics based on the weights along a trajectory, hence visualizing the diversity of functions between the subspace helps understand the difference between Bayesian neural networks and ensembles. We use a representative set of four subspace sampling methods: a random subspace, a Monte Carlo dropout, a diagonal Gaussian approximation, and a low-rank covariance matrix Gaussian approximation. In the descriptions of the methods, let w 0 be the current weight-space position (the weights and biases of our trained neural net) around which we will construct the subspace. Results on CIFAR-10 using two different architectures. For each of these architectures, the left subplot shows the cosine similarity between different solutions in weight space, and the right subplot shows the fraction of labels on which the predictions from different solutions disagree. • Random subspace sampling: We start at an optimized solution w 0 and choose a random directionv in the weight space. We step in that direction by choosing different values of t and looking at predictions at configurations w 0 + tv. We do this for many random directionsv. • Monte Carlo dropout subspace: We start at an optimized solution w 0 and apply dropout with a randomly chosen p keep to it. We do this many times, each time choosing a random p keep, and look at predictions at dropout p keep (w 0). • Diagonal Gaussian subspace: We start at an optimized solution w 0 and look at the most recent iterations of training proceeding it. For each trainable parameter w i, we calculate its mean mean i and standard deviation std(w i). To sample solutions from the subspace, we draw each parameter independently as w i ∼ N (mean i, std i). We repeat this many times and obtain predictions in each. This corresponds to sampling from a normal distribution with a diagonal covariance matrix. • Low-rank Gaussian subspace: We start at an optimized solution w 0 and look at the most recent iterations of training proceeding it. For each trainable parameter w i, we calculate its mean mean i. For a rank-k approximation, we calculate top k principal components of the weight vectors in the most recent iterations of training {p i ∈ R params} k. We sample from a k-dimensional normal distribution and obtain the weight configurations as w ∼ mean Figure 4 shows that functions sampled from a subspace (denoted by colored squares) corresponding to a particular initialization, are much more similar to each other. While some subspaces are more diverse, they still do not overlap with functions from another randomly initialized trajectory. Diversity versus Accuracy plots To illustrate the difference in another fashion, we sample functions from a single subspace and plot diversity (as measured by disagreement between predictions) versus accuracy in Figure 5. Comparing these subspace points (colored dots) to the baseline optima (green star) and the optima from different random initializations (denoted by red stars), we observe that random initializations are much more effective at sampling diverse and accurate solutions, than subspace based methods constructed out of a single trajectory. Figure 4: Results using SimpleCNN on CIFAR-10: t-SNE plots of validation set predictions for each trajectory along with four different subspace generation methods (showed by squares), in addition to 3 independently initialized and trained runs (different colors). As visible in the plot, the subspacesampled functions stay in the prediction-space neighborhood of the run around which they were constructed, demonstrating that truly different functions are not sampled. The diversity score used above quantifies the difference of two functions, by measuring fraction of points on which their predictions differ. We chose this approach due to its simplicity; one could also compute the KL-divergence or other distances between the output probability distributions Let d diff denote the fraction of predictions on which the two functions differ. It is 0 when the two functions make identical class predictions, and 1 when they differ on every single example. To account for the fact that the lower the accuracy of a function, the higher its potential d diff due to the possibility of the wrong answers being random and uncorrelated between the two functions, we normalize this by (1 − a), where a is the accuracy. For a reference function f * of accuracy a * and a function f of accuracy a whose predictions are obtained by randomly perturbing the predictions of f *, the expected fractional difference is, where C is the number of classes. If the function f of accuracy a were entirely independent of f *, then the expected fractional difference would be. Those two limiting behavioursthe function f being derived from f * by a perturbation, and f and f * being completely independent -form the two dashed lines in Figure 5. We refer to Appendix D for further details on the limiting curves. The diversity reached is not as high as the theoretical optimum even for the independently initialized and optimized solutions, which provides scope for future work. Figure 6 shows the radial loss landscape (train as well as the validation set) along the directions of two different optima. The left subplot shows that different trajectories achieve similar values of the loss, and the right subplot shows the similarity of these functions to their respective optima (in particular the fraction of labels predicted on which they differ divided by their error rate). While the loss values from different optima are similar, the functions are different, which confirms that random initialization leads to different modes in function space. Figure 6: Results using MediumCNN on CIFAR-10: Radial loss landscape cut between the origin and two independent optima and the predictions of models on the same plane. We construct a low-loss tunnel between different optima using the procedure proposed by , which is a simplification of the procedures proposed in and. As shown in Figure 7 (a), we start at the linear interpolation point (denoted by the black line) and reach the closest point on the manifold by minimizing the training loss. The minima of the training loss are denoted by the yellow line in the manifolds. Figure 7 (b) confirms that the tunnel is indeed low-loss. In order to visualize the 2-dimensional cut through the loss landscape and the the associated predictions on along a curved low-loss path, we divide the path into linear segments, and compute the loss and prediction similarities on a triangle given by this segment on one side and the origin of the weight space on the other. We perform this operation on each of the linear segments from which the low-loss path is constructed, and place them next to each other for visualization. Figure 8 visualizes the loss along the manifold, as well as the similarity to the original optima. Note that the regions between radial yellow lines consist of segments, and we stitch these segments together in Figure 8. The accuracy plots show that as we traverse along the low-loss tunnel, the accuracy remains fairly constant as expected. However, the prediction similarity plot shows that the low-loss tunnel does not correspond to similar solutions in function space. What it shows is that while the modes are connected in terms of accuracy/loss, their functional forms remain distinct and they do not collapse into a single mode. Our observations in the previous section suggest that subspace-based methods and ensembling should provide complementary benefits in terms of uncertainty and accuracy. To test this, we evaluate the performance of the following four variants using SmallCNN on CIFAR-10: • Baseline: optimum at the end of a single training trajectory. • Subspace sampling: average predictions over the solutions sampled from a subspace. • Ensemble: train baseline multiple times with random initialization and average the predictions. • Ensemble + Subspace sampling: train multiple times with random initialization, use subspace sampling within each trajectory. Figure 8: Results using MediumCNN on CIFAR-10: Radial loss landscape cut between the origin and two independent optima along an optimize low-loss connector and predictions similarity along the same planes. Figures 9(a) and 9(b) show the for low rank Gaussian subspace and diagonal Gaussian subspace respectively. The validate our hypothesis as (i) subspace sampling and ensembling provide complementary benefits, and (ii) the relative benefits of ensembling are higher as it averages predictions over more diverse solutions. Weight averaging within a subspace One could use the mean and diagonal/low-rank variance to approximate each mode of the posterior, however that increases the number of parameters required for each mode. Using just the mean weight for each mode would not increase the number of parameters. proposed stochastic weight averaging (SWA) for better generalization. One could also compute an (exponential moving) average of the weights along the trajectory, inspired by Polyak-Ruppert averaging in convex optimization, (see also for a Bayesian view on iterate averaging). As weight averaging has been already studied by, we do not discuss it in detail. Figure S1 provides an illustration of why these strategies might help with generalization. We use weight averaging (WA) on the last few epochs which corresponds to using the mean of the subspace within each mode. Figure 10 (a) shows that weight averaging achieves better performance within each mode, and ensemble + WA performs as well as ensemble + subspace combination methods, without any additional parameter overhead. Figure 10(b) shows accuracy and Brier score on CIFAR-10, both on the usual test set (corresponding to the intensity = 0 column) as well as on the CIFAR-10-C benchmark proposed which contains corrupted versions of CIFAR-10 with varying intensity values, making it useful to verify calibration under dataset shift . We see that ensembling and weight-averaging provide complementary benefits. WA improves over the vanilla baseline, but combining WA with ensembling over multiple random initializations improves performance further. Figure 9 reports accuracy and Brier score on the usual CIFAR-10 test set as a function of ensemble size. Under dataset shift, it is particular important to have diverse functions to avoid overconfident predictions (as averaging over similar functions would not reduce overconfidence). To illustrate the effect on another challenging dataset, we repeat these experiments on ImageNet using the same ResNet20V1 architecture. Due to computational constraints, we focus mainly on the experiment decoupling the effect of weight averaging vs ensembling. Figure 11 Our show that trajectories of randomly initialized neural networks explore different modes in function space, which explains why deep ensembles with random initializations help. They are essentially orthogonal to each other in the space of weights and very diverse in terms of their predictions. While these modes can be connected via optimized low-loss paths between them, we demonstrate that they correspond to distinct functions in terms of their predictions. Therefore the connectivity in the loss landscape does not imply connectivity in the space of functions. Subspace sampling methods such as weight averaging, Monte Carlo dropout, and various versions of local Gaussian approximations, sample functions that might lie relatively far from the starting point in the weight space, however, they remain in the vicinity of their starting point in terms of predictions, giving rise to an insufficiently diverse set of functions. Using the concept of the diversityaccuracy plane, we demonstrate empirically that these subspace sampling methods never reach the combination of diversity and accuracy that independently trained models do, limiting their usefulness for ensembling. A VISUALIZING THE LOSS LANDSCAPE ALONG ORIGINAL DIRECTIONS AND WA DIRECTIONS Figure S1 shows the loss landscape (train as well as the validation set) and the effect of WA. Figure S1: Loss landscape versus generalization: weights are typically initialized close to 0 and increase radially through the course of training. Top row: we pick two optima from different trajectories as the axes, and plot loss surface. Looking at x and y axes, we observe that while a wide range of radii achieve low loss on training set, the range of optimal radius values is narrower on validation set. Bottom row: we average weights within each trajectory using WA and use them as axes. A wider range of radius values generalize better along the WA directions, which confirms the findings of. Random seed affects both initial parameter values as well the order of shuffling of data points. We run experiments to decouple the effect of random initialization and shuffling; Figure S2 shows shows the . We observe that both of them provide complementary sources of randomness, with random initialization being the dominant of the two. As expected, random mini-batch shuffling adds more randomness at higher learning rates due to gradient noise. We run additional experiments comparing the diversity of solutions found vs their test accuracy on CIFAR-100. CIFAR-100 is an intermediate step between CIFAR-10 and ImageNet, and is overall Figure S2: The effect of random initializations and random training batches on the diversity of predictions. much more challenging to learn than CIFAR-10. Our additional are presented in Figure S3. Solutions obtained by subspace sampling methods described in Section 4 have a worse trade off between prediction diversity (needed for ensembling) and accuracy, compared to independently initialized and trained optima. This is consistent with our on CIFAR-10 in Figure 5. In Figures 5 and S3 we bound our empirical by two theoretically derived curves, limiting the expected trade off between diversity and accuracy in the best and worst case scenarios. The ing functions are presented in the main text in Section 3.2. We will show the detailed derivations here. Given a C-class classification problem and a reference solution with accuracy a *, we would like to obtain a function d diff (a) which gives us the fraction of labels on which another solution disagrees with the reference solution as a function of its accuracy. The best case scenario is when the predicted labels are uncorrelated with the reference solution's labels. On a particular example, the probability that the reference solution got it correctly is a *, and the probability that our solution got it correctly is a. On those examples, the predictions do not differ since they both have to be equal to the ground truth label. The probability that the reference solution is correct on an example while our solution is wrong is a * (1 − a). The probability that the reference solution is wrong on an example while our solution is correct is (1 − a *)a. On the examples where both solutions are wrong (probability (1 − a *)(1 − a)) there are two cases: a) the two solutions agree (an additional factor of 1/(C − 1)) or b) disagree (an additional factor of (C − 2)/(C − 1)). Only the case b) contributes to the fraction of labels on which they disagree. Hence we end up with The other extreme case is when the predictions of our new solution are just the predictions of the reference solution perturbed by perturbations of different strength. Then, the solutions retain a great amount of correlation. Let the probability of a label changing be p. We will consider 4 cases: a) the label of the correctly classified image does not flip (probability a * (1 − p)), b) it flips (probability a * p), c) an incorrectly labelled image does not flip (probability (1 − a *)(1 − p)), and d) it flips (probability (1 − a *)p). The ing accuracy a(p) obtains a contribution a * (1 − p) from case a) and with probability 1/(C − 1) contribution (1 − a *)p from d). Therefore a(p) = a * (1 − p) + p(1 − a *)/(C − 1). Inverting this relationship, we get p(a) = (C − 1)(a * − a)/(Ca * − 1). The fraction of labels on which the solutions disagree is simply p by our definition of p, and therefore d diff (a; a *, C) = (C − 1)(a * − 1a)

represent this text to find a much much longer article about it with say hundreds-thousands of words

We study deep ensembles through the lens of loss landscape and the space of predictions, demonstrating that the decorrelation power of random initializations is unmatched by subspace sampling that only explores a single mode.

represent the tldr for finding the corresponding scientific article

We propose a hypothesis for why gradient descent generalizes based on how per-example gradients interact with each other.