PITE: Multi-Prototype Alignment for Individual Treatment Effect Estimation

2025-11-13

AI Summary

Estimating individual treatment effects (ITE) from observational data remains challenging due to confounding bias and the neglect of local structural information, leading to suboptimal accuracy. To address this, we propose Multi-Prototype Alignment (MPA), a novel method that leverages cluster centroids of similar individuals as prototypes. MPA jointly models intra-cluster local structure and performs fine-grained cross-treatment-group alignment in the latent space, thereby preserving both distributional consistency and individual heterogeneity. Unlike global balancing or instance-level matching approaches, MPA introduces, for the first time, a cross-group local structural alignment mechanism built upon instance-to-prototype matching. Within an end-to-end deep learning framework, clustering, prototype matching, and structural alignment are jointly optimized. Extensive experiments on multiple benchmark datasets demonstrate that MPA significantly outperforms 13 state-of-the-art methods, achieving superior ITE estimation accuracy and robustness.

Technology Category

• Machine Learning: Multi-instance/Multi-view Learning
• Intelligent Robots: State Estimation
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

Abstract

Estimating Individual Treatment Effects (ITE) from observational data is challenging due to confounding bias. Most studies tackle this bias by balancing distributions globally, but ignore individual heterogeneity and fail to capture the local structure that represents the natural clustering among individuals, which ultimately compromises ITE estimation. While instance-level alignment methods consider heterogeneity, they similarly overlook the local structure information. To address these issues, we propose an end-to-end Multi- extbf{P}rototype alignment method for extbf{ITE} estimation ( extbf{PITE}). PITE effectively captures local structure within groups and enforces cross-group alignment, thereby achieving robust ITE estimation. Specifically, we first define prototypes as cluster centroids based on similar individuals under the same treatment. To identify local similarity and the distribution consistency, we perform instance-to-prototype matching to assign individuals to the nearest prototype within groups, and design a multi-prototype alignment strategy to encourage the matched prototypes to be close across treatment arms in the latent space. PITE not only reduces distribution shift through fine-grained, prototype-level alignment, but also preserves the local structures of treated and control groups, which provides meaningful constraints for ITE estimation. Extensive evaluations on benchmark datasets demonstrate that PITE outperforms 13 state-of-the-art methods, achieving more accurate and robust ITE estimation.

Problem

• Estimating Individual Treatment Effects with confounding bias
• Capturing local structure and individual heterogeneity in data
• Aligning multi-prototype distributions across treatment groups

Innovation

• Multi-prototype alignment for treatment effect estimation
• Instance-to-prototype matching captures local structure
• Cross-group prototype alignment reduces distribution shift