Aprendiendo cuándo actuar o rechazar: Protección de modelos de razonamiento agentico para un uso seguro de herramientas en múltiples pasos

Learning When to Act or Refuse: Guarding Agentic Reasoning Models for Safe Multi-Step Tool Use

March 3, 2026

Autores: Aradhye Agarwal, Gurdit Siyan, Yash Pandya, Joykirat Singh, Akshay Nambi, Ahmed Awadallah

cs.AI

Resumen

Los modelos de lenguaje agentes operan en un régimen de seguridad fundamentalmente diferente al de los modelos de chat: deben planificar, invocar herramientas y ejecutar acciones de largo horizonte, donde un solo paso en falso, como acceder a archivos o introducir credenciales, puede causar un daño irreversible. Los métodos de alineación existentes, optimizados en gran medida para la generación estática y la finalización de tareas, fallan en estos entornos debido a la toma de decisiones secuencial, la retroalimentación adversarial de herramientas y el razonamiento intermedio sobreconfiado. Presentamos MOSAIC, un marco de post-entrenamiento que alinea a los agentes para un uso seguro de herramientas multi-etapa, haciendo que las decisiones de seguridad sean explícitas y aprendibles. MOSAIC estructura la inferencia como un bucle de planificar, verificar, luego actuar o rechazar, con un razonamiento de seguridad explícito y el rechazo como acciones de primera clase. Para entrenar sin etiquetas a nivel de trayectoria, utilizamos el aprendizaje por refuerzo basado en preferencias con comparaciones por pares de trayectorias, lo que captura distinciones de seguridad que a menudo pasan desapercibidas para las recompensas escalares. Evaluamos MOSAIC de forma zero-shot en tres familias de modelos, Qwen2.5-7B, Qwen3-4B-Thinking y Phi-4, y en benchmarks fuera de distribución que abarcan tareas dañinas, inyección de prompts, uso benigno de herramientas y fugas de privacidad entre dominios. MOSAIC reduce el comportamiento dañino hasta en un 50%, aumenta el rechazo de tareas dañinas en más de un 20% en ataques de inyección, reduce las fugas de privacidad y preserva o mejora el rendimiento en tareas benignas, demostrando una generalización robusta entre modelos, dominios y configuraciones agentes.

English

Agentic language models operate in a fundamentally different safety regime than chat models: they must plan, call tools, and execute long-horizon actions where a single misstep, such as accessing files or entering credentials, can cause irreversible harm. Existing alignment methods, largely optimized for static generation and task completion, break down in these settings due to sequential decision-making, adversarial tool feedback, and overconfident intermediate reasoning. We introduce MOSAIC, a post-training framework that aligns agents for safe multi-step tool use by making safety decisions explicit and learnable. MOSAIC structures inference as a plan, check, then act or refuse loop, with explicit safety reasoning and refusal as first-class actions. To train without trajectory-level labels, we use preference-based reinforcement learning with pairwise trajectory comparisons, which captures safety distinctions often missed by scalar rewards. We evaluate MOSAIC zero-shot across three model families, Qwen2.5-7B, Qwen3-4B-Thinking, and Phi-4, and across out-of-distribution benchmarks spanning harmful tasks, prompt injection, benign tool use, and cross-domain privacy leakage. MOSAIC reduces harmful behavior by up to 50%, increases harmful-task refusal by over 20% on injection attacks, cuts privacy leakage, and preserves or improves benign task performance, demonstrating robust generalization across models, domains, and agentic settings.