Reducing LLM Hallucinations with Knowledge Graphs and Coarse-to-Fine Highlighting

By Newton Chen and Wayne Chu as part of the Stanford CS224W course project.

Retrieval-Augmented Generation (RAG) is great in theory: you ask an LLM a question, it pulls in relevant documents, and then answers using that evidence. In principle, it should eliminate hallucinations.

But there’s a catch. When the retrieved context becomes too long or cluttered, LLMs can get “lost in the middle.” They have trouble separating the key facts from the noise, which puts us right back where we started: hallucinations.

Imagine giving a student a 50-page textbook chapter five minutes before an exam and saying, “The answer is somewhere in here.” They’ll probably panic and guess. That’s essentially how an LLM behaves when flooded with overly long retrieval results.

Our goal: Hand the LLM a metaphorical highlighter. Instead of treating all context as equally important, we want to pre-process the text and emphasize the crucial parts.

To achieve this, we implemented COFT (Coarse-to-Fine Highlighting), a framework introduced by Lv et al. [1]. COFT leverages the structure of Knowledge Graphs (KGs) to guide the LLM’s attention, ultimately improving factual accuracy.

How COFT Works

COFT functions like an intelligent study assistant. Instead of summarizing (which risks losing details), it highlights.

The process runs in three distinct phases:

1. Recaller: Extracts key concepts from the user’s query and expands them using a Knowledge Graph to uncover related concepts.
2. Scorer: Evaluates which of these concepts appear in the retrieved document and assesses their importance.
3. Selector: Chooses the most relevant concepts and highlights them in the text at varying granularities (words, sentences, or paragraphs).

The Recaller — Grounding with Knowledge Graphs

Why bring Graph ML into the picture? Why not simply match query keywords with strings?

Because language is varied. If a user asks about “Cupertino,” a basic string match won’t surface related concepts like “Apple headquarters” or “Silicon Valley.”

A Knowledge Graph fills this gap by linking entities through meaningful relationships. By treating the query terms as nodes and exploring their one-hop neighbors, we semantically expand the search space beyond literal wording.

For our project, we rely on a large-scale KG (e.g., Wikidata5M). With millions of nodes and edges, efficiency is critical. We don’t need heavy Graph Neural Networks — just extremely fast adjacency lookups.

To achieve this, we store the graph in Compressed Sparse Row (CSR) format using PyTorch tensors, a common approach in Graph ML that enables memory-efficient storage and rapid neighbor retrieval.

# A snippet from the COFT __init__ method
# We assume triplets is a list of (head, relation, tail) tuples
src_list: List[int] = []
dst_list: List[int] = []

# 1. Map string IDs to integer indices
for h, _, t in triplets:
    if h in self.node_to_id and t in self.node_to_id:
        u, v = self.node_to_id[h], self.node_to_id[t]
        src_list.extend([u, v])
        dst_list.extend([v, u])

# 2. Convert to PyTorch Tensors for efficient processing
if len(src_list) > 0:
    src_tensor = torch.tensor(src_list, dtype=torch.long)
    dst_tensor = torch.tensor(dst_list, dtype=torch.long)

    # 3. Create Compressed Sparse Row (CSR) structure
    # Sort by source index to align edges
    perm = torch.argsort(src_tensor)
    src_sorted = src_tensor[perm]
    self.dst_sorted = dst_tensor[perm]  # This holds the neighbors

    # Calculate pointers (indptr) where each source node's edges begin/end
    self.indptr = torch.zeros(self.num_nodes + 1, dtype=torch.long)
    unique_src, counts = torch.unique(src_sorted, return_counts=True)
    self.indptr[unique_src + 1] = counts
    self.indptr = torch.cumsum(self.indptr, dim=0)

With this structure, finding 1-hop neighbors is a highly efficient tensor slicing operation:

# A snippet from the _recall_candidates method:
# Given a starting node index 'idx'
start = int(self.indptr[idx].item())
end = int(self.indptr[idx + 1].item())

# Instant retrieval of all neighbors
if end > start:
    neighbor_indices.update(self.dst_sorted[start:end].tolist())

The Recaller outputs a comprehensive set of surface forms (text representations) for both the original query entities and their graph neighbors.

The Scorer — Measuring Importance

At this stage, we’ve gathered a large pool of candidate entities from the Knowledge Graph. The question now is: Which of these actually matter for the document we retrieved?

COFT answers this using a hybrid scoring strategy that blends statistical signals with semantic understanding:

• TF-ISF (Term Frequency — Inverse Sentence Frequency): This boosts entities that appear often within specific sentences but rarely across the entire document. It helps surface information that is locally important rather than globally common.
• Self-Information (via LLM): A semantic metric. We ask an LLM (e.g., GPT-2): Given the query, how “surprising” or informative is this entity? In information theory, more surprise means more information.

To compute this, we evaluate the negative log-likelihood (NLL) of the LLM generating the entity when conditioned on the query. A lower probability (higher NLL) indicates that the entity carries more informational value.

def _get_entity_relevance(self, query: str, entity: str) -> float:
    """
    Calculates Self-Information I(e|q) in bits.
    I(e) = -log2 P(e|q) = Loss(nats) / ln(2)
    """
    # We ask LLM to calculate the loss of generating the 'answer' (entity)
    # given the 'question' (query).
    loss_tensor = self.llm(question=[query], answer=[entity])

    val = loss_tensor.item() if hasattr(loss_tensor, "item") else float(loss_tensor)

    return val / math.log(2)

The Selector — Coarse-to-Fine Highlighting

We have scored entities. We can’t highlight everything, or the highlight becomes useless.

The Selector applies a dynamic threshold to decide how many entities to keep. This threshold depends on:

• Document Length: Longer documents require stricter filtering so the model isn’t flooded with highlights.
• Model Confidence (Global Loss): If the LLM appears uncertain about the document as a whole (high global loss), we keep more entities to give it stronger guidance.

After selecting the top concepts, COFT highlights them at varying levels of granularity, since different tasks benefit from different scopes of focus:

• Word Level: Fine-grained precision.
• Sentence Level: Balanced context.
• Paragraph Level: Broader topic-level emphasis.

The implementation involves standard regex and string manipulation once the top entities are selected.

# A snippet from the  _select_and_format method:
if granularity == Granularity.SENTENCE:
    # Split reference into sentences
    sentences = re.split(r"(?<=[.!?])\s+", reference)
    new_sents = []
    for s in sentences:
        # If a selected crucial entity exists in this sentence...
        if any(e in s for e in selected_entities):
            # ...highlight the whole sentence with markdown bolding
            new_sents.append(f"**{s}**")
        else:
            new_sents.append(s)
    return " ".join(new_sents)

# Similar logic applies for WORD and PARAGRAPH granularity

Results

Here’s what COFT looks like in practice.

Query: “How does deforestation affect global climate patterns?”

Raw Context: …Forests play a critical role in regulating the Earth’s temperature. When trees are cut down, the carbon stored in them is released into the atmosphere as carbon dioxide. This process has cascading effects beyond carbon emissions. This significantly increases the greenhouse effect, trapping heat and leading to global warming. Additionally, the loss of canopy cover disrupts the water cycle, reducing local rainfall and increasing the risk of droughts. Without trees to anchor the soil, heavy rains can cause severe erosion and landslides…

When this context is processed with COFT at sentence-level granularity, the system draws on the Knowledge Graph to identify high-value concepts such as carbon dioxide, greenhouse effect, global warming, water cycle, and droughts, scoring them as most relevant to the query.

COFT Highlighted Context: …Forests play a critical role in regulating the Earth’s temperature. *When trees are cut down, the carbon stored in them is released into the atmosphere as carbon dioxide. This process has cascading effects beyond carbon emissions. This significantly increases the greenhouse effect, trapping heat and leading to global warming. Additionally, the loss of canopy cover disrupts the water cycle, reducing local rainfall and increasing the risk of droughts.* Without trees to anchor the soil, heavy rains can cause severe erosion and landslides…

Feeding this highlighted version into the generator LLM directs its focus toward the core evidence, reducing the likelihood of overlooking important details or producing hallucinated explanations.

Conclusion

Implementing COFT showcased the strength of combining structured knowledge (via Knowledge Graphs) with unstructured text generation from LLMs. Despite their capabilities, modern LLMs are still vulnerable to noisy or overly long contexts.

By using Graph ML techniques, specifically efficient graph representation and traversal with PyTorch, to ground the retrieval process, we can provide these models with the highlighter pen they need to navigate long contexts accurately.

This approach represents a meaningful step toward more reliable, factual AI systems.

References

[1] Lv, Q., Wang, J., Chen, H., Li, B., Zhang, Y., & Wu, F. (2024). Coarse-to-Fine Highlighting: Reducing Knowledge Hallucination in Large Language Models. In Proceedings of the 41st International Conference on Machine Learning, ICML 2024, Vienna, Austria, pages 33594–33623.