I accidentally turned LLM memory into program analysis

Over the past few months, the author has been experimenting with large language model (LLM) agents, particularly for vulnerability research. These models have shown promise in navigating large codebases, explaining unfamiliar subsystems, and assisting in exploring potential attack surfaces. However, the author noticed a recurring issue: as investigations took longer, the models began to lose track of previously established facts.

This led to incorrect suggestions, forgotten assumptions, and continued reasoning from invalid observations. To address this problem, the author started looking into memory systems for LLMs, which typically involve storing old conversations or observations and retrieving relevant pieces when needed. While this works reasonably well, the author felt that the model should maintain what is currently known.

The author then drew a parallel to program analysis, where facts and rules are used to derive additional facts, and the fixed point of derivations can be calculated. Inspired by this, the author wondered if a similar approach could be applied to LLMs during vulnerability research. The author then introduced Datalog, a declarative logic programming language, which describes facts and rules from which new facts can be derived.

By storing facts and rules in a Datalog database, the model would no longer need to reconstruct its entire state over and over again. This approach allows the database to automatically invalidate conclusions when an observation changes, making the process more efficient and accurate. One challenge the author faced was removing facts from the Datalog database, as simply removing a fact could affect multiple derived conclusions.

To handle this, the author had to keep track of how facts were derived and update their support when changes occur. The author also implemented a feature that allows users to ask the database for the provenance of a conclusion, providing the reasoning behind it. This way, if an observation becomes incorrect, the database can remove the affected conclusions automatically.