Primary Submission Category: Design of Experiments
AllocOT: Constrained Treatment Assignment via Semi-Discrete Optimal Transport
Authors: Qiuran Lyu, Mingxun Wang,
Presenting Author: Mingxun Wang*
Treatment assignment is central to causal inference but often must respect real-world constraints, including adaptive updates, fairness across demographic groups, and budget or capacity limits. Moreover, as the number of treatment arms grows and constraints become complex, generic solvers such as Gurobi can become computationally expensive.
We propose Constrained Allocation via Optimal Transport (AllocOT), a semi-discrete optimal transport framework that casts constrained assignment as transporting the covariate distribution to a finite set of treatment arms. AllocOT gives assignment rules for each unit by minimizing an overall cost; for outcome-driven allocation, the cost is the negative predicted outcome under each arm, so minimizing cost maximizes expected benefit. Practical requirements are imposed as linear quota constraints on aggregate assignment frequencies, for example, fixed arm ratios, subgroup-specific exposure bounds for fairness and balance, and capacity or budget caps for scarce or costly interventions.
We validate AllocOT performance across simulation studies covering varying numbers of arms and constraint regimes. The result shows that AllocOT delivers high-quality allocations with substantially improved scalability and runtime compared with generic optimization solvers, making it a practical approach for large-scale constrained treatment assignment in causal inference.