user: text: Write a concise one-sentence TLDR summary of the scientific paper content below.

Title: Learning Weighted Representations for Generalization Across Designs

Paper content:
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 results 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 (Shimodaira, 2000; 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 (Shimodaira, 2000) 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 (Swaminathan & Joachims, 2015) .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 Zubizarreta, 2015) . Closely related
, matching (Stuart, 2010) 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 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 result. 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 results 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 result about estimating IPMs from finite samples from Sriperumbudur et al. (2009) .Theorem A2 (Estimation of IPMs from empirical
samples (Sriperumbudur et al., 2009) ). 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.

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