How to Shrink Confidence Sets for Many Equivalent Discrete Distributions?
Odalric-Ambrym Maillard, Mohammad Sadegh Talebi
TL;DR
This work studies a set of $K$ discrete distributions $(p_k)_{k\in\mathcal K}$ over a common alphabet that are permutation-equivalent through an unknown canonical $q$ and permutations, and aims to tighten individual confidence sets by exploiting this structure. It introduces a low-complexity algorithm to identify compatible matchings among distributions and construct refined confidence intervals that remain valid while using all samples, supported by finite-time high-probability bounds. The analysis shows that refined confidence sets shrink at rates $O\big(1/\sqrt{\sum_k n_k}\big)$ for points in the support and $O\big(1/\max_k n_k\big)$ outside the support, improving over naive per-distribution estimation and enabling significant gains when $\mathcal K$ is large. The approach is instantiated with surrogate confidence intervals (e.g., KL, Bernstein, empirical Bernstein) and demonstrated on reinforcement learning tasks (RiverSwim), where exploiting permutation-equivalence yields notably tighter estimates and lower regret compared to baselines. The results offer practical guidance on when refinement pays off (in terms of data and structure) and how to balance statistical gains with computational considerations, while outlining extensions to broader automorphism families and applications.
Abstract
We consider the situation when a learner faces a set of unknown discrete distributions $(p_k)_{k\in \mathcal K}$ defined over a common alphabet $\mathcal X$, and can build for each distribution $p_k$ an individual high-probability confidence set thanks to $n_k$ observations sampled from $p_k$. The set $(p_k)_{k\in \mathcal K}$ is structured: each distribution $p_k$ is obtained from the same common, but unknown, distribution q via applying an unknown permutation to $\mathcal X$. We call this \emph{permutation-equivalence}. The goal is to build refined confidence sets \emph{exploiting} this structural property. Like other popular notions of structure (Lipschitz smoothness, Linearity, etc.) permutation-equivalence naturally appears in machine learning problems, and to benefit from its potential gain calls for a specific approach. We present a strategy to effectively exploit permutation-equivalence, and provide a finite-time high-probability bound on the size of the refined confidence sets output by the strategy. Since a refinement is not possible for too few observations in general, under mild technical assumptions, our finite-time analysis establish when the number of observations $(n_k)_{k\in \mathcal K}$ are large enough so that the output confidence sets improve over initial individual sets. We carefully characterize this event and the corresponding improvement. Further, our result implies that the size of confidence sets shrink at asymptotic rates of $O(1/\sqrt{\sum_{k\in \mathcal K} n_k})$ and $O(1/\max_{k\in K} n_{k})$, respectively for elements inside and outside the support of q, when the size of each individual confidence set shrinks at respective rates of $O(1/\sqrt{n_k})$ and $O(1/n_k)$. We illustrate the practical benefit of exploiting permutation equivalence on a reinforcement learning task.
