PaperlayerView exact revision on arXiv

Source-linked AI summary

The Alignment Waltz: Jointly Training Agents to Collaborate for Safety

Jingyu Zhang, Haozhu Wang, Eric Michael Smith, Sid Wang, Amr Sharaf, Mahesh Pasupuleti, Benjamin Van Durme, Daniel Khashabi, Jason Weston, Hongyuan Zhan

arXiv:2510.08240v2cs.CL

TL;DR

LLMs must resist adversarial attacks without overrefusing benign but sensitive requests, while blunt safeguards can block useful content. WaltzRL jointly trains collaborating conversation and feedback agents, using improvement-based rewards and adaptive feedback. Across five diverse datasets, it reduces both unsafe responses and overrefusals, including a large reduction in WildJailbreak safety violations.

  • Problem

    LLMs face a trade-off between unsafe responses to adversarial attacks and overrefusal of benign sensitive prompts, while safeguards can block safe helpful content.

    Source: §1 · p. 1

  • Method

    WaltzRL jointly trains conversation and feedback agents to collaborate, rewarding feedback according to improvements in the conversation agent’s response and applying it adaptively.

    Source: §1 · p. 2

  • Results

    Across five diverse datasets, WaltzRL reduces unsafe responses and overrefusals, including safety violations from 39.0% with the base model to 4.6% with ours on WildJailbreak.

    Source: §1 · p. 2

  • Takeaways & Limitations

    WaltzRL advances the safety and overrefusal Pareto front without degrading general helpfulness, while its feedback agent is adapted to a specific conversation agent.

    Source: §5 · p. 11

    WaltzRL is developed strictly for research; downstream deployments should include careful red-teaming, monitoring, and additional guardrails when needed.

    Source: Ethical Considerations · p. 11

Abstract

from arXiv · show

Harnessing the power of LLMs requires a delicate dance between being helpful and harmless. This creates a fundamental tension between two competing challenges: vulnerability to adversarial attacks that elicit unsafe content, and a tendency for overrefusal on benign but sensitive prompts. Current approaches often navigate this dance with safeguard models that completely reject any content that contains unsafe portions. This approach cuts the music entirely-it may exacerbate overrefusals and fails to provide nuanced guidance for queries it refuses. To teach models a more coordinated choreography, we propose WaltzRL, a novel multi-agent reinforcement learning framework that formulates safety alignment as a collaborative, positive-sum game. WaltzRL jointly trains a conversation agent and a feedback agent, where the latter is incentivized to provide useful suggestions that improve the safety and helpfulness of the conversation agent's responses. At the core of WaltzRL is a Dynamic Improvement Reward (DIR) that evolves over time based on how well the conversation agent incorporates the feedback. At inference time, unsafe or overrefusing responses from the conversation agent are improved rather than discarded. The feedback agent is deployed together with the conversation agent and only engages adaptively when needed, preserving helpfulness and low latency on safe queries. Our experiments, conducted across five diverse datasets, demonstrate that WaltzRL significantly reduces both unsafe responses (e.g., from 39.0% to 4.6% on WildJailbreak) and overrefusals (from 45.3% to 9.9% on OR-Bench) compared to various baselines. By enabling the conversation and feedback agents to co-evolve and adaptively apply feedback, WaltzRL enhances LLM safety without degrading general capabilities, thereby advancing the Pareto front between helpfulness and harmlessness.

1 Introduction

Safety alignment must balance resistance to adversarial attacks with helpful answers to benign but sensitive prompts; blunt safeguards can block useful content and worsen overrefusal. WaltzRL instead trains conversation and feedback agents to improve responses collaboratively, reducing both unsafe outputs and overrefusals across diverse datasets.

1 Introduction

  • Standalone safeguards convert responses deemed unsafe into refusals, which can deprive users of safe helpful information and exacerbate overrefusal on sensitive or dual-use prompts.The introduction notes that safeguards reduce unsafe responses at the cost of higher overrefusal, with a larger negative effect when the unsafeguarded system already has low overrefusal.Source: §1 · pp. 1–3
  • WaltzRL treats safety alignment as a positive-sum collaboration: a feedback agent suggests revisions, while a conversation agent incorporates useful feedback and both agents are updated.Its Dynamic Improvement Reward changes with the conversation agent’s improvement after feedback; adaptive triggering is intended to preserve helpfulness and latency.Source: §1 · p. 2
  • WaltzRL deploys the conversation and feedback agents together at inference, unlike prior multi-agent training approaches that deploy only a single defender.The authors state that an attack must jailbreak both agents to succeed.Source: §1 · p. 2
  • Across five diverse datasets, WaltzRL reduces both safety violations and overrefusals, including safety violations from 39.0% with the base model to 4.6% with ours on WildJailbreak.The introduction also reports overrefusals falling from 45.3% to 9.9% on OR-Bench.Source: §1 · p. 2
  • Inference-time collaboration without RL reduces unsafe and overrefusing responses, but triggers feedback excessively; WaltzRL training further improves safety and overrefusal while preventing over-triggering.The trained system uses feedback adaptively.Source: §1 · p. 3
  • An oracle using template feedback from ground-truth safety and overrefusal labels underperforms WaltzRL, indicating detailed feedback is crucial, especially for turning overrefusals into helpful responses.Source: §1 · p. 3

2 WaltzRL: Training Agents for Collaborative Reasoning

WaltzRL frames safety alignment as collaboration between a conversation agent and a feedback agent, trained to improve responses while avoiding overrefusal. Its protocol, reward design, joint training, and adaptive-feedback stages support this collaboration.

2 WaltzRL: Training Agents for Collaborative Reasoning

  • WaltzRL trains a conversation agent and feedback agent to collaboratively produce safe responses while avoiding overrefusal.The two-agent framework is reinforced through multi-agent RL.Source: §2 · p. 3

2.1 Collaboration Protocol in WaltzRL

  • The agents receive separate, non-competing rewards, and the feedback agent returns safety and overrefusal labels plus feedback that can be applied over multiple rounds.Feedback stops when the response is satisfactory or the maximum number of rounds is reached.Source: §2.1 · pp. 3–4

2.2 Shaping Rewards to Encourage Collaboration

  • The Dynamic Improvement Reward scores feedback by the change in conversation-agent reward after revision, while label and format rewards encourage accurate flags and well-formed JSON.The improvement reward is conditioned on label correctness because otherwise it can dominate and label reward can drop during training.Source: §2.2 · p. 4

2.3 Multi-Agent Reinforcement Learning

  • Each training step collects collaborative rollouts, samples both initial-response and feedback-informed trajectories, and updates the conversation and feedback agents as separate actors.The two agents’ policy-gradient updates can run in parallel.Source: §2.3 · pp. 4–6

2.4 Learning to Give Feedback Adaptively

  • A two-stage training strategy first freezes the conversation agent while training feedback labels and format, then trains both agents while disabling the additive label reward.The improvement reward remains conditioned on label correctness to maintain label accuracy as revision-worthy responses become less prevalent.Source: §2.4 · p. 6

3 Experiments

WaltzRL is evaluated against safety, overrefusal, capability, and efficiency baselines across adversarial and borderline prompts. It leads on average ASR and ORR, while adaptive feedback preserves capability and limits triggering on safe queries.

3.1 Experimental Setup

  • Both agents are initialized from Llama-3.1-8B-Instruct and trained on WildJailbreak adversarial prompts and OR-Bench-80K borderline overrefusal prompts.Training uses at most one feedback round, allowing two conversation-agent responses and one feedback response.Source: §3.1 · p. 6
  • The evaluation covers adversarial safety, benign-prompt overrefusal, instruction following, and general capability using five safety and overrefusal datasets plus capability benchmarks.The safety datasets are WildJailbreak, FORTRESS adversarial harmful, and StrongREJECT; the overrefusal datasets are OR-Bench-Hard-1K and FORTRESS benign.Source: §3.1 · p. 6
  • Baselines include the untrained model, single-model RL, safeguards applied to each, untrained inference-time collaboration, and oracle label-converted feedback.The safeguard converts responses classified as unsafe into refusals; oracle feedback uses ground-truth alignment labels to select a template instruction.Source: §3.1 · p. 7

3.2 Evaluation Results

  • WaltzRL outperforms all baselines on average ASR and ORR across evaluation datasets, while safeguards increase overrefusal.Adding safeguards raises ORR from 8.6% to 14.9% after single-model RL and from 25.7% to 29.8% on baseline responses.Source: §3.2 · p. 7, Table 1 · p. 7
  • Inference-time collaboration reduces both ASR and ORR over the baseline, and WaltzRL training reduces both further.Oracle label-converted feedback reduces ASR from 26.5 to 7.0 but ORR only from 25.7 to 16.6; the authors attribute the limited ORR reduction to the lack of rationale in template feedback.Source: §3.2 · p. 7
  • WaltzRL yields little instruction-following and general-helpfulness degradation despite using no helpfulness prompts during RL.Adaptive feedback is rarely triggered on non-safety AlpacaEval prompts, leading to little degradation of win rate.Source: §3.2 · p. 8, Table 2 · p. 8
  • Feedback triggering falls to 6.7% on general AlpacaEval prompts and remains below 50% on WildJailbreak and OR-Bench.Experiments cap feedback at one round, so the reported latency impact is similar to safeguard models, which prior work considers acceptable for practical deployment.Source: §3.2 · p. 8

Qualitative examples

The ablations support collaborative two-stage training: Stage 2 improves safety and overrefusal outcomes, while staged reward choices balance feedback usefulness and label accuracy. Qualitative examples show feedback can turn an overrefusal into compliance and outline an ideal response.

Qualitative examples

  • Qualitative examples show feedback converting an overrefusal into compliance and giving the conversation agent an outline, sometimes as a quoted ideal response.The response follows outlines generated by the feedback agent.Source: Qualitative examples · p. 8

3.3 Ablations and Analysis

  • All three reward setups produce more improved than worsened responses; setup A slightly underperforms B and C on feedback improvement but achieves the best label accuracy.B retains label conditioning within DIR, while C removes label reward entirely; conditioning DIR on label reward is crucial for maintaining high label accuracy.Source: §3.3 · pp. 8–9
  • The two-stage recipe uses setup A first to learn accurate labels, then setup B to further improve feedback usefulness.Setup A combines DIR, label, and format rewards; setup B removes the additive label reward while retaining label-conditioned DIR.Source: §3.3 · pp. 8–9
  • Freezing the conversation agent throughout training produces worse outcomes than multi-agent co-evolution.This ablation freezes the conversation agent in both stages and optimizes only the feedback agent.Source: §3.3 · p. 9
  • Replacing DIR with the next-round conversation agent’s outcome reward produces a notably worse outcome.The ablation isolates the effect of the feedback agent’s Dynamic Improvement Reward.Source: §3.3 · p. 9
  • Stage 1 trains valid feedback formatting and label prediction, while Stage 2 raises both initial-response and revised-response rewards.The final response reward remains notably higher than the initial response reward at the end of RL training, indicating additional gains from feedback beyond single-model RL.Source: §3.3 · pp. 9–10
  • Removing Stage 2 collaborative training leads to significantly higher ASR and ORR, with similar label accuracy and feedback trigger rate.The authors conclude Stage 2 improves safety and reduces overrefusal while maintaining label accuracy learned in Stage 1.Source: §3.3 · p. 10

4 Related Work

WaltzRL is positioned against prior debate, safeguard, and multi-agent alignment approaches through collaborative positive-sum training and deployment of both agents. The paper distinguishes this from zero-sum competition and single-defender deployment.

Debate for AI safety

Safeguarding LLMs

  • Standalone safeguards offer flexibility as safety standards change, but the paper says they can enhance safety without reducing overrefusal; WaltzRL instead extends deliberation to multi-agent dialogue.The related-work discussion also notes complementary work on helpfulness or constructiveness under safety constraints [Zhang et al., 2025a; Duan et al., 2025; Yuan et al., 2025].Source: §4 · p. 10

Self-play and multi-agent RL

  • Unlike prior two-player alignment games that are zero-sum, WaltzRL uses positive-sum rewards to encourage collaboration and deploys both agents at inference.Liu et al. (2025) and Zheng et al. (2024) deploy only a trained defender LLM [p31 reference marker is absent; no marker attached].Source: §4 · pp. 10–11

5 Conclusion and Future Work

WaltzRL’s positive-sum multi-agent RL shows promising safety and overrefusal results without degrading general helpfulness. The paper trains feedback agents for specific conversation agents and identifies generalist feedback agents as future work.

5 Conclusion and Future Work

  • WaltzRL reports promising movement of the safety–overrefusal Pareto front without degrading general helpfulness.The conclusion characterizes the approach as a positive-sum game in which both agents are rewarded by the same outcome, encouraging collaboration.Source: §5 · p. 11
  • The feedback agent is trained to suit a specific conversation agent; future work could develop generalist feedback agents usable off the shelf with different conversation agents.Source: §5 · p. 11

Ethical Considerations

The paper frames WaltzRL as research-purpose work and cautions that downstream deployment needs additional safeguards. Agent prompts establish the initial collaboration framework used as a baseline and starting point for reinforcement learning.

Ethical Considerations

  • WaltzRL is developed strictly for research; downstream deployment should include careful red-teaming, monitoring, and additional guardrails when needed.The framework focuses on safety alignment through multi-agent reinforcement learning.Source: Ethical Considerations · p. 11

A Agent Initialization and Task Format Details

  • Separate system prompts assign the conversation and feedback agents their roles, establishing a prompting-based collaboration baseline and the starting point for reinforcement learning.The conversation agent integrates useful suggestions while responding to the original query; the feedback agent provides feedback and a high-level safety-guideline summary.Source: A Agent Initialization and Task Format… · p. 17

B.1 Training Data

WaltzRL training uses 20,000 prompts drawn equally from WildJailbreak and OR-Bench, divided evenly between its two training stages.

B.1 Training Data

  • 20,000 prompts comprise 10,000 WildJailbreak adversarial attacks and 10,000 OR-Bench borderline prompts, with 5,000 from each dataset assigned to each training stage.The OR-Bench prompts are deduplicated from OR-Bench-Hard-1K.Source: §B.1 · p. 17

B.2 Training Infrastructure and Hyperparameters

Training uses open-source RL infrastructure and specified hyperparameters, while alignment labels distinguish unsafe responses from refusals to benign prompts.

B.2 Training Infrastructure and Hyperparameters

  • Training uses OpenRLHF and MARTI2, vLLM rollouts, batch size 32, β=0.01, α=0.65, λ=0.25, γ=0.1, clip ϵ=0.2, and learning rate 5·10^-7.Both agents use a maximum prompt length of 4096 and maximum generation length of 512 per turn.Source: §B.2 · p. 17

C Alignment Labels Details

  • WildGuard labels response harmfulness and refusal, while the prompt source supplies prompt harmfulness; these labels yield unsafe and overrefuse statuses, with overrefusal meaning refusal on a benign prompt.A response is unsafe when WildGuard flags harmful content; it is overrefusing when a non-harmful prompt receives a refusal.Source: C Alignment Labels Details · p. 17

D Evaluation Details

The evaluation details define average-based safety and overrefusal metrics, describe general-capability benchmarks, and report that removing mixed trajectory sampling worsens the aggregate outcome.

D Evaluation Details

  • Table 8 defines F1 as the harmonic mean of 1-average ASR and 1-average ORR, and reports a worse outcome when mixed trajectory sampling is ablated.F1 is an aggregate score for balancing helpfulness and safety, with higher values better.Source: Table 8 · p. 18
  • ASR averages harmful-response scores on adversarial prompts, while ORR averages refusal scores on borderline prompts.WildGuard classifies response harmfulness for ASR and refusal for ORR.Source: D Evaluation Details · pp. 17–18
  • General helpfulness and instruction following are evaluated on AlpacaEval 2.0, IFEval, GPQA, MMLU, and TruthfulQA using win-rate, instruction-following, and accuracy measures.AlpacaEval uses length-controlled win rate and win rate; IFEval uses four accuracy variants, while the remaining datasets use multiple-choice accuracy.Source: D Evaluation Details · p. 18

E Ablation Studies Continued

Training on only one of the two trajectory types produces worse outcomes than mixed trajectory sampling, supporting the mixed strategy.

E Ablation Studies Continued

  • Training on only τA or only τB produces worse outcomes than mixed trajectory sampling.The ablation is reported as evidence for the effectiveness of mixed trajectory sampling.Source: E Ablation Studies Continued · p. 18

F Limitations

The evaluation is limited to English datasets, one feedback round, and static adversarial prompts; unsafe responses and overrefusals are not fully eliminated.

F Limitations

  • Experiments use English datasets, leaving performance on other languages unevaluated.The paper identifies adversarial attacks in other languages, including side-channel attacks in low-resource languages, as future evaluation.Source: F Limitations · p. 18
  • Feedback is limited to a maximum of one round because of computation resource constraints.The authors note that more feedback rounds might lead to further improvements.Source: F Limitations · p. 18
  • Evaluation covers static adversarial prompts, not dynamic adaptive attacks.The paper cites prior work showing stronger results from dynamic adaptive attacks and suggests extending evaluation to them.Source: F Limitations · p. 18
  • WaltzRL reduces unsafe responses and overrefusals but does not eliminate them.Source: F Limitations · p. 18

G Qualitative examples

This section presents qualitative examples of feedback revising overrefusals and unsafe content, alongside the prompts guiding both agents. The prompts ask the conversation agent to integrate feedback silently and direct the feedback agent to assess safety and overrefusal separately, suggest fixes, and follow specified discussion and response-format constraints.

G Qualitative examples

  • Feedback converts a refusal of a legitimate income-planning query into general information about blind trusts, offshore accounts, and lawful tax strategies.The response emphasizes compliance with applicable laws and consulting a qualified tax or financial professional.Source: G Qualitative examples · p. 20
  • The feedback agent steers an overrefusal on environmental conservation toward a neutral, informative response presenting different perspectives and alternatives.The conversation agent’s revision discusses conservation’s environmental benefits, economic tradeoffs, and approaches intended to balance development with protection.Source: G Qualitative examples · p. 21
  • For a hacking-related query, feedback distinguishes disallowed instructions from permitted historical context, and the revised answer discusses rainbow-table attacks and countermeasures.The initial answer provided simulation code; the feedback instead directed the agent to avoid facilitating hacking while allowing historical context, critique, or reporting.Source: G Qualitative examples · pp. 21–23

H System Prompts

  • The system prompts ask the conversation agent to answer helpfully and safely while silently integrating useful feedback, and task the feedback agent with constructive guidance that is both safe and not overrefusing.The feedback instructions state that safety and avoiding overrefusal are equally crucial.Source: H System Prompts · pp. 24–25
  • The prompt permits analysis of harmful topics when it does not facilitate harm and requires conversational feedback in a JSON-only format adapted from Llama Guard 3 guidance.The required JSON fields are reasoning, unsafe, overrefuse, and feedback.Source: H System Prompts · p. 25, Figure 6
  • The feedback agent reasons about safety and overrefusal, assigns separate binary labels, and gives only fixes for unsafe or overrefusing responses.The prompt requires labels for both aspects and says not to comment on response quality.Source: H System Prompts · p. 25

Related papers

Loading 2510.08240v2…