🤖 AI

One Retrieval to Cover Them All: Co-occurrence-Aware Knowledge Base Reorganization for Session-Level RAG

This paper proposes a co-occurrence-aware knowledge base reorganization method that significantly improves session-level retrieval coverage and efficiency for enterprise RAG systems by clustering related documents and expanding query candidates, arguing that session-level coverage should replace single-query recall as the primary evaluation metric.

Original authors: Shivam Ratnakar, Yixuan Zhu, Cecilia Cheng, Chaya Vijayakumar

Published 2026-07-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shivam Ratnakar, Yixuan Zhu, Cecilia Cheng, Chaya Vijayakumar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to fix a complex problem with your website, like setting up a new online store. You need to do four things: connect your domain name, pick a design template, set up payments, and configure email.

In a standard AI system (called RAG), the computer acts like a librarian who only knows how to find books that look exactly like the question you ask.

You ask: "How do I connect my domain?"
The librarian finds three articles about domains.
The Problem: The librarian doesn't realize you also need the articles about templates and payments right now. Those articles are on completely different shelves in the library. To get them, you have to ask four separate questions, wait for the librarian to find each one, and hope they get it right every time.

The authors of this paper call this a "blind spot." They found that a standard system only finds about 41% of the information you actually need in one go.

The Solution: A New Way to Organize the Library

The researchers propose a clever trick: reorganize the library based on what people actually do together, not just what the books look like.

Think of it like a grocery store.

Standard System: It puts "Milk" next to "Cheese" because they are both white dairy products (semantic similarity).
This New System: It notices that people who buy "Milk" also almost always buy "Cereal" and "Strawberries" (co-occurrence), even though milk and strawberries don't look alike. So, it moves the Cereal and Strawberries to the same aisle as the Milk.

How It Works (The Simple Steps)

Offline Reorganization (The "Night Shift"): Before anyone asks a question, the system looks at thousands of past user sessions. It sees that when people ask about domains, they often also need help with templates and payments. It groups these unrelated topics together into "clusters" or "neighborhoods."
Online Retrieval (The "Day Shift"): When you ask your question, the system does two things:
- It finds the most obvious answer (like the standard librarian).
- The Magic Step: It immediately looks at the "neighborhood" where that answer lives and grabs the other items in that cluster (the templates and payments), even if they don't look like your question.
The Result: Instead of giving you just the domain articles, it gives you the domain articles plus the template and payment articles in a single answer.

Why This Matters

The paper tested this on a real enterprise knowledge base with over 6,000 articles. Here is what they found:

One Call, More Answers: Instead of needing four separate questions to get all the info, a single question now covers 58% of the user's needs (up from 41%).
Fewer Steps: To get 70% of the information you need, users now need 34% fewer questions. It's like finishing your shopping trip in one trip instead of four.
Smaller Library: Because the system groups things efficiently, it can shrink the effective size of the knowledge base to 20% of its original size while still finding the right answers.
Works Everywhere: This trick worked no matter which "brain" (AI model) they used to understand the words, and it worked across six different types of topics (like payments, email, and design).

The Trade-Off

The authors admit there is a tiny cost. Because they are grabbing a wider net of information, sometimes the very first answer might be slightly less perfect than before (dropping from 96% to 93% precision). However, they argue this is a great trade-off. It is much better to have 58% of the solution in front of you immediately than to have the "perfect" single answer and then have to ask three more questions to get the rest.

What They Didn't Claim

They did not claim this works for one-off questions where there is no follow-up (like a trivia quiz).
They did not claim this makes the AI's final written answer perfect; they only measured if the right documents were found.
They did not use real user logs for the training (due to privacy), but instead used expert-labeled data and simulated user paths to teach the system how to group things.

In short: This paper suggests that for complex, real-world problems, we shouldn't just ask the AI to find the "most similar" document. We should teach the AI to find the "most helpful group" of documents based on how humans actually use them together.

Technical Summary: One Retrieval to Cover Them All

Problem Statement

Current Retrieval-Augmented Generation (RAG) systems are optimized for single-query relevance, retrieving documents that are semantically similar to an isolated user query. However, enterprise users typically interact with knowledge bases (KBs) through sessions: coherent episodes of related questions that span semantically distant parts of the KB.

The paper identifies a critical "session coverage gap": standard RAG systems fail to retrieve the full set of documents a user needs in a single session because those documents often reside in different semantic neighborhoods. For example, a user setting up a website needs articles on domain configuration, template design, and payment processing. While functionally related, these documents are semantically distant in embedding space.

The Blind Spot: Standard retrieval metrics (e.g., Recall@k) optimize for finding the single best document for a single question, ignoring multi-document information needs.
The Consequence: On the WixQA benchmark, a standard RAG system covers only 41% of a user's session-level information need with a single retrieval call ( $k=8$ ). Users must issue multiple follow-up queries to access the remaining knowledge, increasing latency, API costs, and user friction.

Methodology

The authors propose Co-occurrence-Aware Knowledge Base Reorganization, a pre-retrieval approach that restructures the KB offline to align with user behavior patterns rather than just semantic similarity. The method consists of two phases:

1. Offline: Co-occurrence-Aware Clustering

The system learns a new embedding space based on document co-occurrence patterns (documents accessed together in sessions) rather than semantic text similarity.

Sequence Construction: Co-occurrence sequences are built from three sources:
1. Ground-truth: Expert-labeled queries with multiple relevant article IDs.
2. Embedding-neighborhood walks: Random walks through the semantic similarity graph with a high jump probability (40%) to capture cross-neighborhood relationships.
3. QA-driven sequences: Synthetic question-answer pairs and their source document neighbors.
Embedding Learning: Document IDs are treated as tokens, and the sequences as sentences. A Word2Vec (CBOW) model is trained to generate co-occurrence embeddings ( $z_d$ ). This captures functional relatedness: documents needed together in a session are close in $z$ -space, even if they are semantically distant.
Hierarchical Clustering: Agglomerative clustering is applied to the co-occurrence embeddings to group documents into clusters. The target is to reduce the effective KB size to 20% of its original size (e.g., 6,221 documents $\to$ 1,244 clusters).

2. Online: Hybrid Retrieval with Cluster Expansion

At query time, the system performs a hybrid retrieval process:

Vanilla Retrieval: Retrieve the top- $k_v$ (e.g., 3) documents based on standard semantic similarity (using any encoder).
Cluster Expansion: For each retrieved document, gather all other documents within its pre-computed cluster as expansion candidates.
Re-ranking: Re-rank the union of the vanilla results and the expansion candidates using the original query-document similarity to return the final top- $k$ (e.g., 8) results.

This approach is encoder-agnostic (the cluster structure is learned independently of the query encoder) and pre-retrieval (the KB is reorganized once offline).

Key Contributions

Session-Level Evaluation Metrics: The paper introduces Single-Query Session Coverage (the fraction of a session's documents retrieved in one call) and Calls to $\tau$ -Coverage (the number of calls needed to cover a specific fraction of a session). These metrics shift the evaluation focus from single-query recall to holistic session satisfaction.
Co-occurrence-Aware KB Reorganization: A novel method that reorganizes the KB offline using co-occurrence patterns to enable cluster-expanded hybrid retrieval.
Empirical Validation: Demonstration of consistent gains across four different embedding models, six functional domains, and varying session complexities on an enterprise benchmark.

Results

Evaluated on WixQA (6,221 enterprise support articles) and an E-Commerce Support corpus:

Session Coverage: The proposed method raises single-query session coverage from 41.4% (Vanilla RAG) to 59.1% at $k=8$ (+17.7% absolute; 95% CI: [14.1, 20.4]).
Retrieval Efficiency: To reach 70% session coverage, the method reduces the average number of retrieval calls from 4.0 (Vanilla) to 2.6, a 34% reduction.
Encoder Robustness: The method improves coverage across all four tested embedding models (MiniLM, BGE, GTE, E5) with absolute gains ranging from +17.1% to +20.6%.
Domain Generalization: Gains are consistent across six distinct functional domains (e.g., Bookings, Apps, Editor), with improvements ranging from +12.7% to +21.5%.
KB Compression: The clustering reduces the effective KB size to 20% of the original while improving coverage.
Precision Trade-off: First-hit precision (Hits@K) decreases slightly (96.0% to 93.0%), which the authors argue is an acceptable trade-off for significantly broader session coverage.
Boundary Conditions: The method shows no gain on HotpotQA, confirming that the improvement relies on persistent knowledge bases with document reuse across sessions, not on the expansion mechanism itself.

Significance and Claims

The paper argues that session-level coverage, not single-query recall, should be the primary metric for enterprise RAG evaluation. The significance of this work lies in:

Bridging the Semantic Gap: It demonstrates that semantic similarity and user information needs are distinct signals. Documents needed together are not always semantically similar, and relying solely on embedding similarity misses critical context.
Efficiency and Cost: By retrieving the right documents in the first call, the method reduces the number of API calls, latency, and user friction required to resolve complex support issues.
Complementarity: The approach is orthogonal to query rewriting or post-retrieval compression. It can be combined with existing query-side optimizations (like multi-query rewriting) or post-retrieval methods (like EDC2-RAG) for end-to-end session optimization.
Practicality: The offline reorganization is computationally efficient (under 20 minutes for 6k documents) and the online overhead is negligible, making it viable for production deployment.

The authors conclude that while their co-occurrence sequences are currently derived from synthetic and embedding-based sources, the method establishes a framework where co-occurrence patterns can serve as a powerful, complementary signal to semantic embeddings for enterprise RAG systems.