# End-to-End Balancing for Causal Continuous Treatment-Effect Estimation

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 47 references that are in the map.
- First posted on arXiv 27 Jul 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Balancing Covariates via Propensity Score Weighting](https://mlsanity.com/paper/1609.07494/) (2016)
- [Nonparametric methods for doubly robust estimation of continuous treatment effects](https://mlsanity.com/paper/1507.00747/) (2015)
- [VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments](https://mlsanity.com/paper/2103.07861/) (2021)
- [Covariate Balancing Propensity Score by Tailored Loss Functions](https://mlsanity.com/paper/1601.05890/) (2016)
- [Using Wasserstein Generative Adversarial Networks for the Design of Monte Carlo Simulations](https://mlsanity.com/paper/1909.02210/) (2019)
- [Minimal Dispersion Approximately Balancing Weights: Asymptotic Properties and Practical Considerations](https://mlsanity.com/paper/1705.00998/) (2017)
- [Permutation Weighting](https://mlsanity.com/paper/1901.01230/) (2019)
- [Entropy balancing is doubly robust](https://mlsanity.com/paper/1501.03571/) (2015)
- [Nonparametric Estimation of Population Average Dose-Response Curves using Entropy Balancing Weights for Continuous Exposures](https://mlsanity.com/paper/2003.02938/) (2020)
- [Kernel Optimal Orthogonality Weighting: A Balancing Approach to Estimating Effects of Continuous Treatments](https://mlsanity.com/paper/1910.11972/) (2019)
- [Double Robust Representation Learning for Counterfactual Prediction](https://mlsanity.com/paper/2010.07866/) (2020)

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

- [Probabilities of Causation for Continuous and Vector Variables](https://mlsanity.com/paper/2405.20487/) (2024)
- [Beyond the Known: Decision Making with Counterfactual Reasoning Decision Transformer](https://mlsanity.com/paper/2505.09114/) (2025)
- [Causal Estimation of Exposure Shifts with Neural Networks](https://mlsanity.com/paper/2302.02560/) (2023)
- [Bayesian Causal Inference in Doubly Gaussian DAG-probit Models](https://mlsanity.com/paper/2304.05976/) (2023)
- [Removing Spurious Correlation from Neural Network Interpretations](https://mlsanity.com/paper/2412.02893/) (2024)
- [Conditional Front-door Adjustment for Heterogeneous Treatment Assignment Effect Estimation Under Non-adherence](https://mlsanity.com/paper/2505.05677/) (2025)
- [Supervised Learning Objectives for High-Dimensional Treatment Effects](https://mlsanity.com/paper/2605.27281/) (2026)
- [Causal Generalization of Continuous Treatment Effects under Covariate Shift](https://mlsanity.com/paper/2608.19383/) (2026)
- [Weighted Spline-Expanded Networks with Distributional Balancing for Continuous Treatment Effects](https://mlsanity.com/paper/2609.33740/) (2026)

## Abstract (arXiv)

We study the problem of observational causal inference with continuous treatments in the framework of inverse propensity-score weighting. To obtain stable weights, we design a new algorithm based on entropy balancing that learns weights to directly maximize causal inference accuracy using end-to-end optimization. In the process of optimization, these weights are automatically tuned to the specific dataset and causal inference algorithm being used. We provide a theoretical analysis demonstrating consistency of our approach. Using synthetic and real-world data, we show that our algorithm estimates causal effect more accurately than baseline entropy balancing.