PaperlayerView exact revision on arXiv

Source-linked AI summary

Bandits in Prod: Hyperparameter Optimization at Inference Time

Louis Abraham, Tuan-Anh Nguyen, Nicolas Devatine

arXiv:2609.01335v2cs.LGcs.AI

TL;DR

Production systems may need to tune inference-time configurations from noisy live feedback without representative validation data. The paper formulates this as OHPO, introduces IMABO with IMOSS and configurable proposal oracles, and reports better regret than baselines across classical and LLM-agent settings.

  • Problem

    OHPO addresses configuration tuning from live request feedback when representative validation data is unavailable.

  • Method

    IMABO combines a bandit policy for allocating pulls among sampled configurations with an oracle for proposing new ones, using IMOSS as an anytime policy.

  • Results

    IMABO outperformed its baselines on classical machine-learning models and a language-model agent while remaining computationally efficient for online optimization.

  • Takeaways & Limitations

    IMABO supports mixed, conditional, and tree-structured spaces without discretizing coordinates and requests fewer distinct configurations than rounds because its active set grows as t^β.

  • Takeaways & Limitations

    Adaptive sampling can create survivorship bias, making the true performance of configurations favored by unusually strong early rewards difficult to estimate.

Abstract

from arXiv · show

Many production systems can assess a configuration only by using it on live requests and observing noisy feedback. Modern agentic systems are a prominent example, with inference-time choices such as model selection, retrieval depth, prompting strategy, and decoding temperature, yet often with no representative validation data. We formalize this setting as Online Hyperparameter Optimization (OHPO) and cast it as an infinitely many-armed bandit over mixed and conditional search spaces. We introduce IMABO, a general framework that combines any bandit policy for choosing among already sampled configurations with any oracle for proposing new ones. We instantiate it with IMOSS, a restart-free anytime policy whose active set grows as $t^β$, and prove an expected cumulative quantile-regret bound of $O(p_ρ^{-1/β} + T^{(1+β)/2})$, where $β\in(0,1)$ controls active-set growth and $p_ρ$ lower-bounds the probability that a proposed configuration falls in the top-$ρ$ fraction of the search space. We combine IMOSS with three practical oracles: a Tree-structured Parzen Estimator, an incumbent-mutation oracle driven by a per-coordinate bandit, and a pretrained tabular foundation model, all three improving over the uniform random oracle baseline. IMABO outperforms all baselines in terms of regret across diverse OHPO settings, from tuning classical machine-learning models to configuring LLM-based agents. Our implementation is available at https://github.com/Tiime-Software/IMABO.

1 Introduction

OHPO addresses production settings where configurations can be evaluated only through noisy live requests rather than representative validation data. IMABO frames this as an infinitely many-armed bandit and combines adaptive bandit selection with oracle-guided configuration proposals.

  • Problem: OHPO evaluates inference-time configurations through live requests when no representative validation dataset is available.Examples include model choice, decoding temperature, and prompting strategy.
  • Problem: Each configuration is an arm, while each served request is a pull in an infinitely many-armed bandit.Mixed continuous, integer, and categorical spaces can contain far more configurations than available requests.
  • Approach: IMABO pairs a bandit policy over sampled configurations with an oracle that proposes new configurations.The framework requires only that configurations can be scored from live feedback and makes no further search-space assumptions.
  • Approach: IMOSS provides an anytime bandit policy with a sublinear cumulative quantile-regret bound and multiple arm-creation oracles.The concrete variants include Random, TPE, TabPFN, and mutate-KL×PE.
  • Evaluation: IMABO outperforms baselines across classical HPOBench tasks and a real LLM-based question-answering system.The evaluated LLM system uses retrieval, prompting, and model choices in a mixed search space.

2 Related Work

The related work spans offline HPO, infinitely many-armed bandits, and production-oriented online tuning. Prior methods often rely on fixed discretizations, smoothness assumptions, or enumerated arms, while IMABO targets broader mixed spaces.

  • Hyperparameter optimization: Offline HPO methods optimize fixed objectives using grid search, random search, or model-based surrogates such as Gaussian processes, random forests, and TPE.TPE handles continuous and categorical parameters without requiring a metric on the search space.
  • Multi-, continuum-, and infinite-armed bandits: Infinite-armed bandit methods address settings with more arms than pulls, but existing approaches commonly assume nearby arms have similar rewards or sample new arms randomly.The smoothness assumption is unsuitable for categorical choices such as model selection.
  • Bandits & HPO: Offline bandit-HPO methods allocate training budgets across configurations, while D-TTTS uses a fixed reservoir of configurations.These approaches do not match the live-serving setting described here.
  • Implementation: The paper provides code and experiments at the stated public repository.The implementation link is given in the paper.
  • Inference-time tuning of LLM systems: AutoRAG-HP tunes discretized RAG pipelines online without a regret guarantee, whereas BanditSpec optimizes token-level decoding latency.Other related systems treat a bare model identity as an arm from an enumerated pool.

3 Online Hyperparameter Optimization

OHPO models live configuration selection as sequential decision-making over a mixed search space with noisy rewards and no replayable validation set. Because global-optimum competition is infeasible, performance is measured against a reward quantile.

  • Problem formulation: OHPO allows continuous, integer, categorical, and tree-structured hyperparameters, with each committed configuration serving the next request.Rewards lie in [0, 1] and are drawn from an unknown configuration-specific distribution.
  • Problem formulation: The learner receives only the served request’s reward, with no side validation set and no ability to replay past rounds under another configuration.The horizon is unknown.
  • Regret: Performance is compared with a top-ρ quantile of mean reward rather than the global optimum.The quantile is defined relative to a baseline sampling distribution over the search space.
  • Regret: Cumulative ρ-regret sums the reward shortfall below the target quantile across served requests.Simple ρ-regret instead measures the quality of the configuration returned after T rounds.

4 Infinitely Many-Armed Bandit with Oracle

IMABO maintains an active set of configurations, uses a bandit policy to select known arms, and delegates new-arm proposals to an oracle. IMOSS grows the active set sublinearly and supports several practical oracle designs for mixed spaces.

  • IMABO: At each round, IMABO either admits and serves a new oracle-proposed arm or serves a known arm selected by the bandit policy.The algorithm updates arm statistics after observing the live reward.
  • IMABO: IMABO differs from standard infinitely many-armed bandits by using a history-dependent oracle instead of a fixed reservoir to admit new arms.A baseline-sampling oracle recovers the standard reservoir formulation, while history-aware oracles can target promising regions.
  • IMOSS: IMOSS uses the anytime MOSS index with current active-set size K_t and serves the arm with the highest index.Its exploration bonus shrinks with both the arm’s observations and the active-set size.
  • IMOSS: The active set grows as ⌊t^β⌋, trading off discovering promising regions against obtaining reliable estimates for each arm.A warm-up draws N_s = 10 configurations from the baseline distribution before index-based selection takes over.
  • Guarantee: The regret analysis requires only that an oracle proposes a top-ρ configuration with conditional probability at least p_ρ whenever expansion occurs.IMOSS-Random satisfies this condition with p_ρ = ρ.
  • Oracles: IMOSS-TPE fits good and bad density models over whole configurations and proposes candidates with high density ratio.It natively handles continuous, integer, and categorical parameters without a metric on the search space.
  • Implementation: On finite spaces, an oracle proposal that duplicates an active arm is served again rather than admitted as a duplicate.This implementation avoids requerying the oracle until it produces a fresh arm.
  • Oracles: The mutation oracle changes one coordinate of the empirically best arm, selecting coordinates with a KL-UCB bandit and values with a one-dimensional Parzen estimator.The proposal chooses among 24 candidate values and must differ from the current coordinate value.

5 Experiments

The experiments evaluate IMABO across synthetic, discrete, continuous, and LLM-based online hyperparameter optimization settings. IMOSS-based variants generally improve cumulative or online regret through adaptive allocation and learned or mutation-based proposal oracles, while performance depends on search-space structure.

  • Synthetic benchmarks: IMOSS-TPE achieves the most balanced trade-off between simple regret and online average regret across the four synthetic functions.It maintains lower average regret than tree-based baselines on Sin1 and Garland, and outperforms StroquOOL on Rastrigin at T = 10000 with comparable simple regret.
  • Synthetic benchmarks: Hier-MAB-11 excels on separable Sin1 and Rastrigin but degrades on rapidly oscillating Garland and non-separable Gaussian.On Gaussian, its coordinate-wise updates remain near the local mode even though the grid contains a point close to the global mode.
  • Synthetic benchmarks: IMOSS-TPE achieves lower cumulative regret than every fixed-replication TPE-k baseline across all four synthetic functions.The advantage persists against TPE-50 and TPE-70, indicating that adaptive evaluation allocation contributes beyond the number of configurations available to the surrogate.
  • OpenML benchmarks: IMOSS-mutate-KL×PE attains the lowest cumulative regret on all three discrete OpenML HPO tasks, while every IMABO variant beats UCB-AIR.The learned oracles are strongest overall, and proposal-quality diagnostics help explain their advantage.
  • OpenML benchmarks: IMOSS-TPE and IMOSS-mutate-KL×PE obtain the lowest cumulative regret on both continuous OpenML tasks and reach near-zero simple regret.They are essentially tied with StroquOOL on Logistic Regression but clearly ahead on Support Vector Machine.
  • LLM question answering: On HotpotQA, IMOSS-mutate-KL×PE and IMOSS-TabPFN improve over Random on both metrics and online performance, with IMOSS-TabPFN attaining the lowest online average regret.Their coordinate-local proposals suit the near-separable search space, whereas IMOSS-TPE is more dispersed and trails Hier-MAB.

6 Practical Usage

The paper extends IMABO for production settings with delayed, censored feedback and evolving search spaces. Experiments show that delay-aware switching improves online regret by concentrating scarce observations across fewer active configurations.

  • Production challenges: IMABO extensions address delayed or missing rewards and configuration spaces that grow during long-running production searches.The policy continues selecting configurations without waiting for feedback, while schedule resets preserve past rewards when the space expands.
  • Delayed and censored feedback: The delay-aware rule counts pending pulls fractionally using the estimated reward-return probability, allowing expansion and arm indices to reflect expected feedback.A pending pull contributes p of a reward, while pulls exceeding a fixed patience window are dropped from the count.
  • Experimental setup: The LCBench evaluation uses three seven-dimensional mixed search spaces and Bernoulli rewards derived from surrogate validation accuracy.The instances are higgs, APSFailure, and Fashion-MNIST, accessed through the YAHPO-Gym surrogate.
  • Delayed and censored feedback: Across all three LCBench instances, the delay-aware rule sits between the delay-blind variant and the instant-feedback skyline, closing most of the gap to the skyline.The delayed methods share the same censored stream, so their performance difference is attributed to the switching rule.
  • Delayed and censored feedback: The delay-aware rule settles at about 60 active arms and gains tighter estimates by spreading approximately 1,900 observed rewards across fewer arms than the naive rule.The naive rule reaches approximately 100 arms, while both delayed methods observe the same approximately 1,900 rewards; the delay-aware optimizer finishes below UCB-AIR despite instant feedback for UCB-AIR.
  • Evolving search spaces: When the search space changes, only the admission schedule is reset, so newly proposed configurations can use all previously collected rewards.This differs from a true restart, which would discard prior knowledge.

7 Conclusion

The paper frames online hyperparameter optimization as an infinitely many-armed bandit problem and introduces IMABO with IMOSS for mixed, conditional search spaces. Across classical machine-learning and language-model-agent tasks, oracle-based IMABO variants outperform baselines, while adaptive sampling creates an estimation challenge.

  • Conclusion: IMABO combines a bandit policy for allocating pulls with an oracle for suggesting configurations in mixed and possibly conditional search spaces.Its instantiations use Random, TPE, mutate-KL×PE, or TabPFN oracles.
  • Conclusion: IMABO variants outperform their baselines across classical machine-learning models and a language-model agent while remaining computationally efficient for online optimization.The framework operates without discretizing coordinates and includes a cumulative quantile-regret guarantee.
  • Limitations and future work: Adaptive sampling leaves promising configurations heavily observed while most configurations remain weakly estimated, creating potential survivorship bias.The paper identifies uncertainty-aware or debiased estimators for adaptive data collection as a promising direction.
  • Future directions: The modular framework could support contextual bandits, language-model proposal oracles, prompt search, and constrained or multi-objective extensions when the relevant information is available.These are proposed directions rather than evaluated capabilities in the supplied conclusion passage.

A Theoretical Analysis

The theoretical analysis extends MOSS allocation to a growing active set and uses that result to establish the paper’s quantile-regret guarantee. The proof anchors performance to the first top-ρ configuration admitted rather than requiring comparison with an immediately available optimal arm.

  • Lemma 1: Lemma 1 extends the usual MOSS allocation bound to an active set whose size grows over time.A regular growth condition sandwiches the time-varying confidence radius between fixed-K confidence radii.
  • Lemma 1: The proof combines standard MOSS underestimation and overestimation decomposition with a peeling argument to derive a gap-dependent bound and then a distribution-free guarantee O(√K_T T).The lemma remains relative to a reference arm admitted at a possibly random time and active thereafter.
  • Theorem 1: Theorem 1 defines τ_ρ as the first admission time of a top-ρ configuration and uses that configuration as the reference arm for splitting ρ-regret.This connects the growing-active-set allocation result to quantile-based performance analysis.

T + RMOSS

IMOSS grows an active set of configurations while using MOSS to allocate exploitation rounds among sampled arms. Under a top-ρ coverage condition, its cumulative ρ-regret combines discovery, expansion, and exploitation costs.

  • Regret decomposition: The regret proof decomposes cost into discovery before τ_ρ, expansion rounds after τ_ρ, and MOSS exploitation rounds after τ_ρ.The discovery term is controlled by the coverage probability, while the other terms depend on active-set growth and MOSS regret.
  • Regret bound: The expansion contribution is O(T^β), while the exploitation contribution is O(T^(1+β)/2) and dominates because β < 1.The proof bounds expansion rounds by the number of active arms and applies the MOSS lemma to exploitation rounds.
  • Policy structure: IMOSS alternates between requesting new arms and serving active arms through a MOSS index, with active-set size K_t = Θ(t^β).The schedule uses β ∈ (0, 1), and expansion occurs whenever the active set is smaller than the target growth rate.
  • Proof mechanism: Choosing a gap threshold balances the baseline gap cost against confidence-controlled pulls of larger-gap arms.The analysis uses gap ordering, midpoint thresholds, and underestimation/overestimation events to charge large-gap selections.

B Complexity Analysis

The implementation separates IMOSS policy cost from oracle-call cost. Index evaluation is generally dominant, except that TabPFN calls can dominate when active-set growth is sufficiently fast.

  • IMOSS: IMOSS stores O(T^β) arms and uses O(d T^β) memory, while each exploitation round evaluates the index across the active set.The active-set size is K_T = O(T^β), and statistics are maintained incrementally.
  • Random oracle: For fixed candidate-pool and dimension sizes, IMOSS-Random runs in O(T^(1+β)) total time.A random oracle draws one configuration in O(d) per call.
  • TPE oracle: TPE adds sorting, density fitting, candidate sampling, and density-ratio scoring, but its total cost remains asymptotically dominated by index passes for fixed c and d.The resulting IMOSS-TPE complexity is O(T^(1+β)).
  • Mutate-KL×PE oracle: Mutate-KL×PE uses one-coordinate Parzen modeling and is d times cheaper than TPE in density work, while also remaining dominated by index passes.Its proposal step selects the best arm, refreshes a coordinate bandit, and models only the selected coordinate.
  • TabPFN oracle: TabPFN contributes O(E T^(3β)) once the active set exceeds the candidate-pool size, matching index-pass order at β = 1/2 and dominating for larger β.In practice, constant factors from the pretrained network govern wall-clock cost.

C Oracle ablation: which oracle, and when?

The oracle ablation compares four proposal mechanisms under the same active-set schedule and MOSS allocation rule. The design isolates global versus local proposal behavior and tests the oracles across six reward landscapes.

  • Experimental setup: The comparison uses β = 0.5, T = 5000, and 10 seeds, so differences are attributable to the proposal mechanism alone.All four oracles share the same active-set schedule and MOSS allocation rule.
  • Controlled comparison: IMOSS-TPE and IMOSS-mutate-KL×PE share the same Parzen density-ratio ranking but differ between global sampling and local one-coordinate mutation.This controlled comparison isolates where proposals are drawn.
  • Benchmarks: The evaluation covers three discrete random-forest grids and three LCBench instances, yielding six structurally diverse reward landscapes.The two benchmark families are intended to differ widely in landscape structure.

C.1 Landscape structure and surrogate accuracy

Oracle usefulness depends on both reward-landscape structure and prediction accuracy. The random-forest benchmarks range from broad, nearly one-dimensional optima to sparse interaction-dependent optima, while TabPFN prediction error improves during search.

  • Landscape structure: Segment and numerai28.6 are primarily controlled by max_depth, whereas credit-g requires jointly finding large max_depth and max_features.Credit-g therefore has an interaction-dependent high-reward region rather than a single dominant coordinate.
  • Landscape taxonomy: Credit-g has a sparse optimum, segment has a broad near-optimal region, and numerai28.6 has a single-hyperparameter optimum with less uniform rewards.Figures 9 and 10 provide the taxonomy used to interpret oracle performance.
  • Surrogate accuracy: TabPFN’s pool-wide prediction MSE drops by roughly an order of magnitude over the first 2000 iterations on segment and credit-g, while remaining low and flat on numerai28.6.The pool-wide measure evaluates all 100 candidates at each probed exploration step.
  • Surrogate accuracy: Prediction error on the proposed configuration falls by one to two orders of magnitude across all three tasks and ends below pool-wide MSE.This measures accuracy where the oracle actually acts.
  • Surrogate accuracy: On numerai28.6, proposal-focused accuracy keeps improving despite flat pool-wide error because exploited arms become more accurately predicted while near-tied arms dominate residual pool error.The two error views therefore distinguish accuracy on selected configurations from accuracy across the candidate pool.

C.2 Comparing the oracles

The oracle comparison separates regret into the gap to the best admitted arm and the cost of pulling other arms. Learned and local oracles reduce regret mainly by making non-best pulls cheaper, while higgs exposes locality’s limitation.

  • Regret decomposition: The regret decomposition separates the best-arm gap from the cost of the other arms, assigning the first to oracle quality and the second to allocation away from the best admitted arm.An oracle can reduce regret by admitting a better arm or by making the remaining arms cheaper to pull.
  • Overall comparison: All three learned oracles accumulate less regret than uniform admission on every task.The comparison uses cumulative regret at T = 5000, averaged with standard deviations over 10 seeds.
  • Overall comparison: On four of six tasks, nearly all regret comes from the cost of other arms; Fashion-MNIST balances both terms, while higgs is dominated by the best-arm gap.On higgs, no oracle finds a configuration near the optimum.
  • Local versus global proposals: The two local oracles reduce the cost of other arms below global TPE, but neither improves the best-arm gap; they are statistically indistinguishable except on APSFailure.On APSFailure, mutate-KL×PE performs better, while it makes the best-arm gap slightly worse overall.
  • Local versus global proposals: Local proposals make non-best pulls cheaper by packing the admitted set with arms near a good configuration, although MOSS still spends most pulls away from the best arm.Global proposals add arms from across the space, where equivalent wandering is more expensive.
  • Locality limitation: Local mutation fails on higgs because one-parameter changes cannot reach a good region requiring several parameters to be correct simultaneously.TabPFN admits the best arm among the four oracles on this task by modeling the whole reward surface.

D Additional Experimental Details

The HotpotQA RAG experiments vary model, temperature, retrieval depth, and prompt template at inference time. Prompts combine retrieved passages with the question, while repeated configuration queries are cached to avoid duplicate API calls.

  • Search space: The HotpotQA RAG search space includes generation models accessed through OpenRouter, retrieval settings, prompting strategies, and decoding parameters.The search space is summarized in Table 5.
  • Prompt templates: Three system prompt templates are used: few_shot adds examples, zero_shot removes them, and naive keeps only a one-line instruction.The selected template is sent as the system message.
  • Request construction: Each request combines the top-k retrieved passages and HotpotQA query with the prefixes Context: and Question:, respectively.Dense embedding search with all-MiniLM-L6-v2 retrieves passages from the full corpus.
  • Execution: Repeated question–configuration pairs are served from a local cache, preventing identical API calls from being re-issued when experiments are rerun or resumed.Other generation parameters are left unset and default to the provider’s own values.
  • Prompt templates: The prompt examples demonstrate concise answers for factual and yes/no questions using only the supplied context.Rules require short answers and return unknown when the context lacks the answer.
Loading 2609.01335v2…