Advanced RAG 08：使用 Self-RAG 打造高質量、可追溯的 RAG System 原創 精華
編者按：RAG 技術通過檢索并利用外部知識源，能夠較為有效地提升生成內容的準確性和多樣性。然而，經典 RAG 流程也存在一些不足，例如不必要的檢索會浪費計算資源，并可能引入無關內容或錯誤信息，影響生成質量。
本文介紹了 Self-RAG 這一技術，通過引入 Reflection Tokens，語言模型能夠根據具體需求動態決定是否檢索外部知識，從而大大減少了不必要的檢索操作。與此同時，Self-RAG 通過特殊的訓練流程，使生成的內容不僅通順自然，與事實知識相符，甚至還可以追溯知識源。
當然，Self-RAG技術訓練過程相對復雜，在生成階段也融入了諸多特殊機制，在一定程度上增加了推理成本。不過，本文作者也提出了優化 Self-RAG 的一些建議，如簡化 Reflection Tokens 設計、探索不同模型大小的影響等，指明了一些技術發展方向。
作者 | Florian June
編譯 | 岳揚
本文從一個常見的生活場景入手：參加開卷考試。我們通常會采用以下兩種作答策略：
- 方法一：對于熟悉的題目，直接快速作答；對于不熟悉的題目，快速翻閱參考書，找到相關部分，在腦海中整理分類和總結后，再在試卷上作答。
- 方法二：每一個題目都需要參考書本進行解答。先找到相關部分，在腦海中進行整合和總結后，再到試卷上書寫答案。
顯然，方法一更受考生青睞，是首選方法。方法二不僅耗時，還有可能引入無關的或錯誤的信息，導致出現混淆和錯誤，甚至在考生原本擅長的領域也不例外。
然而，方法二是經典的 RAG （檢索->整合->生成）流程[1]，而方法一則代表 Self-RAG 流程[2]，本文將進一步探討這個問題。
01 Overview
圖 1 將 RAG 和 Self-RAG[2] 的主要流程進行了比較：
圖 1：Overview of Self-RAG 。Self-RAG（右）的主要功能為檢索（retrieve）、評判（critique）和生成（generate），使生成的文本內容不僅通順流暢，而且與事實知識相符，并且可以追溯到原始知識源。資料來源：https://arxiv.org/pdf/2310.11511.pdf
Self-RAG 主要由三大步驟組成：
- 按需檢索（Retrieval as Needed）： 當模型需要檢索信息時，例如查詢“美國各州的名字是怎么來的？”(圖 1 右上方）時，模型輸出中會包含一個 [Retrieve] token，表示需要檢索與 Query 相關的內容。相反，當被要求寫 "寫一篇主題為 ‘記憶中深刻的暑假’ 的文章"（圖 1 右下方）時，模型會選擇直接生成答案，無需進行檢索。
- 以并行方式生成內容（Parallel Generation）： 模型會同時使用提示詞（prompt）和檢索到的內容來生成模型輸出。在整個過程中，有三種 reflection tokens （譯者注：用于指示模型在生成過程中應該執行哪些操作、效果如何的token，比如前文的 [Retrieve]。）會顯示檢索內容的相關性如何。
- 內容的評估和選擇： 對步驟 2 中生成的內容進行評估，并選擇最佳文檔段落作為輸出。
請注意，上述模型是經過特殊訓練的，其訓練過程將在本文后續部分進行討論。
02 Reflection Tokens 的介紹
如圖 2 所示，與 RAG 相比，self-RAG 框架的不同之處在于它在生成過程中使用了 reflection tokens 進行更精確的控制。
圖 2：Self-RAG 中使用的四種 reflection tokens 。每種類型都使用多個 tokens 來表示其輸出值（output）。底部三行是三類 critique tokens ，粗體字表示這一類中最理想的 critique tokens 。x、y、d 分別表示輸入（input）、輸出（output）和相關段落（relevant passage）。來源：Self-RAG[2]
一般來說，self-RAG 會做出四種不同的判斷：
- [Retrieve]： 決策過程，決定是否從資源 R 中檢索額外信息。
- [IsREL]： 相關性檢查，確定給定數據 d 是否包含解決問題 x 所需的信息。
- [IsSUP]： 驗證過程，檢查數據 d 中是否有依據支持所生成回復 y 中的內容。
- [IsUSE]： 輸出結果是分數數值（數值范圍 1-5 ），其中 5 代表所生成的回復對問題的解決有用程度最高。
在 RAG 系統中，檢索是一個必須進行的固定步驟，無論條件如何，都要進行檢索。相比之下，self-RAG 引入了 reflective tokens ，使 LLM 更具適應性（adaptable）和智能性（intelligent）。在文本生成過程中，如果 LLM 遇到了需要額外信息支持的不確定區域，它會在遇到 reflective token 時暫停執行文本生成任務，系統會執行一次快速而精確的信息檢索操作，最后 LLMs 會利用新獲取的信息繼續完成當前的文本生成任務。
03 Code Explanation 通過解讀代碼了解 self-RAG
為了直觀地理解 Self-RAG 過程，我們需要首先查看、檢查代碼，然后詳細討論模型的訓練過程。
Self-RAG 這種技術是開源的[3]，知名的開源 Python 庫 Langchain?[4] 和 LlamaIndex 都分別實現了 Self-RAG 功能。本文將以 LlamaIndex 庫中 Self-RAG 的具體技術實現[5]作為參考進行說明。
3.1 環境配置
首先，進行環境配置。
(base) Florian@instance-1:~$ conda create -n llamaindex python=3.11
(base) Florian@instance-1:~$ conda activate llamaindex
(llamaindex) Florian@instance-1:~$ pip install llama-index
(llamaindex) Florian@instance-1:~$ pip install huggingface-hub
(llamaindex) Florian@instance-1:~$ huggingface-cli login安裝完成后，LlamaIndex 的版本信息如下，請確認：
llama-index 0.10.20
llama-index-core 0.10.20.post2下載論文提供的 Llama2-7B 模型，模型大小約為 4.08G。
(llamaindex) Florian@instance-1:~$ huggingface-cli download m4r1/selfrag_llama2_7b-GGUF selfrag_llama2_7b.q4_k_m.gguf --local-dir "YOUR_DOWNLOAD_MODEL_DIR" --local-dir-use-symlinks False
(llamaindex) Florian@instance-1:~$ ls "YOUR_DOWNLOAD_MODEL_DIR"
selfrag_llama2_7b.q4_k_m.gguf3.2 測試代碼
測試代碼如下。首次執行時需要下載 SelfRAGPack[5]。
import os
os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"
from llama_index.core import Document, VectorStoreIndex
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.readers import SimpleDirectoryReader
from pathlib import Path
# Option: download SelfRAGPack
# The first execution requires the download of SelfRAGPack.
# Subsequent executions can comment this out.
from llama_index.core.llama_pack import download_llama_pack
download_llama_pack(
"SelfRAGPack",
"./self_rag_pack")
from llama_index.packs.self_rag import SelfRAGQueryEngine
# The directory where the Llama2 model was previously downloaded and saved.
download_dir = "YOUR_DOWNLOAD_MODEL_DIR"
# Create testing documents
documents = [
Document(
text="A group of penguins, known as a 'waddle' on land, shuffled across the Antarctic ice, their tuxedo-like plumage standing out against the snow."
),
Document(
text="Emperor penguins, the tallest of all penguin species, can dive deeper than any other bird, reaching depths of over 500 meters."
),
Document(
text="Penguins' black and white coloring is a form of camouflage called countershading; from above, their black back blends with the ocean depths, and from below, their white belly matches the bright surface."
),
Document(
text="Despite their upright stance, penguins are birds that cannot fly; their wings have evolved into flippers, making them expert swimmers."
),
Document(
text="The fastest species, the Gentoo penguin, can swim up to 36 kilometers per hour, using their flippers and streamlined bodies to slice through the water."
),
Document(
text="Penguins are social birds; many species form large colonies for breeding, which can number in the tens of thousands."
),
Document(
text="Intriguingly, penguins have excellent hearing and rely on distinct calls to identify their mates and chicks amidst the noisy colonies."
),
Document(
text="The smallest penguin species, the Little Blue Penguin, stands just about 40 cm tall and is found along the coastlines of southern Australia and New Zealand."
),
Document(
text="During the breeding season, male Emperor penguins endure the harsh Antarctic winter for months, fasting and incubating their eggs, while females hunt at sea."
),
Document(
text="Penguins consume a variety of seafood; their diet mainly consists of fish, squid, and krill, which they catch on their diving expeditions."
),
]
index = VectorStoreIndex.from_documents(documents)
# Setup a simple retriever
retriever = VectorIndexRetriever(
index=index,
similarity_top_k=10,
)
model_path = Path(download_dir) / "selfrag_llama2_7b.q4_k_m.gguf"
query_engine = SelfRAGQueryEngine(str(model_path), retriever, verbose=True)
# No retreival example
response = query_engine.query("Which genre the book pride and prejudice?")
# Retreival example
response = query_engine.query("How tall is the smallest penguins?")測試代碼運行結果如下（大部分 llama_cpp 的調試信息已被刪除）：
...
...
Model metadata: {'tokenizer.ggml.add_eos_token': 'false', 'tokenizer.ggml.eos_token_id': '2', 'general.architecture': 'llama', 'llama.rope.freq_base': '10000.000000', 'llama.context_length': '4096', 'general.name': 'LLaMA v2', 'tokenizer.ggml.add_bos_token': 'true', 'llama.embedding_length': '4096', 'llama.feed_forward_length': '11008', 'llama.attention.layer_norm_rms_epsilon': '0.000010', 'llama.rope.dimension_count': '128', 'tokenizer.ggml.bos_token_id': '1', 'llama.attention.head_count': '32', 'llama.block_count': '32', 'llama.attention.head_count_kv': '32', 'general.quantization_version': '2', 'tokenizer.ggml.model': 'llama', 'general.file_type': '15'}
Using fallback chat format: None
llama_print_timings: load time = 4887.53 ms
llama_print_timings: sample time = 11.29 ms / 22 runs ( 0.51 ms per token, 1947.76 tokens per second)
llama_print_timings: prompt eval time = 4887.46 ms / 24 tokens ( 203.64 ms per token, 4.91 tokens per second)
llama_print_timings: eval time = 5883.27 ms / 21 runs ( 280.16 ms per token, 3.57 tokens per second)
llama_print_timings: total time = 10901.84 ms / 45 tokens
Final answer: The book "Pride and Prejudice" is a romantic novel by Jane Austen.
...
...
llama_print_timings: load time = 4887.53 ms
llama_print_timings: sample time = 11.74 ms / 20 runs ( 0.59 ms per token, 1703.29 tokens per second)
llama_print_timings: prompt eval time = 7473.66 ms / 37 tokens ( 201.99 ms per token, 4.95 tokens per second)
llama_print_timings: eval time = 5414.34 ms / 19 runs ( 284.96 ms per token, 3.51 tokens per second)
llama_print_timings: total time = 13076.88 ms / 56 tokens
Input: ### Instruction:
How tall is the smallest penguins?
### Response:
[Retrieval]<paragraph>Penguins consume a variety of seafood; their diet mainly consists of fish, squid, and krill, which they catch on their diving expeditions.</paragraph>
Prediction: [Relevant]The height of the smallest penguin species can vary depending on the species.[No support / Contradictory][Utility:5]
Score: 1.4213598342974367
10/10 paragraphs done
End evaluation
Selected the best answer: [Relevant]The smallest penguin species is the Little Blue Penguin (also known as the Fairy Penguin), which can grow to be around 40 centimeters (16 inches) in height.[Fully supported][Utility:5]
Final answer: The smallest penguin species is the Little Blue Penguin (also known as the Fairy Penguin), which can grow to be around 40 centimeters (16 inches) in height.我們可以看到，第一個 query 不需要執行檢索操作，而第二個 query 已進行了檢索操作并進行了輸出內容的評估。
理解測試代碼的關鍵在于弄清楚類 SelfRAGQueryEngine[6] 的實現，現在我們來深入探討一下這個類。
3.2 類 SelfRAGQueryEngine
首先來看看構造函數[7]，它主要用于使用 llama_cpp 加載 Llama2-7B 模型。
class SelfRAGQueryEngine(CustomQueryEngine):
"""Simple short form self RAG query engine."""
llm: Any = Field(default=None, description="llm")
retriever: BaseRetriever = Field(default=None, description="retriever")
generate_kwargs: Dict = Field(default=None, description="llm generation arguments")
verbose: bool = Field(default=True, description="Verbose.")
def __init__(
self,
model_path: str,
retriever: BaseRetriever,
verbose: bool = False,
model_kwargs: Dict = None,
generate_kwargs: Dict = None,
**kwargs: Any,
) -> None:
"""Init params."""
super().__init__(verbose=verbose, **kwargs)
model_kwargs = model_kwargs or _MODEL_KWARGS
self.generate_kwargs = generate_kwargs or _GENERATE_KWARGS
try:
from llama_cpp import Llama
except ImportError:
raise ImportError(_IMPORT_ERROR_MSG)
self.llm = Llama(model_path=model_path, verbose=verbose, **model_kwargs)
self.retriever = retriever接下來，我們將介紹處理 query 的相關函數[8]。其主要流程如圖 3 所示：
圖 3：query 處理功能的主要流程。圖片由作者提供。
關鍵部分已作注釋，以便讀者更好地理解。
def custom_query(self, query_str: str) -> Response:
"""Run self-RAG."""
# Obtain responses using the Llama2 model.
response = self.llm(prompt=_format_prompt(query_str), **_GENERATE_KWARGS)
answer = response["choices"][0]["text"]
source_nodes = []
# Determine if a retrieval is necessary.
if "[Retrieval]" in answer:
if self.verbose:
print_text("Retrieval required\n", color="blue")
# The step 1 of Figure 1, retrieve as needed.
documents = self.retriever.retrieve(query_str)
if self.verbose:
print_text(f"Received: {len(documents)} documents\n", color="blue")
paragraphs = [
_format_prompt(query_str, document.node.text) for document in documents
]
if self.verbose:
print_text("Start evaluation\n", color="blue")
# Step 2 and 3 in Figure 1, generate in parallel and evaluate
# (the code does not implement parallelism)
critic_output = self._run_critic(paragraphs)
paragraphs_final_score = critic_output.paragraphs_final_score
llm_response_per_paragraph = critic_output.llm_response_per_paragraph
source_nodes = critic_output.source_nodes
if self.verbose:
print_text("End evaluation\n", color="blue")
# Select the paragraph with the highest score and return it.
best_paragraph_id = max(
paragraphs_final_score, key=paragraphs_final_score.get
)
answer = llm_response_per_paragraph[best_paragraph_id]
if self.verbose:
print_text(f"Selected the best answer: {answer}\n", color="blue")
answer = _postprocess_answer(answer)
if self.verbose:
print_text(f"Final answer: {answer}\n", color="green")
return Response(response=str(answer), source_nodes=source_nodes)從代碼中可以看到，圖 1 中的三個步驟都有所體現。然而，LlamaIndex 的 self-RAG 功能并未實現并行化。感興趣的讀者可以查看 self._run_critic 函數了解更多信息，該函數也處理與各種 reflection tokens 相對應的分數。
04 如何訓練 Llama2–7B 模型
我們以前曾多次使用過 Llama2-7B 模型，本文來探討一下如何獲取和訓練該模型。
4.1 訓練目標
使語言模型能夠生成包含 reflection tokens 的文本。
4.2 兩個模型
在訓練過程中，需要使用兩個模型：批判模型（critic model） C 和生成模型（generator model） M。批判模型 C 主要生成 M 所需的已經標注好的用于有監督學習任務的數據（supervision data）。
然而，在推理過程中，只需使用模型 M，不需要批判模型 C。
4.3 批判模型 C
批判模型（critic model）經過訓練可生成 reflection tokens 。使用該模型是為了在任務輸出中能夠離線插入 reflection tokens ，從而更新訓練語料庫（training corpus）。
手動為每個文本段落標注 reflection tokens 的成本非常高。Self-RAG 利用 GPT-4 根據不同 reflection token 的定義、輸入和輸出，為每個 reflection token 分配唯一的特定指令，從而高效地完成數據標注任務。例如，[retrieval] token 會指導模型在進行內容評估時是否需要檢索外部文檔。
獲得訓練數據 D_critic 后，我們就可以基于傳統的、基于條件的語言模型（standard conditional language model）構建用于訓練機器學習模型的目標函數，如下所示：
批判模型 C 可以用任何預訓練的語言模型進行初始化和進一步微調。例如，它可以直接使用與生成模型相同的預訓練模型（如Llama 2 7B）進行初始化。
4.4 生成模型 M
圖 4 顯示了在 Self-RAG 框架中如何收集用于訓練生成模型（Generator）的 supervision data（譯者注：已經標注好的用于有監督學習任務的數據） 。給定一對輸入輸出（x, y），self-RAG 使用檢索器和批判模型對原始的模型輸出 y 進行標注，從而創建 supervision data（譯者注：已經標注好的用于有監督學習任務的數據） 。對于 y 中的每個段落 yt（yt ∈ y）：
圖4：收集生成模型（generator）的訓練數據。圖中的每次條件判斷都是通過批判模型 C 執行的。該圖片由作者提供，靈感來源于 Self-RAG[2] section 3.2.2 。
請注意，圖 4 中的每次條件判斷都是通過批判模型 C 執行的。獲得的訓練數據如圖 5 所示：
圖 5：Self-RAG 訓練示例。左邊的例子不需要進行外部檢索，而右邊的例子需要檢索外部文檔，因此插入了相關文檔段落。來源：Self-RAG[2]。
獲得訓練數據 D_gen 后，我們就可以構建用于訓練語言模型的 standard next-token prediction objective function （譯者注：“Standard”指這是語言模型訓練中最普遍、最基礎的目標函數形式?！癗ext-token prediction”指的是該函數的目標是預測序列中的下一個token（詞元），“Objective function”即目標函數，可能是損失函數等內容。）啦，如下所示：
生成模型 M 需要預測的不僅僅是輸出內容，還包括 reflection tokens 。
05 作者對 self-RAG 的見解和思考
總體來看，self-RAG 為 RAG 技術的增強和優化提供了一種全新的視角。不過，這種技術需要更復雜的訓練流程，并且在內容生成階段（generation phase） 融入了一些特殊機制，不僅生成目標輸出文本，還會生成多種類型的反饋標簽，并在生成時根據這些標簽執行多個判斷操作，這樣會不可避免地會增加推理成本?？赡軙?/span>嚴重影響重視實時性能要求的項目。
此外，該框架還有很大的優化空間。為激發進一步的討論和創新，這里有一些建議：
- 如何優化 reflection tokens 。Self-RAG 設計了四種 reflection tokens 。除了 [Retrieve] token 之外，其他三個（[IsREL]、[IsSUP]、[IsUSE]）都有一定的相似性。可以考慮使用更少的 reflection tokens ，或者設計表示其他語義的 reflection tokens ，這可能是一個值得探索的優化方向。
- 為什么批判模型（critic model）要使用大語言模型（LLM）呢？我認為這可能是由于像[IsUSE]這樣的 token 嚴重依賴于常識性知識。判斷某個回答的有用，是一項較小的模型可能也能夠勝任的任務。然而，這些較小的模型通常僅學習了特定的訓練數據，缺乏全面的知識。因此，使用LLM作為批判模型（critic model）是有道理的。
- 批判模型（critic model）模型大小的選擇。我們已經在 7B 和 13B 的模型上對 Self-RAG 進行了測試，結果非常出色。但是，如果我們切換到較小尺寸的 LLM，比如 3B ，會觀察到哪些差異？同樣，如果我們過渡到使用更大的 LLM ，比如 33B ，我們又能預見到多大的系統性能提升呢？
- 為什么不使用基于人類反饋的強化學習（RLHF）呢？該論文建議在特定任務示例數據（task examples）上訓練目標語言模型。然后使用離線但是批判模型對這些數據進行增強（使用 reflection tokens），與 RLHF 相比，訓練成本大大降低。此外，self-RAG 中的 reflection tokens 使得在推理階段的內容生成可控，而 RLHF 專注于在訓練期間與人類的偏好對齊。不過，論文中沒有包含任何與 RLHF 相關的對比實驗。
06 Conclusion
本文從一個直觀的生活場景（開卷考試）入手，介紹了 self-RAG 技術的基本流程，并輔以代碼解釋。文章還分享了作者的一些見解和思考。
如果你對 RAG（檢索增強生成） 技術非常感興趣，歡迎瀏覽、分享本系列其他文章。 : )
Thanks for reading!
Julian Yip
Multi-Cloud Data Architect | Azure, GCP, Databricks Certified | ML and MLOps Practitioner
END
參考資料
[1]https://medium.com/ai-in-plain-english/a-brief-introduction-to-retrieval-augmented-generation-rag-b7eb70982891
[2]https://arxiv.org/pdf/2310.11511.pdf
[3]https://github.com/AkariAsai/self-rag
[4]https://github.com/langchain-ai/langgraph/blob/main/examples/rag/langgraph_self_rag.ipynb?ref=blog.langchain.dev
[5]https://github.com/run-llama/llama_index/tree/v0.10.20/llama-index-packs/llama-index-packs-self-rag
[6]https://github.com/run-llama/llama_index/blob/v0.10.20/llama-index-packs/llama-index-packs-self-rag/llama_index/packs/self_rag/base.py
[7]https://github.com/run-llama/llama_index/blob/v0.10.20/llama-index-packs/llama-index-packs-self-rag/llama_index/packs/self_rag/base.py#L174
[8]https://github.com/run-llama/llama_index/blob/v0.10.20/llama-index-packs/llama-index-packs-self-rag/llama_index/packs/self_rag/base.py#L245
本文經原作者授權，由 Baihai IDP 編譯。如需轉載譯文，請聯系獲取授權。
原文鏈接：
???https://ai.gopubby.com/advanced-rag-08-self-rag-c0c5b5952e0e????
