adu2021 Alignment Waltz: Collaborative Safety Training
Core Concept
Current LLM safety approaches reject any output containing unsafe content, causing high overrefusal rates on benign but sensitive topics. WaltzRL frames safety as a collaborative game where a conversation agent generates responses and a feedback agent suggests improvements. Both agents improve jointly through RL, reducing unsafe outputs from 39% to 4.6% while lowering overrefusals from 45.3% to 9.9%.
The key insight: Safety and helpfulness aren't opposed—a feedback agent guides the conversation agent toward safe-and-helpful outputs rather than simple rejection.
Architecture Overview
- Dual-Agent Framework: Conversation agent (generates) + feedback agent (improves)
- Dynamic Improvement Reward: Incentivizes acceptance and implementation of feedback
- Adaptive Deployment: Feedback agent engages only when needed
- Collaborative Learning: Both agents benefit from each other's improvements
- Positive-Sum Game: Safety and helpfulness both increase simultaneously
Implementation Steps
1. Design the Conversation Agent
Generates responses while being guided toward safety.
class ConversationAgent:
"""
Agent that generates responses to user queries.
"""
def __init__(self, model):
self.model = model
def generate_response(self, query, feedback=None):
"""
Generate response, optionally conditioned on feedback.
Args:
query: User query
feedback: Optional feedback from feedback agent
Returns:
Generated response
"""
if feedback:
prompt = (
f"User query: {query}\n\n"
f"Feedback to incorporate: {feedback}\n\n"
f"Improved response: "
)
else:
prompt = f"User query: {query}\n\nResponse: "
response = self.model.generate(
prompt,
max_tokens=256,
temperature=0.7
)
return response
def get_response_logprob(self, query, response):
"""Compute log probability of response given query."""
prompt = f"User query: {query}\n\nResponse: {response}"
return self.model.compute_log_prob(prompt)
2. Design the Feedback Agent
Analyzes responses and suggests improvements.
class FeedbackAgent:
"""
Agent that identifies unsafe elements and suggests improvements.
"""
def __init__(self, model, safety_classifier):
self.model = model
self.safety_classifier = safety_classifier
def analyze_response(self, query, response):
"""
Check response for safety issues.
Returns:
(has_issues, issue_description, suggested_fix)
"""
# Classify safety
safety_score = self.safety_classifier.score_safety(response)
if safety_score > SAFETY_THRESHOLD:
return (False, None, None) # Response is safe
# Identify specific issues
issues = self.safety_classifier.identify_issues(response)
return (True, issues, None)
def generate_feedback(self, query, response):
"""
Generate constructive feedback for improvement.
Args:
query: Original user query
response: Response to improve
Returns:
Feedback string (how to improve the response)
"""
has_issues, issues, _ = self.analyze_response(query, response)
if not has_issues:
return None # Response already good
# Generate improvement suggestions
feedback_prompt = (
f"User query: {query}\n\n"
f"Response: {response}\n\n"
f"Safety issues: {issues}\n\n"
f"How to improve this response to address the issues while "
f"maintaining helpfulness: "
)
feedback = self.model.generate(
feedback_prompt,
max_tokens=100,
temperature=0.5
)
return feedback
def get_feedback_logprob(self, query, response, feedback):
"""Compute log probability of feedback."""
prompt = (
f"Response: {response}\n"
f"Feedback: {feedback}"
)
return self.model.compute_log_prob(prompt)
3. Implement Dynamic Improvement Reward
Reward when feedback is incorporated and improves response quality.
class DynamicImprovementReward:
"""
Compute reward for feedback-guided improvement.
"""
def __init__(self, safety_classifier, helpfulness_scorer):
self.safety = safety_classifier
self.helpfulness = helpfulness_scorer
def compute_improvement_reward(self, response, improved_response, feedback):
"""
Reward successful incorporation of feedback.
Args:
response: Original response
improved_response: Response after feedback
feedback: The feedback provided
Returns:
Scalar reward
"""
# Safety improvement: did issue get addressed?
original_safety = self.safety.score_safety(response)
improved_safety = self.safety.score_safety(improved_response)
safety_improvement = improved_safety - original_safety
# Was feedback actually incorporated?
# Check overlap between feedback concepts and improvements
feedback_concepts = extract_concepts(feedback)
improvement_indicators = extract_concepts(improved_response)
incorporation_score = len(
set(feedback_concepts) & set(improvement_indicators)
) / (len(set(feedback_concepts)) + 1e-8)
# Helpfulness should not decrease
original_helpfulness = self.helpfulness.score(response)
improved_helpfulness = self.helpfulness.score(improved_response)
helpfulness_penalty = 0 if improved_helpfulness >= original_helpfulness else -0.2
# Combined reward
improvement_reward = (
0.6 * safety_improvement +
0.3 * incorporation_score +
0.1 * helpfulness_penalty
)
return improvement_reward
4. Implement WaltzRL Training Loop
Train both agents jointly with dynamic interaction.
def train_waltzrl(
conversation_model,
feedback_model,
dataset,
config
):
"""
Multi-agent RL training: conversation + feedback agents.
"""
conv_agent = ConversationAgent(conversation_model)
feedback_agent = FeedbackAgent(
feedback_model,
safety_classifier=config.safety_classifier
)
reward_computer = DynamicImprovementReward(
config.safety_classifier,
config.helpfulness_scorer
)
conv_optimizer = torch.optim.Adam(conversation_model.parameters(), lr=1e-5)
feedback_optimizer = torch.optim.Adam(feedback_model.parameters(), lr=1e-5)
for epoch in range(config.num_epochs):
epoch_loss = 0
for batch_idx, query in enumerate(dataset):
# Step 1: Conversation agent generates initial response
response = conv_agent.generate_response(query)
# Step 2: Check if feedback needed (adaptive)
needs_feedback = should_request_feedback(
query,
response,
config.safety_classifier
)
if needs_feedback:
# Feedback agent generates improvement suggestions
feedback = feedback_agent.generate_feedback(query, response)
if feedback:
# Conversation agent improves response
improved_response = conv_agent.generate_response(
query,
feedback=feedback
)
# Compute reward for improvement
improvement_reward = reward_computer.compute_improvement_reward(
response,
improved_response,
feedback
)
# Train conversation agent: encourage accepting feedback
conv_logprob = conv_agent.get_response_logprob(
query,
improved_response
)
conv_loss = -improvement_reward * conv_logprob
# Train feedback agent: encourage helpful feedback
feedback_logprob = feedback_agent.get_feedback_logprob(
query,
response,
feedback
)
feedback_loss = -improvement_reward * feedback_logprob
# Update both agents
conv_optimizer.zero_grad()
conv_loss.backward()
torch.nn.utils.clip_grad_norm_(
conversation_model.parameters(), 1.0
)
conv_optimizer.step()
feedback_optimizer.zero_grad()
feedback_loss.backward()
torch.nn.utils.clip_grad_norm_(
feedback_model.parameters(), 1.0
)
feedback_optimizer.step()
epoch_loss += (conv_loss + feedback_loss).item()
else:
# No improvement needed
epoch_loss += 0
else:
# Response already safe: small reward for conversation agent
conv_logprob = conv_agent.get_response_logprob(query, response)
conv_loss = -0.1 * conv_logprob # Small positive reward
conv_optimizer.zero_grad()
conv_loss.backward()
conv_optimizer.step()
epoch_loss += conv_loss.item()
if (batch_idx + 1) % 100 == 0:
print(f"Epoch {epoch}, Batch {batch_idx}: "
f"loss={epoch_loss / (batch_idx + 1):.4f}")
return conversation_model, feedback_model
5. Adaptive Feedback Deployment
Only engage feedback agent when necessary.
def should_request_feedback(query, response, safety_classifier, threshold=0.6):
"""
Decide if feedback agent should engage.
Strategy: Only ask for feedback on borderline cases.
Skip if clearly safe or clearly unsafe (latter gets rejection).
"""
safety_score = safety_classifier.score_safety(response)
if safety_score > 0.8:
return False # Clearly safe, no feedback needed
if safety_score < 0.2:
return False # Clearly unsafe, will be rejected anyway
# Borderline (0.2-0.8): feedback can help improve
return True
def deploy_with_adaptive_feedback(
conversation_model,
feedback_model,
query,
safety_classifier
):
"""
Generate response with optional feedback engagement.
"""
conv_agent = ConversationAgent(conversation_model)
feedback_agent = FeedbackAgent(feedback_model, safety_classifier)
# Generate initial response
response = conv_agent.generate_response(query)
# Check if feedback needed
if should_request_feedback(query, response, safety_classifier):
feedback = feedback_agent.generate_feedback(query, response)
if feedback:
# Improve based on feedback
response = conv_agent.generate_response(query, feedback=feedback)
return response
6. Evaluation: Safety and Helpfulness Tradeoff
Measure improvements in both dimensions.
def evaluate_waltzrl(conversation_model, feedback_model, test_dataset):
"""
Evaluate safety and helpfulness metrics.
"""
results = {
"unsafe_rate": 0,
"overrefusal_rate": 0,
"helpful_and_safe": 0
}
unsafe_count = 0
overrefusal_count = 0
helpful_safe_count = 0
conv_agent = ConversationAgent(conversation_model)
feedback_agent = FeedbackAgent(feedback_model, config.safety_classifier)
for query in test_dataset:
# Generate with adaptive feedback
response = deploy_with_adaptive_feedback(
conversation_model,
feedback_model,
query,
config.safety_classifier
)
# Evaluate
is_safe = config.safety_classifier.score_safety(response) > 0.7
is_helpful = config.helpfulness_scorer.score(response) > 0.6
is_overrefusal = len(response) < 10 # Suspiciously short response
if not is_safe:
unsafe_count += 1
if is_overrefusal and should_be_helpful(query):
overrefusal_count += 1
if is_safe and is_helpful:
helpful_safe_count += 1
results["unsafe_rate"] = unsafe_count / len(test_dataset) * 100
results["overrefusal_rate"] = overrefusal_count / len(test_dataset) * 100
results["helpful_and_safe"] = helpful_safe_count / len(test_dataset) * 100
print(f"Unsafe rate: {results['unsafe_rate']:.1f}%")
print(f"Overrefusal rate: {results['overrefusal_rate']:.1f}%")
print(f"Both safe & helpful: {results['helpful_and_safe']:.1f}%")
return results
Practical Guidance
Hyperparameters:
- Feedback decision threshold: 0.6 (0.2-0.8 = request feedback)
- Safety improvement weight: 0.6 (primary goal)
- Incorporation score weight: 0.3 (feedback must be used)
- Helpfulness penalty weight: 0.1 (don't sacrifice helpfulness)
- Learning rate: 1e-5 (conservative for safety)
When to Use:
- Deploying models that interact with users on sensitive topics
- Current approach rejects too many benign queries (overrefusal)
- Want collaborative improvement rather than hard rejection
- Have access to safety and helpfulness classifiers
When NOT to Use:
- Maximally adversarial red-teaming (waltz is cooperative)
- Real-time inference with strict latency (two-agent check adds overhead)
- Models that should never generate unsafe content (use hard filtering)
- Domains where any risk is unacceptable
Reference
Alignment Waltz: Jointly Training Agents for Safety — arXiv:2510.08240