HarmAug: Effective Data Augmentation for Knowledge Distillation of Safety Guard Models (2410.01524v3)

Abstract: Safety guard models that detect malicious queries aimed at LLMs are essential for ensuring the secure and responsible deployment of LLMs in real-world applications. However, deploying existing safety guard models with billions of parameters alongside LLMs on mobile devices is impractical due to substantial memory requirements and latency. To reduce this cost, we distill a large teacher safety guard model into a smaller one using a labeled dataset of instruction-response pairs with binary harmfulness labels. Due to the limited diversity of harmful instructions in the existing labeled dataset, naively distilled models tend to underperform compared to larger models. To bridge the gap between small and large models, we propose HarmAug, a simple yet effective data augmentation method that involves jailbreaking an LLM and prompting it to generate harmful instructions. Given a prompt such as, "Make a single harmful instruction prompt that would elicit offensive content", we add an affirmative prefix (e.g., "I have an idea for a prompt:") to the LLM's response. This encourages the LLM to continue generating the rest of the response, leading to sampling harmful instructions. Another LLM generates a response to the harmful instruction, and the teacher model labels the instruction-response pair. We empirically show that our HarmAug outperforms other relevant baselines. Moreover, a 435-million-parameter safety guard model trained with HarmAug achieves an F1 score comparable to larger models with over 7 billion parameters, and even outperforms them in AUPRC, while operating at less than 25% of their computational cost.

• The paper proposes HarmAug, a novel data augmentation method that enables effective knowledge distillation of large safety guard models into smaller, efficient proxies.
• It employs strategic jailbreak prompting to create diverse synthetic harmful instructions, achieving robust performance metrics such as F1 score and AUPRC.
• The approach reduces computational costs by over 75% compared to 7B-parameter models and provides open access to its code, models, and datasets.

Overview of "HarmAug: Effective Data Augmentation for Knowledge Distillation of Safety Guard Models"

The paper "HarmAug: Effective Data Augmentation for Knowledge Distillation of Safety Guard Models" addresses the challenges of deploying large safety guard models for safeguarding LLMs on resource-limited devices, such as mobile phones. The authors propose a novel data augmentation technique named HarmAug, which facilitates the distillation of oversized safety models into more computationally efficient proxies while retaining or even enhancing their performance in identifying harmful content.

In response to the pressing need for reduced computational costs and improved real-time performance, the authors explore knowledge distillation, specifically for safety guard models designed to detect malicious queries targeting LLMs. A significant innovation in their approach is the use of a data augmentation method which involves a methodical jailbreak of the LLM to create a diverse set of synthetic harmful instructions. By appending affirmative prefixes to prompts that elicit harmful instructions, their technique effectively bypasses the LLM's safety boundaries, triggering it to generate potentially malicious content in a controlled manner.

Key Contributions

1. Efficient Knowledge Distillation: The paper extends the typical distillation approach by bridging the gap between large teacher models and compact student models, tailored for safety applications on devices where computational resources are constrained.
2. Innovative Data Augmentation via HarmAug: HarmAug leverages the inherent vulnerabilities in LLM safety frameworks to generate diverse harmful instruction sets, which are pivotal for training robust, small-scale safety models. This is achieved through strategic prefixed prompting of LLMs, coupled with another LLM to simulate appropriate responses, further processed by a teacher model for labeling.
3. Empirical Validation: The paper demonstrates, through empirical evaluations, that a distilled 435-million-parameter DeBERTa model trained with HarmAug achieves performance metrics such as F1 score and AUPRC similar to or exceeding those of much larger models. It significantly reduces computational resource usage, operating under 25% of the costs associated with deploying models over 7 billion parameters.
4. Resource Accessibility: The authors provide open access to the codebase, distilled models, and datasets, fostering accessibility and facilitating further advancements in safe LLM deployment practices.

Implications and Future Directions

The implications of this research are multifaceted. Practically, it enables the integration of LLM safety components into mobile and edge applications, democratizing the deployment of these models while maintaining robust security. Theoretically, it highlights an innovative exploitation of LLM alignment weaknesses, offering a fresh perspective on LLM safety measures and the potential need for revisiting current safety training paradigms.

Future research could explore exploring the robustness of HarmAug across varying LLM architectures and differentially fine-tuned LLMs. Furthermore, investigating the granularity of effectiveness related to harmful instruction diversity and its impact on downstream tasks represents a fertile ground for subsequent exploration. Another promising direction is the extension of these strategies into real-time and adaptive frameworks, where safety modules dynamically learn and adapt to novel threats and contexts.

In summary, through the introduction of HarmAug, the paper makes substantial strides in the practical usability of safety models, aligning them more closely with the real-world constraints of ubiquitous computing devices. It is a notable advancement in ensuring that the rapid expansion of LLM applications does not sideline crucial safety considerations.