SEAL: Safety-enhanced Aligned LLM Fine-tuning via Bilevel Data Selection

Thank you for the careful review and support. Our response to your questions follows.

1. It is recommended that the authors conduct an ablation study comparing LoRA fine-tuning that incorporates weights in both the Feed Forward Network (FFN) and Multi-Head Self-Attention (MHSA) blocks with their current PEFT approach. Additionally, an analysis discussing how various LoRA configurations affect the safety-performance trade-off would be beneficial.

Thank you for the valuable suggestion. Indeed, the fine-tuning method can play a role in the performance discrepancy between different baselines and models. Due to the limitation of computational resource, in the main paper, we focused our experiments on the effectiveness and computational complexity of SEAL against other baselines. While an ablation study on how LoRA affects the performance of SEAL is certainly interesting and worth considering in the revision.

2. The author used the memory-efficient variant of SEAL by default in the experiments. To provide a more comprehensive understanding of the impact of this choice, I suggest the authors perform a quantitative comparison of memory consumption and performance between the standard and memory-efficient SEAL variants.

Great suggestion! Comparing the performance of the standard and memory efficient algorithm is the most important future direction of this branch of work. As compared to the memory efficient variant, the standard algorithm needs to additionally compute the value function $v(x)$ and its gradient. Currently, we find the standard algorithm difficult to run on our setup when the model size increases beyond 7 billion. This motivated us to propose the memory efficient variant which is more scalable to model size. In order to test standard algorithm, we believe additional implementation side techniques can be used. For example, one may offload the model to CPU when computing the value function $v(x)$ and its gradient. However, offloading can also increase the overall time complexity. It is indeed an interesting future direction, and can potentially further improves the performance of SEAL.

Thank you for the careful review and support. Our response to your questions follows.

1. Could you compare the results of training the selector and LLMs separately? It is intuitive that this method seems to be more effective.

We are sorry for the confusion. LLM fine-tuning and selector fine-tuning are already separate steps in SEAL workflow as described in Section 3.3. Specifically, we first train a data selector via solving the bilevel data selection formulation. Then note that we do not use the output model of the bilevel algorithm, we fine tune the aligned model on the task data filtered by the data selector.

2. Could you further explain the results in Figure 3 and Figure 4. I am confused about why the performance of SEAL on the SLIMORCA dataset is not always superior to the baseline DSIR (Sometimes even worse).

In Figure 3, SEAL outperforms DSIR on SlimOrca test dataset, while in Figure 4, SEAL is slightly outperformed by DSIR on SlimOrca test. The major reason is that performance of SEAL is dependent on the aligned model's architecture and parameters. Specifically, the aligned model determines the initial point of both SEAL's data selection training and the fine-tuning process. Since Figure 4 and 3 are tested on different models, variation in performance is expected. However, it is worth noting that SEAL is able to outperform DSIR in terms of average win rate on the three test datasets across all the models tested. We believe the average win rate on multiple datasets is a better metric than win rate on a single test dataset, as it showcases a model's capability on multiple domains more comprehensively.

3. From the current results, it is difficult to know whether the win rate improvements of SEAL are due to enhanced safety or better performance on downstream tasks. Could you provide more explanation about which datasets focus more on safety evaluation versus performance evaluation? Additionally, can you present more distinguishable results to show the model comparisons in terms of safety and performance, respectively?

We are sorry for the confusion. In this work, we mainly use three test datasets: Anthropic Helpful and Harmless (HH) dataset, which contains a mixture of samples that can be used to evaluate either the helpfulness (instruction-following performance) or the harmlessness (safety) of the model; SlimOrca dataset, which contains instruction prompts and can be used to evaluate the task performance of the model; and HEx-PHI dataset containing red-teaming prompts that can be used to evaluate the safety of the model. Therefore, the model's fine-tuning performance and safety can be directly observed from the figures and plots in the main paper. For example, in Figure 3, we observe that SEAL is able to outperform all baselines on HH and HEx-PHI, and achieves similar performance to DSIR on SlimOrca. While in terms of average win rate on all test datasets, SEAL outperforms all baselines. Thus in this case we conclude that SEAL achieves comparable task performance to other methods while being able to better preserve safety. Overall, SEAL is able to output a stronger model.

Thank you again for the careful review.

Thank you for the careful review. Our response to your questions follows.

1. Overall, the article's approach lacks novelty, both in the idea itself and in the optimization method.

We respectfully disagree. The use scenario of SEAL is LLM fine-tuning, which is different from the vision tasks traditionally performed in previous bilevel works (e.g., [1][2]). Due to the different data format, loss types, model architectures and even evaluation process, it was unclear how the traditional bilevel data re-weighting approach can be used in LLM training, let alone effectively. Due to the immense size of the language models, we find that the existing method in [1][2] are inefficient in LLM tasks. Therefore, we proposed more computationally efficient variants in Algorithm 1&2 where we utilized a fully single-loop structure and saved the computational cost of computing the value function's gradient. In the experiments, we showcase the effectiveness of bilevel LLM data selection through both win rate comparison and quality analysis for the first time. We also tested the transferability of the data selectors, which was not present in previous bilevel works.

[1] Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and reverse gradient-based hyperparameter optimization. arXiv:1703.01785, 2017.

[2] Han Shen and Tianyi Chen. On penalty-based bilevel gradient descent method. ICML, 2023.

2. For all experiments, there was no adjustment for the variance of performance given after randomization of seeds. Especially for the randomized selection method, a large number of replicated experiments are necessary.

For the random selection baselines, the results reported in the paper are averaged over three random seeds. For other baselines, we are conducting experiments to test the variance between runs, and will share the results in future versions.

3. I don't know why the authors didn't design additional but direct baselines, such as directly letting the model be given weights or rankings through the prompt engineering.

Prompting LLMs to output absolute weights or rankings might have some variance and context length issues, while it is interesting to consider efficient implementations of such methods. We will consider adding them in the revisions.

4. The results now lack insight, for example, the HEX-PHI data actually contains 11 categories, and it would be interesting to analyze the different categories in detail and explore the limitations of the sample selection algorithm.

Thank you for the suggestion. In the revised paper, we have added Section (B.5) where the categorized win rate comparison on the Llama-3-8B-Instruct model is provided, Hope this resolves the question.

5. LoRA introduces additional parameters, so it is difficult to directly attribute for performance improvements. Therefore, I believe that full parameter fine-tuning should be used for all methods. Also, the models do not perform well. Sometimes it is not even better than Random.

We have tested full-parameter in Section B.1 on the Pythia-2.8b model. Only for the larger model like Llama-3, we use LoRA to save computational cost. It is worth noting that LoRA is widely used to fine-tune larger models like Llama family [1][2], and it has been observed in these works that the performance degradation is minimal as compared to full-parameter fine-tuning.

Regarding the performance compared to random, it is worth noting that across all models tested, our method consistently outperforms random selection in terms of win rate averaged over three test datasets. Compared to random, SEAL’s average performance gain on LLAMA-3-8B-INSTRUCT is around 8.5% in Table1, and is around 9.8% on MERLINITE-7B in Table 2. On specific test datasets like HEx-PHI, SEAL outperforms random by 10% on both Merlinite-7b and Llama-3-8b, and is only 1% worse than random in either benign tests on Llama-2-7b or small model tests on Pythia-2.8b. While we argue that average test performance should be preferred over performance on one single test dataset, as it more comprehensively demonstrates a model's capability over multiple domains.

[1] CY Hsu, YL Tsai, CH Lin, PY Chen, CM Yu, CY Huang. Safe LoRA: the Silver Lining of Reducing Safety Risks when Fine-tuning Large Language Models. NeurIPS, 2024.
[2] Xiangyu Qi, Yi Zeng, Tinghao Xie, Pin-Yu Chen, Ruoxi Jia, Prateek Mittal, and Peter Henderson. Fine-tuning aligned language models compromises safety, even when users donot intend to! arXiv preprint arXiv:2310.03693,2023.

Thank you again for the careful review.

Thank you for the careful review. Our response to your comments follows.

Weakness: One potential drawback of this paper is it assume the safe dataset is readily available in the supervised finetuning process. With the safe dataset available, a straightforward idea would be directly optimize the model with a combination of safe data and task data.

Our method is designed to filter the unsafe or low-quality samples in the task/fine-tuning data, and output a better task dataset. We prove this point by showing in the experiments that the model fine-tuned on SEAL-selected task dataset is better across several test datasets than the model fine-tuned on the original task dataset. Then combining the safe data with the SEAL-selected task dataset will improve over simply combining the safe dataset and the unfiltered task dataset.

To support this point, within the limited time of the rebuttal period, we have conducted the suggested ablation study.

Due to the time limitation, we conducted these experiments on the Pythia-1b model. The experimental setup is the same as that in Section (B.1). In the table, safety+fine-tuning dataset (x% non-safe) represents fine-tunng on the combination of safety dataset and the fine-tuning dataset which contains of x% unsafe samples, and Safety+SEAL is fine-tuning on the combination of safety dataset and SEAL-selected fine-tuning dataset. The evaluation metric is the win rate against safety+fine-tuning dataset (x% non-safe) (see Section A.2 for description of win rate). The model is considered better if the win percent is higher. It can be observed that SEAL is able to consistently outperform directly combining the safety and task dataset without selection, across different unsafe ratios. This is because SEAL's data selection gets rid of the unsafe samples completely in the fine-tuning process, while directly combining safety data and task data will always suffer from the influence of these unsafe samples.

Q1. Since a safe dataset is already available in the training, why not just train the model with the objective in Equ. (4) to balance safety and task objective, without further effort on doing data selection?

This is because Equ. (4) still involves the unsafe samples, thus optimizing on a combination of safe data and non-selected task data can still lead to sub-optimal performance. In the paper, we show with win rate comparison in Section 4 and quality analysis in Section (B.2)&(B.3) that the SEAL is able to pick out unsafe samples in the task dataset. As a result, training with the filtered task data yields improved performance compared to training with the original task dataset, which is further supported by the added ablation study in our response to the previous question.

Q2. Why does the target domain win rate also improve as we select data based on safety alignment?

The reason is that seal is also able to filter the low-quality samples, e.g., see the quality analysis in Section (B.2) in the appendix, where SEAL filters the samples that have improper target answers. Since fine-tuning on these samples will damage the model's performance in target domain, filtering them out helps the performance.

Thank you again for the careful review. We hope that our response addresses your questions, potentially leading to a higher score. Please let us know if you have further questions.

Thank you for the valuable feedback. Our response to your comments follows.

1. An ablation study varying the proportion of safety-aligned vs non-safety-aligned data in the fine-tuning set to demonstrate SEAL's effectiveness across different scenarios (e.g., varying the proportion of non-safety-aligned data from 0% to 100%); Experiments comparing SEAL (𝛾>0) to a baseline (𝛾=0) across different ratios of safety-aligned to non-safety-aligned data.

Thank you for the great suggestions. Within the limited time of rebuttal period, we have conducted the suggested ablation studies and added the new Section (B.4) in the revised paper. For convenience, we also present the results and discussion below.

We run these experiments on Pythia-1b model with the same setup as that in Section (B.1). In the table, safety+fine-tuning dataset (x% non-safe) represents fine-tuning on the upper-level safety dataset combined with the original fine-tuning dataset with x% of non-safety-aligned samples; Safety dataset only ($\gamma$=0) represents fine-tuning on the upper-level safety dataset only, so it is equivalent to SEAL with $\gamma$=0; and Safety+SEAL($\gamma$>0) is fine-tuning on the mixture of upper-level safety dataset and SEAL-selected fine-tuning dataset. The metric used in this table is the win rate (see Section A.2 for description of win rate) against safety+fine-tuning dataset (x% non-safe) under each non-safe ratio. The model is considered better if the win rate is higher.

It can be observed that when the fine-tuning dataset is completely safe with $0\%$ unsafe ratio, only using the safe dataset ($\gamma$=0) has an disadvantage on the fine-tuning domain (SlimOrca) since the other two baselines additionally utilize the safe fine-tuning dataset. When the unsafe ratio is 10%, SEAL is able to achieve a noticeable average win rate advantage over the baselines. Although $\gamma=0$ improves performance on safety domain (HEx-PHI), it has worse performance on the fine-tuning domain (SlimOrca) than SEAL with $\gamma>0$. This is because $\gamma=0$ fails to utilize the safe subset of fine-tuning dataset, while SEAL is able to select a safe fine-tuning subset relatively well, which can be observed by its performance gain over no-selection. Similar observation can be made when the unsafe ratio increases to 50%, $\gamma=0$ is outperformed by SEAL ($\gamma>0$) on the fine-tuning's target domain (SlimOrca). It is worth noting that our goal is to improve the model's capability on the fine-tuning domain while preserving its safety as much as possible, we can see SEAL helps achieving this goal. Lastly, we would like to clarify that SEAL belongs to the data selection method family. The use scenario for data selection methods assumes at least part of the samples in the original fine-tuning dataset is usable. An example of such a use case is fine-tuning on large datasets that contain useful information but may not be entirely safe, and thus needs filtering. In this case, a 100% unsafe ratio leads to no motivation of fine-tuning at all. Therefore, 100% unsafe ratio is out-of-scope generally for all data selection methods.