Unintended Misalignment from Agentic Fine-Tuning: Risks and Mitigation

TL;DR

The paper investigates unintended safety risks that arise when fine-tuning LLMs for agentic tasks, showing that improved benign-task performance can come at the cost of safety, namely higher harmful task execution and weaker refusal behavior. It introduces Prefix INjection Guard (PING), an automatic, iterative method that generates and selects natural-language prefixes to prepend to agent responses, guiding the model to refuse harmful requests while preserving performance on benign tasks. Across web navigation and code generation benchmarks, PING consistently improves safety compared to baselines and remains compatible with external guardrails, with only minor declines in benign task success. Internal-model analyses using linear probes reveal that PING shifts final-token representations toward safety, providing a mechanistic explanation for the observed improvements. The work underscores the importance of integrating safety considerations into agentic training and deployment, and offers a scalable, architecture-agnostic approach to layered safety in LLM agents.

Abstract

Beyond simple text generation, Large Language Models (LLMs) have evolved into agentic systems capable of planning and interacting with external tools to solve complex tasks. This evolution involves fine-tuning LLMs on agent-specific tasks to enhance their proficiency. However, safety concerns are frequently overlooked during this fine-tuning process. In this work, we show that aligned LLMs can become unintentionally misaligned, leading to a higher likelihood of executing harmful tasks and a reduced tendency to refuse them when fine-tuned to execute agentic tasks. To address these safety challenges, we propose Prefix INjection Guard (PING), a simple yet effective method that prepends automatically generated natural language prefixes to agent responses, guiding them to refuse harmful requests while preserving performance on benign tasks. Specifically, we introduce an iterative approach that alternates between (1) generating candidate prefixes and (2) selecting those that optimize both task performance and refusal behavior. Experimental results demonstrate that PING significantly enhances the safety of fine-tuned LLM agents without sacrificing their effectiveness. PING consistently outperforms existing prompting approaches across diverse benchmarks in both web navigation and code generation tasks. Our analysis of internal hidden states via linear probes reveals that prefix tokens are crucial for behavior modification, explaining the performance gains. WARNING: This paper contains contents that are unethical or offensive in nature.

Figures (13)

• Figure 1: An overview of our work. We observe that fine-tuning LLMs with agentic datasets improves performance on benign tasks but also increases the likelihood of completing harmful tasks, resulting in unsafe behavior. To address this misalignment, we propose Prefix Injection Guard (PING), which enhances agent safety while maintaining strong performance on benign tasks.
• Figure 2: Differences in the initial token probability distribution and responses for harmful tasks in the web domain before and after fine-tuning. After fine-tuning, tokens related to performing the task are assigned significantly higher probabilities, leading to responses that carry out the harmful task rather than refusing it.
• Figure 3: Success and refusal rates for web navigation and code generation tasks across diverse models for open-source models. Different markers indicate different methods. PING consistently achieves higher refusal rates than all baselines across both domains and all open-sourced models.
• Figure 4: Success and refusal rates of code generation tasks for closed-source models, with different markers indicating methods. PING achieves high refusal rates for both.
• Figure 5: Linear probe logit values for each token in the input sequence of a harmful task from WebDojo, evaluated on the Llama-3.1-8B-Instruct agent. The final token logit is lower when the suffix is injected (top), and higher when the prefix is injected before the model response (bottom).