Cross-Treatment Effect Estimation for Multi-Category, Multi-Valued Causal Inference via Dynamic Neural Masking

TL;DR

This work tackles counterfactual causal inference under multi-category, multi-valued treatments by introducing XTNet, a unified neural architecture with BasicNet for baseline effects, EffectNet for cross-treatment interactions, and MaskNet for dynamic masking. It also proposes MCMV-AUCC, a cost-aware evaluation metric tailored to multi-dimensional treatment spaces, and provides theoretical justification for its lower expected error compared to existing metrics. Empirically, XTNet demonstrates superior ranking accuracy and effect estimation on synthetic and real-world data, with ablation studies confirming the contributions of each architectural component and the imbalance loss. A production-scale online A/B test further validates XTNet’s practical impact, showing improved gross merchandise value and treatment-order gains. Overall, the approach offers a scalable, effective solution for estimating complex treatment effects in settings with heterogeneous interventions across multiple categories.

Abstract

Counterfactual causal inference faces significant challenges when extended to multi-category, multi-valued treatments, where complex cross-effects between heterogeneous interventions are difficult to model. Existing methodologies remain constrained to binary or single-type treatments and suffer from restrictive assumptions, limited scalability, and inadequate evaluation frameworks for complex intervention scenarios. We present XTNet, a novel network architecture for multi-category, multi-valued treatment effect estimation. Our approach introduces a cross-effect estimation module with dynamic masking mechanisms to capture treatment interactions without restrictive structural assumptions. The architecture employs a decomposition strategy separating basic effects from cross-treatment interactions, enabling efficient modeling of combinatorial treatment spaces. We also propose MCMV-AUCC, a suitable evaluation metric that accounts for treatment costs and interaction effects. Extensive experiments on synthetic and real-world datasets demonstrate that XTNet consistently outperforms state-of-the-art baselines in both ranking accuracy and effect estimation quality. The results of the real-world A/B test further confirm its effectiveness.

Figures (3)

• Figure 1: Multi-Category, Multi-Valued Treatment Scenario: An illustrative example with two distinct treatment types. The objective is to estimate patient outcomes (e.g., body temperature) under different treatment combinations.
• Figure 2: Current Model Architectures
• Figure 3: Architecture of XTNet for multi-treatment causal effect estimation. Our XTNet consists of three main components: BasicNet, EffectNet, and MaskNet. The BasicNet produces baseline outcome predictions using input features without cross-treatment effects. The EffectNet estimates cross-treatment effects from other treatments, which are then added to the baseline outcomes to obtain the final effect estimation. The MaskNet generates parameter masks for the EffectNet to construct treatment-specific masked networks. This masking mechanism provides flexibility for handling varying numbers of treatment categories.

Theorems & Definitions (12)

• Definition 3.1: Multi-Category CATE
• Definition 5.1: Stochastic Ideal Allocation
• Definition 5.2: Expected Metric Error
• Theorem 5.1: Stochastic Dominance Principle
• Definition 5.3: AUCC Allocation Policy
• Lemma 5.2: RoI Ordering
• Theorem 5.3: Qini vs AUCC
• Definition 5.4: MV-AUCC Allocation Policy
• Theorem 5.4: AUCC vs MV-AUCC
• Definition 5.5: MCMV-AUCC Allocation Policy