KNEE动态分块:高性能RAG的革新算法
DRAG with KNEE 是一种基于膝点剪枝优化的动态检索增强生成算法。它在 Qdrant 中构建分层向量树,结合密集与稀疏嵌入的混合搜索,并采用倒数排名融合(RRF)策略,实现上下文的自适应精准召回——有效突破传统静态 top_k 方法的局限,如上下文溢出或关键信息丢失等问题。
构建分层索引结构
文档按页拆分。每一页由大语言模型(LLM)生成简洁摘要和关键词,这些内容作为树结构中的叶节点。
叶节点以3至5个为一组,合并为父节点。每个父节点继承子节点的关键词,并通过聚合子内容生成自身摘要。该过程持续递归,直至到达根节点(parent_id = -1)。
在 Qdrant 中,数据以扁平化的点列表形式存储,元数据用于过滤。这种设计支持灵活操控分层结构,仅通过负载字段即可实现层级调整。
混合搜索与RRF融合
点比较结合密集与稀疏嵌入,使用倒数排名融合(RRF)计算得分:
$$RRFscore(d \in D) = \sum_{r \in R} \frac{1}{k + r(d)}$$
其中 $D$ 为文档集合,$R$ 表示不同排序结果(密集、稀疏),$k=60$。
搜索算法
分支搜索(Branch Search)
首先定位顶层根节点,然后递归比较父节点与子节点。若子节点得分更高,则深入下一层;否则锁定当前父节点。
def parent_vs_children(query: str, parent: ScoredPoint) -> list[ScoredPoint]:
file_name = parent.payload["file_name"]
parent_id = parent.payload["id"]
child_ids = parent.payload["child_ids"]
all_ids = [parent_id] + child_ids
if len(child_ids) == 0:
console.print(
f"Parent with id {parent_id} is leave!", style="bold underline violet"
)
return []
console.print()
console.print("#" * 20, style="red")
console.print(
f"Children (ids: {child_ids}) [underline]VS[/underline] Parent (id: {parent_id}) [ FILE: {file_name} ]\nFor query: {query[:20]}...",
style="bold violet",
)
with RAGalicClient() as client: # artifacts from a previous life :)
dense_model: str = client.client.embedding_model_name # type: ignore
sparse_model: str = client.client.sparse_embedding_model_name # type: ignore
children_and_parent_filter = Filter(
must=FieldCondition(key="id", match=MatchAny(any=all_ids))
)
sorted_points = client.client.query_points(
collection_name="ragalic",
prefetch=[
Prefetch(
query=Document(text=query, model=dense_model),
using="dense",
limit=len(all_ids),
filter=children_and_parent_filter,
),
Prefetch(
query=Document(text=query, model=sparse_model),
using="sparse",
limit=len(all_ids),
filter=children_and_parent_filter,
),
],
query=FusionQuery(fusion=Fusion.RRF),
query_filter=children_and_parent_filter,
limit=len(all_ids),
).points
children_better_then_parent = []
for point in sorted_points:
if point.payload["id"] == parent_id:
break
children_better_then_parent.append(point)
ids = [child.payload["id"] for child in children_better_then_parent]
console.print(
f"Child ids which is better than parent: {ids if ids else 'N/A'}",
style="italic purple",
)
console.print("#" * 20, style="red")
console.print()
return children_better_then_parent
def branch_search(query: str, num_roots: int = 3) -> list[Chunk]:
roots = find_roots(query=query, num_to_find=num_roots)
final_points = []
points_to_process = roots
while len(points_to_process) > 0:
new_point_to_process = []
for point in points_to_process:
new_points = parent_vs_children(query=query, parent=point)
if len(new_points) == 0:
console.print(
f"--- NEW FINAL POINT (id: {point.payload['id']} | file: {point.payload['file_name']} | pages: {point.payload['page_start']} - {point.payload['page_end']}) ---",
style="bold green",
)
final_points.append(point)
else:
new_point_to_process.extend(new_points)
points_to_process = new_point_to_process
console.print(
f"--- WAS FOUND {len(final_points)} POINTS FOR QUERY: '{query[:20]}...' ---",
style="bold green on white",
)
return prepare_chunks(final_points)
使用方式:
from rag_lib.search import branch_search
results = branch_search(query="your_query", num_roots=3)
for chunk in results:
print(f"{chunk.file_name} (pages {chunk.page_start}-{chunk.page_end}):
{chunk.text[:200]}...")
优点:召回率高。缺点:可能引入无关分支。
束搜索(固定束宽)
- 取当前父节点及其子节点的集合。
- 父节点中胜出者淘汰其子节点;子节点中胜出者淘汰其父节点。
- 幸存者竞争固定束宽。
可有效降低噪声,但存在误删相关段落或误保留无关内容的风险。
自适应束搜索(膝点法)
采用膝点法动态调节束宽。通过识别 RRF 曲线的拐点(下凸区域),判断最优截断位置。
动态剪枝:每轮先执行膝点剪枝,再生成下一轮候选集。可剔除约62%的无关文档,显著提升效率与精度。
— Editorial Team
暂无评论。