RAFT: Adaptación de Modelos de Lenguaje a RAG Específico de Dominio

RAFT: Adapting Language Model to Domain Specific RAG

March 15, 2024

Autores: Tianjun Zhang, Shishir G. Patil, Naman Jain, Sheng Shen, Matei Zaharia, Ion Stoica, Joseph E. Gonzalez

cs.AI

Resumen

El preentrenamiento de modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) en grandes corpus de datos textuales es ahora un paradigma estándar. Al utilizar estos LLMs para muchas aplicaciones posteriores, es común incorporar adicionalmente nuevos conocimientos (por ejemplo, noticias de actualidad o conocimiento específico de un dominio) en el modelo preentrenado, ya sea mediante técnicas basadas en RAG (Retrieval-Augmented Generation) o mediante ajuste fino (fine-tuning). Sin embargo, la metodología óptima para que el modelo adquiera dichos conocimientos sigue siendo una pregunta abierta. En este artículo, presentamos Retrieval Augmented FineTuning (RAFT), una receta de entrenamiento que mejora la capacidad del modelo para responder preguntas en entornos de "libro abierto" dentro de un dominio específico. En RAFT, dada una pregunta y un conjunto de documentos recuperados, entrenamos al modelo para ignorar aquellos documentos que no ayudan a responder la pregunta, a los que llamamos documentos distractores. RAFT logra esto citando textualmente la secuencia correcta del documento relevante que ayudaría a responder la pregunta. Esto, junto con la respuesta estilo cadena de pensamiento (chain-of-thought) de RAFT, mejora la capacidad del modelo para razonar. En RAG específico de dominio, RAFT mejora consistentemente el rendimiento del modelo en los conjuntos de datos PubMed, HotpotQA y Gorilla, presentando una receta posterior al entrenamiento para mejorar LLMs preentrenados en RAG dentro de un dominio. El código y la demostración de RAFT están disponibles en github.com/ShishirPatil/gorilla.

English

Pretraining Large Language Models (LLMs) on large corpora of textual data is now a standard paradigm. When using these LLMs for many downstream applications, it is common to additionally bake in new knowledge (e.g., time-critical news, or private domain knowledge) into the pretrained model either through RAG-based-prompting, or fine-tuning. However, the optimal methodology for the model to gain such new knowledge remains an open question. In this paper, we present Retrieval Augmented FineTuning (RAFT), a training recipe that improves the model's ability to answer questions in a "open-book" in-domain settings. In RAFT, given a question, and a set of retrieved documents, we train the model to ignore those documents that don't help in answering the question, which we call, distractor documents. RAFT accomplishes this by citing verbatim the right sequence from the relevant document that would help answer the question. This coupled with RAFT's chain-of-thought-style response helps improve the model's ability to reason. In domain-specific RAG, RAFT consistently improves the model's performance across PubMed, HotpotQA, and Gorilla datasets, presenting a post-training recipe to improve pre-trained LLMs to in-domain RAG. RAFT's code and demo are open-sourced at github.com/ShishirPatil/gorilla.