Paper: Adapting Retrieval Models to Task-Specific Goals using Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The experiments are quite weak and non-standard. This weakness itself may warrant rejection in a top venue. There are numerous retrieval datasets and strong baselines and it is not clear why the authors selected the datasets (1 public dataset that is not commonly used) and an internal dataset, and pretty much without any sensible baselines. Especially for the recall metric, it is standard so it’s not clear why no standard baselines or datasets are used.

Author response:
Our goal was to choose a dataset where LLM can be reliably used as a reward model. That is why we chose symmetric short-text retrieval tasks where the relevance function can be described easily using LLM.

Reviewer comment (Evaluation):
There are many ways to compose a task that do not optimize recall metric so it is not clear why such task is chosen.

Author response:
Using LLM as a rater requires that the prompted LLM align with the task’s objectives. That is why we choose symmetric, short-text retrieval tasks where the LLM can be reliably used.

Reviewer comment (Novelty):
RL in information retrieval has a very rich literature, and optimizing arbitrary metric is inherited from RL, so the novelty/story from this perspective is limited. The major contribution is really the algorithm mentioned in S1, but the significance of the proposed method is unclear.

Author response:
The novelty is the application to large action spaces. Most RL work in information retrieval focuses on tens of actions, whereas our method can scale to millions of actions.

Reviewer comment (Theory):
How valid is the assumptions made on LLMs? Using LLMs as rater is a promising area, but not solved problem.

Author response:
Using LLM as a rater requires that the prompted LLM align with the task’s objectives. That is why we choose symmetric, short-text retrieval tasks where the LLM can be reliably used.

Reviewer comment (Evaluation):
It is unclear how good the LLMs are at aligning the fine-tuning objective to the objective we are actually aiming to solve, for instance, if the criteria is diversity, etc.

Author response:
Yes, that is correct. Our formulation works whenever the reward for a list of items can be decomposed into a sum/aggregation of individual rewards for each item in the list (e.g., for the novelty reward w.r.t. a pre-existing top-k list). It will not work for diversity reward where two items in the predicted list need to be compared.

Reviewer comment (Theory):
The reframing of the problem as having two action spaces is a bit unclear in terms of how it alleviates the complexity of the setting.

Author response:
The intuition is that the convergence depends quadratically on number of actions and linearly in the number of states. So, our setup reduces the actions while increasing the number of states; the net result is a substantial increase in convergence speed.

Reviewer comment (Theory):
It is unclear how the problem is still cast as an MDP for the theoretical results to hold, and this should be clearly articulated in the main body of the paper.

Author response:
It is a one-step MDP. In each round, a state (<query, item>) is sampled, an action is taken (0/1), and the environment provides the reward. It can be understood as a contextual bandit setup.

Reviewer comment (Novelty):
Once we have an oracle, potentially other avenues of improving the performance appear. For instance, framing the problem as a bandit problem (for example, contextual linear cascading bandit [1,2] where item vectors can be generated with the LLMs), considering the problem an example of Positive-Unlabelled [3] or Active learning, where we can augment the initial dataset based on weaknesses the model uncovers in its own predictions. In light of these possible alternatives, what are the reasons to believe using the REINFORCE algorithm for the fine-tuning step is an impactful approach and not just another approach?

Author response:
Compared to positive-unlabelled or active learning setting, the benefit of RL is that we can optimize any non-differentiable loss. Contextual Bandits do not scale well to millions of actions. Most such approaches optimize over tens of items whereas our problem considers optimization over millions of items.

Reviewer comment (Theory):
Can you better describe the intuition why the convergence speed of the algorithm is substantially increased from framing the problem as having the states being pairs <state, item> and having binary actions?

Author response:
The intuition is that the convergence depends quadratically on number of actions and linearly in the number of states. So, our setup reduces the actions while increasing the number of states; the net result is a substantial increase in convergence speed.

Reviewer comment (Evaluation):
Can you provide an interpretation of how impactful the $0.2\%$ increase in the click yield is?

Author response:
The click yield is an incremental metric: it is measured only for the novel keywords that are introduced by our algorithm but not produced by the production algorithm. The click yield is lower than our expectation but still significant.

Reviewer comment (Experiments):
For QADSM, the reward model seems rather small, especially compared to typical embedding models used for the retrieval. My understanding is that the reward model can/should actually be orders of magnitude larger and more capable to generate best supervision signal for the embedding model training. Any reason authors decided to use a variation of T5-base?

Author response:
We use t5-base since its accuracy was already good on the dataset. It may be even more powerful to use a bigger model.

Reviewer comment (Experiments):
About the use case, I was wondering if authors could share any findings/experiments or thinking on how to apply such technique when new items are being added to the index. It is easy for an offline fixed dataset to build such relevance models or finetune embeddings but how would this work when new items are added? Do we need to retrain the models each time or are we claiming generalization in the learned embeddings?

Author response:
The learned embedding model will have generalization. As long as the model to encode keywords is fixed, the same trained model for queries can be used.  In other words, our method, PGRetrieve, can be considered as an *offline RL* training method. Once trained, the model is used in an identical fashion to a model learnt using supervised learning.

Reviewer comment (Reproducibility):
What's the prompt to get relevance judgement from GPT-3.5?

Author response:
Here’s the prompt used. 'Given the query “q”, is the following query expressing the exact same intent, "k"? Please answer in a single word: Yes/No.' Note that keyword is also presented as a query to the large language model.