AI・機械学習
LLMのメモリをプログラム解析に誤って転換してしまった
I accidentally turned LLM memory into program analysis (pwning.systems)
要約
LLMエージェントを脆弱性リサーチに活用する中で、長時間にわたる調査でLLMが以前の発見や除外されたアプローチを忘れるという問題に直面しました。この問題を解決するため、LLMのメモリシステムをプログラム解析の考え方に応用し、Datalogのような宣言型ロジック言語を用いて、事実とルールに基づいた推論と、事実変更時の影響範囲の自動更新を実現するシステム「Lemmalog」を開発しました。これにより、LLMは自然言語理解に集中し、推論の維持はLemmalogが行うという分業が可能になりました。
全文翻訳
過去数ヶ月間、LLMエージェント、特に脆弱性リサーチのために多くの時間を費やしてきました。それらは大規模なコードベースをナビゲートし、馴染みのないサブシステムを説明し、潜在的な攻撃サーフェスを探索するのに驚くほど役立つようになっています。しかし、調査が数時間かかり始めると、同じ問題に繰り返し直面しました。モデルは、実際に確立したことを徐々に忘れ去ってしまうのです。すでに除外されたアプローチを提案したり、前提が誤りであったことを忘れたり、もはや有効でない観察から自信を持って推論を続けたりすることがありました。明らかに、LLMに何かが間違っていると伝えても、それが依存していたすべてのことを信じなくなるわけではありません :)
私は当初、LLMを複雑な脆弱性リサーチにより役立つものにし、この種の幻覚を減らしたいと考えてメモリシステムを調べ始めました。LLMにメモリを与えるためのソリューションはすでにたくさんあります。通常、これは古い会話や観察をどこかに保存し、それらを埋め込み、モデルが再び必要とするときに最も関連性の高い部分を取得することを含みます。これはかなりうまく機能しますが、私を悩ませる何かがありました。
脆弱性リサーチのセッション中、モデルに私たちが話したことを覚えておいてほしいだけではありません。現在の私たちの知識を維持してほしいのです。調査中に以下のようなことを確立したと想像してください。
攻撃者はobject_aを制御する
object_aはobject_bを指している
object_bはカーネルオブジェクトである
ここから、攻撃者はカーネルオブジェクトを制御できると結論付けることができます。通常のメモリシステムは、これらの観察をすべて保存し、バグの悪用可能性について尋ねるときにそれらを再び取得できます。LLMは同じ結論を導き出します。素晴らしい!
しかし、2時間後にLLDBでobject_aが実際にはobject_bを指しておらず、以前の観察が誤った前提に基づいていたことを発見したとします。その時点で、私たちのメモリには以下のようなものが含まれているかもしれません。
object_aはobject_bを指している
攻撃者はobject_bを制御できる
object_aは実際にはobject_bを指していない
ここで、これらのメモリのサブセットを取得し、LLMがどの結論がまだ有効であるかを正しく把握することを願います。
これは私には少し馴染みのあるものに感じ始めました。これはプログラム解析に似ています。
私が通常行う作業の多くはプログラム解析を含みます。プログラムを解析するとき、通常、プログラムに関する多くの事実と、それらから追加の事実を導き出すルールがあります。例えば、以下のような事実を知っていると想像してください。
calls(foo, bar)
calls(bar, baz)
ある関数が別の関数を呼び出し、その関数自体が3番目の関数に到達できる場合、最初の関数も3番目の関数に到達できるというルールを定義できます。最終的に、プログラムから導き出せるすべてのものを含む固定点を計算します。さらに重要なのは、入力事実のいずれかが変更された場合、すべてを最初からやり直すのではなく、影響を受ける結果のみを更新するための多くのテクニックがあることです。
これは、脆弱性リサーチ中にLLMに求めていたものと全く同じです。観察が変更された場合、モデルにトランスクリプトから調査全体を再構築させ、すべてに気づくことを期待したくありません。影響を受ける結論が自動的に無効になることを望みます。
この問題を見る視点を変えたとき、なぜ私たちはLLMにその状態全体を何度も再構築させているのだろうかと考え始めました。もし私たちがそれを維持したらどうなるだろうか?そして、これが私がLLMのためにDatalogエンジンを書くことになった経緯です :)
Datalog
先に進む前に、Datalogが実際には何であるかを簡単に説明しておくと役立つでしょう。Datalogは宣言型ロジックプログラミング言語です。何かがどのように計算されるかを説明する命令を書く代わりに、事実と、それらから新しい事実を導き出すことができるルールを記述します。例えば、以下の事実を保存できます。
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).
まだ特にエキサイティングなものではありません。しかし、後で以下を発見したとします。
points_to(object_a, object_b).
が間違っていたとします。もしcontrols_kernel_object(attacker)がその事実から導き出されていた場合、どの結論が変更された観察に依存しているかを正確に把握でき、それを自動的に無効にすることができます。これは、古い情報をすべてプロンプトに入れ、LLMに同じことに気づくように頼むよりもかなり優れています。
Lemmalog
これは最終的にLemmalogになりました。基本的な考え方は、LLMが必ずしも自身の知識を維持する責任を負うべきではないということです。代わりに、問題を2つの部分に分割しました。LLMはあいまいな部分を担当します。「LLDBは、解放されたオブジェクトが後で書き込みの宛先として再利用されることを示しています。」
v freed(object_a)
reused_as(object_a, write_target)
そしてLemmalogは決定論的な部分を担当します。
事実
v
ルール
v
導出された事実
これは、LLMが依然として自然言語、ソースコード、デバッガ出力、および調査中に現れる他のすべての厄介な情報を理解する責任を負うことを意味します。LLMはこれに非常に長けています。しかし、その情報が構造化された事実に変換されると、モデルにその結果を繰り返し決定させる必要はなくなります。データベースが代わりにそれを行うことができます。
撤回
私が遭遇した最初の興味深い問題の1つは、事実の削除でした。Datalogデータベースに事実を追加するのは比較的簡単です。新しい事実を追加し、それによって追加の結果を生み出す可能性のあるすべてのルールを評価します。何かを削除するのは少し面倒です。次の例を考えてみましょう。
a.
b.
c :- a.
c :- b.
ここでcには真であるための2つの別々の理由があります。aを削除した場合、bがまだ別の導出を提供しているため、cを単純に削除することはできません。しかし、aとbの両方を削除すると、cも消えるはずです。
これは脆弱性リサーチ中に非常に重要であることが判明しました。なぜなら、結論は複数の観察によって裏付けられる可能性があるからです。例えば、
candidate_3_is_exploitable
は、特定のエクスプロイトプリミティブが機能しないことが判明した場合でも、同じ結果への別の独立したパスがあるため、真のままである可能性があります。したがって、Lemmalogは事実がどのように導出されたかを追跡し、何かが変更されたときにそれらのサポートを更新する必要があります。
都合の良いことに、これはもう1つの有用なプロパティも提供します。なぜ何かが真であるかを尋ねることができます。
なぜ?
AIエージェントを数時間実行して何かを調査し、最終的に結論に至ったと想像してください。
candidate_3_is_exploitable
それは良いですが、私もその理由を知りたいです。Lemmalogはすでに導出された事実の依存関係を追跡しているため、結論の出所を尋ねることができます。例えば、概念的には次のようになります。
candidate_3_is_exploitable
| +-- attacker_controls_pointer
| | | +-- observation_41
| +-- pointer_reaches_target
| +-- observation_57
+-- rule_12
もしobservation_41が後に誤りであることが判明した場合、この結論はもはや有効ではない可能性があることを知っており、データベースもそれを知っているため、影響を受ける結論を削除できます。これは元々、増分評価を正しく機能させるために主に必要でしたが、AIエージェントにそれがなぜそう信じているのかを尋ねることができることが、非常に役立つことも判明しました。
これはまた、LLM支援リサーチで遭遇した、より厄介な失敗モードの1つにも対処しています。モデルが自信を持って「このポインタは攻撃者によって制御されているとすでに確立しています」と言うことがありますが、実際にはそうではありません。結論がLemmalogに存在する場合、どこから来たのかを尋ねることができます。それを支持する出所がない場合、それは主要なものではありません。