Invariant Feature Extraction Through Conditional Independence and the Optimal Transport Barycenter Problem: the Gaussian case

TL;DR

This paper introduces a method for extracting invariant features that predict a response variable without confounding, using optimal transport theory in the Gaussian case, with extensions to non-Gaussian scenarios.

Contribution

It develops a novel feature extraction approach based on conditional independence and optimal transport barycenter, providing a closed-form solution in the Gaussian case.

Findings

The method achieves invariant feature extraction in Gaussian settings.

It provides a closed-form linear feature extractor based on eigenvectors.

Extensions to non-Gaussian and non-linear cases are feasible with minimal modifications.

Abstract

A methodology is developed to extract $d$ invariant features $W=f(X)$ that predict a response variable $Y$ without being confounded by variables $Z$ that may influence both $X$ and $Y$. The methodology's main ingredient is the penalization of any statistical dependence between $W$ and $Z$ conditioned on $Y$, replaced by the more readily implementable plain independence between $W$ and the random variable $Z_Y = T(Z,Y)$ that solves the [Monge] Optimal Transport Barycenter Problem for $Z\mid Y$. In the Gaussian case considered in this article, the two statements are equivalent. When the true confounders $Z$ are unknown, other measurable contextual variables $S$ can be used as surrogates, a replacement that involves no relaxation in the Gaussian case if the covariance matrix $\Sigma_{ZS}$ has full range. The resulting linear feature extractor adopts a closed form in terms of the first $d$ eigenvectors of a known matrix. The procedure extends with little change to more general, non-Gaussian / non-linear cases.