The Alignment Waltz: Jointly Training Agents to Collaborate for Safety
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.
68e87d9195e8e6771df38a45
Jingyu Zhang
68e87d9195e8e6771df38a46
Haozhu Wang
68e87d9195e8e6771df38a47
Eric Michael Smith
68e87d9195e8e6771df38a48
Sid Wang
68e87d9195e8e6771df38a49
Amr Sharaf
68e87d9195e8e6771df38a4a
Mahesh Pasupuleti
68e87d9195e8e6771df38a4b
Benjamin Van Durme
68e87d9195e8e6771df38a4c
Daniel Khashabi
68e87d9195e8e6771df38a4d
Jason Weston
68e87d9195e8e6771df38a4e
Hongyuan Zhan
The Alignment Waltz: Jointly Training Agents to Collaborate for Safety
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.
WaltzRL, a multi-agent reinforcement learning framework, improves LLM safety and helpfulness by collaboratively training a conversation agent and a feedback agent, reducing unsafe responses and overrefusals.
LLMs adversarial attacks safeguard models multi-agent reinforcement learning conversation agent feedback agent Dynamic Improvement Reward DIR WildJailbreak OR-Bench Pareto front