Imparare quando Agire o Rifiutarsi: Proteggere i Modelli di Ragionamento Agente per un Uso Sicuro di Strumenti Multi-Step

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

March 3, 2026

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

cs.AI

Abstract

I modelli linguistici agentici operano in un regime di sicurezza fondamentalmente diverso rispetto ai modelli di chat: devono pianificare, richiamare strumenti ed eseguire azioni a lungo termine in cui un singolo passo falso, come l'accesso a file o l'inserimento di credenziali, può causare danni irreversibili. I metodi di allineamento esistenti, ottimizzati principalmente per la generazione statica e il completamento di compiti, falliscono in questi contesti a causa del processo decisionale sequenziale, del feedback avversario degli strumenti e del ragionamento intermedio eccessivamente sicuro. Introduciamo MOSAIC, un framework di post-addestramento che allinea gli agenti per un utilizzo sicuro e multi-step degli strumenti, rendendo esplicite e apprendibili le decisioni di sicurezza. MOSAIC struttura l'inferenza come un ciclo di pianificazione, verifica, quindi azione o rifiuto, con un ragionamento esplicito sulla sicurezza e il rifiuto come azioni di prima classe. Per addestrare senza etichette a livello di traiettoria, utilizziamo l'apprendimento per rinforzo basato su preferenze con confronti a coppie di traiettorie, che cattura distinzioni di sicurezza spesso trascurate da ricompense scalari. Valutiamo MOSAIC in zero-shot su tre famiglie di modelli, Qwen2.5-7B, Qwen3-4B-Thinking e Phi-4, e su benchmark fuori distribuzione che abbracciano compiti dannosi, prompt injection, uso benigno di strumenti e fughe di dati transdominio. MOSAIC riduce i comportamenti dannosi fino al 50%, aumenta il rifiuto di compiti dannosi di oltre il 20% sugli attacchi di injection, riduce le fughe di dati e preserva o migliora le prestazioni sui compiti benigni, dimostrando una robusta generalizzazione attraverso modelli, domini e contesti agentici.

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.