Online Pandora's Box for Contextual LLM Cascading

Technical Summary: Online Pandora's Box for Contextual LLM Cascading

1. Problem Formulation

The paper addresses the operational challenge of LLM cascading, where an organization must sequentially query a portfolio of Large Language Model (LLM) APIs to generate high-quality content cost-effectively. Unlike standard model routing (selecting a single API per request), cascading involves a sequential search process: the decision-maker (DM) queries APIs one by one, incurring costs and observing outputs, until a stopping condition is met, at which point a single output is selected and deployed.

The authors formalize this as an Online Contextual Pandora's Box problem with a distinct two-phase structure:

1. Query Phase: The DM observes a request context $x_t$ and sequentially queries boxes (APIs). Querying box $a$ reveals a stochastic output $\omega_{at}$ and incurs a cost $c_{at}$. The DM decides whether to continue querying or stop based on observed outputs and costs.
2. Selection Phase: Upon stopping, the DM selects one of the queried outputs to deploy. The DM observes only the downstream reward $r_t$ associated with the selected output, not the rewards of unselected queried outputs.

Key Distinction: This differs from classical Pandora's Box models where opening a box directly reveals its reward. Here, the reward is "output-mediated": the reward depends on the request-output pair $(x_t, \omega_{at})$ via a shared reward function $\mu^*$, while API-specific heterogeneity arises only through the distributions of outputs and costs. The feedback is partial; the DM learns about the reward function only from selected outputs, but learns about reservation indices from all queried outputs.

The goal is to minimize cumulative regret over a horizon $T$ relative to an oracle policy that knows the conditional output distributions, cost functions, and the reward function.

2. Methodology

The authors propose COSMOS (COst-aware Sequential MOdel query and Selection), an online learning policy that avoids estimating full conditional distributions. Instead, it directly models and learns reservation indices and the reward function using parametric structures.

Core Components:

• Parametric Modeling:
  • Reward Function: Modeled as a generalized linear function: $\mu^*(x, \omega) = G(\theta^*{}^\top \phi(x, \omega))$.
  • Reservation Indices: Modeled as generalized linear functions of context: $\sigma^*_a(x) = \Lambda(\rho_a^\top \psi(x))$. These indices are derived from Weitzman's optimal policy, satisfying the moment equation $E[(\mu^*(x, \omega) - \sigma^*_a(x))^+ | x] = E[c_a | x]$.
• Optimism-Based Learning: The policy constructs optimistic upper bounds for both the reward function and the reservation indices.
  • Reward Estimation: Uses a penalized Maximum Likelihood Estimator (MLE) for $\theta^*$ based on deployed outputs, constructing Upper Confidence Bounds (UCB) via a self-normalized martingale inequality.
  • Index Estimation: Uses a Generalized Method of Moments (GMM) approach (specifically an $M$-estimation formulation) to estimate $\rho_a$. The loss function is constructed such that its first-order condition recovers the sample analogue of the Weitzman moment equation. This allows learning the index directly from queried output-cost pairs without estimating the full distribution.
• Regret Decomposition: A key theoretical insight is that under optimism, the per-period regret decomposes into two independent error terms:
  1. The error in estimating the reward for the selected output.
  2. The sum of errors in estimating the reservation indices for all queried boxes.
     This decomposition allows the authors to apply standard UCB-style analysis (elliptical potential arguments) to both components separately.

Algorithm Flow (COSMOS):

1. Observe context $x_t$.
2. Construct optimistic estimates $\hat{\sigma}_{at}$ and $\hat{\mu}_t$.
3. Order boxes by descending $\hat{\sigma}_{at}$.
4. Sequentially query boxes. Stop querying box $k$ if the maximum optimistic reward of queried outputs exceeds the optimistic index of the next box.
5. Select the queried box with the highest optimistic reward estimate.
6. Observe the realized reward and update estimators.

3. Key Contributions

1. Novel Model Formulation: The paper introduces an online contextual Pandora's Box model specifically tailored to LLM cascading. It captures the critical "output-mediated" feedback structure where rewards are only observed for deployed outputs, distinguishing it from classical settings where rewards are revealed upon inspection.
2. Direct Index Modeling: Instead of learning complex conditional distributions of outputs and costs, the authors propose directly modeling the reservation indices using parametric structures. This is achieved via a GMM-based $M$-estimation approach that leverages the moment condition defining the optimal index.
3. Regret Guarantees: The paper provides formal regret bounds for the COSMOS policy.
   • Known Reward Regime: If the reward function is known (e.g., via offline data), the policy achieves a regret bound of $\tilde{O}(A m \sqrt{T})$, where $A$ is the number of APIs and $m$ is the dimension of the index feature map.
   • Unknown Reward Regime: When the reward function must be learned online, the policy achieves a regret bound of $\tilde{O}([d + A(m + \sqrt{dm})]\sqrt{T})$, where $d$ is the dimension of the reward feature map.
   • These bounds improve upon existing contextual Pandora's Box results (e.g., $\tilde{O}(T^{5/6})$) by leveraging local curvature assumptions to achieve the optimal $\tilde{O}(\sqrt{T})$ rate.

4. Results and Assumptions

The theoretical results rely on several regularity conditions:

• Assumption 1 & 2 (Indices): The link function $\Lambda$ is smooth and strictly increasing; the loss function exhibits local curvature around the true index (ensuring the moment equation crosses zero with a bounded slope).
• Assumption 3 & 4 (Reward): The reward function follows a generalized linear model with a strictly increasing link $G$; rewards belong to a canonical exponential family.
• Assumption 5 (Anti-concentration): Ensures that the features of selected outputs (used for reward learning) are sufficiently informative to control the estimation error of indices (which rely on queried outputs). This prevents a "mismatch" between the data used for learning the reward and the data used for learning the indices.

The paper proves that under these conditions, the COSMOS policy achieves dimension-dependent $\tilde{O}(\sqrt{T})$ cumulative regret.

5. Significance and Claims

The paper claims to be the first to formulate LLM cascading as an online contextual Pandora's Box problem with joint query and selection, providing formal regret guarantees. It bridges the gap between heuristic LLM cascading methods (like FrugalGPT) and rigorous online learning theory.

The significance lies in:

• Practical Relevance: It addresses the real-world trade-off between inference cost and output quality in LLM deployment, offering a principled way to manage sequential API queries.
• Theoretical Advancement: It extends the Pandora's Box literature to settings with output-mediated feedback and develops a novel learning algorithm combining GMM and UCB techniques to handle the specific feedback structure of LLM cascading.
• Efficiency: By directly modeling reservation indices rather than full distributions, the approach is computationally more tractable and statistically efficient for high-dimensional contexts.

The authors note that their work assumes a standard sequential querying structure and parametric models, leaving extensions to nonparametric settings, batched querying, and output ensembling as future directions.