DISCRET: Synthesizing Faithful Explanations For Treatment Effect Estimation

Abstract

Designing faithful yet accurate AI models is challenging, particularly in the field of individual treatment effect estimation (ITE). ITE prediction models deployed in critical settings such as healthcare should ideally be (i) accurate, and (ii) provide faithful explanations. However, current solutions are inadequate: state-of-the-art black-box models do not supply explanations, post-hoc explainers for black-box models lack faithfulness guarantees, and self-interpretable models greatly compromise accuracy. To address these issues, we propose DISCRET, a self-interpretable ITE framework that synthesizes faithful, rule-based explanations for each sample. A key insight behind DISCRET is that explanations can serve dually as database queries to identify similar subgroups of samples. We provide a novel RL algorithm to efficiently synthesize these explanations from a large search space. We evaluate DISCRET on diverse tasks involving tabular, image, and text data. DISCRET outperforms the best self-interpretable models and has accuracy comparable to the best black-box models while providing faithful explanations. DISCRET is available at https://github.com/wuyinjun-1993/DISCRET-ICML2024.

Figure 5.

The curve of ATE errors on test split of IHDP by DISCRET

Figure 6.

Frequency of the outcome values on Uganda dataset

Figure 7.

Extracted concepts from Uganda dataset

Figure 1.

Motivating examples from the Uganda dataset. We predict how providing economic aid (the treatment) helps to develop remote regions of the country (the outcome) via satellite images. The task is to estimate the ITE for each sample $x_1$ and $x_2$. DISCRET predicts that, because both images have several indicators of rich soil and urbanization, they will have similar ITE if given aid. Self-interpretable models such as Causal Forest (Athey & Wager, 2019) produce consistent ITE estimates (i.e., samples with same explanations have same model predictions, viz. 3.97 and 3.97), but have poor accuracy ($IT{\hat{E}}_{x_1}\ll IT{E}_{x_1}=4.25$). Black-box models such as TransTEE (Zhang et al., 2022), are accurate but do not produce similar predictions for samples $x_1$ and $x_2$ with similar explanations, when the explanations are sourced from post-hoc explainers such as Anchor (Ribeiro et al., 2018). DISCRET produces both consistent and accurate predictions.

Figure 2.

Illustration of DISCRET on the IHDP dataset, which tracks premature infants. Given a sample $x$, DISCRET synthesizing an explanation $L_{1:k}$ where it iteratively constructs each literal in the explanation. In particular, DISCRET (i) embeds the given sample and any previously generated literals (${\Theta }_0$), (ii) passes the embedding to the feature selection network $\left({\Theta }_1\right)$ to pick a feature, and then (iii) passes the embedding and selected feature to the constant selection network $\left({\Theta }_2\right)$ to get a thresholding constant. The operator is auto-assigned based on the feature and sample. DISCRET executes this explanation on the database to find relevant samples, which are used (i) during training to compute a reward function for ${\Theta }_0,{\Theta }_1$ and ${\Theta }_2$, and (ii) during testing to calculate the ITE.

Figure 3.

Consistency scores (higher is better) for DISCRET and a black-box model (TransTEE) combined with a post-hoc explainer. Our results confirm that DISCRET produces faithful explanations, and importantly, show that post-hoc explanations are rarely faithful, as evidenced by low consistency scores across datasets.

Figure 4.

DISCRET identifies similar samples across diverse datasets – tabular (IHDP), image (Uganda), and text (EEEC). 1) In the first setting, given a tabular sample $x$ describing a premature infant, DISCRET establishes a rule associating extremely underweight (weight ≤ 1.5) infants born to teenage mothers (mom age ≤ 19) with a history of drug use; such groups likely benefit from childcare visits (treatment), and will have highly improved cognitive outcomes. 2) In the second scenario on satellite images, for a sample $x$, DISCRET discerns a rule based on the presence of concepts like “high soil moisture” (reddish-pink pixels) and absence of minimal soil (brown pixels); thus characterizing areas with high soil moisture. DISCRET’s synthesized rule aligns with findings that government grants (treatment) are more effective in areas with higher soil moisture content (outcome) (Jerzak et al., 2023b). 3) Likewise, the text setting aims to measure the impact of gender (treatment) on the mood (outcome). Given a sentence $x$ where the gendered noun (“Betsy”) does not affect the semantic meaning, DISCRET’s rule focuses on mood-linked words in the sentence, i.e., “hilarious”.