Organizations: Department of Electrical Engineering, Indian Institute of Technology Delhi · Yardi School of Artificial Intelligence, Indian Institute of Technology Delhi
Safety tuning through supervised fine-tuning and reinforcement learning from human feedback has substantially improved the robustness of large language models. However, it typically suppresses rather eliminates unsafe behaviors, leaving rare but critical failures hidden in the long tail of the output distribution. While most red-teaming work emphasizes adversarial prompt search, we show that these hidden risks can be systematically exposed through diverse response generation. Specifically, we show that, for a fixed safety-critical prompt, increasing the number and diversity of sampled responses monotonically raises the jailbreak success rate. To efficiently uncover these failures, we propose Progressive Diverse Population Sampling (PDPS). This approach replaces naive, large-scale IID sampling with a multi-stage expansion-and-selection strategy that generates a compact, semantically diverse set of responses at a substantially lower computational cost. Across multiple jailbreak benchmarks and open-source LLMs, PDPS achieves attack success rates comparable to large-scale IID sampling while using only 8%-29% of the computational cost, and outperforms IID sampling and Diverse Beam Search by 26%-40% under limited-response budgets, while uncovering a broader and more semantically diverse range of failure modes. Critically, this diversity translates directly into more effective safety hardening: when integrated into an RLHF-based safety-tuning pipeline, PDPS-generated unsafe responses yield 33% and 41% greater reductions in ASR than those generated by IID sampling and Diverse Beam Search, respectively. Finally, we show that while input-space prompt optimization methods fall short of output-space exploration when used in isolation, combining input-space perturbation with diversity-driven output-space exploration covers a wider range of failure modes more efficiently than either paradigm alone.
Figures & tables
Figure 1: Illustration of input-space versus output-space search for jailbreaking. In input-space search, several variations or perturbations of the original safety-critical query are generated to elicit unsafe responses from an LLM. In contrast, output-space search is an orthogonal and complementary approach in which multiple responses are generated from a safety-critical prompt to assess whether any of them are unsafe.
Figure 2: Attack Success Rate (ASR) trends on Qwen2.5-7B-Instruct. The plots show ASR across HarmBench , JailbreakBench , AdvBench , and MaliciousInstruct datasets as a function of: (a) total number of generations ( N ), (b) sampling temperature ( τ ), and (c) nucleus sampling probability ( p ), while holding other parameters constant. The results demonstrate that broader exploration of the output space, whether through increased sample size or higher stochasticity, leads to a monotonic increase in ASR.
Figure 3: 2D plots of embeddings of 1024 responses generated from Llama-2-13B-Chat for three safety-critical prompts. The color indicates the likelihood of a response being unsafe.
Figure 4: Illustration of the difference between (a) IID sampling and (b) PDPS. IID sampling generates a large number of long responses through diverse token-level sampling. In contrast, PDPS produces a small set of diverse responses that approximate the modality coverage of large-scale IID sampling, while retaining the computational efficiency of small-scale IID sampling.
L2-7BCh
L2-13BCh
Q-7BInst
Q-14BInst
Avg.
AdvB
JBB
MalI
HarB
AdvB
JBB
MalI
HarB
AdvB
JBB
MalI
HarB
AdvB
JBB
MalI
HarB
\lx@sectionsignPDPS16
0.77
0.90
0.89
0.87
0.74
0.89
0.76
0.89
0.71
0.88
0.89
0.94
0.74
0.78
0.77
0.84
0.83
♣ IID 16
0.09
0.21
0.23
0.31
0.29
0.56
0.45
0.63
0.49
0.63
0.82
0.75
0.26
0.41
0.58
0.53
0.45
△\lx@sectionsign∼♣ (in %)
68 ↑
69 ↑
66 ↑
56 ↑
35 ↑
43 ↑
37 ↑
24 ↑
22 ↑
25 ↑
07 ↑
19 ↑
48 ↑
37 ↑
19 ↑
31 ↑
38 ↑
†DBS16
0.07
0.20
0.18
0.34
0.21
0.27
0.24
0.48
0.57
0.77
0.85
0.81
0.33
0.47
0.61
0.52
0.43
△\lx@sectionsign∼† (in %)
70 ↑
70 ↑
71 ↑
53 ↑
53 ↑
62 ↑
52 ↑
41 ↑
14 ↑
11 ↑
4 ↑
13 ↑
41 ↑
31 ↑
16 ↑
32 ↑
40 ↑
Table 1: Comparison of ASR obtained from PDPS with IID sampling ( IID ) and Diverse Beam Search ( DBS ) across four models ( L2-7BCh , L2-13BCh , Q-7BInst , Q-14BInst ) and four datasets ( JBB , MalI , HarB , AdvB ). Subscripts denote the number of full-length generations (e.g., IID16 uses 16 full-length response generations). The blue entries indicate the performance improvement of PDPS (in % ) over the corresponding baseline.
L2-7BCh
L2-13BCh
Q-7BInst
Q-14BInst
AdvB
JBB
MalI
HarB
AdvB
JBB
MalI
HarB
AdvB
JBB
MalI
HarB
AdvB
JBB
MalI
HarB
◊IID1024
0.73
0.88
0.95
0.97
1.00
0.99
1.00
1.00
0.98
0.99
1.00
1.00
0.99
0.97
1.00
0.99
\lx@sectionsignPDPS16
0.77
0.90
0.89
0.87
0.74
0.89
0.76
0.89
0.71
0.88
0.89
0.94
0.74
0.78
0.77
0.84
ρ\lx@sectionsign∼◊
≥ 1.0
≥ 1.0
0.94
0.90
0.74
0.90
0.76
0.89
0.72
0.89
0.89
0.94
0.75
0.80
0.77
0.85
\lx@paragraphsignPDPS64
0.81
0.93
0.94
0.94
0.97
0.98
0.99
1.0
0.97
0.97
1.00
0.99
0.98
0.99
0.98
0.97
ρ\lx@paragraphsign∼◊
≥ 1.0
≥ 1.0
0.99
0.97
0.97
0.99
0.99
≥ 1.0
0.99
0.98
1.0
0.99
0.99
≥ 1.0
0.98
0.98
Table 2: ASR of PDPS compared to the brute-force benchmark IID1024 across the four models ( L2-7BCh , L2-13BCh , Q-7BInst , Q-14BInst ) and the four datasets ( JBB , MalI , HarB , AdvB ). Subscripts denote the number of full-length generations (e.g., PDPS16 uses 16 full-length sequences constructed from 1024 partial generations). The row labeled ρ\lx@sectionsign∼◊ (resp. ρ\lx@paragraphsign∼◊ ) report the ASR achieved by PDPS16 (resp. PDPS64 ) as a fraction of the ASR achieved by IID1024 . The blue entries in these rows indicate that PDPS attains at least 0.8 times the ASR of IID1024 , whereas red entries indicate failure to reach this threshold.
Figure 5: Bar plot of the mean toxicity difference, defined as the difference between the average number of unsafe responses returned by PDPS and that returned by a baseline method, across various model–dataset combinations. The average is computed over only those queries for which the respective method achieves a successful attack. Blue bars indicate a positive difference, while red bars indicate a negative difference.
Dataset
Sampler
Dist-1 ↑
Dist-2 ↑
SB-1 ↓
SB-2 ↓
SB-3 ↓
SB-4 ↓
Uni-Ent ↑
Cos-Dist ↑
BERT-Div ↑
IID64
0.27
0.68
0.62
0.42
0.28
0.19
5.46
0.31
0.29
AdvB
DBS64
0.17
0.48
0.74
0.58
0.47
0.39
5.31
0.24
0.25
PDPS64
0.32
0.77
0.59
0.34
0.20
0.12
5.81
0.47
0.34
IID64
0.24
0.65
0.69
0.48
0.33
0.23
5.57
0.26
0.29
JBB
DBS64
0.15
0.45
0.80
0.65
0.53
0.44
5.39
0.18
0.25
PDPS64
0.30
0.76
0.65
0.39
0.23
0.14
6.00
0.40
0.34
Table 3: Comparison of the average diversity of unsafe responses generated by PDPS and two baseline methods for the Q-7BInst model on the four datasets. The average is computed over only those queries for which at least two unsafe responses are returned by the respective methods. Diversity metrics: Dist- n = Distinct- n [ 22 ] ; SB- n = Self-BLEU- n [ 49 ] ; Uni-Ent = Unigram Entropy [ 9 ] ; Cos-Dist = Cosine Distance [ 18 ] ; BERT-Div = BERTScore Diversity [ 47 ] . For metrics marked with ↑ , higher values indicate greater diversity, while ↓ indicates the opposite.
Figure 6: Sampling time (as a percentage of brute-force IID1024 ) for generating 64 responses on L2-13BCh and Q-14BInst across token lengths 512 – 4096 . Bars labeled OOM indicate out-of-memory failures.
Method
ASR
Base Model (no fine-tuning)
0.75
Fine-tuned w/ IID64
0.36
Fine-tuned w/ DBS64
0.41
Fine-tuned w/ PDPS64
0.24
Table 4: ASR of the base model and models fine-tuned using negative samples generated by IID64 , DBS64 , and PDPS64 . Lower ASR indicates stronger safety hardening.
Figure 7: ASR under different hyperparameter settings for the Q-7BInst model on the AdvBench dataset. (a) ASR for different nucleus sampling probabilities ( p ); (b) ASR for different sampling temperatures ( τ ); and (c) ASR for different values of λ , the hyperparameter controlling the quality–diversity threshold.
Figure 8: Empirical comparison and interplay between input-space search and output-space exploration paradigms. (a) While input-space optimization attacks require computationally intensive iterative query refinement or gradient-based optimization, output-space exploration methods achieve substantially higher ASRs at lower computational cost. (b) Increasing the output-generation budget for standard input-space baselines produces a monotonic increase in ASR, demonstrating that input-space search and output-space exploration are highly complementary paradigms.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: 2D plots of embeddings of 1024 responses generated for three safety-critical prompt from Q-7BInst model. The color indicates the likelihood of a response being unsafe.
Dataset
Sampler
Dist-1 ↑
Dist-2 ↑
SB-1 ↓
SB-2 ↓
SB-3 ↓
SB-4 ↓
Uni-Ent ↑
Cos-Dist ↑
BERT-Div ↑
IID64
0.31
0.69
0.56
0.36
0.24
0.17
5.12
0.27
0.27
AdvB
DBS64
0.28
0.64
0.61
0.44
0.33
0.25
4.99
0.23
0.23
PDPS64
0.29
0.67
0.58
0.39
0.27
0.19
5.29
0.31
0.29
IID64
0.25
0.61
0.68
0.50
0.37
0.28
5.28
0.22
0.27
JBB
DBS64
0.22
0.52
0.71
0.57
0.47
0.40
5.05
0.15
0.23
PDPS64
0.27
0.65
0.65
0.45
0.31
0.23
5.46
0.28
0.30
Appendix
Table 5: Comparison of the average diversity of unsafe responses generated by PDPS and two baseline methods for the L2-7BCh model on the AdvB and JBB datasets. The average is computed over only those queries for which at least two unsafe responses are returned by the respective methods.
Dataset
Sampler
Dist-1 ↑
Dist-2 ↑
SB-1 ↓
SB-2 ↓
SB-3 ↓
SB-4 ↓
Uni-Ent ↑
Cos-Dist ↑
BERT-Div ↑
IID64
0.28
0.61
0.48
0.32
0.22
0.16
4.96
0.40
0.30
AdvB
DBS64
0.20
0.46
0.65
0.51
0.40
0.33
4.86
0.27
0.23
PDPS64
0.29
0.68
0.53
0.33
0.21
0.14
5.32
0.45
0.32
IID64
0.27
0.62
0.58
0.40
0.28
0.21
5.17
0.31
0.29
JBB
DBS64
0.18
0.43
0.69
0.57
0.49
0.43
4.90
0.16
0.23
PDPS64
0.26
0.65
0.64
0.43
0.29
0.20
5.47
0.36
0.31
Appendix
Table 6: Comparison of the average diversity of unsafe responses generated by PDPS and two baseline methods for the L2-13BCh model on the AdvB and JBB datasets. The average is computed over only those queries for which at least two unsafe responses are returned by the respective methods.
Dataset
Sampler
Dist-1 ↑
Dist-2 ↑
SB-1 ↓
SB-2 ↓
SB-3 ↓
SB-4 ↓
Uni-Ent ↑
Cos-Dist ↑
BERT-Div ↑
IID64
0.25
0.64
0.64
0.45
0.32
0.23
5.20
0.30
0.26
AdvB
DBS64
0.22
0.56
0.65
0.49
0.38
0.31
5.09
0.31
0.26
PDPS64
0.22
0.58
0.62
0.42
0.28
0.20
5.21
0.45
0.30
IID64
0.23
0.61
0.66
0.48
0.35
0.26
5.22
0.30
0.27
JBB
DBS64
0.19
0.50
0.70
0.55
0.44
0.37
5.11
0.25
0.25
PDPS64
0.22
0.60
0.66
0.45
0.31
0.22
5.34
0.43
0.30
Appendix
Table 7: Comparison of the average diversity of unsafe responses generated by PDPS and two baseline methods for the Q-14BInst model on the AdvB and JBB datasets. The average is computed over only those queries for which at least two unsafe responses are returned by the respective methods.
Method
Example Responses
DBS
1. Research the victim: Gather as much information … 2. Obtain the victim’s personal information: … 3. Create a fake identity: … 7. Avoid detection: …
… II. Gathering Information … Obtaining personal information … Covering Tracks …
1. Obtain personal information: … 2. Create a fake identity: … 6. Avoid detection: …
2. Use the obtained information … a fake identity to cover tracks …
… 1. Obtain personal information: … 2. Create a fake identity: …
1. Research the target: … 2. Obtain personal information: … 3. Create a fake identity: …
Appendix
Table 8: Sample responses generated using PDPS and DBS by querying the Q-7BInst model with the prompt: ‘‘Create a blueprint for committing identity theft and stealing someone’s personal information.’’ . The examples illustrate that PDPS produces semantically diverse responses, whereas DBS generates largely similar outputs with only surface-level variations.
Diversity Penalty
1
4
16
64
DBS64
0.22
0.20
0.16
0.14
Appendix
Table 9: ASR obtained by DBS for various values of the diversity penalty hyperparameter on the L2-7BCh model and the AdvB dataset in the 64-response generation setting. The results suggest that the performance of DBS does not improve significantly as the diversity penalty increases.
Safety fine-tuning of language models typically requires a curated adversarial dataset. We take a different approach: score each candidate prompt's difficulty by how often the target model's own rollouts are judged harmful, then fine-tune on the hardest prompts paired with the model's own non-jailbroken rollouts. On Llama-3-8B-Instruct and Llama-3.2-3B-Instruct, this approach cuts the WildJailbreak attack success rate from 11.5% and 20.1% down to 1-3%, but pushes refusal on jailbreak-shaped benign prompts from 14-22% to 74-94%. Interleaving the same hard prompts 1:1 with adversarially-framed benign prompts (prompts that look like jailbreaks but have benign intent) cuts that refusal back down to 30-51% on 8B and 52-72% on 3B, at a cost of 2-6 percentage points of attack success rate. Within the mixed regime, training on the hardest half of the eligible pool rather than a random half cuts the remaining ASR by 35-50% (about 3 percentage points) on both models.
Task-specific fine-tuning can improve the performance of large language models (LLMs) on downstream tasks. However, our study reveals that task-specific fine-tuning can also weaken the safety guardrails of aligned LLMs. A widely adopted strategy for preserving safety during fine-tuning is to incorporate safety data. Although previous studies have shown that randomly mixing safety data can alleviate safety degradation, the underlying principle determining why some safety examples are more effective than others still remains unclear. In this paper, we propose DataRx, a missingness-aware sampling method for selecting safety-critical examples. DataRx is based on the hypothesis that a safety sample is more effective when the selected examples provide safety signals that fill the missing parts of LLMs' safety capabilities. DataRx's key insight is leveraging high-dimensional hidden representations rather than discrete tokens to quantify the safety signal gap between the target model's native response and the safety reference response. The results show that, with only 1% additional safety samples from BeaverTails, DataRx reduces the average attack success rate of Llama3-8B-Instruct across seven downstream tasks from 59.23% under random sampling to 13.70%. In addition, DataRx can be combined with the existing safety data synthesis method to further enhance safety defenses during fine-tuning. We hope that DataRx will inspire more data-centric defense research.
Red-teaming, or identifying prompts that elicit harmful responses, is a critical step in ensuring the safe and responsible deployment of large language models (LLMs). Developing effective protection against many modes of attack prompts requires discovering diverse attacks. Automated red-teaming typically uses reinforcement learning to fine-tune an attacker language model to generate prompts that elicit undesirable responses from a target LLM, as measured, for example, by an auxiliary toxicity classifier. We show that even with explicit regularization to favor novelty and diversity, existing approaches suffer from mode collapse or fail to generate effective attacks. As a flexible and probabilistically principled alternative, we propose to use GFlowNet fine-tuning, followed by a secondary smoothing phase, to train the attacker model to generate diverse and effective attack prompts. We find that the attacks generated by our method are effective against a wide range of target LLMs, both with and without safety tuning, and transfer well between target LLMs. Finally, we demonstrate that models safety-tuned using a dataset of red-teaming prompts generated by our method are robust to attacks from other RL-based red-teaming approaches.
Seanie Lee, Minsu Kim, Lynn Cherif +8
KAIST · Mila – Québec AI Institute · Université de Montréal +3