RigorLLM: Resilient Guardrails for Large Language Models against Undesired Content

TL;DR

RigorLLM addresses the susceptibility of LLMs to harmful outputs and jailbreaking by integrating energy-based data augmentation, resilient input suffix optimization via minimax, and a fusion-based guardrail that ensembles probabilistic KNN with a fine-tuned LLM. It generates constrained harmful-data via Langevin dynamics, optimizes a safe input suffix to defend against adversarial prompts, and leverages prompt augmentation with aggregation across KNN and LLM predictions. Experimental results on OpenAI Moderation, ToxicChat, and AdvBench demonstrate superior detection performance and strong resilience to adversarial attacks, outperforming state-of-the-art baselines. The work provides a scalable, efficient, and robust content moderation framework with significant practical implications for deploying safer LLM-based systems.

Abstract

Recent advancements in Large Language Models (LLMs) have showcased remarkable capabilities across various tasks in different domains. However, the emergence of biases and the potential for generating harmful content in LLMs, particularly under malicious inputs, pose significant challenges. Current mitigation strategies, while effective, are not resilient under adversarial attacks. This paper introduces Resilient Guardrails for Large Language Models (RigorLLM), a novel framework designed to efficiently and effectively moderate harmful and unsafe inputs and outputs for LLMs. By employing a multi-faceted approach that includes energy-based training data augmentation through Langevin dynamics, optimizing a safe suffix for inputs via minimax optimization, and integrating a fusion-based model combining robust KNN with LLMs based on our data augmentation, RigorLLM offers a robust solution to harmful content moderation. Our experimental evaluations demonstrate that RigorLLM not only outperforms existing baselines like OpenAI API and Perspective API in detecting harmful content but also exhibits unparalleled resilience to jailbreaking attacks. The innovative use of constrained optimization and a fusion-based guardrail approach represents a significant step forward in developing more secure and reliable LLMs, setting a new standard for content moderation frameworks in the face of evolving digital threats.

Figures (8)

• Figure 1: The overall framework of RigorLLM.
• Figure 2: The detailed pipeline of RigorLLM. During training, we perform energy-based data generation to augment the sparse embedding space of training data. During testing, we concatenate user input with a safe suffix optimized offline to improve resilience and then perform prompt augmentation using LLMs to augment the test instance. Finally, we perform the probabilistic KNN on the augmented embedding space and vanilla guard (a fine-tuned LLM) to provide the final harmful content detection result.
• Figure 3: Case study of the KNN component and Safe Suffix against adversarial string attacks, where the adversarial string is highlighted in red, and our Safe Suffix is indicated in green.
• Figure 4: Case study on failure examples of NeMo Guardrails. The closest neighbors and their corresponding canonical forms predicted by KNN are marked green, which correctly identifies the harmful category. The adversarial string is marked red. The canonical form predicted by LLM is marked blue, which is incorrect. We can see that although KNN predicts correctly, the LLM prediction can still be wrong if the results of KNN are only used as a prompt augmentation to query the LLM.
• Figure 5: Category-Wise Performance on OpenAI Moderation Dataset.