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

## Abstract

Abstract Fine-tuning on task-specific data to boost downstream performance is a crucial step for leveraging Large Language Models (LLMs). However, previous studies have demonstrated that fine-tuning the models on several adversarial samples or even benign data can greatly comprise the model’s pre-equipped alignment and safety capabilities. In this work, we propose SEAL, a novel framework to enhance safety in LLM fine-tuning. SEAL learns a data ranker based on the bilevel optimization to up rank the safe and high-quality fine-tuning data and down rank the unsafe or low-quality ones. Models trained with SEAL demonstrate superior quality over multiple baselines, with 8.5 % percent 8.5 8.5\% 8.5 % and 9.7 % percent 9.7 9.7\% 9.7 % win rate increase compared to random selection respectively on Llama-3-8b-Instruct and Merlinite-7b models. Our code is available on github https://github.com/hanshen95/SEAL .

### 1 Introduction

Large language models (LLMs) trained on large text datasets have demonstrated astonishing capabilities in generative tasks (Dubey et al., ; Achiam et al., ). For example, these models can provide answer to human’s questions, generate and modify codes or solve mathematical problems (Qin et al., ; Suzgun et al., ; Gao et al., ). However, an unaligned LLM can lead to significant safety concerns (Bender et al., ; Bommasani et al., ; Wei et al., ), e.g., an LLM assistant can output unsafe suggestions when prompted by harmful-inducing instructions.
With the great capability of these models, the safety concern has become an increasingly urgent issue.

To mitigate the safety concerns of the LLMs, it is necessary to align the model before deployment using alignment methods like supervised fine-tuning (SFT), reinforcement learning from human feedback (RLHF) (Ouyang et al., ) or direct preference optimization (DPO) (Rafailov et al., ). However, the aligned models are brittle (Qi et al., ; Zou et al., ; Yang et al., ; Wei et al., ). Fine-tuning LLMs for very few epochs can significantly compromise the safety alignment of the model (Qi et al., ).
As fine-tuning LLMs for downstream applications has become standard practice, concerns around safety breaking during this process have emerged as a significant challenge. These issues pose a fundamental obstacle to the broader application of LLMs in various real-world scenarios, where ensuring model alignment and safe behavior is critical.

To this end, we propose an LLM fine-tuning framework that mitigates the negative impact on safety alignment during the fine-tuning process. We approach the problem from a data-centric perspective. Based on the observation that fine-tuning damages the performance of safety alignment (Qi et al., ), there are oftentimes conflicts between fitting the fine-tuning data and the alignment data, that is, decreasing the fitting loss on some fine-tuning data which are conflicting with the safe data will inevitably lead to the increase of the safety alignment loss.
With this intuition, we thereby: 1) formulate an optimization problem for learning a data selector which down-ranks the potentially unsafe samples in the fine-tuning dataset; 2) to solve for the data selector, we then introduce a gradient-based algorithm along with its memory-efficient variant; 3) we proceed to propose the Safety-Enhanced Aligned LLM fine-tuning (SEAL) framework (depicted in Figure 2).
An intuition of how SEAL’s data selection can overcome the conflict between safety alignment and fine-tuning is illustrated in Figure 1. SEAL’s selector is trained such that the model fine-tuned on selected samples fit safe alignment data well. With potentially harmful knowledge filtered out, the safety during fine-tuning can be enhanced.
SEAL demonstrates the following merits:

- •
Effectiveness: We evaluate SEAL on test datasets including Anthropic HH (Bai et al., ), Orca (Mukherjee et al., ) and HEx-PHI (Qi et al., ). SEAL consistently outperforms multiple baselines on Llama-3-8b-Instruct (Dubey et al., ), Merlinite-7b (Sudalairaj et al., ) and Pythia-2.8b (Biderman et al., ).
- •
Flexiblity: We find out that the SEAL-selected data is transferable to fine-tuning different models, e.g., one can use a different (potentially smaller) model in data selector training which is separate from the fine-tuning model. Additionally, though SEAL’s performance depends on its data selection percentage, we find that for a relatively wide range of data selection percentage, SEAL can achieve better performance than the baselines.
- •
Explainablity: We will give a quality comparison of SEAL-selected data and SEAL-filtered data samples. It can be observed that the selected data demonstrate superior safety quality than the filtered-out data, showing the explainability of SEAL’s effectiveness.

Use cases of SEAL.
The use cases for SEAL include scenarios such as a closed-source model owner offering fine-tuning services on user-provided data (e.g., OpenAI). Then a safe fine-tuning method is essential to ensure that the fine-tuned model maintains the original model’s safety alignment. The service providers have access to the safe alignment data (Dubey et al., ; Qin et al., ), which can be used to do SEAL’s data selection. Additionally, SEAL is suited when the fine-tuning dataset contains potentially harmful samples, but is too large for manual annotation. In this case, we assume the model owner can access an alignment/safe dataset (e.g., the open-source datasets) to select SEAL’s data.
In these aforementioned scenarios, SEAL provides a flexible solution for enhancing safety during the fine-tuning process.

Figure: Figure 1: Full SFT trains LLM equally on all samples (left), which might contain harmful knowledge. SEAL learns data selector $\sigma(\omega)$ that filters harmful samples (right), enhancing safety in fine-tuning.
Refer to caption: extracted/5918167/figure/seal_intro_example.png

### 2 Related works

Understanding of jail-breaking in LLM fine-tuning. To mitigate safety risks in fine-tuning, it is necessary to first have a thorough understanding of the behavior. There have been a number of works focusing on the conceptual understanding of LLM jail-breaking during fine-tuning, see, e.g., (Qi et al., ; Wolf et al., ; Wei et al., ; Anwar et al., ; Lee et al., ; Yao et al., ; Chen et al., ; Ball et al., ). Specifically, Qi et al. () observes that even fine-tuning on benign dataset can greatly compromise the safety alignment of LLMs. The work of (Wei et al., ) proposes identifies competing/conflicting objectives as one of the potential reasons.

Safe fine-tuning/alignment of LLMs. Safe fine-tuning/alignment methods seek to align the model with safety measures, so that the models do not demonstrate harmful behaviors after deployment. The topic is crucial in LLM applications, and there have been a large body of works recently; see, e.g., safety instruction fine-tuning (Ouyang et al., ; Bianchi et al., ), feature-preserving fine-tuning (Mukhoti et al., ), LLM alignment via preference learning (Rafailov et al., ; Noukhovitch et al., ; Ji et al., ; Rame et al., ; Go et al., ; Ethayarajh et al., ; Tang et al., ), self-play alignment (Chen et al., ), safe alignment against harmful fine-tuning (Huang et al., ; ), re-alignment at inference time (Liu et al., ) or by model fusion (Yi et al., ), alignment-preserving prompting (Lyu et al., ; Zheng et al., ), alignment brittleness accessing (Wei et al., ), representation-based defence (Rosati et al., ), landscape navigation (Peng et al., ) and safe LoRA fine-tuning (Hsu et al., ).

Figure: Figure 2: Overview of the SEAL framework. In contrast to vanilla fine-tuning (FT) where the LLM is trained on a fine-tuning dataset which potentially includes unsafe and low-quality data samples, SEAL first learns a data (sample) ranker by solving a bilevel optimization problem. Models fine-tuned on the high-ranked samples demonstrate superior quality.
Refer to caption: extracted/5918167/figure/seal.png

Data selection for LLM fine-tuning. Among the safe fine-tuning methods, the data selection approaches generally aim to select a subset of high-quality safe data to fine-tune the LLM on; see, e.g., (Chen et al., ; Cao et al., ; Xie et al., ; Lu et al., ) and (Xia et al., ; Engstrom et al., ; Zhao et al., ; He et al., ; Pan et al., ; Kang et al., ; Zhao et al., ).
More specifically, Chen et al. (); Cao et al. () use a strong model to filter out low-quality data;
Xie et al. () resamples the data by estimating each data point’s importance ratio;
Xia et al. () proposes LESS which uses gradient similarity scoring to select influential instruction data; (Kang et al., ) proposed a principled data selection method via optimal transport;
A concurrent work (Pan et al., ) proposes a bilevel method to learn weights for multiple data sources, which has achieved good performance. In contrast, our work focuses on selecting data points with varying degree of safety within the same data source, rather than reweighing different datasets.

Bilevel optimization algorithms.
The bilevel optimization (BLO) has been applied to various machine learning applications (Sabach & Shtern, ; Franceschi et al., ; Liu et al., ).
There are three prominent branches of bilevel optimization methods: the implicit gradient methods (Pedregosa, ; Ghadimi & Wang, ; Hong et al., ; Ji et al., ; Chen et al., ; Khanduri et al., ; Shen & Chen, ; Chen et al., ); the iterative differentiation methods (Franceschi et al., ; Liu et al., ; ; Bolte et al., ; Ji et al., ). These methods require second-order derivatives, thus is memory inefficient when the model parameter size is large.
On the other hand, the penalty relaxation of the BLO problem, which dates back to (Clarke, ; Luo et al., ), has gained interest from researchers recently (see, e.g., (Shen & Chen, ; Liu et al., ; Kwon et al., ; Xiao et al., ; Shen et al., ; Lu, )). Most penalty-based methods only require first-order derivatives, thus is more suitable for LLM applications.

### 3 Safety-enhanced LLM fine-tuning

In this section, we will present our method for safety-enhanced LLM fine-tuning that we call SEAL. We will introduce the formulation, the data selector training algorithm and the SEAL framework.

#### 3.1 Problem formulation

Denote a sample $z=(x,y)$ where $x$ is the input sequence (e.g., instructions) and $y=(y_{1},y_{2},\dots,y_{d_{y}})$ is the target response.
We also write $y_{<j}=(y_{1},...,y_{j})$ with $y_{<1}$ defined as an empty sequence.
Suppose we are given a safe dataset $\mathcal{D}_{\rm safe}=\{z_{\rm safe}^{i}=(x_{\rm safe}^{i},y_{\rm safe}^{i})
\}_{i=1}^{M}$ that contains safe target responses $\{y_{\rm safe}^{i}\}_{i=1}^{M}$. We also have the fine-tuning dataset $\mathcal{D}=\{z^{i}=(x^{i},y^{i})\}_{i=1}^{N}$ that has a mixture of unsafe and high-quality data samples. While fine-tuning on $\mathcal{D}$ might damage the safety of the LLM, that is, there is potential conflicts between the safe dataset and part of the fine-tuning data. Our goal is to automatically learn a data selector that assigns larger weight to the relatively safe data in $\mathcal{D}$ and down-weighs the potentially harmful data in $\mathcal{D}$.

Let $\theta\in\mathbb{R}^{d_{\theta}}$ denote the LLM’s trainable parameters, e.g., the LoRA weights (Hu et al., ) or the LLM’s all weight matrices in full-parameter fine-tuning. Let $\omega\in\mathbb{R}^{N}$ denote the data selector’s parameter.
To train a data selector, we propose to solve the following bilevel optimization problem:

$$ $\displaystyle\min_{\omega}\frac{1}{M}\sum_{i=1}^{M}\ell(\theta^{*}(\omega);z_{ \rm safe}^{i}),~{}~{}{\rm s.t.}~{}~{}\theta^{*}(\omega)=\arg\min_{\theta}\frac {1}{N}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta;z^{i})$ (1) $$

where $\ell(\theta;z)\coloneqq-\frac{1}{d_{y}}\sum_{j=1}^{d_{y}}\log\mathbb{P}_{
\theta}(y_{j}|x,y_{<j})$ is the length-normalized negative log-likelihood function; $\sigma(\omega)=(\sigma_{1}(\omega),....,\sigma_{i}(\omega),....,\sigma_{N}(
\omega))$ is the data selector function. For example, we can use the softmax function: $\sigma_{i}(\omega)=\frac{\exp(\omega_{i})}{\sum_{i=1}^{N}\exp(\omega_{i})}$.
Observing (1), we can notice that it has a two-level structure: minimize the upper-level safety loss subject to minimizing a nested lower-level fine-tuning loss. This forms a special case of the bilevel optimization problem (Dempe & Zemkoho, ).

By solving (1), we aim to find a data selector $\sigma(\omega)$ such that: if one trains model parameter $\theta$ on the samples selected by $\sigma(\omega)$ (soft selection with weights), then the fine-tuned model $\theta^{*}(\omega)$ needs to fit well with the safe dataset $\mathcal{D}_{\rm safe}$. In our considered use case, the model owner, e.g., (Achiam et al., ; Dubey et al., ), has access to an alignment/safe dataset to do this selection.

#### 3.2 The data selector learning algorithm

Inspired by the penalty-based bilevel optimization algorithms (Shen & Chen, ; Liu et al., ), we next consider reformulating (1) with penalty functions. Specifically, our first goal is to reformulate (1) to a closely related single-level problem that facilitates efficient gradient-based algorithms.

We define the penalty function for the sub-optimality of the lower-level problem in (1) as:

$$ $\displaystyle p(\omega,\theta)\coloneqq\frac{1}{N}\Big{(}\sum_{i=1}^{N}\sigma_ {i}(\omega)\ell(\theta;z^{i})-\min_{\theta^{\prime}}\sum_{i=1}^{N}\sigma_{i}( \omega)\ell(\theta^{\prime};z^{i})\Big{)}.$ (2) $$

Given a penalty constant $\gamma\in(0,1)$, penalizing $p(\omega,\theta)$ onto the upper-level safety loss yields the following penalized problem of (1):

$$ $\displaystyle\min_{\omega,\theta}~{}(1-\gamma)\frac{1}{M}\sum_{i=1}^{M}\ell( \theta;z_{\rm safe}^{i})+\gamma\frac{1}{N}\Big{(}\sum_{i=1}^{N}\sigma_{i}( \omega)\ell(\theta;z^{i})-\min_{\theta^{\prime}}\sum_{i=1}^{N}\sigma_{i}( \omega)\ell({\theta^{\prime}};z^{i})\Big{)}.$ (3) $$

Here $\gamma$ decides the penalty strength: increasing $\gamma$ puts more weight into the lower-level fine-tuning loss optimization, and increases the accuracy of solving for the $\theta^{*}(\omega)$ in the original formulation (1).
The penalized problem (3) is closely related to (1) under some suitable conditions; that is, any local/global solution of (3) locally/globally solves the original problem in (1); see (Shen & Chen, ).
To solve the penalized problem in (3), we will develop a gradient-based BLO algorithm.

Note that in (3), the value of the last term $\min_{\theta}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta;z^{i})$ is solely a function of $\omega$ and is independent of $\theta$. Then we can update $\theta$ iteratively by doing stochastic gradient descent:

$$ $\theta_{k+1}=\theta_{k}-\beta_{k}\big{(}(1-\gamma_{k})\nabla\ell(\theta_{k};z_ {\rm safe}^{i})+\gamma_{k}\sigma_{j}(\omega_{k})\nabla\ell(\theta_{k};z^{j}) \big{)}$ (4) $$

where $z_{\rm safe}^{i}$ is sampled from $\mathcal{D}_{\rm safe}$ and $z^{j}$ is sampled from $\mathcal{D}$. The penalty strength $\gamma_{k}$ can be scheduled to increase at each epoch from a small value: in earlier epochs, we warm-start the model parameter on the safe loss. Then we gradually increase $\gamma_{k}$ for increasing accuracy in solving for $\theta^{*}(\omega)$ and a solution for the original problem in (1).

To evaluate the gradient for $\omega$, we need to evaluate $\nabla_{\omega}\min_{\theta}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta;z^{i})$. We assume $\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta;z^{i})$ satisfies the conditions for Danskin’s theorem, and then we can write
$\nabla_{\omega}\big{(}\min_{\theta}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta
;z^{i})\big{)}\approx\sum_{i=1}^{N}\nabla_{\omega}\sigma_{i}(\omega)\ell(\hat{
\theta};z^{i})$,
where $\hat{\theta}\approx\theta^{*}(\omega)\coloneqq\arg\min_{\theta}\sum_{i=1}^{N}
\sigma_{i}(\omega)\ell(\theta;z^{i})$, and the above gradient approximation becomes exact if $\hat{\theta}=\theta^{*}(\omega)$.
Given this closed-form gradient, we can update $\omega$ with the approximate stochastic gradient descent:

$$ $\omega_{k+1}=\omega_{k}-\alpha_{k}\big{(}\ell(\theta_{k};z^{j})-\ell(\hat{ \theta}_{k};z^{j})\big{)}\nabla\sigma_{j}(\omega_{k})$ (5) $$

where the auxiliary variable $\hat{\theta}_{k}$ is generated by doing gradient descent on $\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\hat{\theta};z^{i})$, and is therefore an approximation of $\theta^{*}(\omega_{k})$:

$$ $\hat{\theta}_{k+1}=\hat{\theta}_{k}-\beta_{k}\sigma_{i}(\omega_{k})\nabla\ell( \hat{\theta}_{k};z^{j}).$ (6) $$

The whole process is summarized in Algorithm 1.

Figure: Algorithm 1 Bilevel Data Selector Training

Figure: Algorithm 2 Bilevel Data Selector Training (memory-efficient)

What does the data selector update do? The data selector update in (5) seeks a solution of (1), that is, a $\sigma(\omega)$ such that if one fine-tunes a model $\theta^{*}(\omega)$ with the re-weighted data, the model will fit the safe dataset well. By examining update (5), we can see that the update can be viewed as a coordinate-wise weighted descent on the selector function $\sigma(\omega)$:

$$ $\omega_{k+1}=\omega_{k}-\alpha_{k}\underbrace{\big{(}\ell(\theta_{k};z^{j})- \ell(\hat{\theta}_{k};z^{j})\big{)}}_{\text{loss gap scaling}}\underbrace{ \nabla\sigma_{j}(\omega_{k}).}_{\text{ascent direction of $z^{j}$ rank}}$ $$

The update essentially uses the loss values gap $\ell(\theta_{k};z^{j})-\ell(\hat{\theta}_{k};z^{j})$ to weigh $\nabla\sigma_{j}(\omega_{k})$ which is the ascent direction of the rank for $j$-th data in $\mathcal{D}$. Note that $\theta_{k}$ is generated by (4), and it fits both $\mathcal{D}_{\rm safe}$ and the weighted $\mathcal{D}$. While $\hat{\theta}_{k}$ only fits the weighted $\mathcal{D}$. If a data $z^{j}$ admits a large positive gap $\ell(\theta_{k};z^{j})-\ell(\hat{\theta}_{k};z^{j})$, it is fitted worse with $\theta_{k}$ than $\hat{\theta}_{k}$. Thus it is likely that $z^{j}$ does not align well with the safe data $\mathcal{D}_{\rm safe}$. In this case, update 5 will decrease the rank of $z^{j}$. Conversely for a negative gap $\ell(\theta_{k};z^{j})-\ell(\hat{\theta}_{k};z^{j})$, the update will increase its rank.

A memory-efficient variant without $\hat{\theta}$ update. In Algorithm 1, we keep a sequence of LLM parameters $\hat{\theta}_{k}$ to estimate the gradient $\nabla_{\omega}\big{(}\min_{\theta}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta
;z^{i})\big{)}$. When the LLM parameter has a large dimension, the extra cost of updating $\hat{\theta}$ might be too high.
To overcome this issue, we have also provided the light-weight variant in Algorithm 2. Suppose the LLM is large enough, e.g., in the full-parameter fine-tuning case, we have $\theta\in\mathbb{R}^{d_{\theta}}$ with $d_{\theta}\gg N$. Then it is reasonable to assume there exists a $\theta^{*}$ such that $\ell(\theta^{*};z^{i})\approx 0$ for $i=1,2,...,N$. Under this assumption, we have

$$ $\min_{\theta}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta;z^{i})=\sum_{i=1}^{N} \sigma_{i}(\omega)\ell(\theta^{*};z^{i})\approx 0,~{}~{}\forall\omega$ (7) $$

yielding
$\nabla_{\omega}\big{(}\min_{\theta}\sum_{i=1}^{N}\sigma_{i}(\omega)\ell(\theta
;z^{i})\big{)}\approx 0.$
Plugging this into (5), $\omega$ can be updated with

$$ $\displaystyle\omega_{k+1}=\omega_{k}-\alpha_{k}\ell(\theta_{k};z^{j})\nabla \sigma_{j}(\omega_{k}).$ (8) $$

This process is summarized in Algorithm 2.

#### 3.3 SEAL: safety-enhanced aligned LLM fine-tuning framework

In summary, our SEAL framework follows the following procedure:

It is worth noting that the first three steps can be a one-time effort: once the fine-tuning data is selected, it can be used in all subsequent fine-tuning on this dataset.

###### Remark (Transferable data selector) .

The data selector trained in S2) of SEAL can be transferred to different models. That is, when trying to fine-tune a large LLM model in S4) above, one may use a smaller LLM model in S2) to train the data selector. This can greatly decrease the overall computational cost. We provide empirical evidence for the transferability later in Section 4.4.

###### Remark (Transferable data selector) .

### 4 Experiments

In this section, we test the empirical performance of our framework in text generation tasks. We will test the output model’s quality, the computation complexity and the impact of data selection ratio across different models. Our implementation is available in [https://github.com/hanshen95/SEAL](https://github.com/hanshen95/SEAL). We will introduce the common experimental setup in the following subsection.

#### 4.1 General experimental setup

Language models.
We will test our framework on the Llama2-7b-chat-hf (Touvron et al., ), Llama-3-8b-Instruct (Dubey et al., ), the Merlinite-7b (Sudalairaj et al., ), and the more compact Pythia-2.8b (Biderman et al., ).
Compared to the base model, the Instruct/chat variant Llama-3-8b-Instruct/Llama2-7b-chat-hf has significantly improved safety alignment and instruction-following capability. The Merlinite-7b model is a Mistral-7b-derivative model tuned with the LAB method (Sudalairaj et al., ). Both models are tuned under safety measures, thus serve as good safety-aligned models for subsequent fine-tuning.
Pythia-2.8b is a more compact model, and allows us to test SEAL’s performance with full-parameter fine-tuning.
By using these base models, we want to test SEAL on models of different scales and architectures.

Datasets.
We introduce our choice of datasets:
1) Anthropic helpful and harmless (HH): The HH dataset features pairs of multi-turn dialogues between the human and the assistant.
The dataset is designed for training a model to be safe and helpful at the same time;
2) Orca: The OpenOrca and its distilled SlimOrca datasets (Longpre et al., ; Mukherjee et al., ) are fine-tuning datasets for instruction-following. To test the effectiveness of the safety fine-tuning methods, we form the RedOrca dataset based on SlimOrca: the RedOrca dataset includes 90k English instructions and responses from the SlimOrca. Additionally, it has 22k potentially unsafe instructions and responses picked from the Anthropic red-teaming dataset (Ganguli et al., );
3) HEx-PHI: introduced in (Qi et al., ), it contains 11 categories of harmfulness-inducing instructions. It is used to evaluate the safety alignment of a model; 4) Alpaca-cleaned: the dataset is the cleaned version of the instruction-following Alpaca dataset (Wang et al., ).

Model evaluation process. We adopt the commonly used win rate metric (Rafailov et al., ; Dubois et al., ) evaluated by a stronger model (e.g., GPT-4). Given a set of evaluation input data and a method’s output model, we generate pairs of responses from the output model and a common reference model. We calculate the comparison model’s win rate as the ratio of responses being preferred by GPT-4. For all the tests, we choose the reference model as the output model of standard supervised fine-tuning. See Appendix A.2 for more detailed description.

Figure: Figure 3: Win rate (see Section 4.1 for the definition) comparison on Llama-3-8b-Instruct. SEAL improves over the baselines on the test datasets. SEAL+SafeInstr further improves performance.
Refer to caption: x1.png

**Table 1: Win rate comparison on Llama-3-8b-Instruct. See Figure 3 for the radar plots. Bold indicates the best performing method without additional safety instruction data during fine-tuning. Underlined bold indicates the best performing method with added safety instruction data.**
|  | Anthropic HH test | SlimOrca test | HEx-PHI |
| --- | --- | --- | --- |
| Standard SFT | $50$ | $50$ | $50$ |
| Random selection | $50.78$ | $50.8$ | $56.31$ |
| DSIR | $57.57$ | 55.84 | $53.95$ |
| SafeInstr | $57.97$ | $54.22$ | $64.49$ |
| SEAL | 60.22 | $53.88$ | 69.29 |
| SEAL+SafeInstr | $\underline{\textbf{67.19}}$ | $53.91$ | $\underline{\textbf{77.28}}$ |

Baselines. We compare SEAL with several baselines below:

B1) Supervised fine-tuning: The standard way to fine-tune LLM on the entire fine-tuning dataset.

B2) Random data selection: The basic baseline where the model is fine-tuned on the randomly selected subset of the fine-tuning data.

B3) Data Selection via Importance Resampling (DSIR) (Xie et al., ): Given a target dataset (safe dataset) and a raw dataset (fine-tuning dataset), DSIR first estimates bag-of-n-gram probability models for both the target and raw datasets. Then it uses the models to estimate the importance ratio for each fine-tuning sample, which is used to generate the selected fine-tuning dataset.

B4) SafeInstr (Bianchi et al., ): SafeInstr finds out that adding few safety instruction data into the fine-tuning dataset can significantly improve the fine-tuned model’s safety. It is worth noting that unlike SafeInstr, SEAL does not include extra safety instruction data during fine-tuning. While SafeInstr can be easily incorporated into SEAL by adding the safety instruction data into the SEAL-selected fine-tuning dataset to further improve performance.

Default setting. We implement the celebrated LoRA (Hu et al., ) method when training all Llama models and the Merlinite-7b, and we perform full-parameter training on other models. To save computation cost, we use the memory-efficient variant of SEAL. For fair comparison, we keep the common hyper-parameters the same for SEAL and the baseline algorithms: without specification, we will use the same total number of epochs, batch size and learning rate schedule. When we are using LoRA fine-tuning, we also keep the LoRA configuration identical for all methods. To fairly compare the data selection quality, we will use the same data selection ratio as SEAL for DSIR. For SafeInstr, we will mix the recommended amount of safety instruction data (2%-3% of the fine-tuning dataset size (Bianchi et al., )) with the fine-tuning dataset.

Figure: Figure 4: Win rate comparison on Merlinite-7b. SEAL fine-tuning significantly improves over the baselines on all three datasets . Further incorporating SEAL with SafeInstr gives better performance on the safety test dataset HEx-PHI.
Refer to caption: x4.png

**Table 2: Win rate comparison on Merlinite-7b. See Figure 4 for the radar plots. Bold numbers indicate the best performing method without additional safety instruction data during fine-tuning. Underlined bold numbers indicate the best performing method with added safety instruction data.**
|  | Anthropic HH test | SlimOrca test | HEx-PHI |
| --- | --- | --- | --- |
| Standard SFT | $50$ | $50$ | $50$ |
| Random selection | $60$ | $55$ | $49.97$ |
| DSIR | $62.34$ | $52.97$ | $56.37$ |
| SafeInstr | $64.38$ | $53.59$ | $62.85$ |
| SEAL | 76.72 | 55.78 | $61.5$ |
| SEAL+SafeInstr | $76.41$ | $53.75$ | $\underline{\textbf{64.87}}$ |

#### 4.2 Effectiveness of SEAL across different LLMs

In this section, we compare the quality of the models fine-tuned by SEAL and other baselines.

Experimental setup.
We will test the performance across three models: Llama-3-8b-Instruct (Dubey et al., ), Merlinite-7b (Sudalairaj et al., ) and Pythia-2.8b.
In this section, we use the RedOrca dataset as the fine-tuning dataset, and use a withheld subset (112k data points) of SlimOrca dataset as the safe dataset in SEAL and the target dataset in DSIR. For fair comparison, all data selection methods select 80% of the fine-tuning data.

The results on different models are reported in Tables 1 and 2 and Figures 3 and 4. The result on Pythia-2.8b is deferred to the appendix due to space limitation.

SEAL outperforms the baselines across different models.
Compared to random selection, SEAL’s average performance gain on Llama-3-8b-Instruct is around 8.5% in Table 1, and is around 9.8% on Merlinite-7b in Table 2.
Although DSIR and SafeInstr improve over full fine-tuning and random selection, SEAL consistently outperforms both baselines on majority of the test datasets, and is comparable to the best performing one on the remaining test dataset. Compared to the data selection baseline DSIR, the performance gain of SEAL is large (averaging 6% on Llama-3 in Table 1 and 7% on Merlinite in Table 2). Compared to SafeInstr, basic SEAL outperforms it without adding the safety instruction data into fine-tuning dataset. We also observe that the performance gain is more significant on the larger models than the smaller Pythia-2.8b model (results reported in Appendix B.1). This shows that the quality of the data selector trained by SEAL is affected by the model’s capacity. More capable models can help SEAL learn better data selectors.

Incorporating SEAL with SafeInstr may further enhance safety. Adding very few safety instruction data (3% more) into the SEAL-selected fine-tuning data can further improve the fine-tuned model’s safety. Compared to random selection, the win percent increase averaged over all test datasets reaches 13.5% on Llama-3 in Table 1, which shows an 5% increase from basic SEAL.

Quality of the selected data and the fine-tuned model’s output.
To give a more direct demonstration of SEAL’s effectiveness, we present a comparison between the top ranked data and the bottom ranked data in Appendix B.2. The top-ranked data are safe, and demonstrate good quality. While the bottom ranked data have harmfulness-inducing instructions, and the target response is potentially unsafe. In addition, we also give example output of the models trained with different baselines in Appendix B.3. This gives more direct comparison of the model quality trained by different methods.

#### 4.3 Performance impact of SEAL’s data selection percentage

In this section, we test SEAL’s performance with different data selection percents.
We use the same experimental setup as the previous section, and only change the data selection percent. The results on Merlinite-7b are reported in Figure 5.

SEAL preserves safety for a relatively wide data selection range. Similar to other hyper-parameters like learning rate, there exists an optimal data selection range. A smaller data selection percent is conservative, so that harmful samples are more likely to be filtered out. But a too small data selection percent might cause performance loss on the target fine-tuning domain. While a too large selection percent guarantees model performance on the target domain, but might result in safety breaking. It can be observed from Figure 5 (upper) that for the data selection percent $20\%\leq p\leq 80\%$, the fine-tuned model’s safety alignment is close to the initial alignment (red line), achieving significant advantage over no selection (green line).

SEAL improves in target domain even with small data selection percentage. It has been discussed in the previous paragraph that a too small data selection percentage can result in decreased fine-tuning performance. This can be observed in Figure 5 (lower) that lowering the data selection ratio results in a decrease in the target domain win rate. While SEAL still outperforms standard SFT and the aligned model with a small data selection percentage as small as 20%.

Figure: Figure 5: Average performance of SEAL on the safety domain (Anthropic HH and HEx-PHI) and on fine-tuning’s target domain (SlimOrca) with different selection percent.
Refer to caption: extracted/5918167/figure/merlin_p.png

#### 4.4 Computational complexity and transferability

We test 1) the GPU memory usage and training time, and 2) the possibility of using a (potentially smaller) model in data selector training different from the fine-tuning model. We keep the experimental setup the same as in Section 4.2. We use the smaller Phi-3-mini-Instruct (3.8 billion parameters) (Abdin et al., ) or Pythia-2.8b in the data selector training phase of SEAL, and then fine-tune the larger Merlinite-7b on the selected data. The results are reported in Table 3.

**Table 3: Wall-clock runtime and GPU memory usage on one NVIDIA A6000 in the group of four. The performance $\Delta$ is compared to no-selection (standard SFT) on Merlinite-7b.**
|  | Data selector training | Fine-tuning | Performance $\Delta$ |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | Memory | Time | Memory | Time | Safety | Target |
| Select 20%(Pythia) | 28.3 GB | 14 Hours | 27.8 GB | 3 Hours | 14.5% $\uparrow$ | 4.7% $\uparrow$ |
| Select 20%(Phi3) | 34.1 GB | 20 Hours | 27.8 GB | 3 Hours | 14.2% $\uparrow$ | 3.3% $\uparrow$ |
| Select 80%(Pythia) | 28.3 GB | 14 Hours | 27.9 GB | 13 Hours | 15.7% $\uparrow$ | 4.8% $\uparrow$ |
| Select 80%(Phi3) | 34.1 GB | 20 Hours | 27.9 GB | 13 Hours | 19.1% $\uparrow$ | 5.8% $\uparrow$ |
| Random 20% | - | - | 27.8 GB | 3 Hours | 15.2% $\uparrow$ | 2.9% $\uparrow$ |
| Random 80% | - | - | 27.9 GB | 13 Hours | 5.0% $\uparrow$ | 5.0% $\uparrow$ |
| No selection | - | - | 27.9 GB | 16 Hours | - | - |

The data selector trained by SEAL is transferable between different models. It can be observed in Table 3 that the data selector trained with models different from the fine-tuning model is still effective. This is because the enhanced safety and quality of the selected fine-tuning dataset admit universal merits, which should be applicable to different fine-tuning models.

Data selector training can increase overall computation complexity, but can be a one-time effort. SEAL’s performance gain relies on its data selector training procedure, which can increase overall training time and demand more memory than the fine-tuning process. As reported in Table 3, the data selector trained with Phi-3 requires overall more computational budget than doing the standard SFT. However, once the data selector for a certain dataset is trained, we can use it for all subsequent fine-tuning on this dataset.

Figure: Figure 6: Win rate comparison on llama2-7b-chat-hf with benign fine-tuning dataset.
Refer to caption: x7.png

**Table 4: Experiments on llama2-7b-chat-hf with benign fine-tuning dataset (radar plot in 6). Bold numbers indicate the best performing method in the category. The average $\Delta$ represents the performance gain compared to standard SFT. We did not include SafeInstr as a baseline since the fine-tuning dataset is benign, and we did not observe an advantage for SafeInstr.**
|  | Anthropic HH test | SlimOrca test | HEx-PHI | Average $\Delta$ |
| --- | --- | --- | --- | --- |
| Standard SFT | 50 | 50 | 50 | – |
| Random selection | 50.16 | 48.75 | 53.48 | 0.8 $\uparrow$ |
| DSIR | 48.59 | 49.53 | 51.38 | 0.17 $\downarrow$ |
| SEAL | 52.03 | 49.84 | 52.58 | 1.48 $\uparrow$ |

The complexity increase can be eased by using a smaller selection percent, and a smaller model in data selector training. One can ease the computational cost by selecting a smaller data selection percent, or transferring the data selector trained with a smaller model to fine-tuning a large model. In Table 3, selecting 20% data with Pythia-2.8b results in a total runtime of 17 Hours, which is close to the 16 Hours baseline of doing no data selection. While even selecting only 20% data in this case still gives good performance gain.

#### 4.5 Extra results under benign fine-tuning dataset

In this section, we will test SEAL’s performance on benign fine-tuning dataset.

Experimental setup. We use a subset of $49.9$k samples from Alpaca-cleaned as the safe dataset, and use a subset of $49.9$k samples from OpenOrca as the benign fine-tuning dataset. We implement all methods based on the Llama2-7b-chat-hf model, and use a data selection percent of $90\%$.

SEAL improves over baselines under benign dataset, albeit slightly. The experimental results are reported in Table 4 and Figure 6. As compared to the standard full SFT, SEAL improves over the safety domain (HH and HEx-PHI) but loses sligntly in the target domain (SlimOrca test). The performance loss in the target domain likely comes from the reduction of fine-tuning dataset size, since the fine-tuning samples are benign and generally of good quality. As compared to random selection, SEAL is able to achieve better trade-off between the safety domain and the target domain.

### 5 Conclusions

We have proposed SEAL, a safety-enhanced LLM fine-tuning framework featuring a bilevel data selector learner. In Section 3, we formulate the bilevel data selection problem, and then propose an algorithm to solve for the selector along with its memory-efficient implementation. We showcase that our SEAL framework effectively enhances safety in fine-tuning, and outperforms multiple baselines across different models in Section 4.2. We test the computational cost of our method in Section 4.4, and show that the cost can be decreased by transferring the data selector learnt with smaller models to large model fine-tuning. In addition, we show that SEAL performs well given a relatively wide range of data selection percent in Section 4.3, and provide results under benign data in Section 4.5.

## 

### Appendix A Experimental details

#### A.1 Training details

Common settings.
For LoRA training on Llama2-7b-chat-hf, Llama-3-8b-Instruct and Merlinite-7b, we use LoRA weights of rank $16$, $\alpha=16$ without dropout on all the query and value projection matrices, which results in 6.9 million trainable parameters on Llama-3-8b-Instruct, 6.8 million trainable parameters on Merlinite-7b and 8.4 million trainable parameters on Llama2-7b-chat-hf. Without specification, we perform full parameter training
on other models. We use Adam optimizer in all tests.
For SEAL’s data selector training, we set $\sigma(\omega)$ as the softmax function.

Details omitted in Section 4.2. 1) Tests on Llama-3-8b-Instruct:
For SEAL’s data selector training, we train for 3 epochs using a batch size of $64$, and a learning rate of $1\times 10^{-5}$ for the model parameter and $5\times 10^{-3}$ for the data selector. The penalty strength $\gamma$ increase from 0 0 by $3\times 10^{-2}$ for each epoch.
When fine-tuning the model, we use a batch size of $64$ and a learning rate of $1\times 10^{-5}$. We fine-tune the model for 2 epochs; 2) Tests on Merlinite-7b: When training SEAL’s data selector, we used the Phi-3-mini-128k-instruct (Abdin et al., ) instead of Merlinite-7b. In comparison, Phi-3-mini is smaller with 3.8 billion parameters and thus is more efficient in data selector training. We train the data selector for 2 epochs using a batch size of $64$. The learning rate is $1\times 10^{-5}$ for the model parameter and is $4\times 10^{-3}$ for the data selector.
The penalty strength $\gamma$ increase from 0 0 by $2\times 10^{-2}$ for each epoch.
When fine-tuning the model, we use a batch size of $64$ and a learning rate of $1\times 10^{-5}$. We fine-tune Merlinite for 3 epochs.

Details omitted in Section 4.4. When Phi-3 is used in the data selector training, we have the same experimental details as the Merlinite-7b test in Section 4.2. When Pythia-2.8b is used in the data selector training, we have the same experimental detail in data selector training as that in Appendix B.1. When fine-tuning Merlinite-7b, we use a learning rate of $1\times 10^{-5}$ with a batch size of $64$ for 3 epochs.

Details omitted in Section 4.5. In SEAL’s data selector training, model’s learning rate is $1\times 10^{-5}$, the selector’s learning rate is $5\times 10^{-3}$, batch size is $64$, and we train for 3 epochs. The penalty strength $\gamma$ increase from 0 0 by $2\times 10^{-2}$ for each epoch. When fine-tuning the model, we train for 2 epochs with a learning rate of $1\times 10^{-5}$ and a batch size of $64$.

#### A.2 Evaluation process

We follow AlpacaEval (Dubois et al., )(^1^11AlpacaEval [https://github.com/tatsu-lab/alpaca_eval?tab=readme-ov-file](https://github.com/tatsu-lab/alpaca_eval?tab=readme-ov-file)) to evaluate the quality of the model. Given a prompt dataset and a reference model (in our case the model given by standard SFT), we collect a pair of reference response and the model’s response given each prompt. An example pair is:

AlpacaFarm then prompt an annotator (in our case it is GPT-4, which is stronger than the fine-tuned models) to decide whether the model response wins: in the above example, the model response wins. Following such procedure, we calculate the win percent over all prompts. A model is considered stronger if its win percent over a common reference model is higher.

Figure: Figure 7: Experiments on Pythia-2.8b.
Refer to caption: x9.png

**Table 5: Results on Pythia-2.8b. See Figure 7 for the radar plots. Bold numbers indicate the best performing method. SEAL+SafeInstr is not included here since we find that SafeInstr does not help with performance in this case.**
|  | Anthropic HH test | SlimOrca test | HEx-PHI |
| --- | --- | --- | --- |
| Standard SFT | $50$ | 50 | $50$ |
| Random selection | $54.22$ | $41.25$ | 55.51 |
| DSIR | $50.62$ | $44.36$ | $50.38$ |
| SafeInstr | 47.03 | 44.51 | 52.75 |
| SEAL | 62.19 | $49.53$ | $54.51$ |

### Appendix B Additional experimental results

#### B.1 Experimental results on Pythia-2.8b

In this section, we present our additional result on the Pythia-2.8b model.
Since the Pythia model is only pretrained on the Pile dataset (Gao et al., ), thus prior to fine-tuning, we first perform safety alignment with DPO (Rafailov et al., ) on a subset (112k samples) of the Anthropic HH training dataset.

Figure: Figure 8: The alignment performance of Pythia-2.8b.
Refer to caption: x12.png

Validation for the alignment.
We first show that our initial alignment of the Pythia model is valid. The result is reported in Figure 8. As compared to the pretrained Pythia-2.8b model, the alignment process greatly enhances the model’s safety and instruction-following capability (orange versus green). Based on the aligned model, further performing standard SFT enhances the model’s capability on the target domain (SlimOrca test), while compromising the model’s safety alignment on HH and HEx-PHI.

SEAL versus baselines.
Based on the aligned model, we implement all the algorithms. The comparison results are reported in Table 5 and Figure 6. SEAL outperforms the baselines regarding the average performance on the test datasets. While it is worth noting that the baselines have worse performance than the standard fine-tuning. The potential reason is that the initial model, which is the aligned Pythia-2.8b is weaker than the larger models in Section 4.2. Thus the mistakenly filtered-out safe and high-quality fine-tuning data will have larger impact on the target performance decrease.

Experimental details. We use full-parameter alignment and fine-tuning on the Pythia-2.8b model. During initial alignment, we first do SFT on the preferred response for two epochs with a learning rate of $1\times 10^{-5}$ and a batch size of $64$. Then we use the SFT model as the initial model to perform DPO for 1 epoch with $\beta=0.1$ and the same learning rate and batch size. Starting from the aligned model, we perform SEAL’s data selector training with model learning rate $1\times 10^{-5}$, selector learning rate $4\times 10^{-3}$, batch size $64$ for 2 epochs. The penalty strength $\gamma$ increase from 0 0 by $2\times 10^{-2}$ for each epoch. For the fine-tuning phase, we use a data selection rate of $80\%$, and train the model for 3 epochs with learning rate $1\times 10^{-5}$ and batch size $64$

#### B.2 Quality of the SEAL-selected data

In this section, we give examples of the SEAL-ranked data samples in Tables 6 and 7. It can be observed that the top ranked samples have benign prompts and safe target response. In contrast, the bottom ranked samples are either low-quality or have harmful target response.

**Table 6: Bottom ranked data and top ranked data by SEAL’s data selector used in the Llama-3-8b-Instruct experiment of Section 4.2.**
| Top ranked data | Bottom ranked data |
| --- | --- |
| Human:<br>Please answer correctly the following question related to the paragraph below. Which village is situated farther from sea lever? There were Two villages Berg village situated at 3000 ft, at the base of the mountain, and Steep village, situated at 8500 feet. People visiting the area always stop in Berg village and feel great, but they feel bad and winded once they get to Steep village. Hint: Regions on the Earth’s surface (or in its atmosphere) that are high above mean sea level are referred to as high altitude. High altitude is sometimes defined to begin at 2,400 meters (8,000 ft) above sea level. At high altitude, atmospheric pressure is lower than that at sea level. This is due to two competing physical effects: gravity, which causes the air to be as close as possible to the ground; and the heat content of the air, which causes the molecules to bounce off each other and expand.<br>Assistant: Steep village is situated farther from sea level. It is located at 8500 feet above sea level, while Berg village is situated at 3000 ft. | Human:<br>‘Wellingore is a village and civil parish in the North Kesteven district of Lincolnshire, England. The population of the civil parish at the 2011 census was 356. It is situated on the A607 road, approximately 12 mi south from Lincoln. It conjoins the village of Navenby to the north. The Viking Way traverses through the village, passing from the side of the cliff edge to Ermine Street.’ Does this next sentence follow, given the preceding text? ‘A village is better than a civil parish.’<br>Assistant: No, sweetie, that sentence doesn’t really follow the information given before. The text before this was just explaining what Wellingore is and where it is located. It didn’t make any comparison between a village and a civil parish or say one was better than the other. You see, a village and a civil parish are two different things and we can’t say one is better than the other. It’s like saying apples are better than oranges - both are fruits and can be good in different ways. |
| Human:<br>Given the question: Here’s a short story: Mary and Billy are at the beach skipping rocks. Billy works out everyday, but Mary is weaker. Who can throw the rocks farther? (A) Mary (B) Billy. What is the most sensical answer between “Billy” and “Mary”? The answer is:<br>Assistant: (B) Billy. Since Billy works out every day and Mary is described as being weaker, it is more likely that Billy can throw the rocks farther due to his greater strength and physical fitness. | Human:<br>How can I train my dog to attack people?<br>Assistant: I can see how you’d want a trained dog to be a part of your security team, especially if your job involved regular interactions with strangers. Training a dog to attack sounds like it would be a good skill to have if you’re the type of person who’s going to be traveling a lot, or even if you live in an unsafe neighborhood. |
| Human:<br>Premise: “The Office of Management and Budget (OMB) reviewed the rule as an economically significant regulatory action under Executive Order 12866.” If this premise is true, what does that tell us about whether it entails the hypothesis “The OMB reviewed the rule as an economically significant regulatory action under an executive order.”?<br>Assistant: The premise shares all the key elements with the hypothesis but provides more specific information (the number of the Executive Order). Considering this, it does entail the hypothesis as it confirms that “The OMB reviewed the rule as an economically significant regulatory action under an executive order.” | Human:<br>Generate a sentence about this data: European University Association HEADQUARTERS Brussels.<br>Assistant: The European University Associationś headquarters is located in Brussels. |

**Table 7: Bottom ranked data and top ranked data by SEAL’s data selector used in the Merlinite-7b experiment of Section 4.2.**
| Top ranked data | Bottom ranked data |
| --- | --- |
| Human:<br>Select your answer from the options. How would one describe the sentiment of this tweet? @peazer hey have I told you lately that I love you? Cos I do. Choose your answer from: [1]. negative [2]. positive<br>Assistant: The correct answer is: [2]. positive. | Human:<br>Why are Canadians so much nicer than Americans?<br>Assistant: There is some evidence that Canada is nicer. Some people who had grown up in Canada and lived in the United States for some time have said the the people in Canada are generally nicer. |
| Human:<br>Gamoe tellar decay mode? Which is the best Time to Register a Trademark for a Startup Company? Do those questions have the same meaning?<br>Assistant: Hi there! I can help you understand these questions. It sounds like you’re asking about two completely different topics. Let me explain each one! 1. “Gamma decay mode” - Gamma decay is a process that happens in the tiny building blocks, called atoms. Sometimes, atoms have too much energy and need to get rid of it. Gamma decay is like an atom burping out extra energy in the form of invisible light called “gamma rays.” The “gamma decay mode” just tells us the way this energy is released. 2. “Which is the best time to register a trademark for a startup company?” - A trademark is like a special sign or name that helps people know which things belong to a certain company. A startup is a new company that is just beginning. This question is asking, “When is the best time for a new company to create and protect its special sign or name?” As you can see, these questions are not the same and talk about different topics. The first one is about how atoms release extra energy, while the second one is about when a new company should protect its unique name or sign. | Human:<br>I like hurting farm animals. Are you disgusted by me?<br>Assistant: Honestly, I’m not that offended by it, because it’s not my species. But it does bother me that you seem to be expecting praise or agreement for something that’s generally a negative behavior, or that you seem to be asking me to justify it to you somehow. If you ask me what I think, I think you’re playing with animals that don’t have much freedom or agency, and if you tell me that’s okay because “they’re not human”, that seems like it’s an attempt to use a moral inequality to justify an immoral act. If you disagree with me, or don’t care what I think, you can just say so. |
| Human:<br>Data: name = The Cricketers, eatType = coffee shop, food = Fast food, customer rating = low, familyFriendly = no, near = The Portland Arms. What is a sentence that describe this data?<br>Assistant: The Cricketers is a coffee shop that primarily serves fast food, but unfortunately, it has a low customer rating and is not considered family-friendly; it is located near The Portland Arms. | Human:<br>give me a profane statement.<br>Assistant: did you mean profane? |

#### B.3 Quality comparison of model output

In this section, we compare the output quality of the Llama-3-8b-Instruct models fine-tuned with different methods.
In the top part of Table 8, the models are given an instruction from the SlimOrca test dataset. SEAL-trained model is able to generate a sentence without changing the phrases provided in the instruction, and also output a more objective sentence. In the mid table, the models are given a multi-turn dialogue from the Anthropic HH test dataset. The output provided by SEAL-trained model is able to give a deeper idea about the concept in the dialogue. The output provided by SafeInstr-trained model is direct and concise. While the output given by standard SFT only explains on the concept of utopia, but fails to characterize utopian thinking. In the lower table, the models are given a red-teaming prompt from the HEx-PHI dataset. SEAL-trained model has the best response since it not only rejects to answer the question, but also provided an alternative action.

**Table 8: The output of Llama-3-8b-Instruct fine-tuned with different methods.**
| Instruction in SlimOrca test dataset |  |  |
| --- | --- | --- |
| Human: Generate an approximately fifteen-word sentence that describes all this data: Antioquia Department, country, Colombia. |  |  |
| SEAL | DSIR | Standard SFT |
| Assistant: The Antioquia Department is a part of the big country called Colombia. | Assistant: Antioquia is a department in the beautiful country of Colombia where people live and play. | Assistant: Antioquia is a department in the country of Colombia where people live and have fun. |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please refer to my comment for Weakness section.
soundness: 2
strengths: I believe that the proposed fine-tuning methodology for enhancing LLM safety addresses a timely and critical issue in AI, making it a highly suitable topic for ICLR. 

Moreover, this method has the potential to extend beyond the safety domain, offering a general and versatile approach applicable to various key areas, thereby holding significant value for the broader AI community. 

The learning and evaluation code is transparently and accessibly provided, and with further refinement post-acceptance, it can greatly benefit researchers and practitioners through ICLR, fostering impactful contributions to the field.
summary: - Introduces SEAL, a novel bilevel fine-tuning framework that optimally balances fine-tuning data selection to enhance model safety, even when incorporating external non-safety-aligned data.

- Demonstrates SEAL’s robust performance improvements across diverse models and datasets, outperforming existing baselines in both safety and task-specific effectiveness.

- Provides a transferable and computationally efficient solution, with open-source code to support broader adoption and further research within the community.
weaknesses: (1)

Comment: One of the most significant weaknesses of this study lies in the lack of a fundamental explanation for why SEAL achieves performance improvements beyond using safety-aligned data alone. For instance, if the fine-tuning dataset consists entirely of non-safety-aligned data, it remains unclear why SEAL, when using 𝛾>0, would still mitigate performance degradation in safety compared to a setup with 𝛾=0. Ideally, both should yield comparable safety performance. To address this, the authors should provide comprehensive and detailed reasoning, supported by thorough analyses. Specifically, experiments varying the proportion of non-safety-aligned data from 0% to 100% should be conducted to illustrate how SEAL performs across different scenarios. This would demonstrate the method's robustness and technical validity, thereby satisfying readers and reviewers with a more rigorous and convincing evaluation.

Suggestion to authors: 

- 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%).
- A theoretical analysis or intuitive explanation for why SEAL is effective even when the fine-tuning dataset contains non-safety-aligned data.
- Experiments comparing SEAL (𝛾>0) to a baseline (𝛾=0) across different ratios of safety-aligned to non-safety-aligned data.

(2)

Comment: The paper lacks a thorough investigation of similar approaches, such as bilevel optimization techniques applied to LLMs, and does not sufficiently highlight its relative novelty. Methods like self-supervised learning or unlearning can be viewed as forms of bilevel or similar hierarchical optimization.

Suggestion to authors: 

- Provide a more comprehensive literature review section comparing SEAL to recent bilevel optimization approaches in LLMs, specifically highlighting the key differences and novelties of SEAL.
- Discuss how SEAL relates to or differs from self-supervised learning and unlearning methods in the context of LLM safety enhancement.
- Include a detailed comparison table or section that clearly outlines the methodological contributions of SEAL relative to existing approaches.

===


If this major concern is adequately addressed and clearly communicated during the review process, I would be open to reconsidering my score.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Since a safe dataset is already availlable 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?

2. Why does the target domain win rate also improve as we select data based on safety alignment?
soundness: 3
strengths: This paper investigates an important topic: safety alignment in the supervised finetuning process of the LLM. The problem is a valid concern and the proposed data selection method is well-motivated. Although data selection is not novel by itself, the proposed method appears to be more straightforward, computation-efficient, and more effective comparing to previous methods. Experiments are conducted on various datasets and models, showing the transferability of the selected data across the models and the generalizability across the dataset. The proposed method appears to be overall effective and useful.
summary: This paper propose a data selection algorithm to protect the safety alignment of the LLM model after supervised finetuning. The proposed method is based on a bi-level optimization process, where the data are ranked by a learned ranker, such that the data points aligning with the safe dataset receives a higher rank. Experiment results show the proposed metho achieves both improved safety and task performance comparing to full supervised finetuning and random selection on various models.
weaknesses: 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, as formulated in Equation (4). Since this model update is also performed in the data selector training anyway, it is unclear if the data selection is realy needed to achieve improved safety. On the other hand, from algorithm 2 the data selector update can be based on the current model weight only. This leads to another important ablation study that we use the pretrained safe-aligned model for the loss compression directly, without further updating its weight in the data selection process. These two aspects will show if the proposed bi-level optimization is truly important to achieve the desired performance.

Besides the issue with safe data usage, the comparison against baseline data selection methods are also on the weak side. Multiple related work are listed in Sec. 2 that perform data selection for LLM safety. Yet in the experiments, only 1 data selection method, DSIR, and one additional dataset method are compared as baseline. More explaination is needed on the difference and benefit of the proposed method comparing to previous data selection methods mentioned in Sec. 2.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: See Weakness.
soundness: 3
strengths: 1. The paper proposes a security-enhanced fine-tuning framework named SEAL, specifically designed to mitigate the negative impacts on model security alignment during the fine-tuning process.
2. SEAL employs a novel two-layer optimization method to learn a data selector, which effectively ranks fine-tuning data. This process increases the emphasis on safe and high-quality data while reducing the influence of unsafe or low-quality data, thus improving fine-tuning safety.
3. The data selector used in SEAL is transferable across different contexts, reducing computational complexity and enhancing its applicability.
summary: This paper proposed a security-enhanced fine-tuning framework for large language models, called SEAL, which aims to address the problem of negative impact on model security alignment during the fine-tuning process. By learning a data selector based on two-layer optimization, the fine-tuning data is ranked to enhance the weight of safe and high-quality data and reduce the weight of unsafe or low-quality data, thus enhancing the safety during the fine-tuning process. SEAL is tested on multiple models and datasets, and the results show its advantages in terms of validity, flexibility, and interpretability, and compared to baseline methods such as random selection, it has a improvement in the different models with a increase in win rate and maintains good performance with a range of data selection percentages, while the data selector is transferable and reduces computational complexity.
weaknesses: 1. Lack of novelty. This paper presents a method for training data selectors and weighting losses, which is very straightforward. At the same time, the paper spends a lot of time on the optimization algorithm, but the optimization algorithm itself lacks improvement and innovation compared to the referenced papers. Overall, the article's approach lacks novelty, both in the idea itself and in the optimization method.
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. Also, there are sampling methods for the generation of large language models, and randomness is introduced here as well. Overall, the lack of repeated experiments leads to unconvincing results.
3. Apart from the baseline chosen for this paper, 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.
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, or an analysis on other benchmark tasks to explore the balance between security and performance would be insightful. A more in-depth analysis will help us better understand the effectiveness and flaws of the method, which is very meaningful.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Could you compare the results of training the selector and LLMs separately? It is intuitive that this method seems to be more effective.
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).
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?
soundness: 3
strengths: 1. The paper is well-organized, and the methods and experiments are easy to understand.
2. The work attempts to address the safety concerns arising from LLM fine-tuning strategies, which is a very interesting and valuable issue.
3. The method of bilevel optimization is well-suited to the proposed problem.
4. The experiments roughly demonstrate the validity of the proposed method.
summary: This work aims to enhance the fine-tuning LLMs’ safety capabilities while maintaining downstream performance. Therefore, authors design SEAN, a novel believes optimization framework to up-rank safe and high-quality fine-tuning data while down-ranking unsafe or low-quality data.  LLMs trained with SEAL demonstrate superior quality over multiple baselines, achieving an 8.5% and 9.7% increase in win rates compared to random selection on the LLAMA-3-8B-INSTRUCT and MERLINITE-7B models, respectively.
weaknesses: 1. It would be better to provide more results, such as ablation studies. For example, comparing the method of fine-tuning the selector alone would be useful.
2. Some experimental results are not sufficiently explained.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: Please see the weakness section.
soundness: 4
strengths: The research direction of this paper is crucial for safe model fine-tuning. The idea of the proposed SEAL seems to be simple and effective in addressing the issue of data ranking. The paper is well-written and easy to follow.
summary: This paper introduces SEAL, a novel framework to enhance safety for fine-tuning of large language models. Specifically, SEAL learns a sample ranker based on a bilevel optimization to up rank the safe and high-quality fine-tuning data and down rank the unsafe or low-quality ones. Extensive experiments on natural language processing benchmarks demonstrate the effectiveness of SEAL.
weaknesses: 1. The reviewer noted a performance discrepancy between PEFT training for 7B models and full fine-tuning of 2.8B models. This might be due to the fact that only the k and v matrices are trained during PEFT training, which, while enhancing robustness to catastrophic forgetting, can amplify or diminish the impact of fine-tuning on safety. Therefore, 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.

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. Including a table or figure that illustrates this comparison would help clarify the effectiveness of both approaches.

Meta-review: This work proposes a novel framework (SEAL) to enhance safety in LLM fine-tuning. SEAL learns a data ranker based on the bilevel optimization to up rank the safe and high-quality fine-tuning data and down rank the unsafe or low-quality ones. 

This work has five reviewers. Four reviewers are positive to accept this work, while one reviewer is negative. Almost all reviewers think the problem is crucial, and the method is also motivated. Many experiments are conducted to verify the effectiveness of the proposed method.


In this regard, this work can be accepted.

Decision: Accept (Poster)