Sample-Adaptivity Tradeoff in On-Demand Sampling

AI summary

Overview

• Research area: Machine learning theory — specifically multi-distribution learning (MDL), on-demand sampling, and the theory of adaptivity in learning algorithms.
• Technical level: Advanced. The paper is written for a theoretical machine learning audience and relies on VC dimension, PAC learning, boosting, regret-minimization (Hedge), and convex optimization over the simplex.
• Scope: The paper characterizes how many rounds of adaptive sampling are needed to reach a given sample complexity in multi-distribution learning, both in the realizable and agnostic settings, and introduces a general abstract framework (OODS) that captures the tradeoff.

What This Paper Is About

Multi-distribution learning gives a learner sampling access to k unknown distributions D₁, …, D_k and asks for a single predictor that minimizes worst-case error across all of them. Prior algorithms reach near-optimal sample complexity but need a large number of adaptive sampling rounds — sometimes as many rounds as samples. This paper asks how much adaptivity is actually necessary: it quantifies the tradeoff between the number of rounds an algorithm uses and the total number of samples it must draw, and it builds a general framework for reasoning about that tradeoff.

Key Contributions

1. A tight sample-adaptivity tradeoff for realizable MDL. The paper proves that an r-round algorithm can achieve sample complexity roughly dk^{Θ(1/r)}/ε, and that a k^{Ω(1/r)} overhead is unavoidable, giving a smooth tradeoff between sample complexity and round complexity.
2. A near-optimal few-round algorithm for agnostic MDL. The paper gives a min{Õ(√k), O(k log k)}-round algorithm achieving the optimal Õ((d+k)/ε²) sample complexity, improving over previous algorithms whose round complexity was at least poly(1/ε).
3. A new framework: Optimization via On-Demand Sampling (OODS). OODS abstracts the sample-adaptivity tradeoff as maximizing an unknown concave function f over the simplex Δ^{k−1}, where the algorithm can only query f inside a "trust region" (a box or an ellipsoid) determined by a cap vector that can grow each round. The paper establishes nearly tight upper and lower bounds on round complexity in this setting, and shows that OODS algorithms transfer to agnostic MDL.
4. Algorithm-independent lower bounds. The paper proves poly(k) lower bounds on the round complexity of OODS, implying that achieving sub-polynomial round complexity for agnostic MDL would require techniques that bypass the inherent hardness of the OODS formulation.

Main Findings

• Realizable upper bound (Theorem 1): Algorithm 1 is an r-round (ε, δ)-PAC algorithm for realizable MDL with sample complexity O(k^{2/r} log k · d/ε + k log(k) log(k/δ)/ε). At r = log k this recovers the near-optimal Õ((d+k)/ε) bound; at a small constant r (e.g., r = 4) it gives Õ((d√k + k)/ε).
• Constant rounds suffice for a large improvement: The paper notes that a small constant number of rounds — "say, 3 rounds (!)" — achieves an Õ(d√k/ε) sample complexity, a significant improvement over the Ω(dk/ε) required with no adaptivity (r = 1). Remark 1 states that using boost-by-majority or recursive boosting achieves Õ((d√k + k)/ε) using exactly 3 rounds.
• Realizable lower bound (Theorem 2): For every r = O(log k) and sufficiently large d, every r-round algorithm for realizable MDL has sample complexity Ω(dk^{1/r}/(r log²k)). This nearly matches the dk^{2/r} log k term of Theorem 1, up to a poly(r, log k) factor and a factor of 2 in the exponent of k^{Θ(1/r)}.
• Two-round lower bound intuition: For r = 2 the lower bound is Ω(d√k). The hard instance uses linear functions over F₂^d with difficulty levels consisting of one copy of Θ(d), √k copies of Θ(d/√k), and k − √k − 1 copies of Θ(d/k). A three-round learner can succeed with O(d) sample complexity by sampling Θ(d/k) from each distribution, then Θ(d/√k) from the √k + 1 remaining suspects, then Θ(d) on the last distribution; a two-round learner must skip a step and incurs Ω(d√k).
• Agnostic upper bound (Proposition 1): There is a min{Õ(√k), O(k log k)}-round MDL algorithm with sample complexity Õ((d+k)/ε²).
• Round complexity of LazyHedge (Propositions 2 and 3): The box version of LazyHedge takes at most O(k log k) rounds; the ellipsoid version takes at most Õ(√k) rounds. The proof for the box version tracks a "culprit" index whose historical high must grow by a factor of C, so each index can be the culprit at most O(log_C k) times. The ellipsoid proof splits cap updates into Type I (a coordinate reaches a historical high > 1/√k) and Type II, and uses Lemma 3 of ZZC+ (24), which states that Σ_i max_t w_i^(t) ≤ O(log⁸(k/(εδ))) = Õ(1) with probability 1 − δ.
• Round complexity in perspective: A straightforward implementation of the ZZC+ (24) approach needs T = Θ((log k)/ε²) sampling rounds; the new Õ(√k) round complexity is lower when ε ≪ 1/k^{1/4}.
• Lower bounds in OODS: The paper proves poly(k) lower bounds on the round complexity of OODS, showing that getting to sub-polynomial round complexity would require fundamentally new techniques.

Methodology in Plain English

For the realizable setting, the authors run a variant of AdaBoost with a specific notion of margin over the k distributions. In each round, a base PAC learner is called on a mixture q_t of the k distributions to produce a predictor h_t with low error on that mixture; Markov's inequality then implies h_t has error above a threshold τ on only a small weighted fraction of distributions. The mixture weights are then updated in the AdaBoost style based on whether each distribution's error exceeds τ/2, and the final predictor is a majority vote over the r predictors produced. Choosing the margin θ = r/(2 log k) and the probability p so that the AdaBoost product term falls below 1/k guarantees that at least a 1/2 + θ/2 fraction of the r predictors have error at most τ on every distribution, which yields final error (1 + 1/θ)τ = ε.

For the realizable lower bound, the authors construct a hard instance based on linear functions over F₂^d. The target classifier is a uniformly random linear function, and the k distributions are uniform over randomly chosen subspaces whose dimensions (the "difficulty levels") are set to Θ(d), Θ(d/√k), and Θ(d/k) in the proportions described above; these subspaces are chosen to be linearly independent, so samples from one distribution reveal nothing about the target on another. The argument shows every r-round algorithm must "skip" a step and incur a k^{1/r} overhead.

For the agnostic setting, the authors start from the approach of ZZC+ (24): maintain k datasets, run Hedge for T = Θ((log k)/ε²) iterations with step size η = Θ(ε), and in each iteration do an ERM step on the current mixture followed by a Hedge update using fresh error estimates. They modify this into LazyHedge, which maintains a cap vector w̄ as a proxy for dataset sizes and only draws extra samples when the current weight vector falls outside the observable region defined by the cap. Two observable regions are considered: a box {w : w_i ≤ w̄_i} and an ellipsoid {w : Σ w_i²/w̄_i ≤ 1}.

The OODS framework generalizes this: instead of sampling, the algorithm maximizes an unknown concave f over the simplex using a first-order oracle restricted to the observable region, and it can expand the cap each round. The sample overhead Σ w̄_i^{(r)} proxies the sample complexity, and the number of cap expansions proxies the round complexity.

Why This Matters

The paper settles a question left open across a series of results (BHPQ (17), CZZ (18), NZ (18), HJZ (22), AHZ (23), ZZC+ (24), Pen (24)): whether a constant number of adaptive rounds, or one independent of the accuracy ε, could recover the sample complexity of fully adaptive algorithms. It answers yes for the realizable case (with a smooth, tight tradeoff) and shows a poly(k)-round algorithm suffices in the agnostic case, while proving that sub-polynomial round complexity would need to overcome the inherent hardness of OODS. The OODS abstraction is offered as a tool of independent interest that captures most existing MDL algorithms.

Real-world relevance, as framed by the framework's motivating applications:

• Federated learning, where data is spread across many clients and the learner must decide how much to sample from each client.
• Multi-task learning and domain adaptation, where the goal is one predictor that performs well across heterogeneous data sources.
• Fair and robust machine learning, where worst-case error across subpopulations is the objective.
• Dataset selection and data collection pipelines, where adaptively reweighting or augmenting the training set improves downstream accuracy.

Industry relevance: adaptivity forces sequential data collection, which slows end-to-end training and limits parallelism and scalability. Knowing that a constant number of rounds (indeed 3 rounds in the realizable case) can be nearly as sample-efficient as full adaptivity means practitioners can batch data collection and still retain most of the statistical benefit.

Future Directions

• Closing the exponent gap in the realizable tradeoff: the upper bound has k^{2/r} where the lower bound has k^{1/r}, so a factor-of-2 gap in the exponent of k remains at intermediate values of r.
• Can the poly(k) round complexity barrier in OODS and agnostic MDL be broken? The paper's lower bounds suggest sub-polynomial round complexity would require fundamentally new techniques that bypass the inherent hardness of OODS; identifying whether such techniques exist is left open.
• Extending OODS to other observable regions and objectives: the paper analyzes box and ellipsoid trust regions, and the framework abstracts the sample-adaptivity tradeoff broadly enough to potentially capture other algorithm families.
• Transferring the OODS toolkit to other sequential decision problems, such as the batched multi-armed bandit setting (JJNZ (16), GHRZ (19), JYT+ (24), JZZ (25)) and adaptivity-query tradeoffs in submodular optimization, which the related work section lists as parallel lines of study.

Target Audience

Theoretical machine learning researchers working on learning theory, PAC learning, multi-distribution learning, and the role of adaptivity in algorithm design. It will also interest researchers in bandits and online learning, in convex optimization with restricted oracle access, and in federated or multi-task learning who need to reason about how many rounds of adaptive data collection are worth the statistical savings. Readers without a background in VC dimension, boosting, and regret-minimization will find the formal results difficult to follow, though the high-level tradeoff statements are accessible.

Authors’ abstract

We study the tradeoff between sample complexity and round complexity in on-demand sampling, where the learning algorithm adaptively samples from $k$ distributions over a limited number of rounds. In the realizable setting of Multi-Distribution Learning (MDL), we show that the optimal sample complexity of an $r$-round algorithm scales approximately as $dk^{Θ(1/r)} / ε$. For the general agnostic case, we present an algorithm that achieves near-optimal sample complexity of $\widetilde O((d + k) / ε^2)$ within $\widetilde O(\sqrt{k})$ rounds. Of independent interest, we introduce a new framework, Optimization via On-Demand Sampling (OODS), which abstracts the sample-adaptivity tradeoff and captures most existing MDL algorithms. We establish nearly tight bounds on the round complexity in the OODS setting. The upper bounds directly yield the $\widetilde O(\sqrt{k})$-round algorithm for agnostic MDL, while the lower bounds imply that achieving sub-polynomial round complexity would require fundamentally new techniques that bypass the inherent hardness of OODS.