Efficient and Sound Probabilistic Verification for AI Agents

2026-06-18

AI Summary

This work addresses the challenge of verifying probabilistic safety policies in AI agents, which existing runtime monitoring approaches struggle to handle due to their reliance on deterministic strategies and inability to account for uncertainty and correlations. The paper proposes a novel verification framework based on distributionally robust optimization (DRO), introducing this technique for the first time into Datalog-style probabilistic policy verification. By eschewing assumptions of predicate independence, the method provides rigorous upper bounds on policy violation probabilities. Integrating DRO with probabilistic logical reasoning and Datalog’s formalism, the approach achieves a superior trade-off between safety and utility while maintaining strict guarantees on violation probability bounds. Empirical evaluations on benchmarks involving terminal and tool-calling agents demonstrate significant improvements over state-of-the-art methods.

Abstract

Securing AI agents that operate in complex digital environments has become a critical need, and runtime monitoring approaches that formulate and enforce policies expressed in a formal language like Datalog offer a promising solution. However, existing approaches are restricted to deterministic policies. In many practical applications of AI agents, there is a need to enforce security policies in the face of ambiguity, leading to probabilistic predicates or state transitions (for example, a declassifier or Personally Identifiable Information (PII) detector that has some failure probability on each invocation). Furthermore, in many such applications, one cannot easily make the independence assumptions necessary to invoke prior work on probabilistic inference in Datalog. We address this by introducing a sound and efficient framework for such verification based on distributionally robust optimization, computing sound upper bounds on the probability of policy violation regardless of possible correlations between predicates. On standard benchmarks for terminal and tool calling agents, we demonstrate that our approach outperforms prior art and improves the security-utility trade-off while ensuring rigorous bounds on the probability of policy violation.

Problem

• probabilistic verification
• AI agents
• security policies
• distributional robustness
• Datalog

Innovation

• probabilistic verification
• distributionally robust optimization
• Datalog
• AI agents
• policy enforcement