The Alignment Waltz: Jointly Training Agents to Collaborate for Safety
Authors: name: Jingyu Zhang, name: Haozhu Wang, name: Eric Michael Smith, name: Sid Wang, name: Amr Sharaf, name: Mahesh Pasupuleti, name: Benjamin Van Durme, name: Daniel Khashabi, name: Jason Weston, name: Hongyuan Zhan
Organization: AI at Meta

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.