Rewriting the search plan mid-query gives PAI-2 an 18% lift

Large language models have an odd superpower: they can speak so confidently that it sometimes seems they actually know. The trouble is that confidence and knowledge are not the same thing.

When memory alone is not enough: teaching an LLM to look for facts in a graph instead of a text

Large language models have an odd superpower: they can speak so confidently that it sometimes seems they actually know. The trouble is that confidence and knowledge are not the same thing. That gap opens widest where a question takes several reasoning steps, fact-checking and careful handling of context. And that is exactly where most retrieval-augmented generation systems break: the model answers beautifully, but not always to the point.

The authors of the paper on PersonalAI 2.0, or PAI-2, offer a fairly practical response to this. They take an LLM, attach external memory in the form of a knowledge graph and — the important part — make the system build a search plan instead of simply pulling the nearest chunks, revise that plan as it goes, and walk the graph iteratively. What you get is not "retrieve a piece of text, generate an answer" but a small investigation: what exactly needs to be established, which entities matter, where to go next, whether there is enough evidence yet or the plan should be rewritten.

This matters not because it promises some magical universal memory for AI, but because it shows that the step after ordinary RAG probably runs in the direction of structured memory and controlled retrieval.

What the authors actually propose

At the center of the paper is the PAI-2 architecture, a question-answering system that searches a memory graph. The graph holds knowledge extracted from documents: entities, statements, relations and episodic records. An LLM works on top of that graph, not as a know-it-all but as a dispatcher for the search.

The core idea is simple: a hard question is rarely worth answering with a single query. Instead, PAI-2 breaks the task into stages:

1. cleans up and sharpens the question where needed;

2. determines whether it is really several independent sub-questions;

3. builds an initial search plan;

4. extracts named entities from the current step;

5. matches them to nodes in the graph;

6. generates auxiliary probe queries;

7. traverses the graph outward from those anchor points;

8. filters for the relevant triples;

9. summarizes what it has found;

10. checks whether that is enough to answer;

11. and if not, refines the plan and keeps going.

So the system does not just pull out "the most similar thing." It runs the search piece by piece, with a checkpoint in between: are we even heading the right way?

That is the main difference from most GraphRAG variants, where retrieval tends to stay fairly static. In PAI-2 it is adaptive. If halfway along the route it turns out the original plan was too vague or missed an important entity, the system can rewrite it.

Why this matters at all

Ordinary retrieval-augmented generation has a fundamental weakness: it works well when the answer sits right next to the query, in one or two chunks of text. The moment you have to connect several documents, reconstruct a chain of entities or pin down an intermediate fact, quality starts to sag. You get hallucinations, fragmentary answers, or a confident "no data available" when the data is in fact sitting in the store.

A knowledge graph is convenient as external memory in exactly this sense: it stores the links between objects explicitly. But a graph alone is not enough either. Search the nodes without a strategy and the system either drowns in noise or walks straight past the route it needed.

What makes PAI-2 interesting is that it combines two worlds:

For personal assistants, enterprise agents, search over internal knowledge bases and generally any application where a factual error is expensive, this is a very down-to-earth and useful idea.

How the search works: not one query but a series of guesses

Methodologically the paper reads like engineering work with a great deal of practical detail. The authors clearly were not trying to invent "a new theory of intelligence" — they are building a system that can be broken into legible modules and tested piece by piece.

The mechanics of the auxiliary queries are especially curious. Once entities have been pulled out of a plan step and matched to graph nodes, the system writes several sharper phrasings of the question, each resting on its own combination of the nodes it found. That helps it probe different parts of the graph even when what it needs to know about an object is not stored in one place but scattered across several subgraphs.

The triples it retrieves are then not simply handed to the model as a list. Answers are generated for each auxiliary query first, and only then aggregated into a single piece of knowledge for the current plan step. After that the LLM decides: answer now, or keep searching.

One detail worth noting: the authors show separately that the plan refinement mechanism is not decoration but a genuinely working component. Switch it off and quality drops noticeably. The logic is clear enough — plenty of questions cannot be answered correctly from the first rough route. You need the ability to change course mid-way.

What the experiments showed

PAI-2 was tested on six well-known question-answering datasets: Natural Questions, TriviaQA, HotpotQA, 2WikiMultihopQA, MuSiQue and DiaASQ. The opponents were LightRAG, RAPTOR, HippoRAG 2 and the system's own previous version, PAI-1.

The headline metric is a score from an LLM judge. That is a fairly standard approach by now in tasks where plain string matching says little about whether an answer is right. The authors separately checked how reliable that scoring is against human annotation and got a high correlation, so the evaluation method here looks conscientious enough.

In short, the results look like this:

Across benchmarks the picture is not quite even. The system looks best on TriviaQA, 2WikiMultihopQA and MuSiQue — precisely where multi-hop retrieval matters most. On Natural Questions, by contrast, PAI-2 trails HippoRAG 2 by a visible margin. The authors are honest about the likely reasons: questions there are often poorly normalized, written in lowercase, and the expected answer type is not always obvious. For a system that likes explicit entities and tidy planning, that is an uncomfortable regime.

The strongest ablation: planning really does work

The most convincing part of the paper is probably not the final table of numbers but the ablation on plan refinement. The authors take a question about the directors of two films to show why a static plan is not enough.

The system initially produces steps along the lines of "who directed the first film," "who directed the second film," "what country is the first film's director from," "what country is the second film's director from." The problem is that the last two steps are too vague: without substituting the actual names, the search easily misses.

With plan refinement switched on, the system substitutes the intermediate facts it has already found: instead of an abstract "director of the second film" there is a specific person. And only then does the graph hand back what was needed.

This sounds almost banal, but small things like this are often where the difference lies between "an architecture on a slide" and a system that actually works. Planning here is not a pretty word but a mechanism for sharpening the query as knowledge accumulates.

The other thing that matters: the quality of the memory itself

Systems like this have two levels of quality. The first is how well the agent searches. The second is how sound the memory it walks through actually is. If the graph is built badly, no amount of clever retrieval will save it.

The authors evaluate their graph construction pipeline separately on the MINE-1 benchmark, which measures how much of the factual information in the source text survives into the graph. And here PAI-2 does very well: 89% information retention, noticeably higher than Wikontic and KGGen in the comparison they report.

This may matter just as much as the question-answering scores. External memory for an LLM dies in two ways: when it is hard to search, and when things are written into it badly. Judging by the paper, PAI-2 is trying to close both problems at once.

Limitations the authors do not keep quiet about

Part of what makes the work strong is that the authors do not pretend they have solved everything.

First, the system is slow. By their own measurements PAI-2 is roughly twice as slow as PAI-1: about 1.5 minutes per question on average. Tolerable for research, already a challenge for production use. The bottlenecks are predictable — LLM inference, vector search and graph traversal.

Second, in its current form the graph suffers from duplicates at the level of meaning. If the same thought is phrased two different ways, several near-identical subgraphs can appear. Hence the need for several auxiliary queries per step.

Third, ambiguous entities, time and statement types remain a weak spot. The authors themselves suggest where to go next: label facts, opinions and predictions better, model time intervals explicitly, fix a vocabulary of predicates, and add semantic deduplication.

The takeaway

PAI-2 is not another paper in the "we bolted a graph onto an LLM" genre. Its central claim is subtler: a reliable answer takes more than having memory — you have to be able to think your way through it along a route. Break the question apart, build a plan, pin down the entities, launch several search trajectories, check whether the facts gathered so far are sufficient, and rewrite the strategy when they are not.

Which is why the paper is interesting not only to people building GraphRAG systems but to anyone working with agents, enterprise search and external memory for LLMs. It shows that the next level of quality comes not from "an even bigger model" but from organizing retrieval more rigorously.

The main conclusion is this: structured memory without planning is too dumb, and an LLM without external memory is too sure of itself. PAI-2 tries to join the best of both worlds — and, judging by the results, does it rather convincingly.