Table of Contents
Fetching ...

Ultra-marginal Feature Importance: Learning from Data with Causal Guarantees

Joseph Janssen, Vincent Guan, Elina Robeva

TL;DR

This work introduces ultra-marginal feature importance (UMFI), which uses dependence removal techniques from the AI fairness literature as its foundation, and proves that UMFI satisfies axioms for feature importance methods that seek to explain the causal and associative relationships in data.

Abstract

Scientists frequently prioritize learning from data rather than training the best possible model; however, research in machine learning often prioritizes the latter. Marginal contribution feature importance (MCI) was developed to break this trend by providing a useful framework for quantifying the relationships in data. In this work, we aim to improve upon the theoretical properties, performance, and runtime of MCI by introducing ultra-marginal feature importance (UMFI), which uses dependence removal techniques from the AI fairness literature as its foundation. We first propose axioms for feature importance methods that seek to explain the causal and associative relationships in data, and we prove that UMFI satisfies these axioms under basic assumptions. We then show on real and simulated data that UMFI performs better than MCI, especially in the presence of correlated interactions and unrelated features, while partially learning the structure of the causal graph and reducing the exponential runtime of MCI to super-linear.

Ultra-marginal Feature Importance: Learning from Data with Causal Guarantees

TL;DR

This work introduces ultra-marginal feature importance (UMFI), which uses dependence removal techniques from the AI fairness literature as its foundation, and proves that UMFI satisfies axioms for feature importance methods that seek to explain the causal and associative relationships in data.

Abstract

Scientists frequently prioritize learning from data rather than training the best possible model; however, research in machine learning often prioritizes the latter. Marginal contribution feature importance (MCI) was developed to break this trend by providing a useful framework for quantifying the relationships in data. In this work, we aim to improve upon the theoretical properties, performance, and runtime of MCI by introducing ultra-marginal feature importance (UMFI), which uses dependence removal techniques from the AI fairness literature as its foundation. We first propose axioms for feature importance methods that seek to explain the causal and associative relationships in data, and we prove that UMFI satisfies these axioms under basic assumptions. We then show on real and simulated data that UMFI performs better than MCI, especially in the presence of correlated interactions and unrelated features, while partially learning the structure of the causal graph and reducing the exponential runtime of MCI to super-linear.
Paper Structure (39 sections, 14 theorems, 31 equations, 15 figures, 1 table, 3 algorithms)

This paper contains 39 sections, 14 theorems, 31 equations, 15 figures, 1 table, 3 algorithms.

Key Result

Theorem 3.1

Suppose that the data $(F,y)$ comes from a multivariate Gaussian distribution and that $\nu(\cdot)$ is positively linearly related to $I(Y;\cdot)$, then we can ensure that $U^{F,y}_{\nu}$ satisfies (i) the elimination axiom and (ii) the redundant information invariance and duplication symmetry axiom

Figures (15)

  • Figure 1: Results for the experiments on simulated data from Subsection \ref{['sec:SimulatedData']}. Feature importance scores are shown as a percentage of the total for each of $x_1$ to $x_4$ from $100$ replications. Results are shown for marginal contribution feature importance (MCI), ultra-marginal feature importance with linear regression (UMFI_LR), and ultra-marginal feature importance with pairwise optimal transport (UMFI_OT).
  • Figure 2: The full causal graph generating the data for the blood relation simulation experiment in Section \ref{['sec:Blood_relation_exp']}. Blood related vertices to the response $Y$ (blue) are coloured in red. $S$ and $X_4$ are directly causally related to $Y$, whereas $X_3$ is related to $Y$ via the common ancestor $S$.
  • Figure 3: Median feature importance scores provided by (a) MCI, (b) UMFI with linear regression, and (c) UMFI with pairwise optimal transport, for each gene in the BRCA dataset after 200 iterations. Genes colored in blue are known to be associated with breast cancer while genes colored in grey are random permutations of randomly selected genes, which we assume to be unassociated with breast cancer. The first and third quantiles of the scores are visualized as error bars for each gene.
  • Figure 4: Computation time for a single iteration of each method including: MCI (dark red), MCI with the soft 2-size-submodularity assumption (pink), UMFI_OT (light blue), and UMFI_LR (dark blue), plotted against the number of processed features.
  • Figure 5: PI-diagrams taken from griffith2014quantifying for $I(Y;F)$ when $|F|=2$ (left) and $|F|=3$ (right). Magenta represents unique information, redundant information is colored with yellow, and synergistic information is in cyan. The starred regions represent a single region.
  • ...and 10 more figures

Theorems & Definitions (28)

  • Definition 1
  • Definition 2
  • Theorem 3.1
  • proof
  • Theorem A.1: Symmetry of conditional mutual information yeung2002first
  • Theorem A.2: Chain rule for mutual information yeung2002first
  • Theorem A.3: Supermodularity under independence
  • proof
  • Theorem A.4: Data processing inequality
  • proof
  • ...and 18 more