SAGE_A_Framework_for_RAG_解读
摘要
RAG在特定corpus QA任务上有 成熟应用,然而 大量反例 仍存在。它们不仅仅是LLMs的局限性导致,相反,它们主要由检索不准确信息导致,检索不准的原因是 1 现在RAG技术分块不考虑语义。2 丢失必要信息 和 过多噪声 之间的权衡。
本文提出一个RAG框架 SAGE。1 解决分块问题上,提出训练一个语义分块模型,能将corpus分成语义完整块。2 设计 一个chunk挑选算法,根据相似度的下降速度(delta)动态选择块。3 为进一步确保精度,提出让LLMs评估检索的块是excessive or lacking. 实验证明很牛逼。
INTRUDUCTION
bg of problem – RAG: enhance a generation model’s ability answer questions. With the rise of LLMs, RAG demonstrated proficiency in QA.
Limitions overview:
L1:
challenge of addressing L1:
L2:
challenge of addressing L2:
approach overview:
contributions:
PRELIMINARIES(预备知识)
RAG system phases:
- vector DB construction: corpus is devided into chunks. Convert chunks into vector representations (embedding chunks). Stored these vectors.
- Convert the received question into vector. (embedding question). Caculate similarity with chunks in vector DB and then identifying the top N chunks which serve as the context for an LLM.
- The given question and context fixed into a prompt template forming a complete prompt. Feed the prompt to an LLM. It generates a response as the final answer.
Cost of LLM inference:
The matrix is the money required to obtain ansers from an LLM. Typically, RAG use LLM provided by LLM providers. The charges based on the volume of input and output tokens. OpenAI’s GPT-4 charge a lot.
I: input tokens; O: output tokens; c_i: cost per input token; o_i: cost per output token
Cost efficeincy Metric in RAG:
Acc: metrics used to eval the quality of an answer( e.g. F1-Score, BELU-1, Accuracy); Cost : money calcuated by the preceding equation.

(用网图链接有水印-.-)
OVERVIEW
(A) Vector Database Building: W
A given corpus is split into paragraphs with ‘\n’ initially (isn’t shown in the figure).
- “1 - 2”: All those paragraphs are then segmented by a pretrained model into short but semantically complete chunks as T.
- “3 - 4”: Embed chunks T into vectors.
- “5”: Store vetors in a vector DB.
(B) Retrieval:
A specific question is given.
- “1 - 2”: The given question is embedding into a vector.
- “3 - 4”: Query the vector DB for topN chunk vectors that most align with question’s vector, thereby retrieving the topN chunks.
- “5 - 6”: Use a reranking model to rerank these N chunks into a new order (scored N chunks).
- “7 - 8”: Select the top-ranking K chunks that come before the point where a significant dip in scores is ovserved
(C) Generation based on LLM:
Given the specific question and topK chunks from “(B) Retrieval”.
- “1”: Craft a prompt incorporating the question and the topK chunks. This process is tailed to the question’s type.
- “2 - 3”: Feed the prompt into an LLM to procure the response.
- “4 - 5”: Organize three things (question, topK chunks, answer from “2 - 3”) into a feedback prompt and then submitting it to an LLM. Generating a response of the score of the answer and an assessmen(new K) of the chunks.
- If the score is enough, execute “6 - 7”. Else, execute right-side “6” to adjust the parameter topK.
- If the loop has been executed three times, execute “6 - 7” only.
SEMANTIC SEGMENTATION
High-level Idea
corpus segmentation is crucial. Ineffective segmentation results in semantically incomplete chunks leading to irrelevant retrieval (demonstrated in 8. INSIGHTS OF RAG TASKS).
Two methods predominantly employed: “fixed-size” and “fixed-size with sentence complete”.
- “fixed-size” leads to incomplete sentence in chunks undermining coherence of chunks.
- “fixed-size with sentence complete” breaks continous sentences seperated into different chunks distorting the meaning of a chunk (e.g. unclear references).
- Increasing the fixed length can mitigate the issue of chunks having incomplete meanings while causing each chunk to contain an excessive amount of information which ruins the similarity caculation and brings more noise to LLM.
- While employing a LLM for segmentation is a conceivable approach, it proves to be costly (demonestrated in 7. EXPERIMENT) and slow.
We propose to develop a lightweight and effective segmentation model that ensure that each chunk is complete.
Model Construction
Segmentation model incorporates an embedding model and a Multi-layer Perceptron([pərˈseptrɑn]) (MLP) model.
- An Embedding model utilizing AnglE method, embeds two sentences into two vectors x1 and x2.
- Feature augmentation receives x1 and x2 and then subtracct and mutiply them yielding x1-x2 and x1*x2.
- A MLP model takes x1, x2, x1-x2 and x1*x2 to produce a score which determines wether to split or not.
- The decision to make x1-x2 or x1*x2 are based on the observation that embedding still reflect semantic difference (cite Sentence-bert).
Model Training

- First, gather a substantial amount of semantically segmented training data. Recommend Wikipedia dataset where passages have been semantically segmented into paragraphs.
- Construct triplet training data <s1, s2, label>: If s1 and s2 are consecutive([kən’sekjətɪv]) within the same paragraph, label = 1 representing they should be grouped into the same chunk. Otherwise, label = 0 suggesting a break.
- The procedure are detailed in the figure above. Keep fedding triplets to learn the parameters of the embedding model and MLP model utilizing SDG.
Model Inference
- Each two sentences are put into the segmentation model to gain a score. If < threshold (e.g., 0.5), then a split occur. Conversly, they retained within the same chunk.
- Segmentation Model runs in parallel processing corpus swiftly. (“For instance, gathering all pairs of sentences within a corpus and organize them into multiple batches, each with a size of 512.” batches是不是下面的粗粒度chunks?)
Corpus Segmentation
- Give a corpus, segment it into coarse-grained([ˌkɔːrsˈɡreɪnd]) chunks with length of approximately l (I call it as “large-fixed-size with sentence complete”).
- Subsequently, segmentation model evaluates each pair of adjacent sentences in these coarse-grained chunks.
- (So the boarder of coarse-grained chunks is ingored? I think that’s fine because each chunk is large enough.)
GRADIENT-BASED CHUNK SELECTION

High-level Idea
- Advanced RAG system typically leverage a reranking model to assess first-stage retrieval result (N chunks) with a given question. Subsequently, with the relevance scores assigned by the reranking model, top K chunks are selected as context for QA.
- The idea of using a large K to avoid missing key chunks will lead to an overload of irrelevant chunks, which may confuse LLMs (illustrated in EXPERIMENTS).
- Figure above shows two general cases, in which if chunks preceding a sharp decline on score are selected can easily get the correct answer (demonstrated in EXPERIMENTS).
Gradient based Selection
With top K reranked chunks C_s and their corresponding scores S, preset minimum number min_k and threshold of gradient g.
1 2 3 4 5 6 7 8 9# chks S min_k g def func(chks, S, min_k, g): add_chks = [] for i in range(min_k, K+1): if S[i] > S[i-1]/g: # didn't drop too much add_chks.append(chks[i]) else: break return add_chks(论文里的伪代码有错误:“score = C[k − 1];” 对不上 原文: “if the score of a chunk exceeds 1/g times the score of its predecessor”)
LLM SELF-FEEDBACK
Feedback Loop
After feeding context and question to LLM for final answer, we need to assess the suitability of context and quality of current answer before presented to user finally. By prompt engineer, LLM’s assessment consist of a evaluation score (1-10) and a context adjustment flag (-1/1).
Upon receiving score and flag, if the score is higher than a preset threshold e.g. 9, the current answer is qualified for user, else adjust the parameter min_k in chunks selection: -1 implies reducing min_k by one, 1 implies increasing min_k by one.
EXPERIMENTS
Experimental Setup
Large language models: GPT3.5 turbo, GPT4, GPT4-o-mini, UnifiedQA-3B.
Datasets:
- QuALITY: This dataset comprises multiple-choice questions based on articles around 5,000 tokens each, featuring a challenging portion, QuALITYHARD, with questions most annotators got wrong under time constraints.
- QASPER: