RAGV iz: Diagnose and Visualize Retrieval-Augmented Generation (2024)

Abstract: Retrieval-augmented generation (RAG) combines knowledge from domain-specific sources into large language models to ground answer generation. Current RAG systems lack customizable visibility on the context documents and the model’s attentiveness towards such documents. We propose RAGViz, a RAG diagnosis tool that visualizes the attentiveness of the generated tokens in retrieved documents. With a built-in user interface, retrieval index, and Large Language Model (LLM) backbone, RAGViz provides two main functionalities: (1) token and document-level attention visualization, and (2) generation comparison upon context document addition and removal. As an open-source toolkit, RAGViz can be easily hosted with a custom embedding model and HuggingFace-supported LLM backbone. Using a hybrid ANN (Approximate Nearest Neighbor) index, memory-efficient LLM inference tool, and custom context snippet method, RAGViz operates efficiently with a median query time of about 5 seconds on a moderate GPU node. Our code is available at https://github.com/cxcscmu/RAGViz. A demo video of RAGViz can be found at https://youtu.be/cTAbuTu6ur4.

## introduction large language models (llms), such as gpt-4 (ope, 2024), have revolutionized the field of artificial intelligence with their impressive language understanding and generation capabilities developed through extensive pretraining on large-scale textual data. a key limitation of using pretrained llms for zero-shot answer generation is their lack of access to domain-specific knowledge, as these models rely solely on parametric memory. the fixed knowledge derived from parametric memory often leads to hallucinations. to address this issue, @xcite introduces retrieval-augmented generation (rag), a technique that leverages retrieval mechanisms to incorporate non-parametric memory, typically derived from documents retrieved from domain-specific data stores. various systems have been developed to deliver rag services. for instance, openai assistants (openai, 2024) and pinecone assistant @xcite are "chat-with-your-files" products that use retrieved documents as context for a chatbot. while these rag systems offer state-of-the-art performance in grounded answer generation, they lack explainability regarding the efficacy of the context documents they use to produce those answers. some existing tools have been developed to improve language model explainability, such as bertviz @xcite , an open-source python tool that provides attention visualizations for transformer models. although such tools effectively analyze input token importance, they lack a customizable approach for analyzing the interaction between retrieved context documents and language generation. in this paper, we propose ragviz, a diagnostic tool designed to analyze llm attention mechanisms on the retrieved documents that provide context to ground llm answer generation. ragviz's novelty lies in its focus on the interaction between the retrieval pipeline and the language model. ragviz offers attention visualizations based on different levels of scoring: both cumulative attention scores on documents and individual token attention scores selected by the user. along with document toggling, ragviz enables users to qualitatively assess the effectiveness of retrieved documents and determine whether they contribute to hallucinations. ragviz's system primarily relies on cpu nodes, with the exception of a gpu node that hosts the llm inference server. the system entry point is a web node that hosts the frontend as static content and routes queries to the main cpu node. this node forwards the query to worker nodes for document retrieval, builds the context, and sends the request to the gpu node for llm inference. the generated answer and associated attention scores are then returned as an http response to the frontend. ragviz achieves efficiency through its distributed architecture and optimized llm inference, partitioning large datasets across multiple nodes for parallel processing and faster retrieval. it uses fast inference libraries for low-latency llm output generation. additionally, ragviz is customizable, allowing integration with any retrieval pipeline or attention-based language model architecture supported by huggingface @xcite , offering flexibility for diverse research needs. ## ragviz features and use cases this section first examines the innovative features of ragviz and outlines its key benefits. then, a few potential use cases are explored to demonstrate how ragviz can be valuable to researchers and domain experts. ## system architecture this section introduces ragviz's system architecture and its query pipeline. the system has four main components: the ann (approximate nearest neighbor) index for dense retrieval, the backend server, the llm inference server, and the frontend user interface. these components are implemented separately to allow for configurability. ragviz's system is originally designed for use with a job scheduler like slurm @xcite . ## experiment this section introduces the chosen configurations of ragviz's system demonstration and presents efficiency evaluations. ## conclusion ragviz is a powerful diagnostic tool for analyzing and improving rag pipelines by providing detailed visualizations of attention mechanisms at various levels. its attention-driven insights help users better understand the relationship between retrieved documents and language model outputs, making it invaluable for identifying hallucinations and enhancing retrieval efficacy. as an open-source tool under the mit license, ragviz is available for research and development. we plan to support custom models in the future, allowing users to evaluate their own language models within the rag pipeline. additionally, we aim to improve usability by containerizing services for more efficient deployment and resource management. we will also unify the llm inference process to use one inference library, leading to further improvements in speed and resource utilization.