🤖 AI

Knowledge Graph Representations for LLM-Based Policy Compliance Reasoning

This paper presents an agentic framework that constructs knowledge graphs from AI policy documents to enhance LLM-based reasoning and compliance verification, demonstrating that KG augmentation improves performance across various reasoning tasks and that LLM-discovered schemas can match or exceed formal ontologies.

Original authors: Wilder Baldwin, Sepideh Ghanavati

Published 2026-05-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Wilder Baldwin, Sepideh Ghanavati

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 solve a massive, confusing puzzle where the pieces are rules written in three different languages (the EU AI Act, NIST guidelines, and OWASP standards). You have a very smart assistant (a Large Language Model, or LLM) who has read a lot of books but doesn't have these specific rulebooks memorized perfectly. If you ask the assistant a question, it might guess the answer based on what it "thinks" it knows, but it could get the details wrong or make things up (hallucinate).

This paper introduces a new way to help that smart assistant. Instead of just asking the assistant to remember, the researchers built a digital filing cabinet (a Knowledge Graph) specifically for these AI rules.

Here is how the system works, broken down into simple steps:

1. The "Scan and Sort" Team (Construction)

First, the system takes the long, boring policy documents and breaks them into small, manageable chunks. It then uses an AI agent to read each chunk and pull out the important "facts" (like specific risks, rules, or obligations) and the connections between them.

Think of this like a librarian who reads a 500-page rulebook and creates a set of index cards.

The "Strict" Librarian (AIRO Schema): This librarian follows a very strict, pre-made list of categories (like "Risk," "Threat," "Consequence"). They force every fact into these specific boxes.
The "Creative" Librarian (Open Schema): This librarian doesn't use a pre-made list. Instead, they look at the text and create their own descriptive labels (like "Data Poisoning Risk" or "Transparency Rule") that fit the content best.

2. The "Detective" (Retrieval)

When you ask a question (e.g., "Does this code violate the EU AI Act?"), the system doesn't just ask the assistant to guess. It acts like a detective:

The Direct Path: For simple questions (like "What is a high-risk AI system?"), the detective quickly finds the exact index card and reads the quote.
The Agent Path: For complex questions (like "How do the EU rules and NIST rules work together?"), the detective walks through the filing cabinet, following the connections between cards to gather all the evidence needed to build a complete answer.

3. The "Judge" (Evaluation)

The researchers tested this system with five different AI models (from very powerful ones to smaller, cheaper ones). They asked 42 different types of questions, ranging from simple lookups to complex cross-rule reasoning.

They compared three scenarios:

No Help: The AI answers from its own memory.
Strict Filing Cabinet: The AI uses the "Strict" librarian's cards.
Creative Filing Cabinet: The AI uses the "Creative" librarian's cards.

What They Found (The Results)

The Filing Cabinet Always Helps: Adding the knowledge graph improved the answers for every AI model tested. The AI became more accurate and, crucially, could provide exact quotes from the rules to prove its answer, rather than just guessing.
The "Creative" Librarian Was Surprisingly Good: You might think a strict, pre-made list (the formal ontology) would be better. However, the "Creative" librarian, who made up their own labels on the fly, performed just as well, and sometimes even better, than the strict one. This suggests that letting the AI describe things in its own words can be more effective than forcing it into rigid boxes.
Size Matters (For the Detective): The system worked great for the "big brain" AI models. However, for the smaller, less powerful models, the "Agent Path" (where the detective walks through the graph) actually made things worse. It was like giving a complex map to a child who doesn't know how to read; they got lost. For smaller models, it was better to just give them the direct answer without the complex walking-around step.

The Big Takeaway

The paper proves that if you want an AI to follow complex safety rules, you shouldn't just rely on its memory. You need to give it a structured, searchable map of those rules. Interestingly, you don't need a perfectly pre-planned map; a map that the AI builds itself based on the text works just as well. However, you have to match the complexity of the map-reading task to the intelligence of the AI doing the reading.

1. Problem Statement

As AI features are rapidly integrated into software, regulatory frameworks (e.g., EU AI Act, NIST AI RMF, OWASP Top 10) impose complex, overlapping, and evolving compliance requirements. Practitioners struggle to navigate these cross-referencing obligations manually. While Large Language Models (LLMs) are increasingly used for compliance tasks, they suffer from:

Hallucinations: Inability to reliably cite specific policy clauses.
Outdated Knowledge: Parametric knowledge often lags behind new regulations.
Reasoning Limitations: Difficulty in performing multi-hop reasoning or cross-policy inference.

The paper addresses the need for an automated system that can ground LLM reasoning in authoritative policy text using Knowledge Graphs (KGs) to improve accuracy, traceability, and compliance verification.

2. Methodology

The authors propose an agentic framework that constructs KGs from AI policy documents and uses adaptive retrieval to answer questions. The system operates as a Model-Context-Protocol (MCP) server with a four-stage pipeline:

A. Document Chunking (Scan-and-Review)

Instead of fixed token counts, a chunking agent uses a two-phase process to segment policy text while preserving logical structure:

Scan: A sliding window prompts an LLM to identify natural split boundaries (e.g., end of an Article).
Review: The plan is checked; if chunks are too large (>4,000 chars), the agent proposes further splits.

B. Knowledge Graph Extraction

An extraction agent processes chunks to generate typed entities and relations. The study compares two ontology schemas:

Closed Schema (AIRO): Grounded in the formal AI Risk Ontology (AIRO). It uses a predefined vocabulary (7 entity types, 6 relation types) derived from the EU AI Act and ISO standards.
Open Schema (Emergent): The LLM assigns descriptive, snake-case labels (e.g., technical_risk, regulatory_obligation) without a predefined vocabulary, normalizing them post-extraction.

Cross-Policy Alignment: Entities from different sources (EU, NIST, OWASP) are linked via Corresponds To edges using embedding similarity (cosine $\ge$ 0.70) and string similarity.

C. Adaptive Retrieval & Answer Synthesis

A routing agent classifies questions as "Direct" or "Agent" based on complexity:

Direct Path (Simple Tasks): Embeds the question, retrieves top-5 entities, expands 1-hop neighbors, and synthesizes an answer.
Agent Path (Complex Tasks): Uses a ReAct loop (max 7 steps) with tools for keyword search, semantic search, neighbor expansion, and path finding to traverse the graph for multi-hop reasoning.
Synthesis: The final LLM call receives evidence entities (serialized as relations), raw chunk text, in-context learning (ICL) examples, and task-specific instructions.

D. Evaluation Setup

Datasets: 42 policy QA tasks across three frameworks (EU AI Act, NIST AI RMF, OWASP Top 10).
Task Types: Six categories ranging from T1 (Entity Lookup) to T6 (Cross-Policy Reasoning).
Models: Five LLMs across three tiers (Frontier: gpt-5-mini, gpt-4.1-mini; Mid-range: nemotron:30b, gpt-oss:20b; Small: granite4:micro).
Metrics: Heuristic scoring (word overlap, set-match) and an LLM-as-Judge (gpt-4.1-mini) evaluating accuracy, completeness, and relevance.

3. Key Contributions

Agentic Policy Framework: A complete end-to-end pipeline for constructing and querying policy-derived KGs, exposed as an MCP server for integration.
Ontology Comparison: A rigorous empirical comparison between a formal, standards-derived ontology (AIRO) and an LLM-discovered open schema.
Cross-Policy Reasoning Benchmark: Introduction of Task Type T6, challenging models to map concepts across independently authored frameworks (e.g., linking OWASP risks to EU AI Act requirements).
Model Capability Thresholds: Identification of a threshold where agentic graph traversal benefits larger models but degrades performance for smaller models.

4. Key Results

KG Augmentation Improves Performance: KG-augmented LLMs outperformed the "No-Context" baseline across all five models.
- Judge Score Gains: Improvements ranged from +0.17 to +0.55.
- Best Performer: gpt-oss:20b saw the largest judge improvement (+0.55 with Open KG).
Schema Performance:
- The Open (Emergent) Schema matched or exceeded the formal AIRO schema in judge scores for 3 out of 5 models.
- The open schema's flexibility allowed for richer, more descriptive entity names that aligned better with expected answer vocabularies.
Task-Specific Insights:
- T1/T3 (Lookup/Attributes): Showed the strongest gains (+0.11 to +0.40) because KGs provided verbatim policy quotes that LLMs cannot generate from parametric memory alone.
- T5 (Compliance Checking): KGs transformed answers from hedging ("Partially compliant") to definitive, cited judgments ("Yes, citing Article 50").
- T6 (Cross-Policy): Showed the smallest gains. Success depended heavily on vocabulary overlap between policies; models struggled when terms differed (e.g., "data poisoning" vs. "training data bias").
Model Capability Threshold:
- Large Models: Benefited from the Agent Path (multi-hop traversal), improving complex reasoning.
- Small Models (e.g., granite4:micro): Suffered on complex tasks (T4, T6) when using the agent path, showing performance regression. They performed best with the Direct Path only.

5. Significance and Future Work

Practical Utility: The study demonstrates that grounding LLMs in KGs is essential for regulatory compliance, particularly for providing citable evidence and reducing hallucinations.
Ontology Flexibility: Contrary to the assumption that formal ontologies are superior, an LLM-discovered schema can be equally or more effective for QA tasks due to better semantic alignment with natural language queries.
Strategic Retrieval: The findings suggest that retrieval strategies must be conditioned on model capability. Smaller models should bypass complex agent loops in favor of direct retrieval to avoid degradation.
Future Directions: The authors plan to extend the pipeline to source code analysis (traversing code graphs alongside policy KGs) and explore advanced entity alignment methods (beyond embedding similarity) to solve the cross-policy vocabulary mismatch bottleneck.

Limitations: The evaluation is limited to English-language Western frameworks, uses a small question set (42 items), and relies on automated scoring without human developer user studies.