# Revisiting Counterfactual Regression through the Lens of Gromov-Wasserstein Information Bottleneck

arXiv: https://arxiv.org/abs/2405.15505 · HTML: https://mlsanity.com/paper/2405.15505/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 51 references that are in the map.
- First posted on arXiv 24 May 2024.

## Builds on (most cited first, 22 of 22)

- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [A Brief Review of Domain Adaptation](https://mlsanity.com/paper/2010.03978/) (2020)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [A Survey on Causal Inference](https://mlsanity.com/paper/2002.02770/) (2020)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations](https://mlsanity.com/paper/2002.04083/) (2020)
- [Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition](https://mlsanity.com/paper/1707.02641/) (2017)
- [Wasserstein Dependency Measure for Representation Learning](https://mlsanity.com/paper/1903.11780/) (2019)
- [Convergence Rate of Frank-Wolfe for Non-Convex Objectives](https://mlsanity.com/paper/1607.00345/) (2016)
- [Learning Individual Causal Effects from Networked Observational Data](https://mlsanity.com/paper/1906.03485/) (2019)
- [Counterfactual Representation Learning with Balancing Weights](https://mlsanity.com/paper/2010.12618/) (2020)
- [Learning Weighted Representations for Generalization Across Designs](https://mlsanity.com/paper/1802.08598/) (2018)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [Optimal Transport for Treatment Effect Estimation](https://mlsanity.com/paper/2310.18286/) (2023)
- [Regularized Optimal Transport and the Rot Mover's Distance](https://mlsanity.com/paper/1610.06447/) (2016)
- [InfoOT: Information Maximizing Optimal Transport](https://mlsanity.com/paper/2210.03164/) (2022)
- [Optimal transport weights for causal inference](https://mlsanity.com/paper/2109.01991/) (2021)
- [Deep Treatment-Adaptive Network for Causal Inference](https://mlsanity.com/paper/2112.13502/) (2021)
- [Robustness-enhanced Uplift Modeling with Adversarial Feature Desensitization](https://mlsanity.com/paper/2310.04693/) (2023)
- [Causal Effect Estimation using Variational Information Bottleneck](https://mlsanity.com/paper/2110.13705/) (2021)
- [Reliable Estimation of Individual Treatment Effect with Causal Information Bottleneck](https://mlsanity.com/paper/1906.03118/) (2019)

## Built on by (most cited first, 4 of 4)

- [Proximity Matters: Local Proximity Enhanced Balancing for Treatment Effect Estimation](https://mlsanity.com/paper/2407.01111/) (2024)
- [Orthogonal Representation Learning for Estimating Causal Quantities](https://mlsanity.com/paper/2502.04274/) (2025)
- [Multiple Treatments Causal Effects Estimation with Task Embeddings and Balanced Representation Learning](https://mlsanity.com/paper/2511.09814/) (2025)
- [SGWIB:Sliced Gromov-Wasserstein Information Bottleneck for Video Highlight Detection](https://mlsanity.com/paper/2609.13966/) (2026)

## Abstract (arXiv)

As a promising individualized treatment effect (ITE) estimation method, counterfactual regression (CFR) maps individuals' covariates to a latent space and predicts their counterfactual outcomes. However, the selection bias between control and treatment groups often imbalances the two groups' latent distributions and negatively impacts this method's performance. In this study, we revisit counterfactual regression through the lens of information bottleneck and propose a novel learning paradigm called Gromov-Wasserstein information bottleneck (GWIB). In this paradigm, we learn CFR by maximizing the mutual information between covariates' latent representations and outcomes while penalizing the kernelized mutual information between the latent representations and the covariates. We demonstrate that the upper bound of the penalty term can be implemented as a new regularizer consisting of i) the fused Gromov-Wasserstein distance between the latent representations of different groups and ii) the gap between the transport cost generated by the model and the cross-group Gromov-Wasserstein distance between the latent representations and the covariates. GWIB effectively learns the CFR model through alternating optimization, suppressing selection bias while avoiding trivial latent distributions. Experiments on ITE estimation tasks show that GWIB consistently outperforms state-of-the-art CFR methods. To promote the research community, we release our project at https://github.com/peteryang1031/Causal-GWIB.