Intelligent Knowledge Base Search Tool using Large Language Model and Graph Neural Network

ABSTRACT Within many organizations, a vast number of communications, memos, reports and documents have been accumulated in internal servers. Efficiently discovering relevant entries can reduce time spent addressing organizational needs such as personnel skills matching or anomaly resolution. However, per organization, information retrieval on these disparate data types can be challenging, as systems must be designed for their domain while accounting for unstructured and inconsistent datasets. Traditional querying via search terms often requires relevancy tuning by subject matter experts which makes it difficult to build retrieval systems. We argue that development of retrieval systems can be simplified and enhanced by embedding data with Large Language Models (LLMs), organizing information in a Knowledge Graph (KG) structure, and further encoding their relational features through a Graph Neural Network (GNN). One of the major challenges of using GNNs for information retrieval is optimizing negative edge selection. Training GNNs requires a balanced ratio between positive and negative edges however the space of negative edges is exponentially larger than positive edges. In this work, we extend the LLM-GNN hybrid architecture by applying ensemble voting on a set of trained LLM-GNNs. Preliminary results have shown modest improvement on our personnel-document matching tasks. This work contributes to a developmental effort that aims to help engineers and scientists find new research opportunities, learn from past mistakes, and quickly address future needs. Keywords: Neural document retrieval, Search tool, Knowledge graph, Large Language Model, LLM, graph neural network, GNN, GPT, BERT 1. INTRODUCTION Engineering design/test documents and historical records contain important details of lessons learned that could be relevant to upcoming design/test decisions to avoid past mistakes. However, retrieving and extracting the correct information typically requires subject matter experts (SME) who can construct effective queries against thousands of documents. In the absence of SMEs, junior engineers require robust communication channels or perform a more exhaustive search effort, often using keywords that produce irrelevant results. As a result, engineering tasks such as anomaly resolution take longer than necessary to complete, despite potential solutions existing within the organization’s knowledge base. Another challenge is the inconsistency across documents. While the documents we work with are structured, some fields are often unfilled. Of the fields that are filled, authors complete them differently, potentially leading to confusion when attempting to reference them. This task naturally invites neural information retrieval.1 Large Language Models (LLM) have been demonstrated to understand language and can extract core concepts in text irrespective of the writing style. Furthermore, Graph Neural Networks (GNN) can encode relationships between data, making it a beneficial backend for recommendation systems.2 We designed a novel intelligent search tool centered around a LLM-GNN hybrid architecture that can accept documents or loosely constructed queries to match information from a knowledge base based on anomaly topics and semantic textual meaning. We employ a minimal required effort training approach which notably surpasses the performance of supervised state-of-the-art models. Our knowledge base design leverages organizational metadata, allowing us to employ training strategies without labor intensive data labeling. Moreover, our architecture is lightweight and can run completely offline, rendering it ideal for tasks that would benefit from sensitive enterprise data. We demonstrate that downstream tasks can be effectively structured around simple edge prediction inference. Consequently, this makes our architecture easy to train and versatile across a variety of tasks.