Rank-Learner: Orthogonal Ranking of Treatment Effects

TL;DR

The paper addresses ranking individuals by their treatment effects $\tau(x)$ from observational data, a task where treatment effects are not directly observed and nuisance functions must be estimated. It proposes Rank-Learner, a Neyman-orthogonal two-stage method that directly optimizes a pairwise ranking objective using smooth surrogate targets with a tunable smoothing parameter $\kappa$, and learns a scoring function $g$ that preserves the $\tau(x)$ ordering. The authors prove the orthogonality of their loss to nuisance estimation errors and characterize the population minimizers, showing $g^0(x)=\tfrac{1}{\kappa}\,\tau^0(x)+c$, which emphasizes ranking when $\kappa$ is small. Empirically, Rank-Learner outperforms standard CATE estimators and non-orthogonal rankers across synthetic and semi-synthetic benchmarks, is robust to limited overlap, and scales via subsampling of training pairs. This work provides practitioners with a model-agnostic, two-stage approach to ranking treatment effects with strong theoretical guarantees and practical performance benefits.

Abstract

Many decision-making problems require ranking individuals by their treatment effects rather than estimating the exact effect magnitudes. Examples include prioritizing patients for preventive care interventions, or ranking customers by the expected incremental impact of an advertisement. Surprisingly, while causal effect estimation has received substantial attention in the literature, the problem of directly learning rankings of treatment effects has largely remained unexplored. In this paper, we introduce Rank-Learner, a novel two-stage learner that directly learns the ranking of treatment effects from observational data. We first show that naive approaches based on precise treatment effect estimation solve a harder problem than necessary for ranking, while our Rank-Learner optimizes a pairwise learning objective that recovers the true treatment effect ordering, without explicit CATE estimation. We further show that our Rank-Learner is Neyman-orthogonal and thus comes with strong theoretical guarantees, including robustness to estimation errors in the nuisance functions. In addition, our Rank-Learner is model-agnostic, and can be instantiated with arbitrary machine learning models (e.g., neural networks). We demonstrate the effectiveness of our method through extensive experiments where Rank-Learner consistently outperforms standard CATE estimators and non-orthogonal ranking methods. Overall, we provide practitioners with a new, orthogonal two-stage learner for ranking individuals by their treatment effects.

Figures (6)

• Figure 1: Two-stage learners for ranking treatment effects.Left: Confounded observational data with unobserved treatment effects. Center: First-stage estimation of nuisance functions (i.e., response surfaces and propensity score). Right: Second-stage objectives: our Rank-Learner vs. standard CATE estimation. Here, our proposed Rank-Learner directly optimizes a pairwise, Neyman-orthogonal ranking objective that targets the relative ordering of treatment effects. In contrast, standard CATE estimation optimizes a pointwise regression objective to recover treatment effect magnitudes, which is harder than necessary for ranking.
• Figure 2: Overview of Rank-Learner. In the first stage, we estimate nuisance functions (response surfaces and propensity score) via cross-fitting. In the second stage, we subsample training pairs, construct the pseudo labels, and learn a scoring function by minimizing the orthogonal ranking objective. Both stages can be instantiated with arbitrary machine learning models. At inference time, we rank individuals directly by their learned scores $\hat{g}(x)$.
• Figure 3: Synthetic benchmark (pair subsampling). Test AUTOC (mean $\pm$ s.e. over five seeds) of Rank-Learner as a function of the number of sampled training pairs per epoch ($n=1{,}000$ with $n^2=10^6$ possible training pairs). Higher is better. The horizontal axis shows the fraction of pairs used (log).
• Figure 4: Synthetic benchmark (overlap sensitivity). Test AUTOC (mean $\pm$ s.e. over five seeds) as a function of overlap (decreasing left to right) for $n=500$. Higher is better. Overlap is varied by changing treatment assignment, remaining components of the synthetic data-generating process are kept fixed.
• Figure 5: Synthetic benchmark (pair subsampling). Test AUTOC (mean $\pm$ s.e. over five seeds) of Rank-Learner as a function of the number of sampled training pairs per epoch ($n=1{,}000$ with $n^2=10^6$ possible training pairs). Higher is better. The horizontal axis shows the fraction of pairs used (log).

Theorems & Definitions (3)

• Theorem 5.1: Neyman-orthogonality
• Theorem 5.2: Minimizers of the orthogonal loss
• proof : Proof of Theorem \ref{['thm:orthogonality']}