Lemmalog: Giải pháp biến trí nhớ LLM thành hệ thống phân tích chương trình

Trong vài tháng qua, tôi đã thử nghiệm khá nhiều với các LLM agent, đặc biệt là trong lĩnh vực nghiên cứu lỗ hổng bảo mật.

Chúng đang trở nên cực kỳ hiệu quả trong việc điều hướng các codebase lớn, giải thích các hệ thống con xa lạ và hỗ trợ khám phá các bề mặt tấn công tiềm năng. Tuy nhiên, khi một cuộc điều tra kéo dài vài giờ, tôi liên tục gặp phải cùng một vấn đề: mô hình dần mất dấu những gì chúng ta đã thực sự xác lập.

Nó có thể đề xuất một hướng tiếp cận mà chúng ta đã loại bỏ từ trước, quên mất rằng một giả định đã được chứng minh là sai, hoặc tự tin tiếp tục suy luận từ một quan sát không còn giá trị. Rõ ràng, việc nói với một LLM rằng điều gì đó là sai không có nghĩa là nó sẽ ngừng tin vào tất cả những thứ phụ thuộc vào điều đó:)

Ban đầu, tôi bắt đầu tìm hiểu về các hệ thống bộ nhớ vì muốn làm cho các LLM hữu ích hơn trong việc nghiên cứu lỗ hổng bảo mật phức tạp và giảm thiểu kiểu ảo giác này.

Tất nhiên, đã có rất nhiều giải pháp để cung cấp bộ nhớ cho LLM. Thông thường, việc này bao gồm lưu trữ các cuộc hội thoại hoặc quan sát cũ ở đâu đó, nhúng chúng (embedding), rồi truy xuất những phần liên quan nhất bất cứ khi nào mô hình cần.

Cách này hoạt động khá tốt, nhưng có điều gì đó khiến tôi băn khoăn.

Trong một phiên nghiên cứu lỗ hổng bảo mật, tôi không chỉ muốn mô hình ghi nhớ những gì chúng ta đã nói.

Tôi muốn nó duy trì những gì chúng ta hiện đang biết.

Hãy tưởng tượng rằng trong quá trình điều tra, chúng ta xác lập được những điều sau:

Từ đây, chúng ta có thể kết luận rằng kẻ tấn công có thể kiểm soát một kernel object.

Một hệ thống bộ nhớ thông thường có thể lưu trữ tất cả các quan sát này và truy xuất lại bất cứ khi nào chúng ta hỏi về khả năng khai thác của lỗi. Sau đó, LLM sẽ tự đưa ra kết luận tương tự.

Tuyệt vời!

Tuy nhiên, giả sử hai giờ sau, chúng ta phát hiện trong LLDB rằng object_a thực tế không trỏ đến object_b, và quan sát trước đó của chúng ta dựa trên một giả định sai lầm.

Tại thời điểm đó, bộ nhớ của chúng ta có thể chứa những thứ như:

Bây giờ, chúng ta truy xuất một tập hợp con các bộ nhớ này và hy vọng rằng LLM sẽ xác định đúng những kết luận nào vẫn còn hiệu lực.

Điều này bắt đầu trở nên quen thuộc với tôi.

Nó trông giống như phân tích chương trình (program analysis).

Rất nhiều công việc tôi thường làm liên quan đến phân tích chương trình.

Khi phân tích một chương trình, chúng ta thường có một loạt các sự kiện (facts) về chương trình đó và một số quy tắc (rules) để suy ra các sự kiện bổ sung từ chúng.

Ví dụ, hãy tưởng tượng chúng ta biết:

Chúng ta có thể định nghĩa một quy tắc rằng nếu một hàm gọi một hàm khác, mà bản thân hàm đó có thể dẫn đến hàm thứ ba, thì hàm đầu tiên cũng có thể dẫn đến hàm thứ ba.

Cuối cùng, chúng ta tính toán một điểm cố định (fixed point) chứa mọi thứ có thể suy ra từ chương trình. Quan trọng hơn, nếu một trong các sự kiện đầu vào thay đổi, có rất nhiều kỹ thuật để chỉ cập nhật các kết quả bị ảnh hưởng thay vì chạy lại mọi thứ từ đầu.

Đây cũng chính xác là điều tôi muốn từ một LLM trong quá trình nghiên cứu lỗ hổng bảo mật.

Nếu một quan sát thay đổi, tôi không muốn mô hình phải tái cấu trúc toàn bộ cuộc điều tra từ bản ghi và hy vọng nó nhận ra tất cả các hệ quả. Tôi muốn các kết luận bị ảnh hưởng tự động trở nên không hợp lệ.

Khi nhìn vấn đề từ góc độ này, tôi bắt đầu tự hỏi tại sao chúng ta lại bắt LLM phải tái cấu trúc toàn bộ trạng thái của nó lặp đi lặp lại.

Nếu chúng ta chỉ duy trì trạng thái đó thì sao?

Và đó là cách tôi viết một Datalog engine cho các LLM:)

Datalog

Trước khi tiếp tục, có lẽ nên giải thích ngắn gọn Datalog thực sự là gì.

Datalog là một ngôn ngữ lập trình logic khai báo. Thay vì viết các chỉ dẫn mô tả cách tính toán một thứ gì đó, chúng ta mô tả các sự kiện và quy tắc để từ đó có thể suy ra các sự kiện mới.

Ví dụ, chúng ta có thể lưu trữ các sự kiện sau:

Và sau đó định nghĩa quy tắc sau:

Từ các sự kiện hiện có, engine có thể suy ra:

Chưa có gì quá thú vị.

Tuy nhiên, giả sử sau đó chúng ta phát hiện ra rằng:

là không chính xác.

Nếu controls_kernel_object(attacker) được suy ra từ sự kiện đó, chúng ta biết chính xác kết luận nào phụ thuộc vào quan sát vừa thay đổi, và chúng ta có thể tự động vô hiệu hóa nó.

Điều này tốt hơn đáng kể so với việc đưa tất cả thông tin cũ vào một prompt và cầu nguyện rằng LLM sẽ nhận ra điều tương tự.

Lemmalog