TableRAG: Million-Token Table Understanding with Language Models

Contribution

The contribution of the paper is relatively small. The author mainly builds upon the original Row Column Retrieval and further reduces the retrieval cost by Scheme Cell Retrieval. The original work has included column selection, but this paper only further reduces the returned row to only return specific cells.

We thank the reviewer for the constructive feedback. We would like to emphasize that the existing Row-Column Retrieval requires encoding all rows and columns, which is infeasible for large-scale tables without significant truncation. Our work directly addresses this critical limitation and is significantly differentiated in 3 aspects:

1. New Benchmarks: We introduce two real-world large-scale TableQA benchmarks, ArcadeQA and BirdQA, filling a critical gap in the field. These benchmarks enable rigorous evaluation of TableQA systems on tables with millions of tokens, pushing the boundaries of research in this domain.
2. Scalable Framework: We propose TableRAG, a comprehensive framework that combines novel techniques such as query expansion, frequency-based filtering, and schema-cell retrieval. This framework enables LM-based TableQA to scale to tables with millions of tokens, a significant advancement over existing methods.
3. Rigorous Evaluation: We provide thorough complexity analysis and extensive empirical studies, demonstrating TableRAG's effectiveness and efficiency. Our work paves the way for broader applications of TableQA in real-world scenarios involving large tables.

We will make these clearer in the final version.

Baselines

The baseline of the experimental comparison is relatively weak, and the advantage of finer grained retrieval is obvious.

To the best of our knowledge, TableRAG represents the first dedicated effort to address LM-based large-scale table understanding. We are eager to include any specific baselines the reviewer may recommend that we might have overlooked. While we acknowledge the effectiveness of existing methods on smaller tables, our analysis and studies clearly demonstrate their limitations in scaling due to context length constraints. Furthermore, we have shown that previous attempts at scalable retrieval methods for tabular data encountered inherent limitations. Our carefully designed TableRAG framework overcomes these challenges, achieving superior performance and scalability, as validated through comprehensive complexity analysis and extensive ablation studies.

To better understand TableRAG’s performance compared against state-of-the-art baselines, perform experiments on WikiTableQA with GPT 3.5 as shown below:

In this benchmark, the entire table content often fits within the context window of the language models, which is not representative of real-world large-scale industrial scenarios.

Nevertheless, the results clearly indicate that TableRAG maintains its superior performance even on smaller tables. It's important to highlight that the baseline methods require that LLMs process the entire table, thus limiting their scalability for large tables.

Generalizability

We appreciate the reviewer's question regarding the generalizability of our schema-cell retrieval method. It's important to note that our work focuses on general TableQA, which often involves complex understanding and reasoning beyond simple table manipulation, as shown in the TabFact example in the general response. This distinguishes our approach from Text2SQL methods, which have a narrower scope.

The generalizability of TableRAG is inherently tied to the capabilities of the underlying solver agent. TableRAG can be easily integrated with any downstream LLM-based solver agents. In our implementation, we utilize ReAct, a powerful LLM-based solver with demonstrated effectiveness and generalizability in TableQA tasks [1, 2]. TableRAG's primary role is to efficiently and accurately provide the necessary information to the solver, enabling it to perform the required reasoning and extract relevant details.

1. Rethinking Tabular Data Understanding with Large Language Models (Liu et al., NAACL 2024)
2. ReAcTable: Enhancing ReAct for Table Question Answering (Zhang et al., VLDB 2024)

Evaluation on more models

We follow reviewer’s suggestion and include results obtained using Mistral Nemo, as shown in the general response. The results demonstrate that TableRAG remains the most effective table prompting technique when working with Mistral Nemo, consistent with our previous findings using GPT 3.5 and Gemini 1.0 Pro.

Ablation study

Thanks for the insightful feedback. The accuracy improvement on BirdQA ranges from 0.7% to 2%, which is in line with the improvements observed in prior work such as Binder[3] and Dater[4].

We observe an even more substantial accuracy boost on ArcadeQA, ranging from 2.3% to 6.6%. This is likely attributable to the larger scale of the input in this dataset. As detailed in Table 5, ArcadeQA has approximately twice the average number of columns and cells compared to BirdQA, suggesting that cell retrieval is particularly beneficial when handling larger tables. Nonetheless, both schema and cell retrieval components enhance performance across all tested datasets, and they contribute orthogonal advantages to the overall improvement.

3. Binding Language Models in Symbolic Languages (Cheng et al., ICLR 2023)
4. Large Language Models are Versatile Decomposers: Decomposing Evidence and Questions for Table-based Reasoning (Ye et al., SIGIR 2024)

Figure 8

Thanks for pointing out the typo. The variable leather_wallet_prices should indeed be wallet_prices as obtained in the previous step. We will correct this in the final version.

Description column

Thank you for raising this important point about handling long text columns within our encoding budget. We acknowledge this limitation and designed TableRAG to mitigate its impact in the following ways:

1. Prioritization of Categorical Columns: Our frequency-based encoding prioritizes columns with more distinct categorical values, ensuring that critical information for indexing is processed first. This helps reduce the likelihood of long text columns exhausting the entire budget.
2. Solver Agent's String Matching Capability: Even if the retrieval doesn't cover all cells in long text columns like "description" or "address," our solver agent (ReAct in our case) can still perform basic string matching within those columns to identify relevant information. As demonstrated in Figure 8, the solver agent can detect the existence of the “description” column through retrieval and then use basic string matching techniques (e.g., df['description'] .str.contains('Wallet')) to pinpoint the precise cell within those columns.
3. Empirical Evidence from the Study of Encoding Budget: The ablation study in Figure 6 shows that TableRAG is robust to changes in the encoding budget. We varied the encoding budget from 100 to 10,000 and found that TableRAG's performance remained near optimal on both ArcadeQA and BirdQA. In contrast, Row-Column Retrieval is sensitive to the encoding budget, with performance significantly dropping when the encoding budget is either increased or decreased. This study validates the earlier statement that even if not all relevant cells are retrieved, the solver agent can still obtain the necessary information through programs using the retrieved information in most cases.

While we recognize that there might be more optimized ways to handle long text columns, our current work prioritizes the implementation and evaluation of query expansion and schema-cell retrieval, which we believe are foundational for scaling LM-based TableQA. Addressing the specific challenge of long text columns will be a focus of our future research.

False alignment

Thank you for raising these important questions about the robustness of our retrieval and the alignment between schema and LLM knowledge.

Handling Query Expansion Mismatches: In TableRAG, we employ embedding-based retrieval, which inherently captures semantic similarities. Even if the LLM-generated keywords don't perfectly match the schema or cell values, the retrieval process will still return the most semantically relevant ones. This provides a degree of robustness against potential mismatches in query expansion.

Schema Alignment and LLM Generalizability: TableRAG, having access to both the query and the retrieved context, is expected to leverage its understanding of language and world knowledge to determine which information is pertinent and how to combine it to solve the problem at hand. While we don't explicitly enforce schema alignment within the LLM, we rely on its inherent ability to reason and connect information, even when the schema terms might not have been directly encountered during pre-training.

Baselines

Papers like CodeS [1] have proposed using BM25 retrieval to find cell values, which would serve as an excellent baseline for comparison. Additionally, schema linking—the process of identifying the correct rows and columns—is well-explored. Approaches proposed in studies like TaBERT [2] could also be utilized as baselines.

We thank the reviewer for the valuable references. It's important to note that our work focuses on general TableQA, which often involves complex understanding and reasoning beyond simple table manipulation. This distinguishes our approach from Text2SQL methods, which have a narrower scope. While CodeS primarily focuses on Text2SQL, its BM25 retrieval approach offers a potential enhancement to our TableRAG framework. Thus, we implement BM25 and hybrid retrieval as a replacement for embedding-based retrieval to better understand the impact of different retrieval methods.

The results demonstrate that TableRAG with embedding-based retrieval still outperforms BM25 and hybrid retrieval consistently.

Regarding TaBERT and other LM-based methods, their effectiveness is indeed constrained by the context window of language models, making them less suitable for large-scale tables. Nevertheless, we recognize the benefit of comparing TableRAG against representative baselines, and we evaluate TableRAG on additional WikiTableQA benchmark with GPT 3.5 as shown below:

In this benchmark, the entire table content often fits within the context window of the language models, which is not representative of real-world large-scale industrial scenarios.

Nevertheless, the results clearly indicate that TableRAG maintains its superior performance even on smaller tables. It's important to highlight that the baseline methods require that LLMs process the entire table, thus limiting their scalability for large tables.

Limitation of retaining most frequent distinct cell values

Given an encoding budget, the strategy of retaining only the most frequent distinct cell values seems problematic, especially for queries that require data from columns containing names, addresses, etc. Heuristic-based methods, such as those using syntactic similarity with edit distance, can effectively filter most of the cell values and might be more appropriate.

We appreciate the reviewer highlighting this potential limitation of the cell value filtering strategy. We were indeed aware of this trade-off during the design process and opted for LM-based retrieval with an encoding budget for the following reasons:

1. Semantic Understanding vs. Speed: While heuristic-based methods offer computational efficiency, they often fail to capture subtle semantic relationships that are crucial for accurate retrieval in complex TableQA tasks. We have presented results from replacing embedding-based retrieval with BM25 and a hybrid approach above. While BM25 and the hybrid approach excel at retrieving all cells from the table, their inferior semantic understanding capabilities result in poorer overall performance.
2. Retrieval as Guidance, not the Final Answer: Even with a limited encoding budget, our retrieval mechanism effectively guides the solver agent (PyReAct in our case) to the relevant columns. As demonstrated in Figure 8, the solver agent can detect the existence of the “description” column through retrieval and then use basic string matching techniques (e.g., df['description'] .str.contains('Wallet')) to pinpoint the precise cell within those columns. Thus, while retrieval may not always return all relevant cells, it remains crucial for directing the solver's attention to the right areas of the table.
3. Empirical Evidence from the Study of Encoding Budget: The ablation study in Figure 6 shows that TableRAG is robust to changes in the encoding budget. We varied the encoding budget from 100 to 10,000 and found that TableRAG's performance remained near optimal on both ArcadeQA and BirdQA. In contrast, Row-Column Retrieval is sensitive to the encoding budget, with performance significantly dropping when the encoding budget is either increased or decreased. This study validates the earlier statement that even if not all relevant cells are retrieved, the solver agent can still obtain the necessary information through programs using the retrieved information in most cases.

We will make these clearer in the final version.

Novelty

Many papers about RAG reveals that the accuracy of the evidence is crucially important for the performance, and the idea of this paper is to is another practice in the TableQA domain.

Thank you for highlighting the importance of evidence accuracy in RAG systems. While this principle is well-established for unstructured data, its application to structured data like tables remains an open question. Prior retrieval-based methods for table understanding haven't consistently outperformed non-retrieval methods, primarily due to the challenges of efficient retrieval from tabular data. Our work addresses this gap by introducing a novel cell-schema retrieval method. Additionally, since existing benchmarks only focus on small tables, we contribute two large-scale TableQA datasets to facilitate further research in this area.

Evaluation on open models

The paper only evaluates the model on closed models gpt-3.5 and gemini, but fails to evaluate on open models like Mistral, Llama etc. Why not evaluate on more models? Why not evaluation on commonly used TableQA benchmarks like WikiTableQA?

We follow reviewer’s suggestion and include results obtained using Mistral Nemo, the latest model with a 128K context length from Mistral, as shown below.

The results demonstrate that TableRAG remains the most effective table prompting technique when working with Mistral Nemo, consistent with our previous findings using GPT 3.5 and Gemini 1.0 Pro. Furthermore, even with its 128K context length, Mistral Nemo struggles to handle large tables in ArcadeQA and BirdQA.

Evaluation on WikiTableQA

Why not evaluation on commonly used TableQA benchmarks like WikiTableQA?

In our paper, we have evaluated TabFact, a commonly used TableQA benchmark, and extended it to studying the behavior of TableRAG at different scales. In addition, we follow the reviewer's suggestion and perform experiments on WikiTableQA with GPT 3.5 as shown below:

In this benchmark, the entire table content often fits within the context window of the language models, which is not representative of real-world large-scale industrial scenarios.

Nevertheless, the results clearly indicate that TableRAG maintains its superior performance even on smaller tables. It's important to highlight that the baseline methods require that LLMs process the entire table, thus limiting their scalability for large tables.