CUVET: A Partitioning Approach for Continuous Treatment Assignment At Scale

Treatment assignment problems arise wherever limited budget must be allocated to heterogeneous users, with applications ranging from personalized recommendations to online advertising and healthcare. In such settings, individuals exhibit heterogeneous responses to different treatments, making it essential to learn cost-aware personalized treatments. This paper introduces the Cost per Unit Value Equalization Tree (CUVET) algorithm, a novel treatment assignment approach that partitions the user space. Under a diminishing-returns (power-law) assumption, it solves the treatment assignment problem by equalizing the marginal cost per unit value across each user group. This leads to a closed-form cost-aware treatment assignment solution, making it particularly suited for large-scale applications such as digital advertising. We also release CUVET-policy, a $87$-million-impression public benchmark derived from real-world industrial A/B tests, providing an open-source evaluation framework for decision-focused learning. On both CUVET-policy and the public MT-LIFT dataset, CUVET significantly improves baselines' total value by +1\% and +12.5\% respectively, satisfying budget constraints. The dataset is available on https://huggingface.co/datasets/anonadata/CUVET-policy.