Technical Summary: Amortizing Causal Sensitivity Analysis via Prior-Data Fitted Networks

1. Problem Statement

Causal inference from observational data relies on the untestable assumption of no unobserved confounding. Causal sensitivity analysis addresses this by providing bounds on causal effect estimates (e.g., the Conditional Average Potential Outcome, CAPO) under varying assumptions about the strength of unobserved confounding, parameterized by a sensitivity level $\Gamma$.

Existing methods for computing these bounds are inherently per-instance procedures. For any change in the dataset, causal query, treatment arm, or sensitivity level $\Gamma$, a new, computationally expensive optimization must be instantiated. This limits the practical application of sensitivity analysis, often relegating it to a final check rather than a systematic tool for exploring how conclusions vary across assumptions. The core challenge is to replace these repeated, costly optimizations with a fast, amortized inference mechanism that can handle diverse datasets and queries without retraining.

2. Methodology

The authors propose an amortized approach using Prior-Data Fitted Networks (PFNs), a class of foundation models that perform in-context learning. The methodology consists of three main components:

2.1. Amortized Prediction Task

Instead of solving an optimization problem at test time, the authors train a PFN to learn a mapping from a dataset $D_N$, a causal query $(x, a)$, and a sensitivity level $\Gamma$ directly to the posterior predictive distributions (PPDs) of the lower and upper sensitivity bounds ($\theta^-$ and $\theta^+$).
$$D_N, x, a, \Gamma \longrightarrow \text{PPD}(\theta^-), \text{PPD}(\theta^+)$$
This allows the model to predict bounds for new inputs via a single forward pass, bypassing the need for per-instance optimization.

2.2. General Prior-Data Construction (Label Generation)

A critical hurdle in training PFNs for sensitivity analysis is that, unlike point-identified causal inference, the sensitivity bounds are not directly available from the data-generating process (DGP); they must be computed via non-trivial optimization. The authors introduce a general prior-data construction applicable to the class of Generalized Treatment Sensitivity Models (GTSMs), which includes the Marginal Sensitivity Model (MSM), $f$-sensitivity models, and Rosenbaum's sensitivity model.

The construction involves:

1. Sampling SCMs: Synthetic Structural Causal Models (SCMs) are sampled from a prior $\Pi_S$ to generate observational datasets.
2. Lagrangian Scalarization: To generate training labels ($\theta^-, \theta^+$) for these synthetic DGPs, the authors formulate the bound computation as a multi-criteria optimization problem. They seek to trade off between minimizing/maximizing the causal effect and minimizing the violation of the sensitivity model constraint.
   • The objective is scalarized using a Lagrange multiplier $\lambda$:
     $$\mathcal{L}^\uparrow_\lambda(P) = Q(P(Y(a)|x)) - \lambda \Delta_{x,a}(P)$$
     $$\mathcal{L}^\downarrow_\lambda(P) = -Q(P(Y(a)|x)) - \lambda \Delta_{x,a}(P)$$
   • By sweeping $\lambda$ from high to low, the method traces the Pareto frontier of solutions in the $(\Gamma, \theta)$ plane.
3. Warm-Start Optimization: To efficiently compute the Pareto frontier, the authors employ a warm-start strategy where the solution for a given $\lambda$ initializes the optimization for the next $\lambda$. Theoretical analysis (Theorem 1) shows that under standard convexity and linearity assumptions, this sweep recovers the full Pareto frontier.
4. Dimensionality Reduction: The optimization over full distributions is reduced to a tractable optimization over a single conditional latent distribution $P(U | x, A \neq a)$ using a reparameterization construction, avoiding the need for two-stage neural density estimation required by previous GTSM methods.

2.3. Foundation Model Training

The PFN is trained on the synthetic prior-data pairs (datasets, queries, $\Gamma$) and the computed sensitivity bound labels. The model architecture uses a transformer with two output heads (for lower and upper bounds), each modeling a Gaussian mixture distribution to provide uncertainty-aware predictions. A soft monotonicity regularizer is applied to ensure the predicted bounds behave correctly with respect to $\Gamma$ (i.e., bounds widen as $\Gamma$ increases).

3. Key Contributions

1. Prior-Data Label Construction: A general scheme for generating training labels for sensitivity bounds using Lagrangian scalarization. This approach avoids model-specific analytical derivations and theoretically recovers the full Pareto frontier under convexity conditions.
2. First Foundation Model for Causal Sensitivity Analysis: The training of a PFN that acts as a foundation model for in-context learning of sensitivity bounds. It directly predicts PPDs for bounds given a dataset, query, and sensitivity level.
3. Empirical Validation: Demonstration that the amortized approach approximates sensitivity bounds accurately across various settings (including the Marginal Sensitivity Model) with test-time computation orders of magnitude faster than per-instance methods.

4. Results

• Accuracy: On held-out synthetic datasets, the PFN closely matches analytical bounds for the Marginal Sensitivity Model (MSM). The posterior predictive intervals are well-calibrated, with empirical coverage rates (e.g., 90.3% for the upper bound) close to nominal levels.
• Efficiency: The amortized approach replaces expensive per-instance optimizations with a single forward pass. The median inference time for a batched forward pass is approximately 2.9 seconds, compared to the repeated optimization required by existing methods.
• Warm-Start Efficacy: Experiments show that the warm-start optimization strategy for label generation reduces regret by a factor of 1.90 compared to cold-start methods while being significantly faster.
• Uncertainty Quantification: The model outputs full posterior predictive distributions, providing uncertainty-aware estimates rather than just point predictions.

5. Significance and Claims

The paper claims to present the first foundation model for causal sensitivity analysis. Its primary significance lies in shifting the paradigm from per-instance, computationally intensive sensitivity analysis to an amortized, in-context learning framework.

• Practical Impact: By enabling fast evaluation across many causal queries and sensitivity assumptions, the approach makes sensitivity analysis a viable tool for systematic decision-making rather than a final, isolated check.
• Generality: The prior-data construction is not tied to a specific sensitivity model, making the approach broadly applicable to the class of GTSMs.
• Limitations: The authors modestly note that performance depends on the coverage of the synthetic prior relative to test-time data. The theoretical construction relies on convexity of the sensitivity divergence and linearity of the causal query. Currently, the implementation focuses on binary treatments and scalar outcomes, and label generation for models without closed-form bounds remains an optimization-based approximation subject to Monte Carlo noise.

In summary, this work provides a principled method to "amortize" the computational cost of causal sensitivity analysis, enabling scalable and systematic assessment of unobserved confounding in observational studies.