PaperlayerView exact revision on arXiv

Source-linked AI summary

Toward Optimal LLM Alignments Using Two-Player Games

Rui Zheng, Hongyi Guo, Zhihan Liu, Xiaoying Zhang, Yuanshun Yao, Xiaojun Xu, Zhaoran Wang, Zhiheng Xi, Tao Gui, Qi Zhang, Xuanjing Huang, Hang Li, Yang Liu

arXiv:2406.10977v1cs.CLcs.AI

TL;DR

Pre-collected prompts can be tedious to gather and may miss scenarios where LLMs most need improvement. The paper proposes an iterative game in which an adversarial agent challenges a defensive agent, and reports convergence to a Nash equilibrium alongside improved generalization in safety settings.

  • Problem

    Pre-collected prompts can be difficult to gather comprehensively and may omit scenarios where LLMs most need improvement.

    Source: §1 · p. 1

  • Method

    An adversarial LLM iteratively generates prompts exposing a defensive LLM’s weaknesses, and the defensive LLM adapts its responses using reward-model feedback.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    The iterative process converges to a Nash equilibrium, and safety experiments report enhanced generalization and thorough training for both agents.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    In harmful-input and jailbreak settings, the adversarial agent also serves as an adaptive red-teaming partner that generates challenging prompts for alignment.

    Source: §1 · p. 2

    The approach requires training two separate LLM agents, and the work focuses primarily on safety-related tasks.

    Source: §6 · p. 10

Abstract

from arXiv · show

The standard Reinforcement Learning from Human Feedback (RLHF) framework primarily focuses on optimizing the performance of large language models using pre-collected prompts. However, collecting prompts that provide comprehensive coverage is both tedious and challenging, and often fails to include scenarios that LLMs need to improve on the most. In this paper, we investigate alignment through the lens of two-agent games, involving iterative interactions between an adversarial and a defensive agent. The adversarial agent's task at each step is to generate prompts that expose the weakness of the defensive agent. In return, the defensive agent seeks to improve its responses to these newly identified prompts it struggled with, based on feedback from the reward model. We theoretically demonstrate that this iterative reinforcement learning optimization converges to a Nash Equilibrium for the game induced by the agents. Experimental results in safety scenarios demonstrate that learning in such a competitive environment not only fully trains agents but also leads to policies with enhanced generalization capabilities for both adversarial and defensive agents.

1 Introduction

Pre-collected and rule-based prompt strategies may miss the prompts an aligning LLM struggles with, motivating a more dynamic approach. The paper frames alignment as an iterative two-player game and reports that it converges to a Nash equilibrium while improving both agents’ generalization in safety scenarios.

1 Introduction

  • Traditional prompt expansion methods use difficulty, paraphrase, or self-instruction rules, but do not adapt to prompts the aligning LLM struggles to answer.Static prompt datasets may also saturate performance as reward-model discernment is lost.Source: §1 · p. 1
  • The framework iteratively has an adversarial LLM generate prompts that expose a defensive LLM’s weaknesses, then has the defensive LLM improve its responses to those prompts.The paper describes this competitive interaction as fostering a natural curriculum of increasing complexity for both agents.Source: §1 · p. 2
  • Diversity constraints based on BLEU scores and sentence embeddings prevent the adversarial agent from prematurely converging on a narrow set of prompts, expanding vulnerability coverage.Source: §1 · p. 2
  • The iterative adversarial alignment process converges to a Nash equilibrium, and experiments in harmful-input and jailbreak settings report enhanced generalization and thorough training for both agents.The adversarial agent also serves as an adaptive red-teaming partner that continuously generates challenging prompts.Source: §1 · p. 2

2 Preliminary

This section introduces RLHF’s stages and notation, then motivates dynamically adapting its prompt set to the defensive model’s current weaknesses. It also identifies limits of existing human-written or rule-based prompt construction.

2 Preliminary

  • The section defines a prompt as a token sequence x and an autoregressive response y, with X and Y denoting the sets of possible prompts and responses.Source: §2 · p. 2
  • RLHF optimizes a policy on a specific prompt set using reward-model feedback, aiming to maximize rewards without drifting too far from the supervised fine-tuned model.The workflow comprises supervised fine-tuning, reward modeling, and RL-based policy optimization.Source: §2 · pp. 2–3
  • Prompt-set quality and diversity affect real-world effectiveness, and each optimization cycle should target the policy’s current weaknesses.Human-written and rule-based prompts cannot ensure comprehensive coverage and adaptivity.Source: §2 · p. 3

3 Game-theoretical Preference Optimization (GPO)

Game-theoretical Preference Optimization (GPO) aligns two LLM agents through repeated adversarial prompting and defensive adaptation. It defines their objectives, applies the game to safety, and analyzes equilibrium convergence and the role of diversity.

3 Game-theoretical Preference Optimization (GPO)

  • GPO iteratively pairs an adversarial agent that exposes defensive weaknesses with a defensive agent that adapts to those prompts.The adversary minimizes response reward while maximizing diversity; the defender maximizes reward on the adversary's prompt distribution.Source: §3 · p. 3, §3.1 · p. 4

3.1 A two-agent game framework for alignment

  • The framework models the agents as a zero-sum game: the adversary minimizes response reward minus diversity reward, while the defender maximizes it.The response reward comes from the reward model, and the diversity reward is higher for less-common prompts.Source: §3.1 · pp. 3–4
  • In practice, PPO updates both agents iteratively with a KL regularizer between current and previous policies for training stability.The defender trains on prompts from the adversary's previous-round distribution; the adversary also prioritizes diverse prompts to avoid narrow focus.Source: §3.1 · p. 4

3.2 Application of two-agent alignment in improving LLM safety

  • The safety application uses adversarial red-teaming to generate attack prompts and trains the defensive agent to withstand them.The response-quality reward is the safety level of the output for a given input, typically estimated by a toxicity classifier.Source: §3.2 · pp. 4–5, §3.2.1 · p. 5

3.2.2 Diversity rewards

  • Prompt diversity averages negative SelfBLEU for surface-form variation and embedding similarity for semantic variation.The approach compares each generated prompt with previously generated attack prompts.Source: §3.2.2 · p. 5

3.3 Theoretical analysis

  • Theoretical analysis shows that Algorithm 1 finds an O(T^-1/2)-approximate Nash equilibrium in T iterations, with convergence to Nash equilibrium in the infinite-iteration limit.The guarantee applies to the average policies of the theoretical version of the algorithm.Source: §3.3 · p. 6
  • Without the entropy regularizer, the adversary converges to a point distribution on the prompt with the lowest expected defensive reward, despite the same theoretical guarantee.The paper also analyzes an entropy-based diversity variant under mild assumptions.Source: §3.3 · p. 6

4 Experiments

The experiments evaluate the two-player approach on defensive safety, adversarial attack generation, and jailbreak scenarios. The reported comparisons favor GPO-line methods over RLHF or single-round red-team methods, while analyses examine diversity-reward intensity and decoding temperature.

4 Experiments

  • The experiments assess safety using three datasets and examine adversarial prompt generation, including jailbreak conversion and attack effectiveness.The evaluation spans general conversation and jailbreak scenarios.Source: §4 · p. 6, §4.1 · p. 8

4.1 Main results

  • GPO-line methods improve defensive safety over RLHF, with lower ASR and higher safe reward; adding diversity rewards further enhances performance.Table 1 reports safety evaluations across three datasets of harmful prompts.Source: Table 1 · p. 7, §4.1 · p. 7
  • GPO-line adversarial agents show stronger attacks than single-round red-team LLMs, producing diverse prompts effective across three target models.The evaluation averages results across Llama-2-7b-chat, vicuna-7b-v1.5, and a model trained with standard RLHF.Source: Table 2 · p. 7, §4.1 · p. 8

4.2 Analysis and discussion

  • Across diversity-reward intensities, the method outperforms RLHF, while a moderate intensity is more effective.The agents train alternately; reward intensity affects generated-prompt harmfulness and, in turn, defensive safety.Source: §4.2 · pp. 8–9
  • Across decoding temperatures from 0.1 to 1 in steps of 0.1, the method is more robust on safety tasks than SFT.SFT safety declines rapidly at very low or high temperatures.Source: §4.2 · p. 9

4.1 Main results

  • In jailbreak evaluation, GPO-line methods outperform other alignment methods on prompts from Salad-Bench’s Attack Enhanced subset.Training and test sets are split by attack method, with less common attack types reserved for testing.Source: §4.1 · p. 8, Table 3 · p. 9

5 Related Work

Prior alignment work includes reward-model-based reinforcement and preference learning, as well as methods that learn directly from preference data. Safety alignment also uses red teaming and model hardening, while this paper frames attacker and defender optimization as a two-player game.

Safety Alignment

  • Figure 3 reports more stable alignment performance for the proposed method than SFT under temperature sampling.Source: Figure 3 · p. 9

LLM Alignment

  • RLHF uses reinforcement learning to maximize reward-model signals, while RSO, RRHF, and RAFT also rely on reward-model scoring; DPO, SLiC-HF, and IPO learn from preferences directly.The passage says direct preference methods are susceptible to out-of-distribution data [24] [51] [44].Source: §5 · p. 9

Safety Alignment

  • Safety alignment commonly uses iterative red teaming and model hardening, but prior red-team LMs target static models and MART relies on supervised fine-tuning [14].The passage says MART has difficulty balancing attacker and defender capabilities.Source: §5 · p. 9

6 Conclusion, Limitation and Future Work

The conclusion presents alignment as an iterative two-player game in which adversarial prompts expose defensive weaknesses and the defensive agent improves its responses. The authors report enhanced generalization for both agents and identify safety-task scope and the need to train two agents as limitations.

6 Conclusion, Limitation and Future Work

  • The framework iteratively pairs an adversarial agent that generates diverse, challenging prompts with a defensive agent that improves its responses.The adversary seeks to uncover the defensive LLM’s weaknesses.Source: §6 · pp. 9–10
  • The approach enhances both agents’ generalization capabilities and is evaluated in harmful-input and jailbreak scenarios.The conclusion also reports that the method ensures thorough training.Source: §6 · p. 10
  • The solution requires training two separate LLM agents and was prototyped primarily on safety-related tasks.The authors identify helpfulness and mathematical reasoning as future domains to investigate.Source: §6 · p. 10

A Theoretical Analysis

This section completes the theoretical analysis by establishing notation for distributions over a set and mappings from another set to such distributions.

A Theoretical Analysis

  • For a non-empty set Z, ∆(Z) denotes distributions on Z, while ∆(Z | Z′) denotes mappings from Z′ to distributions on Z.Source: A Theoretical Analysis · p. 14

A.1 A Theoretical Analysis of Algorithm 2

The analysis of Algorithm 2 uses average policies and defines each agent’s regret as a gap against its best response. It then uses closed-form policy updates and assumptions about optimization to analyze those regrets.

A.1 A Theoretical Analysis of Algorithm 2

  • Algorithm 2’s theoretical analysis evaluates average policies over T iterations rather than the final-iteration policies.The analysis also ignores optimization error and assumes the agents attain the specified maxima and minima.Source: §A.1 · p. 14
  • The analysis defines defensive and adversarial regret as each learned average policy’s performance gap against the corresponding best response.The defensive regret compares J(π†, bµT) with J(bπT, bµT); adversarial regret compares J(bπT, bµT) with J(bπT, µ†).Source: §A.1 · p. 14
  • A Gibbs-distribution lemma supplies closed-form updates for the defensive and adversarial agents, which are then used to rewrite their regrets.The passage identifies the defensive update and adversarial update as equations (3.2) and (3.3), respectively.Source: §A.1 · pp. 14–15

A.2 A Theoretical Analysis of the Diversity Reward

The diversity-reward analysis studies an iteration-dependent reward that favors prompts distinct from those generated in the previous iteration. Its convergence result is stated under a lower-bound condition on supported prompt probabilities.

A.2 A Theoretical Analysis of the Diversity Reward

  • The case study sets the adversary’s diversity reward to log(µt−1(x)), making its negative expected reward the cross entropy between µ and µt−1.This reward encourages prompts distinct from the previous iteration’s prompts and serves as a proxy for the practical diversity reward.Source: §A.2 · p. 16
  • Algorithm 3 is the theoretical two-agent algorithm with an entropy regularizer, and its analysis optimizes an objective at every iteration.Source: §A.2 · p. 16
  • Theorem A.4 states a convergence result for Algorithm 3’s average policies under Assumption A.3 and appropriately chosen β and η.Assumption A.3 lower-bounds µt(x) for each prompt x on its support; the passage does not provide the theorem’s bound.Source: §A.2 · p. 16

A.3.1 Proof of Lemma A.1

The appendix gives fragments of proofs for auxiliary optimization results: one states where a maximum is attained, another concludes Lemma A.2, and a final lemma establishes equivalence of maximin and minimax objectives.

A.3.1 Proof of Lemma A.1

  • The proof of Lemma A.1 states that the maximum is attained at p = q.Source: §A.3.1 · p. 17

A.3.2 Proof of Lemma A.2

  • The proof of Lemma A.2 combines bounds from its two terms to conclude the lemma.The displayed proof fragments invoke KL divergence, Pinsker’s inequality, and Hölder’s inequality.Source: §A.3.2 · p. 18

A.4 Auxiliary Lemmas

  • The maximin and minimax objectives are equivalent in the paper’s setup.The proof invokes a minimax theorem and notes that J(π, µ) is linear in both arguments.Source: §A.4 · p. 18

B.1 Evaluation metrics

This section describes the evaluation measures for safety and attack-prompt diversity, including a case study of diversity rewards, and reports checklist items on societal impacts, asset documentation, participant instructions, and risks to human labellers.

B.1 Evaluation metrics

  • Llama-Guard [21] classifies outputs using first-token safe or unsafe probabilities above 0.5, with ASR measuring unsafe outputs; defenders favor lower ASR and higher safe reward, while attackers favor the reverse.The safe reward is the probability of the “safe” token, and the unsafe reward is the probability of the “unsafe” token.Source: §B.1 · p. 19
  • Attack-prompt diversity averages normalized SelfBLEU scores for n-grams n = 2–5 and embedding diversity across test-set data, with higher values indicating greater diversity [52] [36].SelfBLEU measures text diversity and sentence embeddings measure semantic diversity.Source: §B.1 · p. 19

B.2 Hyperparameters

  • GPO trains the defensive agent for 200 steps and the adversarial agent for 400 steps per iteration, with the defender training first and the agents alternating thereafter.The diversity reward coefficient is 5, and RLHF’s total training steps were constrained to match the proposed method.Source: §B.2 · pp. 19–20

B.3 Supervised Fine-tuning of Adversarial Agent

  • For general-dialogue attacks, adversarial-agent SFT used 20,000 paraphrased samples, while jailbreak SFT used 2,239 Attack Enhanced examples to teach jailbreak generation.The general-dialogue data were rewritten from Anthropic red-teaming attempts; the jailbreak samples came from Salad-Bench [23].Source: §B.3 · p. 20

B.4 Reinforcement Learning for Adversarial Agent

  • The adversarial agent’s PPO phase used 40,000 rewritten harmful data points as reinforcement-learning inputs.Source: §B.4 · p. 21

B.5 Case Study

  • With the diversity reward, the adversarial agent generates more diverse and aggressive prompts; without it, it tends to repeat one pattern, while diverse prompts elicit more reasonable defensive responses.The case study contrasts repeated “dispose of a body” attacks with more varied prompts and responses rather than simple refusals.Source: §B.5 · p. 21, Table 4 · p. 21

NeurIPS Paper Checklist

  • The authors report statistical significance for experiments except those requiring human labeling, for which they report one labeling procedure because of its high cost.Source: NeurIPS Paper Checklist · p. 25
  • The authors describe their work as advancing language-model safety and report no foreseeable negative societal impacts.They say their approach downweights unreliable prompt-response pairs during policy optimization to better align models with user preferences.Source: NeurIPS Paper Checklist · p. 26
  • The checklist answers No to whether new assets are well documented and documentation is provided alongside them.Source: NeurIPS Paper Checklist · p. 27
  • The authors say full instructions for the crowdsourcing experiments are included in the appendix.Source: NeurIPS Paper Checklist · p. 28
  • The checklist answers No to the question about describing participant risks and IRB approval, with the justification that human labellers faced no risks.Source: NeurIPS Paper Checklist · p. 28

Related papers

Loading 2406.10977v1…