Test-Time RAG: Enhancing Long Context Understanding in LLMs with Retrieval-Augmented Mechanisms

We’d like to thank the reviewer for their detailed questions and feedback. We’ve addressed questions and concerns below.

• (Question #1) For each ablation and experiment, we include baselines which encapsulate two primary data points to benchmark TTRAG against. Our baseline simulates naively stuffing the full context window with information to test the model’s standalone capabilities for processing long context, while CE (Filtering only) effectively simulates traditional search and retrieval methods for selecting relevant context and including that in the prompt. We demonstrate in our current set of experiments that TTRAG as a whole outperforms solely CE (Filtering only), baselines, and other individual components of TTRAG on relevant benchmarks. For clarity, we refrain from calling CE (Filtering only) as standard RAG because we assume apriori a test-time availability of context. Therefore, we cannot pre-index context in the traditional, standard RAG paradigm. With respect to time consumption, we present results below demonstrating comparable latency between traditional RAG paradigms and the components of TTRAG (on HotpotQA with 1 million token context):

1. Conditional Embeddings vs Traditional RAG | model | average latency (s) | 95th percentile (s) | 99th percentile (s) | |-------|---------------------|-----------------|-----------------| | GPT-3.5-Turbo (with Conditional Embeddings) | 1.283 | 1.876 | 3.742 | | Llama-2-7B (with Conditional Embeddings) | 3.753 | 6.327 | 11.963 | | GPT-3.5-Turbo (Traditional RAG) | 1.076 | 1.539 | 3.315 | | Llama-2-7B (Traditional RAG) | 3.339 | 4.941 | 10.628|
2. Query Rewriting vs. Without Query Rewriting | model | accuracy (%) | average latency (s) | 95th percentile (s) | 99th percentile (s) | |-------|------------|---------------------|-----------------|-----------------| | GPT-3.5-Turbo (with Query Rewriting) | 45.376 | 3.208 | 4.826 | 7.702 | | GPT-3.5-Turbo (without Query Rewriting)| 41.181 |1.283 | 1.876 | 3.742 |

• (Weakness #1) We have improved writing clarity in our most recent revision of our paper. With respect to the specific line highlighted, we are highlighting that a core difference between encoder-decoder models and decoder-only models is the application of attention to tokens, where bidirectional attention effectively creates token representations that encapsulate contextual information from tokens earlier in the sequence and later in the sequence, as opposed to masked attention which is applied only to prior tokens in the sequence.
• (Weakness #2) To clarify, for performance considerations in all of our experiments, we first retrieve the top 10 documents using a generic embedding model like text-embedding-3-small, and then conditionally embed these top 10 documents with the query and retrieve a smaller subset of relevant documents from this ‘pre-filtered’ set. The value-add of conditional embeddings is in capturing fine-grained semantic information, so this application of conditional embeddings accomplishes 2 things: (1) It is highly performant as we always conditionally embed on 10 documents, which introduces very minimal latency and memory. (2) Highlights that CE delivers on capturing detailed semantics. Because each document retrieved in the top 10 set is highly semantically related to the query, the observed gains from applying CE on this pre-filtered set imply that CE is able to differentiate more detailed, relevant context within a set of documents that are all initially viewed as related to the query. Conditional embeddings boasts comparable time consumption when compared to traditional RAG methods, while leading to accuracy gains (please refer to our response for 'Question #1' for this data).
• (Weakness #3) Runaway results for query rewriting are prevented in two ways: First, in the case that the answer is not found in the retrieved documents after all iterations of rewrites (number of iterations until stopping is fixed to 3 in our experiments but this can be treated as a hyperparameter), query rewriting defaults to just fetching the top k (another hyperparameter fixed to 10) documents from similarity search with the original un-rewritten query, lower bounding the performance of query rewriting with normal retrieval. If the answer is found and we want to use the retrieved documents from the rewritten query, at the answer generation stage, retrieved documents from the new query are passed in along with the original query. This prevents the LLM from generating irrelevant responses while maintaining access to information found with the rewritten query.
• (Weakness #4) We have addressed these concerns surrounding figure display in our revised iteration of our paper.

We’d like to thank the reviewer for their constructive feedback and time. We address them below.

• (Question #1) An example to demonstrate the connection between personalized AI and long-context is the following: Consider receiving an email with sensitive information from a coworker, with a large attachment of many pages. Generating a response to this email involves the following considerations: 1) long-context processing and understanding of the email and its attachments 2) an inability to pre-index the attachment, because it was just received and we need to understand it on the fly and it contains private information. Use-cases like this highlight the importance of processing context at test-time. This builds motivation for the rest of the paper, as one unique quality of test-time processing, is that we have simultaneous access to user intent (query) and relevant context, which we then exploit via conditional embeddings, query rewriting, and conditional compute. The long-context benchmarks demonstrate that these techniques perform well at scale (to tie it back to the example: in cases where the attachment in the email is millions of tokens).
• (Question #2) Although not mutually exclusive with incremental indexing, TTRAG focuses on scenarios where we do not have the ability to evaluate and retain the embeddings ahead of time.
• (Question #3 and #4) We have updated our paper and methodology to discuss the complete flow in test time RAG and include results for the collective benefits of the entire pipeline when stitched together. Please refer to table 1, and an updated figure 1 which illustrates how TTRAG works as a comprehensive system. For a brief summary (assume a question-answering setting):

1. We receive a block of context and question.
2. Conditional compute routes the query to one of the compute categories (in this case question-answering)
3. Generate conditional embeddings, select top k documents
4. Use query rewriting to rewrite query if answer isn’t found
5. Re-generate conditional embeddings with rewritten query
6. Repeat steps 4 and 5 until stopping criterion for total rewrites is hit or answer is found (whichever comes first)
7. Generate response

• (Weakness #6) We have modified the paper to demonstrate our experiments from a ‘top-down’ approach, where we start by presenting the performance gains of TTRAG as a whole, then ablating on individual components to demonstrate their individual contributions to performance. Additionally, for each ablation and experiment we include baselines which encapsulate two primary data points to benchmark TTRAG against. Our baseline simulates naively stuffing the prompt with information, while CE (Filtering only) effectively encapsulates traditional search and retrieval methods for selecting relevant context. We demonstrate in our benchmarks that TTRAG outperforms these baselines.

We would like to first thank the reviewer for their detailed feedback and insightful questions. We have addressed questions and weaknesses below.

• (Question #1) We utilize DistillBERT and extract the last hidden state of the encoder portion. We do not introduce a new architecture for generating conditional embedding. We chose to utilize the bidirectional encoder present in BERT because each token can update its representation based on all tokens in the sequence (whether they appear before or after it in the sequence). This advantage is not present if we just extracted the last hidden state of a decoder-only model like Llama-2.
• (Question #2) Averaging is done primarily because at test-time, we don’t have access to information to determine a learned weightage of tokens. We assume we receive the tokens at test time, therefore it is not possible in this setting to learn a weightage on the fly for token ‘importance’. Because we don’t know the prior significance of each token of context, this necessitates taking an average of unmasked tokens.
• (Question #3) CE Filtering Only applies retrieval over embeddings generated by a pretrained embedding model. CE as a whole generates conditional embeddings and then carries out retrieval over these conditional embeddings (no pretrained embedding model). The flow of CE Filtering Only is as follows: generate embeddings of query and documents with pretrained embedding model (we utilize text-embedding-3-small) and run similarity search between query and documents to fetch top k similar documents. CE as a whole does the following: apply CE Filtering Only first to get a pre-filtered set of 10 documents, generate conditional embeddings for query and each document in this pre-filtered set by passing through bidirectional encoder, extracting last hidden state, masking tokens, and averaging across sequence length dimension, and then run similarity search between query and documents to fetch top similar documents amongst the pre-filtered set of documents.
• (Question #4) To clarify, for performance considerations in all of our experiments, we first retrieve the top 10 documents using a generic embedding model like text-embedding-3-small, and then conditionally embed these top 10 documents with the query and retrieve a smaller subset of relevant documents from this ‘pre-filtered’ set. The value-add of conditional embeddings is in capturing fine-grained semantic information, so this application of conditional embeddings accomplishes 2 things: (1) It is highly performant as we always conditionally embed on 10 documents, which introduces very minimal latency and memory. (2) Highlights that CE delivers on capturing detailed semantics. Because each document retrieved in the top 10 set is highly semantically related to the query, the observed gains from applying CE on this pre-filtered set imply that CE is able to differentiate more detailed, relevant context within a set of documents that are all initially viewed as related to the query. Conditional embeddings boasts comparable time consumption when compared to traditional RAG methods, while leading to accuracy gains. We present empirical results below for the time consumption of conditional embeddings below (on HotpotQA with 1 million token context), which demonstrate comparable latencies between traditional RAG methods and conditional embeddings:

1. Conditional Embeddings vs Traditional RAG | model | average latency (s) | 95th percentile (s) | 99th percentile (s) | |-------|---------------------|-----------------|-----------------| | GPT-3.5-Turbo (with Conditional Embeddings) | 1.283 | 1.876 | 3.742 | | Llama-2-7B (with Conditional Embeddings) | 3.753 | 6.327 | 11.963 | | GPT-3.5-Turbo (Traditional RAG) | 1.076 | 1.539 | 3.315 | | Llama-2-7B (Traditional RAG) | 3.339 | 4.941 | 10.628|

• (Question #6) We present an empirical result displaying the latency of query rewriting. Although there is an increase in latency, we observe substantial accuracy gains.

1. Query Rewriting vs Without Query Rewriting | model | accuracy (%) | average latency (s) | 95th percentile (s) | 99th percentile (s) | |-------|------------|---------------------|-----------------|-----------------| | GPT-3.5-Turbo (with Query Rewriting) | 45.376 | 3.208 | 4.826 | 7.702 | | GPT-3.5-Turbo (without Query Rewriting)| 41.181 |1.283 | 1.876 | 3.742 |

• (Weakness #5) Because the focus of TTRAG is to tackle long-context settings rather than re-ranking specifically, we present experiments and references to compare how well current LLMs can naively process long context at test time with TTRAG. With respect to query rewriting, we outline algorithm specifics including stop criterion for iterations in Algorithm 1 in our revised paper. We establish that stop criterion and the number of documents retrieved per iteration can be treated as hyperparameters. In our experiments, we retrieve the top 3 documents during each rewrite iteration with the stopping criterion to be the minimum of either 3 rewrites or when the answer is found in the top 3 retrieved documents. For latency of the pipeline, we display results for conditional embeddings and query rewriting below (on HotpotQA with 1 million token context), demonstrating comparable latencies between TTRAG components and traditional RAG:

1. Conditional Embeddings vs Traditional RAG | model | average latency (s) | 95th percentile (s) | 99th percentile (s) | |-------|---------------------|-----------------|-----------------| | GPT-3.5-Turbo (with Conditional Embeddings) | 1.283 | 1.876 | 3.742 | | Llama-2-7B (with Conditional Embeddings) | 3.753 | 6.327 | 11.963 | | GPT-3.5-Turbo (Traditional RAG) | 1.076 | 1.539 | 3.315 | | Llama-2-7B (Traditional RAG) | 3.339 | 4.941 | 10.628|
2. Query Rewriting vs. Without Query Rewriting | model | accuracy (%) | average latency (s) | 95th percentile (s) | 99th percentile (s) | |-------|------------|---------------------|-----------------|-----------------| | GPT-3.5-Turbo (with Query Rewriting) | 45.376 | 3.208 | 4.826 | 7.702 | | GPT-3.5-Turbo (without Query Rewriting)| 41.181 |1.283 | 1.876 | 3.742 |

• (Weakness #6) We outline our methodology for establishing baselines and experiments as well as any auxiliary data preprocessing steps for reproducibility. For reference, there are 2 key considerations to generalize TTRAG experiments to other datasets:

1. Creating a large enough knowledge base to simulate retrieval over long context (in datasets where context provided per question is not sufficiently large): This is accomplished by extracting context for across all questions in the dataset, concatenating them together until the context block is large enough, and then ensuring that we only ask questions whose answer is contained within this context block. This ensures that we fairly test model’s generation capabilities as the golden documents are guaranteed to be present in the context.
2. Baseline against ‘naively’ filling the context window of current LLMs: The purpose of this baseline is to establish how well a given LLM can utilize information present across a completely filled context window.

Because most current models cannot support multi-million token context lengths, we circumvent this in our experiments by retrieving the top 10 documents (200 tokens per document), then backfilling the context window with randomly retrieved documents from the large context block for the dataset (which is generated in step 1), and then shuffling the documents in the context window. These steps are agnostic to the dataset and are used in our experiments with datasets such as Natural Questions and HotpotQA. For datasets such as QuALITY and QASPER, where substantial context is provided per question, we do not do consideration (1) and just utilize consideration (2) to establish baselines. TTRAG demonstrates performance gains across a variety of datasets with varying structure and type.

We’d like to thank the reviewer for spending time to provide us with valuable feedback. Below, we address your questions and concerns.

• (Weakness #1) We present query rewriting as a component of TTRAG which is enabled by access to both the specific query and context at test-time. A key differentiator between the query rewriting process in TTRAG and the paper highlighted by the reviewer [1] is that query rewriting in TTRAG is driven by an out-of-the-box LLM rewriting the query based on information present in the retrieved documents during each iteration, rather than utilizing a small language model trained via reinforcement learning from LLM feedback [1]. In short our approach does not utilize a training scheme and rewrites queries with information found in documents while [1] utilizes a trainable scheme to do so. Because query rewriting is only a singular component of TTRAG, our motivation section focuses on highlighting current methods of long-context processing (as TTRAG is applied in long context settings). Our comment addressing ‘Weakness #2’ provides more context and motivation for the connection between personalized AI, long-context, and how it relates to TTRAG.

• (Weakness #2) An example to demonstrate the connection between personalized AI and long-context is the following: Consider receiving an email with sensitive information from a coworker, with a large attachment of many pages. Generating a response to this email involves the following considerations: 1) long-context processing and understanding of the email and its attachments 2) an inability to pre-index the attachment, because it was just received, we need to understand it on the fly and it contains private information. Use-cases like this highlight the importance of processing context at test-time. This builds motivation for the rest of the paper, as one unique quality of test-time processing, is that we have simultaneous access to user intent (query) and relevant context, which we then exploit via conditional embeddings, query rewriting, and conditional compute. The long-context benchmarks demonstrate that these techniques perform well at scale (to tie it back to the example: in cases where the attachment in the email is millions of tokens).

• (Weakness #3) Within our benchmarks we include long-context datasets spanning 1 million to 9.6 million tokens. We additionally include ablations on datasets with smaller token contexts per question that can fit within the context window of most modern LLMs to demonstrate that our method performs better than naively filling the entire context window of the LLM. We refer to the sparse vector embeddings that can be generated by BM25 when listing BM25 as the ‘embedding model’.

• (Weakness #4) The motivation behind conditional compute with respect to long context is that certain questions with long context, such as needle in a haystack tasks, can be determined very quickly without needling to explicitly process the entire block of context which is both costly and time consuming. The majority of questions will fall under question and answering, which defaults to the full TTRAG pipeline but in cases that fall in the compute categories we’ve enumerated, we observe latency gains without losing out on performance. At its core, conditional compute relies on an LLM router which classifies the compute category of a query. If the query falls into a category other than default question answering, a pre-implemented compute method is utilized, else we default to the full TTRAG pipeline. For a visual representation of this in action, please refer to figure 1 in our revised paper.

[1] 'Query Rewriting for Retrieval-Augmented Large Language Models' https://arxiv.org/abs/2305.14283

Dear Reviewer,

Thank you again for your insightful review. It has greatly improved our paper, particularly with demonstrating experimental reproducibility and ablations that highlight the low-latency of our proposed TTRAG approach. As the public discussion phase is nearing its conclusion, we wanted to follow up to see if there are any additional questions, concerns, or points that we could clarify or address to further assist the your review. We are more than happy to provide any additional experiments, ablations or clarifications you might need.

Thank you!