# MOSIC: Model-Agnostic Optimal Subgroup Identification with Multi-Constraint for Improved Reliability

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 71 references that are in the map.
- First posted on arXiv 29 Apr 2025.

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

- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [A Survey on Bias and Fairness in Machine Learning](https://mlsanity.com/paper/1908.09635/) (2019)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [Bandits with Knapsacks](https://mlsanity.com/paper/1305.2545/) (2013)
- [Global Convergence of ADMM in Nonconvex Nonsmooth Optimization](https://mlsanity.com/paper/1511.06324/) (2015)
- [Balancing Covariates via Propensity Score Weighting](https://mlsanity.com/paper/1609.07494/) (2016)
- [Convergence Rate of Frank-Wolfe for Non-Convex Objectives](https://mlsanity.com/paper/1607.00345/) (2016)
- [More Efficient Policy Learning via Optimal Retargeting](https://mlsanity.com/paper/1906.08611/) (2019)
- [Causal Rule Sets for Identifying Subgroups with Enhanced Treatment Effect](https://mlsanity.com/paper/1710.05426/) (2017)
- [Reliable Off-Policy Learning for Dosage Combinations](https://mlsanity.com/paper/2305.19742/) (2023)
- [GradTree: Learning Axis-Aligned Decision Trees with Gradient Descent](https://mlsanity.com/paper/2305.03515/) (2023)
- [ADMM for Multiaffine Constrained Optimization](https://mlsanity.com/paper/1802.09592/) (2018)
- [Selecting optimal subgroups for treatment using many covariates](https://mlsanity.com/paper/1802.09642/) (2018)
- [An evaluation framework for comparing causal inference models](https://mlsanity.com/paper/2209.00115/) (2022)
- [Individualized treatment rules under stochastic treatment cost constraints](https://mlsanity.com/paper/2201.06669/) (2022)
- [CAPITAL: Optimal Subgroup Identification via Constrained Policy Tree Search](https://mlsanity.com/paper/2110.05636/) (2021)
- [CURLS: Causal Rule Learning for Subgroups with Significant Treatment Effect](https://mlsanity.com/paper/2407.01004/) (2024)

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

- [Learning Subgroups with Maximum Treatment Effects without Causal Heuristics](https://mlsanity.com/paper/2511.20189/) (2025)

## Abstract (arXiv)

Current subgroup identification methods typically follow a two-step approach: first estimate conditional average treatment effects and then apply thresholding or rule-based procedures to define subgroups. While intuitive, this decoupled approach fails to incorporate key constraints essential for real-world clinical decision-making, such as subgroup size and propensity overlap. These constraints operate on fundamentally different axes than CATE estimation and are not naturally accommodated within existing frameworks, thereby limiting the practical applicability of these methods. We propose a unified optimization framework that directly solves the primal constrained optimization problem to identify optimal subgroups. Our key innovation is a reformulation of the constrained primal problem as an unconstrained differentiable min-max objective, solved via a gradient descent-ascent algorithm. We theoretically establish that our solution converges to a feasible and locally optimal solution. Unlike threshold-based CATE methods that apply constraints as post-hoc filters, our approach enforces them directly during optimization. The framework is model-agnostic, compatible with a wide range of CATE estimators, and extensible to additional constraints like cost limits or fairness criteria. Extensive experiments on synthetic and real-world datasets demonstrate its effectiveness in identifying high-benefit subgroups while maintaining better satisfaction of constraints.