Towards Representation Learning for Weighting Problems in Design-Based Causal Inference

1. Problem Statement

The paper addresses the challenge of design-based weighting in causal inference. In many scenarios (e.g., prospective cohort studies, survey weighting, or the weighting step of doubly robust estimators), researchers must estimate weights to reweight a source distribution ($P$) to match a target distribution ($Q$) without access to outcome data ($Y$).

• The Core Difficulty: Optimal weights (the Radon-Nikodym derivative $dQ/dP$) depend on the unknown Data Generating Process (DGP). Standard approaches either assume a specific parametric form for the propensity score (risking misspecification) or learn representations using outcome data (which violates the "design-based" constraint).
• The Gap: Existing methods that learn representations (e.g., via neural networks) often rely on strong assumptions about the outcome model or lack theoretical guarantees on the bias introduced when the chosen representation is misspecified. There is no principled way to assess the quality of a representation solely based on covariates ($X$) before outcomes are observed.

2. Methodology

The authors propose a framework that separates the learning of a representation ($\phi(X)$) from the calculation of weights, ensuring the representation learning phase remains outcome-agnostic.

A. Theoretical Decomposition of Bias

The authors decompose the total bias of a weighted estimator into three components:

1. Bias wrt Representation: The error due to imperfect balancing of the representation $\phi(X)$ itself.
2. Chosen Weights Bias: The error if the weights depend on $X$ but not solely on $\phi(X)$. (Theoretically shown to be zero if weights are optimized based on $\phi$).
3. Confounding Bias: The error introduced because $\phi(X)$ fails to preserve the unconfoundedness of the original covariates $X$. This is the critical term representing information loss.

They derive that the Confounding Bias is bounded by the product of the outcome model's magnitude and a new metric called the Balancing Score Error (BSE):
$$|\text{Confounding Bias}| \leq \|\tilde{Y}\|_{L^2(P)} \cdot \text{BSE}_{P,Q}(\phi)$$
Where $\text{BSE}_{P,Q}(\phi)$ measures the $L^2$ distance between the true density ratio $dQ/dP$ and its projection onto the representation $\phi(X)$. Crucially, BSE does not depend on the outcome $Y$; it depends only on the marginal distributions of $X$ in $P$ and $Q$.

B. Learning the Representation (The Algorithm)

To minimize the BSE without observing outcomes, the authors leverage the AutoDML loss (from Chernozhukov et al., 2022).

1. Objective: Instead of trying to estimate the true weights directly, they learn a representation $\phi(X; \theta_\phi)$ and a scalar function $g(\cdot; \theta_g)$ such that $g(\phi(X))$ approximates the density ratio $dQ/dP$.
2. Loss Function: They minimize the AutoDML loss:
   $$L_{P,Q}(v) = E_P[v(X)^2] - 2E_Q[v(X)]$$
   This loss can be estimated using only samples from $P$ and $Q$ (covariates only).
3. Neural Network Architecture: They implement this using a neural network where:
   • The input is $X$.
   • An intermediate hidden layer serves as the representation $\phi(X)$.
   • The final layer outputs the scalar function $g(\phi(X))$.
   • The network is trained to minimize the AutoDML loss.
4. Weighting Step: Once $\phi(X)$ is learned, it is plugged into a standard weighting algorithm (e.g., Kernel Optimal Matching or Kernel Mean Matching) to find the final weights $w$. This step minimizes the Integral Probability Metric (IPM) between the weighted source and target distributions in the representation space.

3. Key Contributions

1. Quantification of Information Loss: The paper provides the first rigorous decomposition of bias in design-based weighting that explicitly quantifies the "Confounding Bias" caused by using a suboptimal representation.
2. Balancing Score Error (BSE): They introduce the BSE as a measurable proxy for representation quality that is independent of outcome data. This allows researchers to select or learn representations based purely on covariate balance properties.
3. End-to-End Estimation Procedure: They propose a novel algorithm (Algorithm 1) that learns a flexible, non-linear representation via neural networks using only covariate data, then applies standard weighting methods on that representation.
4. Theoretical Guarantees: They prove that minimizing the BSE (via the AutoDML loss) provides a bound on the overall bias for any outcome model, without requiring the outcome model to be correctly specified or known.

4. Experimental Results

The method was evaluated on three datasets: IHDP (simulated), News (high-dimensional text), and TBI (Traumatic Brain Injury, transportability).

• Metrics: Joint Bias (JB) on Average Treatment Effect (ATE) and Transportability tasks.
• Baselines: Compared against standard Kernel Optimal Matching (KOM) on raw covariates, Propensity Score (PS) weighting, Principal Component Analysis (PCA), and Neural Score Matching (NSM).
• Findings:
  • The proposed method ("Ours + Energy/Linear") consistently outperformed or matched state-of-the-art baselines.
  • On the high-dimensional News dataset, using the learned low-dimensional representation significantly sped up computation while reducing bias compared to using raw covariates.
  • The method was particularly robust in Transportability tasks (TBI dataset), where standard linear kernels failed, but the learned representation with energy balancing succeeded.
  • It outperformed the "Neural Network Head" (using the final layer of the network as weights directly), demonstrating that separating representation learning from weight optimization is beneficial.

5. Significance and Implications

• Design-Based Rigor: The approach strictly adheres to the "design-based" paradigm by ensuring no outcome data is used to construct the weights or the representation. This is crucial for pre-registering studies and avoiding p-hacking.
• Generalizability: The framework is generic; the learned representation can serve as a pre-processing step for any weighting method (e.g., Entropy Balancing, IPW, KOM).
• Handling Misspecification: Unlike traditional propensity score modeling which fails if the functional form is wrong, this method learns a representation that minimizes the worst-case bias bound (BSE), offering robustness against model misspecification.
• Future Directions: The authors suggest this framework could be extended to multimodal data (text, images, tabular) and augmented estimators, bridging the gap between deep learning representation learning and classical causal inference theory.

In summary, this paper provides a theoretically grounded, practical solution for learning optimal covariate representations for causal weighting when outcome data is unavailable, effectively minimizing the bias introduced by dimensionality reduction or feature selection.