# Learning Generalized Gumbel-max Causal Mechanisms

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 35 references that are in the map.
- First posted on arXiv 11 Nov 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Wasserstein GAN](https://mlsanity.com/paper/1701.07875/) (2017)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Stochastic Optimization for Large-scale Optimal Transport](https://mlsanity.com/paper/1605.08527/) (2016)
- [A* Sampling](https://mlsanity.com/paper/1411.0030/) (2014)
- [Counterfactual Off-Policy Evaluation with Gumbel-Max Structural Causal Models](https://mlsanity.com/paper/1905.05824/) (2019)
- [Backpropagation through the Void: Optimizing control variates for black-box gradient estimation](https://mlsanity.com/paper/1711.00123/) (2017)
- [Smooth and Sparse Optimal Transport](https://mlsanity.com/paper/1710.06276/) (2017)
- [Woulda, Coulda, Shoulda: Counterfactually-Guided Policy Search](https://mlsanity.com/paper/1811.06272/) (2018)
- [Sample-Efficient Reinforcement Learning via Counterfactual-Based Data Augmentation](https://mlsanity.com/paper/2012.09092/) (2020)
- [Differential Properties of Sinkhorn Approximation for Learning with Wasserstein Distance](https://mlsanity.com/paper/1805.11897/) (2018)
- [Structured Optimal Transport](https://mlsanity.com/paper/1712.06199/) (2017)
- [Learning Wasserstein Embeddings](https://mlsanity.com/paper/1710.07457/) (2017)

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

- [Counterfactual Identifiability of Bijective Causal Models](https://mlsanity.com/paper/2302.02228/) (2023)
- [Counterfactual (Non-)identifiability of Learned Structural Causal Models](https://mlsanity.com/paper/2301.09031/) (2023)
- [Counterfactual Inference of Second Opinions](https://mlsanity.com/paper/2203.08653/) (2022)
- [Counterfactual Token Generation in Large Language Models](https://mlsanity.com/paper/2409.17027/) (2024)
- [Gumbel Counterfactual Generation From Language Models](https://mlsanity.com/paper/2411.07180/) (2024)
- [Counterfactual Analysis in Dynamic Latent State Models](https://mlsanity.com/paper/2205.13832/) (2022)
- [Evaluation of Large Language Models via Coupled Token Generation](https://mlsanity.com/paper/2502.01754/) (2025)
- [Deep Counterfactual Estimation with Categorical Background Variables](https://mlsanity.com/paper/2210.05811/) (2022)
- [Understanding What Affects the Generalization Gap in Visual Reinforcement Learning: Theory and Empirical Evidence](https://mlsanity.com/paper/2402.02701/) (2024)
- [Counterfactual Influence in Markov Decision Processes](https://mlsanity.com/paper/2402.08514/) (2024)
- [Curiosity in Hindsight: Intrinsic Exploration in Stochastic Environments](https://mlsanity.com/paper/2211.10515/) (2022)
- [Towards Computationally Efficient Responsibility Attribution in Decentralized Partially Observable MDPs](https://mlsanity.com/paper/2302.12676/) (2023)
- [Robust Counterfactual Inference in Markov Decision Processes](https://mlsanity.com/paper/2502.13731/) (2025)
- [Causal Graph Discovery from Self and Mutually Exciting Time Series](https://mlsanity.com/paper/2301.11197/) (2023)
- [Off-Policy Selection for Initiating Human-Centric Experimental Design](https://mlsanity.com/paper/2410.20017/) (2024)
- [Learning to Control Coupled-Dynamics Environments with Joint Markov Decision Processes](https://mlsanity.com/paper/2608.22765/) (2026)

## Abstract (arXiv)

To perform counterfactual reasoning in Structural Causal Models (SCMs), one needs to know the causal mechanisms, which provide factorizations of conditional distributions into noise sources and deterministic functions mapping realizations of noise to samples. Unfortunately, the causal mechanism is not uniquely identified by data that can be gathered by observing and interacting with the world, so there remains the question of how to choose causal mechanisms. In recent work, Oberst & Sontag (2019) propose Gumbel-max SCMs, which use Gumbel-max reparameterizations as the causal mechanism due to an intuitively appealing counterfactual stability property. In this work, we instead argue for choosing a causal mechanism that is best under a quantitative criteria such as minimizing variance when estimating counterfactual treatment effects. We propose a parameterized family of causal mechanisms that generalize Gumbel-max. We show that they can be trained to minimize counterfactual effect variance and other losses on a distribution of queries of interest, yielding lower variance estimates of counterfactual treatment effect than fixed alternatives, also generalizing to queries not seen at training time.