Technical Summary: Progressive Content Refinement with Decaying Reward Joint LinUCB

Problem Statement

While iterative refinement techniques have significantly enhanced Large Language Model (LLM) performance, existing methods face critical limitations. Traditional approaches, such as Self-Refine, often rely on static or heuristically chosen prompts, leading to suboptimal performance as refinement strategies vary in effectiveness across contexts. Furthermore, Multi-Armed Bandit (MAB) frameworks applied to LLM refinement typically assume stationary reward distributions. This assumption ignores the "saturation effect," where repeatedly applying the same refinement prompt yields diminishing returns or even negative impacts over time. This phenomenon, known as "reward decay," leads to the over-exploitation of strategies that have lost their utility. Existing solutions like Rotting Bandits attempt to model this decay but suffer from inefficiencies due to reliance on initial round-robin sampling for parameter estimation. Additionally, traditional contextual bandit approaches (e.g., LinUCB) often treat arm features as disjoint, failing to leverage the semantic similarity between different prompts, which exacerbates the "cold start" problem for new prompts.

Methodology

The authors propose DR-LinUCB (Decaying Reward Joint LinUCB), a novel contextual bandit algorithm designed to explicitly model reward decay and jointly learn arm values within a unified framework.

Core Components

1. Decaying Reward Model:
   The observed reward $r$ for an arm $i$ after being pulled $N_i$ times is modeled as:
   $$r = \tilde{\mu}_i \cdot \exp(-N_i \cdot f_i)$$
   Here, $\tilde{\mu}_i = e_i^\top \theta$ represents the initial expected reward (linearly modeled by a global parameter $\theta$), and $f_i = \max(0, e_i^\top w_f)$ is the decay rate determined by a global decay parameter $w_f$. The arm $i$ is characterized by an embedding $e_i$.

2. Joint Parameter Learning via Expectation-Maximization (EM):
   DR-LinUCB employs an EM algorithm to simultaneously estimate the global reward parameters ($\theta$) and decay parameters ($w_f$).

   • E-Step: Estimates the undecayed reward $\tilde{\mu}_i$ from the observed decayed reward $r_t$ using current decay parameters.
   • M-Step: Updates $\theta$ using standard linear ridge regression on the estimated undecayed rewards. The decay parameters $w_f$ are updated by transforming the exponential decay model into a linear regression problem via logarithmic transformation: $\log(\tilde{\mu}_i) - \log(r_t) = N_i \cdot (e_i^\top w_f)$.

3. Arm Selection with Decay:
   The algorithm selects the arm $i(t)$ that maximizes a modified Upper Confidence Bound (UCB) criterion:
   $$i(t) = \arg\max_{i} \left( e_i^\top \hat{\theta} + C \sqrt{e_i^\top A^{-1} e_i} \right) \cdot d_i$$
   where $d_i = \exp(-N_i \max(0, e_i^\top w_f))$ is the decay factor. This criterion balances exploration (via the UCB term) and exploitation while actively penalizing arms that have been over-exploited and are likely decaying.

4. Prompt Embedding and Generation:
   Refinement prompts are treated as arms and embedded into a high-dimensional space using a fine-tuned MINILM model, followed by dimensionality reduction via PCA. The framework supports two arm generation strategies:

   • ArmGenerator: A self-refinement process where the LLM generates feedback and new prompts based on history.
   • EvoArmGenerator: An evolutionary strategy using EvoPrompt to explore a diverse solution space through mutation and crossover.

Key Contributions

The paper outlines three primary technical contributions:

1. DR-LinUCB Algorithm: A novel bandit algorithm that explicitly incorporates a reward decay model for each arm, enabling adaptive exploration-exploitation in dynamic LLM refinement environments.
2. Joint Learning with Embeddings: A method that embeds refinement prompts as contextual features, assuming a unified arm model. This allows for informed initial parameter estimation for new prompts, mitigating the "cold start" problem common in disjoint LinUCB frameworks.
3. Efficient EM-Based Learning: The use of an Expectation-Maximization algorithm to simultaneously learn contextual and decay parameters, circumventing the inefficient initial exploration phases (e.g., round-robin sampling) found in prior Rotting Bandit approaches.

Experimental Results

The method was evaluated on two benchmarks: GSM8K (math reasoning) and Sentiment Reversal (style transfer), using ChatGPT-3.5-turbo and ChatGPT-4o.

• Performance Gains: DR-LinUCB consistently outperformed strong baselines, including Single Call, Random Exploration, JointLinUCB (without decay), Self-Refine, and REx.
  • On GSM8K with ChatGPT-3.5, DR-LinUCB achieved a success rate of 0.790, significantly surpassing Self-Refine (0.687) and REx (0.530).
  • On Sentiment Reversal with ChatGPT-4o, DR-LinUCB achieved a perfect score of 1.000, outperforming all other methods.
• Ablation Insights:
  • Reward Decay: Comparisons with JointLinUCB (which lacks decay modeling) demonstrated that reward decay is crucial for preventing over-exploitation. In some cases (e.g., ChatGPT-4o on Sentiment Reversal), JointLinUCB performed worse than Random Exploration, highlighting the necessity of decay modeling.
  • Convergence: The inclusion of decay accelerated convergence, with significant performance improvements observed within a limited number of iterative steps.
  • Arm Generation: Feedback-driven prompt exploration (ArmGenerator) generally proved more robust than evolutionary strategies for the specific tasks tested, though the latter provided valuable diversity for ablation studies.

Significance and Claims

The paper claims that DR-LinUCB offers a principled approach to optimizing LLM iterative refinement in dynamic environments where the utility of strategies diminishes over time. By explicitly modeling the "saturation effect," the method effectively balances the exploration of novel prompts with the exploitation of high-performing ones, avoiding the pitfalls of over-exploitation.

The authors note that while the method achieved a perfect score on the Sentiment Reversal task, this is attributed to the task's relative simplicity and the robustness of the evaluation metrics (pretrained sentiment models and LLM judges), rather than a universal capability to solve all complex benchmarks. The paper acknowledges limitations regarding computational cost due to iterative LLM calls and sensitivity to hyperparameters, suggesting future work should focus on cost reduction and adaptive tuning mechanisms.

Ultimately, the work demonstrates that integrating reward decay modeling into a joint contextual bandit framework is essential for mitigating over-exploitation and optimizing the iterative refinement process, leading to more reliable and higher-quality LLM outputs.