Aprendendo Quando Agir ou Recusar: Protegendo Modelos de Raciocínio Agêntico para Uso Seguro de Ferramentas em Múltiplas Etapas

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

Resumo

Os modelos de linguagem agentes operam em um regime de segurança fundamentalmente diferente dos modelos de chat: eles devem planejar, chamar ferramentas e executar ações de longo prazo, onde um único passo em falso, como acessar arquivos ou inserir credenciais, pode causar danos irreversíveis. Os métodos de alinhamento existentes, amplamente otimizados para geração estática e conclusão de tarefas, falham nesses cenários devido à tomada de decisão sequencial, ao feedback adversário de ferramentas e ao raciocínio intermediário excessivamente confiante. Apresentamos o MOSAIC, uma estrutura de pós-treinamento que alinha agentes para o uso seguro de ferramentas em múltiplas etapas, tornando as decisões de segurança explícitas e aprendíveis. O MOSAIC estrutura a inferência como um loop de planejar, verificar, depois agir ou recusar, com raciocínio de segurança explícito e recusa como ações de primeira classe. Para treinar sem rótulos a nível de trajetória, usamos aprendizado por reforço baseado em preferências com comparações pareadas de trajetórias, o que captura distinções de segurança frequentemente perdidas por recompensas escalares. Avaliamos o MOSAIC *zero-shot* em três famílias de modelos, Qwen2.5-7B, Qwen3-4B-Thinking e Phi-4, e em *benchmarks* fora da distribuição abrangendo tarefas nocivas, injeção de *prompts*, uso benigno de ferramentas e vazamento de privacidade entre domínios. O MOSAIC reduz o comportamento nocivo em até 50%, aumenta a recusa de tarefas nocivas em mais de 20% em ataques de injeção, reduz o vazamento de privacidade e preserva ou melhora o desempenho em tarefas benignas, demonstrando generalização robusta entre modelos, domínios e configurações 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.