Technical Summary: Generator-Mediated Bandits (GAMBITTS)

Problem Formulation

This paper addresses a specific class of sequential decision-making problems where the agent's action does not directly determine the environment's reward. Instead, the agent selects an action (e.g., a query or prompt), which triggers a stochastic generator (such as a Large Language Model) to produce a treatment (e.g., a text response). The environment then reacts to this generated treatment, producing a reward.

The authors formalize this as the Generator-Mediated Bandit (GAMBIT) problem. In this framework:

1. Action/Treatment Split: The agent chooses an action $A_t$ (query), but the environment receives a stochastic treatment $G_t$ (generated output).
2. High-Dimensionality: The treatment space $G$ (e.g., natural language text) is high-dimensional, making direct modeling intractable.
3. Latent Embedding: The effect of the treatment on the reward $Y_t$ is assumed to occur through a lower-dimensional sufficient embedding $Z^*$. Since $Z^*$ is unobserved, the agent must rely on a working embedding $Z = h(G)$.
4. Stochastic Mediation: Unlike standard contextual bandits where actions map directly to rewards, here the action influences the reward only indirectly through the stochastic treatment distribution. Standard bandit algorithms fail to exploit the structure of the treatment distribution, leading to inefficient learning.

Methodology: GAMBITTS

The authors propose Generator-Mediated Bandit–Thompson Sampling (GAMBITTS), a framework that explicitly models both the treatment generation process and the reward generation process. The approach utilizes a two-stage sampling procedure to capture uncertainty across the intervention process.

Core Components

The framework models two conditional distributions:

1. Treatment Model: $Z_t \sim f_1(z; A_t, X_t, \theta_1)$, representing the distribution of the treatment embedding given the action and context.
2. Reward Model: $E[Y_t | Z_t, X_t] = m_2(Z_t, X_t; \theta_2)$, representing the expected reward given the treatment embedding and context.

Algorithmic Variants

The paper introduces two primary instantiations of GAMBITTS:

1. Fully Online GAMBITTS (foGAMBITTS):

   • Updates both the treatment model ($\theta_1$) and reward model ($\theta_2$) sequentially using observed data.
   • Suitable when the agent lacks access to a simulator or offline data for the generator.
   • Selects actions by integrating over the posterior distributions of both parameters: $E[Y | a, x] = \int m_2(z, x; \theta_2) f_1(z; a, x, \theta_1) dz$.

2. Partially Online GAMBITTS (poGAMBITTS):

   • Leverages access to a generator simulator or offline data to learn the treatment distribution ($f_1$) independently of user interactions.
   • The treatment model is pre-trained (or estimated empirically) and fixed, while only the reward model ($\theta_2$) is updated online.
   • This decomposition allows for more efficient online adaptation by offloading the learning of the complex treatment distribution.

Ensemble-Based Extensions

To handle nonlinear reward relationships and avoid restrictive parametric assumptions, the authors extend both variants using ensemble sampling (following Lu and Van Roy, 2017). These variants (ens-foGAMBITTS and ens-poGAMBITTS) maintain an ensemble of reward models to approximate the posterior, enabling the use of flexible models like Multilayer Perceptrons (MLPs).

Key Contributions

1. Framework Definition: The formalization of the GAMBIT problem, distinguishing it from standard bandits and causal bandits by focusing on the specific mediation structure where actions yield stochastic, high-dimensional treatments.
2. Algorithmic Instantiations: The development of GAMBITTS algorithms (fully online, partially online, and ensemble-based) that explicitly model the treatment-reward split.
3. Theoretical Guarantees:
   • Derivation of Bayesian regret bounds for both foGAMBITTS and poGAMBITTS.
   • Decomposition of regret sources into treatment uncertainty and reward uncertainty.
   • Proof that GAMBITTS achieves sharper regret bounds than standard Thompson sampling when the treatment embedding dimension $d$ is significantly smaller than the product of context and action space sizes ($d \ll \sqrt{CK}$).
   • Demonstration that poGAMBITTS with linear rewards achieves regret scaling of $\tilde{O}(\sigma_2 d \sqrt{T})$, independent of the number of arms $K$ in the linear setting, provided the treatment model is sufficiently accurate.
4. Empirical Validation: Simulation studies demonstrating that GAMBITTS outperforms standard contextual and non-contextual Thompson sampling, particularly in settings with large action spaces and when leveraging observed treatments to estimate expected rewards.

Results

The authors evaluated GAMBITTS using a simulation study modeled on the 2023 Intern Health Study (IHS), where an agent selects prompts for an LLM to generate mental health support messages.

• Performance: GAMBITTS consistently achieved lower cumulative regret compared to standard Thompson sampling (both contextual and non-contextual).
• Scalability: In experiments varying the number of action arms ($K$), the regret of the linear poGAMBITTS algorithm remained stable, whereas standard methods degraded as $K$ increased.
• Misspecification Sensitivity: The study highlighted that GAMBITTS performance is sensitive to the quality of the working treatment embedding ($Z$). When the embedding is poorly correlated with the true reward-relevant dimension, performance degrades and can fall below standard methods. However, flexible reward models (ensemble-based) helped mitigate some of these issues.
• Noise Impact: The advantage of GAMBITTS over standard methods was shown to increase with higher reward noise ($\sigma_2$), as the ability to generalize across arms via the treatment model becomes more critical.

Significance and Claims

The paper claims that GAMBITTS provides a necessary theoretical and algorithmic bridge between classical bandit learning and modern generative AI applications. By explicitly modeling the stochastic nature of generator-mediated interventions, the framework allows agents to:

• Leverage Treatment Information: Utilize the observed treatment (e.g., the generated text) to accelerate policy learning, rather than discarding this information as noise.
• Generalize Across Actions: Exploit shared structure in the treatment space to learn efficiently even when the number of possible prompts (actions) is large.
• Support Adaptive mHealth: Offer a viable framework for Just-in-Time Adaptive Interventions (pJITAIs) that use LLMs for real-time content generation, a setting where fixed intervention libraries are insufficient.

The authors maintain that while the framework assumes a fixed generator and specific embedding structures, it opens new avenues for personalization in domains like marketing, education, and healthcare where treatments are generated dynamically rather than selected from a static library. They note that future work could address online embedding learning, generator fine-tuning, and safety constraints, but the current work establishes the foundational regret bounds and empirical efficacy of the proposed approach.