Adaptively Grouped Contextual Bandits for Heterogeneous Human-AI Decision Making with Conformal Prediction Sets

Personalizing AI decision support for heterogeneous human decision-makers remains a key challenge. We study a collaboration workflow where AI provides a reduced prediction set via conformal prediction and the human makes the final decision based on the set. We formulate this personalization problem as a contextual bandit, where individual and task features form the context, candidate significance levels $\alpha$ serve as arms, and the optimal prediction-set size varies across contexts. To address large arm spaces and high-dimensional contexts, we introduce the Adaptively Grouped Contextual Bandit (AGCB) framework, which avoids global function approximation by exploiting two Human-AI structural assumptions: continuity and monotonicity. Continuity enables information sharing across nearby contexts and decisions, and drives a data-driven Zooming Mechanism that balances intra-group estimation error against inter-group approximation bias. Monotonicity converts each observation into directional counterfactual information over the $K$ candidate $\alpha$ values, reducing the arm-dependence factor from polynomial to logarithmic in $K$. Together, these mechanisms yield minimax-optimal dependence on the learning horizon $T$ for both cumulative and simple regret objectives. Empirical results confirm that AGCB achieves the strongest overall performance across most heterogeneous, data-scarce settings.

AI systems often help people make decisions by showing a short list of possible answers, such as possible diagnoses, labels, or financial options. A key challenge is deciding how long this list should be. A shorter list is easier to use, but it may leave out the correct answer. A longer list is safer, but it can also make the decision harder by giving the person too many choices. This problem is especially difficult because different people and different tasks may need different amounts of AI assistance. Our method learns how to personalize these lists. It uses two simple ideas about human-AI collaboration. First, people with similar backgrounds facing similar tasks often respond in similar ways to similar AI suggestions. Our method uses this to group similar situations together, while creating smaller groups when the data show important differences. Second, when the correct answer is still included, making the list shorter can help people decide more accurately by removing unnecessary options. Our method uses this pattern to learn from the list that was actually shown, as well as from nearby list lengths that could have been shown. Experiments on both simulated data and real human decision data show that our method gives more effective personalized AI support than existing approaches, especially when there is limited data and users behave differently. Overall, this work aims to make AI decision support more personalized, reliable, and useful in practice.