POQD: Performance-Oriented Query Decomposer for Multi-vector retrieval

Thank for your comments. Our responses are below:

• About the additional training data:

Indeed, in our experiments, we did not include additional training data to train POQD throughout the paper. During the training process, we train POQD using the same set of training samples as those baseline methods.

• About the baseline method:

As explained in Section 5.1, for S-QD, it aims to train a sequence-to-sequence model to generate decomposed sub-queries. Such training process has been used in many prior works such as (Zhou et al., 2022; Zhu et al., 2023; Guo et al., 2022; Wu et al., 2024). We thus reuse the code from these papers to train this query decomposition model. For U-QD, we directly use the released unsupervised model for performing query decomposition. For ICL-QD, since it performs query decomposition by prompting LLMs. We thus leverage the prompts from prior works to decompose queries through in-context learning. For ICLF-QD, we follow (Qi et al.) to add retrieval scores as feedback in the prompts for in-context learning.
The experimental results of these baseline methods are not taken from prior papers. As explained in Appendix C, S-QD, U-QD, ICL-QD and ICLF-QD haven’t been evaluated for decomposing queries for retrieval in the context of RAG. We thus implement these methods in our experiments and obtain these numbers by ourselves.

• About the theoretical analysis:

See our response to reviewer EhPc (About the theoretical analysis)

• About the contribution of the paper:

See our response to reviewer EhPc (About the contribution of the paper)

• About the image retrieval performance of ColBert:

Yes, as mentioned in line 319-321, in the image QA experiments, we leverage the pre-trained CLIP model as the underlying image retrieval model for all methods including ColBert. The reason why Colbert performs worse is that the original version of Colbert has not been applied in image retrieval yet. To mitigate this, we further included the results of Colpali [1] below, which is a variant of ColBert adapted for image retrieval. The results indicate that our method still outperforms Colpali although it is adapted for image retrieval.

• About ColBert text QA results:

See our response to reviewer Nopx (Regarding the retrieval results in Table 1)

• About the number of iterations for step 1:

We configure the number of iterations as 5 by default. If we can find a prompt $p^{new}$ which can lead a training loss smaller than that with $p^{new}$ by up to $\alpha$ within 5 iterations, then we terminate this loop early. Otherwise, we terminate the entire process of updating prompts in Algorithm 2.

• About Algorithm 1:

Some details of this algorithm, particularly the while loop's convergence condition, are deferred to Section 4.3. The loop terminates if the loss with an updated prompt is at least $\alpha$ lower than the initial loss (best_L in Algorithm 2, equivalent to Algorithm 1's initial loss) within 5 iterations. In other words, convergence occurs when any updated prompt reduces the loss by at least $\alpha$ compared to the initial training loss. We will clarify this later.

• About the comparison against the query rewrite strategy:

We further compare POQD against the strategy of performing query rewrite (we follow [2]) on WebQA (text). The following results again show better performance of POQD:

[1] Manuel et al. "Colpali: Efficient document retrieval with vision language models."ICLR 2024.

[2] Ma et al, Query Rewriting for Retrieval-Augmented Large Language Models, arxiv

We would thank your comments. Our responses are below:

• About qualitative examples:

Thanks for pointing out this issue. Indeed, in figure 1, we leverage one example to show the differences between ColBERT, ICL-QD and POQD. We further expand it by reporting the decomposed sub-queries (both before and after filtering step) generated by other baseline methods as follows (the retrieved figures are shown in this link).

Before filtering, the sub-queries from baselines contain many irrelevant tokens, as analyzed below:

1. S-QD: Trained on StrategyQA with human-annotated sub-queries. However, when applied to other datasets, distribution shifts introduce irrelevant tokens.
2. U-QD: Its candidate sub-queries are collected from datasets (e.g., SQuAD), leading to irrelevant tokens in generated sub-queries.
3. ICL-QD/ICLF-QD: Their LLM-generated sub-queries may contain hallucinated irrelevant tokens.

Unlike our method, baselines lack iterative refinement using downstream feedback, causing them to either miss key tokens (e.g., U-QD drops "Hong Kong") or retain unimportant ones (e.g., S-QD keeps "type" and "what"). This leads to incorrect image retrieval and answers. We will add this analysis in the revision.

• About the prompts produced by POQD:

We run Algorithm 1 for four steps and show the prompts and generated prefixes at each step are included in this link. The LLM optimizer indeed produces variations of query decomposition instructions, searching for variants that minimize the retrieval system’s training loss.

• About the filtering step:

As noted in lines 198-201, we only filter irrelevant tokens to ensure query decomposition correctness. For LLM-based methods (ICL-QD, ICLF-QD), retrieval and QA accuracy remain nearly identical with or without POQD's filtering, as shown in our ManyModalQA (image) ablation studies below.

We also removed the filtering step from our method and retested on ManyModalQA (image). The results are shown below, which show that both retrieval and QA accuracy get reduced by 2%-5%, but still higher than baseline methods.

• About the theoretical analysis:

See our response to reviewer EhPc (About the theoretical analysis).

• About the example in Section 2.1:

As explained in Appendix B, image retrieval in image QA involves segmenting images into patches and encoding each of them. Due to the maxsim operation, the similarity between a sub-query and an image depends on the most similar patch to this sub-query. In Section 2.1’s example, CLIP considers Patch A black-filled most similar to "kong," while the ground-truth image’s Patch B buildings is less similar. Since "kong" refers to a gorilla-like monster (see Wiki), the Patch A yields unrealistically high similarity, ranking Lee Kuan Yew’s image above the ground truth. We will clarify this in the revision.

• About query and document encodings: Yes, we encode sub-queries using pre-trained models like Sentence-BERT. For documents, we encode each sentence similarly for all methods except ColBERT to ensure fair comparison (lines 311-321). ColBERT inherently tokenizes queries and documents, so we use its default setup. But we also test it on WebQA by processing documents like other methods. The top-2 retrieval accuracy (48.92) and QA score (59.71) under this setup are worse than ColBERT's default (61.15 and 60.79 resp.).

• About the retrieval time in Figure 6: The retrieval time on WedQA is much smaller than the generator time. You may mean the longer retrieval time on StrategyQA, which as noted in lines 307-311, is due to multi-hop QA requiring repeated reasoning and retrieval per question. We will add this explanation to Section 5.2.

Thanks for your comments. Our responses are below:

• About the contribution of the paper:

While we don't propose a novel LLM-based optimization method, our key contribution is recognizing the need to optimize query decomposition for multi-vector retrieval—a critical factor in improving retrieval-based systems. To address this, we adapt an LLM-based optimization strategy for query decomposition and jointly train it with the generator model. We believe this work fits the scope of ICML's application track.

• About the training algorithm:

While our alternative training algorithm may seem straightforward, it could incur higher training overhead compared to optimizing only the RAG generators—a concern for practical applications. However, our theoretical analysis (lines 281-297) shows that with proper hyper-parameters, the added overhead is minimal without compromising optimality, as empirically confirmed in Figure 4. Thus, our algorithm effectively balances efficiency and effectiveness, as highlighted in Contribution 2.

• About the theoretical analysis:

While our theoretical analysis assumes strong convexity, it also holds under the weaker Polyak-Łojasiewicz star (PL*) condition [2]. Prior work [1] suggests that pre-trained, over-parameterized LLMs fine-tuned with GaLore, a variant of LoRA, likely satisfy the PL* condition. Since the model $\Theta$ in the paper is a pre-trained LLM fine-tuned with GaLore, the analysis in [1] applies. Assuming the loss function $L(\Theta;p)$ satisfies the $\mu$-PL* condition (instead of $\mu$-strong convexity), a variant of Lemma A.1 yields: $L(Θ_k; p) − L(Θ^∗; p) <= (1-\mu/L)^k(L(Θ_0; p) − L(Θ^∗; p))$. Similarly, a modified Theorem 4.4 gives: $L(Θ^*(p^{old}); p^{old}) − L(Θ^*(p^{new}); p^{new}) >= \alpha - (1-\mu/L)^{\tau} M$. These adjustments preserve our conclusion that Algorithm 2 balances efficiency and effectiveness with proper hyper-parameters. We will include this refined analysis.

• About the inconsistent retrieval results:

See our response to reviewer Nopx (Regarding the retrieval results in Table 1)

• About the text QA results:

Indeed, we report both image QA and text QA results for WebQA in Table 1,2. The text QA results on MultiModalQA and ManyModalQA below show that our method consistently outperforms baselines:

• About implementation details:

1. U-QD: We reuse the code from (Perez et al., 2020).
2. ICL-QD & ICLF-QD: Both follow straightforward principles—ICL-QD uses manual prompts for in-context learning, while ICLF-QD adds relevance scores to retrieved documents (see Section 5.1). These prompts will be included in the revision.
3. We used the GPT-4 as the prompt optimizer and query decomposer (details in revision).
4. Parameters: Default $\alpha=0.02$; retrieved items set to 1 (image QA) and 2 (text QA).

• Ablation studies on WebQA (text):

We first varied the value of alpha and tracked its effect on the training loss is varied across the training process, which is visualized here. If alpha is too large (say alpha=0.05), POQD struggles to find a suitable $p^{new}$ in Algorithm 1, thus causing the underfitting issue. In contrast, if alpha is too small (say alpha=0.01), POQD converges much slower than our default with alpha=0.02. Hence, alpha=0.02 can balance the convergence speed and the final performance.

Ablation on the generator model (from Llama to Qwen2.5):

Ablation on the query decomposer (from Llama to GPT4).

Ablation on the retrieval model (from sentence-bert to Roberta):

Ablation on the number of retrieved items:

These results show the superiority of POQD under various conditions.

[1] Liu, X. H. et al. "On the Optimization Landscape of Low Rank Adaptation Methods for Large Language Models." ICLR 2023. [2] Liu, C. et al. "Loss Landscapes and Optimization in Over-Parameterized Non-Linear Systems and Neural Networks." Appl. Comput. Harmon. Anal.

We would thank your comments. You can find our responses to your comments below:

• Regarding the retrieval results in Table 1:

We admit that these results look confusing. Indeed, the discrepancy between Table 1 and Table 2 can arise from the fact that we report Top-20 and Top-100 retrieval accuracy while we only retrieve the most relevant image for image QA and the two most relevant documents for text QA. Hence, if we report Top-1 and Top-2 retrieval accuracy instead (as shown below), we can see that our method always performs better than the state-of-the-art:

• Regarding the theoretical analysis:

We agree that proving the convergence of POQD is trivial. However, the main message that we want to deliver in Theorem 4.4 (and the following explanations between line 290 and 297) is to analyze under what conditions $p^{new}$ would lead to a better training loss at convergence than that of $p^{old}$. This is crucial in demonstrating that the loss $L(\theta;p)$ is indeed optimized with respect to $p$. Otherwise, it is likely that with the finally derived prompt $p^*$ by Algorithm 2, the converged training loss $L(\theta^*(p^*); p^*)$ is even worse than that with the initial random prompt $L(\theta^*(p^{init}); p^{init})$. This thus invalidates the optimality of $p^*$ derived by Algorithm 2. So to guarantee that this solution is optimal, we thus need to make sure that updating $p^{old}$ to $p^{new}$ can lead to sufficiently large training loss reduction (by $\alpha$) within 5 iterations in Algorithm 1 (which will be clarified in the description of Algorithm 2 in the revision). Otherwise, we terminate updating the prompt.

• Regarding the possibly unstable training issue:

We agree that without identifying a $p^{new}$ that can reduce the training loss by $\alpha$, the training process can be unstable. So as long as we cannot find such a $p^{new}$ within 5 iterations in Algorithm 1, which will be further clarified near line 254 to 257 in the revision, we break the while loop in Algorithm 1 and stop updating $p$ any more. The empirical results in Section 5 indeed demonstrate the performance advantage of our method with this strategy. We would love to elaborate more on this point in the revision.

• Regarding the relationship between the prompt scores and loss:

As we point out in line 180 in the right column, the training loss $L(Θ; p)$ is viewed as the score

• Regarding the plots of training loss:

We included the plot of training loss in this link, which shows no spike during the entire training process of Algorithm 2 (including the prompt update process). As mentioned above, we only update $p^{old}$ to $p^{new}$ if the training loss is reduced by at least $\alpha$ within 5 steps (as explained in the response to reviewer PnGe). Otherwise, we terminate the training process. This can thus guarantee the smooth reduction of training loss.