I accidentally turned LLM memory into program analysis

I accidentally turned LLM memory into program analysis

我无意中将大语言模型(LLM)的记忆变成了程序分析

Over the past few months I have been playing around quite a bit with LLM agents, particularly for vulnerability research. They are becoming surprisingly good at navigating large codebases, explaining unfamiliar subsystems and helping explore potential attack surfaces. 在过去的几个月里,我一直在尝试使用大语言模型(LLM)智能体,特别是在漏洞研究方面。它们在导航大型代码库、解释不熟悉的子系统以及帮助探索潜在攻击面方面表现得越来越出色。

However, once an investigation starts taking a few hours, I kept running into the same problem: the model would slowly lose track of what we had actually established. It might suggest an approach that we had already ruled out, forget that an assumption turned out to be false, or confidently continue reasoning from an observation that was no longer valid. Obviously, telling an LLM that something is wrong does not necessarily mean that it will stop believing all of the things that depended on it :) 然而,一旦调查持续几个小时,我就会不断遇到同一个问题:模型会逐渐忘记我们已经确立的事实。它可能会建议一种我们已经排除的方法,忘记某个假设已被证明是错误的,或者自信地基于一个不再有效的观察结果继续推理。显然,告诉 LLM 某件事是错的,并不一定意味着它会停止相信所有依赖于该错误前提的结论 :)

I initially started looking into memory systems because I wanted to make LLMs more useful for complex vulnerability research and reduce this type of hallucination. There are of course already plenty of solutions for giving LLMs memory. Usually this involves storing old conversations or observations somewhere, embedding them, and then retrieving the most relevant pieces whenever the model needs them again. This works reasonably well, but there was something about it that bothered me. 我最初研究记忆系统是因为我想让 LLM 在复杂的漏洞研究中更有用,并减少这类幻觉。当然,目前已经有很多为 LLM 提供记忆的解决方案。通常,这涉及将旧的对话或观察结果存储在某处,对其进行向量化(embedding),然后在模型需要时检索最相关的片段。这效果还算不错,但总有一些地方让我感到困扰。

During a vulnerability research sesh, I don’t just want the model to remember what we said. I want it to maintain what we currently know. Imagine that during an investigation we establish the following: attacker controls object_a, object_a points to object_b, object_b is a kernel object. From this, we may conclude that the attacker can control a kernel object. A normal memory system could store all of these observations and retrieve them again whenever we ask about the exploitability of the bug. The LLM then figures out the same conclusion. Great! 在漏洞研究过程中,我不希望模型仅仅记住我们说了什么,我希望它能维护我们当前已知的知识。想象一下,在调查过程中我们确立了以下事实:攻击者控制 object_a,object_a 指向 object_b,object_b 是一个内核对象。由此,我们可以得出结论:攻击者可以控制一个内核对象。普通的记忆系统可以存储所有这些观察结果,并在我们询问漏洞的可利用性时将其检索出来。然后 LLM 会得出同样的结论。太棒了!

However, suppose that two hours later we discover in LLDB that object_a does not actually point to object_b, and that our previous assumption was based on a wrong assumption. At that point our memory may contain something like: object_a points to object_b, attacker can control object_b, object_a does not actually point to object_b. Now we retrieve some subset of these memories and hope that the LLM correctly figures out which conclusions are still valid. This started to feel a little familiar to me. This looks like program analysis. 然而,假设两小时后我们在 LLDB 中发现 object_a 实际上并不指向 object_b,并且我们之前的观察是基于错误的假设。此时,我们的记忆中可能包含:object_a 指向 object_b,攻击者可以控制 object_b,object_a 实际上并不指向 object_b。现在我们检索这些记忆的一个子集,并希望 LLM 能正确判断哪些结论仍然有效。这让我感到有些熟悉。这看起来就像程序分析。

A lot of the work I normally do involves program analysis. When analysing a program, we usually have a bunch of facts about the program and some rules that derive additional facts from them. For example, imagine we know: calls(foo, bar), calls(bar, baz). We could define a rule stating that if one function calls another function, which itself can reach a third function, then the first function can reach the third function as well. Eventually we calculate a fixed point containing everything we can derive from the program. 我平时做的很多工作都涉及程序分析。在分析程序时,我们通常拥有一堆关于程序的事实,以及一些从这些事实中推导出额外事实的规则。例如,假设我们知道:calls(foo, bar), calls(bar, baz)。我们可以定义一条规则:如果一个函数调用了另一个函数,而后者又能到达第三个函数,那么第一个函数也能到达第三个函数。最终,我们计算出一个不动点(fixed point),其中包含了我们可以从程序中推导出的所有内容。

More importantly, if one of our input facts changes, there are plenty of techniques for updating only the affected results instead of rerunning everything from scratch. This is also exactly what I wanted from an LLM during vulnerability research. If an observation changes, I don’t want the model to reconstruct the entire investigation from a transcript and hopefully notice all of the consequences. I want the affected conclusions to become invalid automatically. 更重要的是,如果我们的输入事实发生变化,有很多技术可以只更新受影响的结果,而不是从头开始重新运行所有内容。这正是我在漏洞研究中对 LLM 的期望。如果观察结果发生变化,我不希望模型从记录中重建整个调查过程,并寄希望于它能注意到所有的后果。我希望受影响的结论能自动失效。

When looking at the problem from this perspective, I started wondering why we were making the LLM reconstruct its entire state over and over again. What if we just maintained it? And this is how I somehow ended up writing a Datalog engine for LLMs :) 从这个角度看问题时,我开始思考:为什么我们要让 LLM 反复重建其整个状态?如果我们只是维护它会怎样?这就是我最终为 LLM 编写了一个 Datalog 引擎的原因 :)

Datalog

Datalog

Before we continue, it is probably useful to briefly explain what Datalog actually is. Datalog is a declarative logic programming language. Instead of writing instructions describing how something should be calculated, we describe facts and rules from which new facts can be derived. 在继续之前,简要解释一下 Datalog 是什么可能很有用。Datalog 是一种声明式逻辑编程语言。我们不是编写指令来描述如何计算某件事,而是描述事实和规则,并从中推导出新的事实。

For example, we could store the following facts: controls(attacker, object_a). points_to(object_a, object_b). kernel_object(object_b). And then define the following rule: controls_kernel_object(Attacker) :- controls(Attacker, ObjectA), points_to(ObjectA, ObjectB), kernel_object(ObjectB). From our existing facts, the engine can therefore derive: controls_kernel_object(attacker). 例如,我们可以存储以下事实:controls(attacker, object_a). points_to(object_a, object_b). kernel_object(object_b). 然后定义以下规则:controls_kernel_object(Attacker) :- controls(Attacker, ObjectA), points_to(ObjectA, ObjectB), kernel_object(ObjectB). 根据我们现有的事实,引擎因此可以推导出:controls_kernel_object(attacker)。

Nothing particularly exciting yet. However, suppose we later discover that: points_to(object_a, object_b). was incorrect. If controls_kernel_object(attacker) was derived from that fact, we know exactly which conclusion depends on the observation that just changed, and we can automatically invalidate it. This is considerably nicer than putting all of the old information into a prompt and asking an LLM to hopefully notice the same thing. 目前还没什么特别令人兴奋的。然而,假设我们后来发现 points_to(object_a, object_b) 是错误的。如果 controls_kernel_object(attacker) 是从该事实推导出来的,我们就能确切地知道哪些结论依赖于刚刚改变的观察结果,并可以自动使其失效。这比把所有旧信息放入提示词中并祈祷 LLM 能注意到同样的事情要好得多。

Lemmalog

Lemmalog

This eventually turned into Lemmalog. The basic idea is that an LLM should not necessarily be responsible for maintaining its own knowledge. Instead, I split the problem into two parts. The LLM handles the fuzzy part: “LLDB shows that the freed object is later reused as the destination of the write.” | v freed(object_a) reused_as(object_a, write_target). And Lemmalog handles the deterministic part: facts | v rules | v derived facts. 这最终演变成了 Lemmalog。其基本思想是:LLM 不一定非要负责维护自己的知识。相反,我将问题分为两部分。LLM 处理模糊的部分:“LLDB 显示被释放的对象随后被重用为写入目标。” | v freed(object_a) reused_as(object_a, write_target)。而 Lemmalog 处理确定性的部分:事实 | v 规则 | v 推导出的事实。

This means that the LLM is still responsible for understanding natural language, source code, debugger output and all the other messy information that appears during an investigation. LLMs happen to be quite good at this. But once that information has been converted into structured facts, we no longer need the model to repeatedly determine all of its consequences. The database can do that instead. 这意味着 LLM 仍然负责理解自然语言、源代码、调试器输出以及调查过程中出现的所有其他杂乱信息。LLM 在这方面恰好非常擅长。但一旦这些信息被转换为结构化事实,我们就不再需要模型反复确定其所有后果。数据库可以代劳。

Retractions

撤回(Retractions)

One of the first interesting problems I ran into was removing facts. Adding facts to a Datalog database is relatively straightforward: add the new fact and evaluate any rules which may now produce additional results. Removing something is a little more annoying. Take the following example: a. b. c :- a. c :- b. Here c has two separate reasons. 我遇到的第一个有趣的问题是删除事实。向 Datalog 数据库添加事实相对简单:添加新事实并评估现在可能产生额外结果的任何规则。删除某些内容则稍微麻烦一些。以以下示例为例:a. b. c :- a. c :- b. 这里 c 有两个独立的推导原因。