W-RAG: Source-Aware Retrieval for Enterprise Document Generation from Heterogeneous Knowledge Bases

Technical Summary: W-RAG – Source-Aware Retrieval for Enterprise Document Generation

Problem Statement
Retrieval-Augmented Generation (RAG) has improved the factual grounding of Large Language Models (LLMs) by incorporating external knowledge. However, existing RAG pipelines rely on a global ranking mechanism that assumes evidence from multiple heterogeneous repositories can be evaluated using a single similarity function. This assumption holds for open-domain retrieval but fails in enterprise document generation. In corporate settings, knowledge bases (KBs) such as policies, regulations, technical specifications, and departmental guidelines serve distinct roles. Global ranking tends to favor larger, denser, or lexically dominant sources, suppressing smaller but critical repositories. This leads to unbalanced context windows where the generated document lacks necessary evidence from specific domains, resulting in incomplete drafts that fail to meet complex enterprise requirements.

Methodology: W-RAG Framework
The authors propose W-RAG, a source-aware retrieval framework designed to regulate evidence composition across heterogeneous KBs. The system operates in three distinct stages:

1. Ontology-Guided Candidate Retrieval:
   Instead of relying solely on semantic similarity, W-RAG uses a domain ontology to extract structured topics from the input query. It retrieves candidate text chunks from each KB using a hybrid scoring function:
   $$R(c_{jk}) = f_{sim}(q, c_{jk}) + \lambda f_{ont}(T, c_{jk})$$
   where $f_{sim}$ measures semantic similarity (cosine similarity of embeddings) and $f_{ont}$ measures alignment with the extracted ontology topics. This ensures candidates are not only lexically relevant but thematically aligned with the target document's requirements.

2. Source Weight Estimation:
   Rather than merging all candidates into a single global rank, W-RAG estimates a specific weight ($w_j$) for each KB. This is computed in real-time by comparing the input topic set against the KB's topic index. The score combines:

   • Topical Overlap: The primary driver, measuring how well the KB matches input topics.
   • Document Availability: A weak prior favoring KBs with broader coverage.
   • Category Bonus: A reward for coarse domain alignment.
     These scores are normalized to determine the percentage of the total context budget ($T_{total}$) each KB should contribute.

3. Weighted Context Construction:
   The estimated weights are converted into token budgets ($T_j$) for each KB. The system selects the top-ranked chunks from each locally ranked candidate set until the specific budget for that KB is exhausted. These chunks are merged into a Single Context Window (SCW), ensuring the final prompt reflects a balanced composition of sources rather than the dominance of a single repository.

Key Contributions

1. Problem Formulation: The paper identifies a critical failure mode in existing RAG systems: global ranking across heterogeneous KBs fails to preserve the source composition necessary for complete enterprise document generation.
2. Methodological Innovation: W-RAG introduces a framework combining ontology-guided retrieval, local ranking within each KB, and source-level weighting to regulate evidence composition.
3. New Dataset and Benchmark: The authors introduce a novel dataset for retrieval-grounded enterprise document generation. It includes 100 requirement documents across four types (Policies, Product Launch, M&A Due Diligence, Academic Programs) and four domains (AI/Digital Media, Healthcare, Climate/Sustainability, Finance). The dataset is paired with four heterogeneous KBs (Corporate, Scientific, Regulatory, Market/News) to test multi-source evidence synthesis.

Experimental Results
The authors evaluated W-RAG against standard Vanilla RAG and an Ontology-Guided RAG (OG-RAG) baseline using GPT-4.1 as the generator and Claude Sonnet 4.6 as an LLM-as-a-judge. Key findings include:

• Performance Gains: W-RAG significantly outperformed baselines. Relative to OG-RAG, it improved Requirement Satisfaction by 58.1% and Context Relevancy by 39.1%. Compared to Vanilla RAG, the improvements were 109.2% and 96.4%, respectively.
• Source Fidelity: The study measured "Realized Knowledge Base Fidelity" to determine if the generated document reflected the intended distribution of sources. W-RAG demonstrated strong alignment between prescribed retrieval weights and the actual source attribution in the generated text, with Jensen-Shannon Divergence scores indicating moderate-to-strong agreement.
• Retrieval vs. Generation: The results indicate that high retrieval quality alone does not guarantee high document quality; the composition of evidence from multiple sources is critical for satisfying enterprise requirements.

Significance and Claims
The paper claims that standard RAG assumptions are poorly matched to the task of enterprise document generation. By explicitly managing how retrieved context is distributed across knowledge bases, W-RAG enables more faithful and complete document generation. The authors posit that in multi-source enterprise settings, the "evidence composition at the document level" is as important as "ranking quality at the chunk level." The work suggests that moving beyond global similarity ranking to source-aware allocation is necessary for reliable synthesis of policies, contracts, and technical documentation.

Limitations
The authors acknowledge that their evaluation relies on a relatively small, synthetic dataset which may not fully capture the scale and noise of real-world enterprise environments. Additionally, the evaluation metrics involve partial subjectivity, and the approach has not yet been validated in open-domain or less structured retrieval environments.