Technical Summary: Efficient and Sound Probabilistic Verification for AI Agents

Problem Statement

The deployment of Large Language Model (LLM) based AI agents in complex digital environments introduces critical security risks, particularly when agents possess programmatic access to system interfaces like terminals and file systems. While runtime monitoring via reference monitors offers a solution for enforcing security policies, existing approaches are limited to deterministic policies. In practical scenarios, agents often rely on upstream components (e.g., PII detectors, redaction tools, or classifiers) that produce probabilistic outputs with inherent failure rates.

Current deterministic monitors handle this ambiguity by applying local thresholds to convert probabilistic outputs into binary decisions. This discretization creates a fragile trade-off between security and utility:

• High thresholds risk under-blocking harmful actions (false negatives).
• Low thresholds risk over-blocking safe actions (false positives).
• Independence Assumptions: Standard probabilistic inference techniques (e.g., Weighted Model Counting) often assume that base facts are mutually independent. In agentic settings, this assumption is frequently invalid; for instance, a redaction tool failing on one file may be perfectly correlated with its failure on another due to shared content or tool behavior. Relying on independence can lead to severe underestimation of risk during multi-step tool chaining.

The paper addresses the need for a verification framework that preserves continuous probabilistic context, handles arbitrary correlations between predicates without assuming independence, and provides rigorous, sound upper bounds on policy violation probabilities in real-time.

Methodology

The proposed framework models agent execution as a distributionally robust optimization problem over an execution trajectory. The methodology consists of four main stages:

1. Probabilistic Datalog Modeling

The agent's execution trace is compiled into a Datalog program where:

• Base Facts ($F$): Represent initial environmental observations (e.g., "file A contains PII") parameterized by marginal probability bounds $[\ell_f, u_f]$ derived from empirical tool failure rates.
• Rules ($R$): Encode state transition logic (e.g., data propagation, merging, declassification) as intensional rules.
• Derivation Graph: The execution trace is translated into a Directed Acyclic Graph (DAG) where nodes represent facts and edges represent rule applications. The authors prove that any Datalog policy execution can be rewritten such that every derived fact depends on at most two parent facts (arity $\le 2$), allowing the logic to be expressed via multilinear polynomials.

2. Exact Optimization Formulation (Theoretical Bound)

To compute the worst-case probability of a policy violation, the authors formulate an exact linear program (LP) over the space of joint probability measures $\mu$.

• Objective: Maximize the probability of reaching an unsafe state (e.g., sensitive data reaching an external recipient).
• Constraints:
  • Structural Consistency: The measure $\mu$ must assign zero probability to states violating the deterministic transition logic of the Datalog rules.
  • Marginal Bounds: The measure must respect the known lower and upper bounds of the initial base facts.
  • Distributional Robustness: Crucially, the formulation makes no assumptions about the joint distribution (independence, positive/negative correlation) unless explicitly specified. It searches for the "worst-case" joint distribution consistent with the marginals.
• Limitation: This exact LP is exponentially sized ($2^{|V_F|}$) and computationally intractable for real-time monitoring.

3. Semidefinite Programming (SDP) Relaxation

To achieve runtime efficiency, the authors introduce a polynomially-sized Semidefinite Programming (SDP) relaxation.

• Moment Tracking: Instead of optimizing over the full joint distribution, the SDP tracks only the first and second-order moments (expectations and pairwise correlations) of the state distribution.
• Linearization: Non-convex quadratic terms arising from covariance constraints are linearized using McCormick envelopes over the marginal intervals.
• Correlation Classes: The framework supports explicit correlation constraints (Positively Correlated, Negatively Correlated, Independent) between specific predicates, allowing the optimization to tighten bounds when domain knowledge is available.
• Soundness: The authors prove that the SDP relaxation yields a sound upper bound ($R^*_{SDP} \ge R^*$) on the true execution risk. This ensures the monitor never under-blocks a violation, preserving security guarantees even if it occasionally over-blocks (reducing utility).

4. Runtime Enforcement

The computed upper bound is compared against a global safety threshold $\tau$. If the estimated risk exceeds $\tau$, the reference monitor blocks the proposed tool call; otherwise, it allows execution.

Key Contributions

1. Probabilistic Verification Model: Identification of security risks in deterministic engines for ambiguous environments and the introduction of a model for probabilistic verification that avoids premature binarization.
2. Exact Risk Formalization: Modeling multi-step agent trajectories via Datalog derivation graphs and formalizing the computation of worst-case policy violation risk as an exact optimization problem over joint measures.
3. Sound SDP Relaxation: Introduction of a polynomial-sized SDP relaxation that tracks second-order moments to efficiently approximate risk. The paper formally proves this relaxation is a strict upper bound on execution risk, ensuring soundness without independence assumptions.
4. Empirical Validation: Evaluation across terminal agent benchmarks (Intercode-NL2Bash, ATBench) and side-channel vulnerability tasks (Praline), demonstrating an improved security-utility trade-off compared to prior art.

Results

The framework was evaluated against three baselines: Praline (a state-of-the-art Datalog inference engine with correlation handling), Monte Carlo sampling (assuming strict independence), and a Deterministic engine with local thresholds.

• Security-Utility Trade-off: The SDP approach consistently outperformed baselines in terms of the Area Under the Curve (AUC) for the utility-security trade-off.
  • vs. Monte Carlo: Monte Carlo achieved high utility but failed on security (high false negatives) because independence assumptions underestimated worst-case risk, especially when inputs were positively correlated.
  • vs. Praline: While Praline achieved perfect security, it suffered from solver timeouts and loose bounds on complex tasks, leading to excessive false positives (low utility). The SDP approach maintained perfect security while achieving higher utility and lower latency on terminal benchmarks.
  • vs. Deterministic: Deterministic engines showed poor performance in both metrics, often failing to block unsafe actions or blocking safe ones due to information loss from thresholding.
• Latency: The SDP relaxation demonstrated polynomial scalability. On Intercode and ATBench benchmarks, SDP was significantly faster (e.g., ~221ms vs ~1015ms for Praline on Intercode). On complex side-channel tasks with many correlation constraints, SDP latency increased but remained predictable, avoiding the exponential tail latencies and timeouts associated with exact SMT-based reasoning.
• Correlation Sensitivity: Experiments showed that as input correlations increased, the independence assumption (Monte Carlo) led to a sharp drop in utility (increased false positives) because it overestimated risk in correlated settings where the true worst-case bound was tighter. The SDP framework adapted to these constraints effectively.

Significance and Claims

The paper claims to provide a paradigm shift from deterministic to probabilistic verification for AI agents. Its primary significance lies in bridging the gap between soundness (rigorous security guarantees) and tractability (real-time performance) in the presence of ambiguity.

• Robustness: By utilizing distributionally robust optimization, the framework ensures security guarantees hold regardless of unknown correlations between predicates, a critical requirement for real-world agent deployments where tool failures may be coupled.
• Practicality: The SDP relaxation enables the enforcement of complex, multi-step security policies with continuous probabilistic inputs without the computational intractability of exact inference or the safety risks of independence assumptions.
• Foundation: The work establishes a scalable foundation for securing autonomous agents operating in ambiguous environments, moving beyond heuristic defenses to formal, probabilistic enforcement.

The authors acknowledge limitations, specifically that risk bounds may degrade (converge to 1.0) over very long execution horizons with high degrees of fact merging, and that the approach relies on pre-defined tool semantics, which can be challenging for arbitrary scripts generated by agents. However, the framework is presented as a significant step toward rigorous, scalable agent security.