Secure AI workflows by enforcing rules in diffusion language models
Only What You've Signed up For: Secure Workflows with Constrained Diffusion Agents
Machine LearningCryptography and Security
Summary
As AI systems get smarter and handle more important tasks, it becomes very important to stop hackers from tricking them with bad instructions. The paper shows a new way to keep AI safe by using a special type of AI model called diffusion language models, which can check and correct their answers step by step. The researchers created a framework called CoDA that ensures the AI always follows certain rules while it is generating text. Tests show CoDA makes it much harder for attackers to succeed without making the AI worse or slower.
What this means in practice
- •For ai product developers: Build safer AI assistants that resist malicious prompt tricks during complex text tasks using built-in generation constraints.$Commercial implications: Enables creation of AI-assistant products with stronger security against prompt manipulation, improving trust for end users.
- •For enterprise security teams: Deploy AI workflow systems that enforce custom task rules to reduce risk from prompt injection attacks in sensitive environments.
Authors
Michael Cardei, Anjun Hu, Chenrui Ma, Francesco Pinto, Chawin Sitawarin, Chulin Xie, Ferdinando Fioretto
Abstract
As agentic AI systems take on increasingly complex and consequential tasks, securing them against prompt injection becomes critical. Existing defenses often rely on detecting or reasoning about injected malicious instructions and can therefore remain vulnerable to adaptive attacks that optimize the injection itself. Diffusion language models (dLLMs), an emerging framework for language generation, offer a new opportunity for security enforcement: by iteratively refining their sequence wide predictions, they enable broader-scope constraint enforcement throughout generation. This paper introduces Policy-Constrained Diffusion Agents (CoDA), a framework that enforces task dependent security policies directly within the denoising process of diffusion language models. At each denoising step, CoDA projects token distributions onto a policy-compliant set, maintaining constraint satisfaction throughout generation. We evaluate CoDA across state-of-the-art diffusion language model families under both static and adaptive prompt-injection attacks. Results demonstrate substantial reductions in attack success rate with competitive task utility and practical inference overhead.