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, CUVET solves the within-cohort allocation problem by equalizing the marginal cost per unit value across each user group, yielding a closed-form cost-aware treatment assignment suited for large-scale industrial deployment. We also release CUVET-policy, an 86.7-million-impression public benchmark derived from real-world industrial A/B tests, providing an open-source evaluation framework for decision-focused learning. Across a multi-method benchmark suite, the Bayesian variant of CUVET achieves the largest cost-feasible value uplift on the public MT-LIFT dataset +11.1% under LP, vs. at most +1.3% for all baselines, while on CUVET-policy, where treatment effects are sub-percent, CUVET variants are the most efficient and strongest cost-feasible methods.