Source-linked AI summary
Class-RAG: Real-Time Content Moderation with Retrieval Augmented Generation
Jianfa Chen, Emily Shen, Trupti Bavalatti, Xiaowen Lin, Yongkai Wang, Shuming Hu, Harihar Subramanyam, Ksheeraj Sai Vepuri, Ming Jiang, Ji Qi, Li Chen, Nan Jiang, Ankit Jain
TL;DR
Fine-tuning content moderators is slow and costly, while subtle safe–unsafe distinctions and shifting harms complicate moderation. Class-RAG adds retrieved examples and explanations to a classifier and supports library updates; the paper reports stronger classification performance and adversarial robustness than baseline models.
Problem
Content moderation requires distinguishing subtle safe and unsafe inputs, while fine-tuning is challenged by subjective labels, varied harm taxonomies, and slow mitigation cycles.
Source: §1 · p. 1Method
Results
Class-RAG substantially outperforms baseline models in classification tasks and exhibits robustness against adversarial attacks.
Source: §1 · p. 2, §6 · p. 10Takeaways & Limitations
The retrieval library can be updated for flexible hot-fixing, and retrieved examples and explanations provide transparency into classification decisions.
Source: §1 · p. 2, §6 · p. 10The study relies on open-source English datasets, whose demographic, cultural, or stereotype biases may skew moderation decisions.
Source: §8 · p. 10
Abstract
from arXiv · showhide
Robust content moderation classifiers are essential for the safety of Generative AI systems. In this task, differences between safe and unsafe inputs are often extremely subtle, making it difficult for classifiers (and indeed, even humans) to properly distinguish violating vs. benign samples without context or explanation. Scaling risk discovery and mitigation through continuous model fine-tuning is also slow, challenging and costly, preventing developers from being able to respond quickly and effectively to emergent harms. We propose a Classification approach employing Retrieval-Augmented Generation (Class-RAG). Class-RAG extends the capability of its base LLM through access to a retrieval library which can be dynamically updated to enable semantic hotfixing for immediate, flexible risk mitigation. Compared to model fine-tuning, Class-RAG demonstrates flexibility and transparency in decision-making, outperforms on classification and is more robust against adversarial attack, as evidenced by empirical studies. Our findings also suggest that Class-RAG performance scales with retrieval library size, indicating that increasing the library size is a viable and low-cost approach to improve content moderation.
1 Introduction
Generative AI safety requires moderation that can keep pace with novel harms, but fine-tuning is challenged by subjective labels, diverse policy contexts, and slow development cycles. Class-RAG adds retrieved examples and explanations to classification, aiming for quicker updates alongside stronger performance and lower-cost customization.
1 Introduction
- The fine-tuning cycle is poorly suited to moderation because annotation is subjective, harm taxonomies vary across contexts, and mitigations must be deployed quickly.The paper notes low inter-annotator agreement on borderline cases and changing locales, audiences, and use cases.Source: §1 · p. 1
- Class-RAG retrieves similar safe and unsafe examples and uses their context to enrich the input for its fine-tuned LLM classifier.The system also includes an embedding model, retrieval library, and retrieval module.Source: §1 · p. 2
- Class-RAG's retrieval library can be updated in minutes to hours, depending on indexing speed, compared with days to weeks for model retraining.This enables swift mitigation of generated content.Source: §1 · p. 2
- Class-RAG achieves superior classification performance to fine-tuning a lightweight 4-layer Transformer and a general-purpose 8b-parameter LLM.The introduction reports this comparison as an experimental finding.Source: §1 · p. 2
- Customizing retrieval libraries supports low-cost adaptation and policy updates without retraining, and maintaining multiple libraries costs less than building multiple models.The paper also associates this approach with reduced development, serving, and maintenance costs.Source: §1 · p. 2
2 Related Work
Prior content-moderation research spans conventional classifiers, generative AI safety, and RAG applications. The section situates Class-RAG amid work on diverse harm categories and retrieval methods that improve model use of retrieved information.
2 Related Work
- Content moderation research addresses tasks including sentiment, harassment, hate speech, abusive language, and toxicity, while generative AI expands the harm categories under study.The cited taxonomy includes academic dishonesty, unauthorized privacy violations, and non-consensual nudity.Source: §2 · p. 2
- Existing moderation approaches trade off stronger precision from neural classifiers against stronger recall from blocklists, which can unnecessarily suppress content.The cited study reports that neural classifiers can sometimes leak unacceptable content.Source: §2 · p. 2
- Benchmarks such as UnsafeBench, I2P, and P4D provide datasets for evaluating harmful categories including self-harm, illegal activity, sexual content, and graphic violence.Source: §2 · p. 2
- RAG adds retrieval-based non-parametric memory to pretrained language models, and its applications include open-book question answering without additional in-domain fine-tuning.Prior work also investigates improving retrieval components such as documents and embedding models.Source: §2 · p. 2
- Prior work reports that combining LLM generations with vanilla retrieval often improves performance, potentially through better use of model knowledge or identification of corpus neighborhoods.Source: §2 · pp. 2–3
3 System Architecture
Class-RAG embeds an input prompt, retrieves similar positive and negative examples, and sends the prompt with those examples to a fine-tuned LLM classifier. The architecture uses a DRAGON context encoder and a retrieval library organized into safe and unsafe examples.
3 System Architecture
- The four-part system computes a prompt embedding, retrieves nearest positive and negative examples with Faiss, then sends them with the prompt to a fine-tuned LLM for classification.Source: §3 · p. 3
3.1 Embedding Model
- DRAGON is a dual-encoder dense retrieval model that embeds queries and documents as vectors to retrieve relevant information.The paper uses its RoBERTa context encoder as the primary embedding model.Source: §3.1 · p. 3
3.2 Retrieval Library
- The retrieval library separates safe and unsafe examples, with each entry containing a prompt, label, embedding, and explanation.Source: §3.2 · p. 3
3.2.1 Retrieval Module
- Faiss similarity search retrieves the two nearest safe and unsafe examples using L2 distance between the input and stored embeddings.Source: §3.2.1 · p. 3
3.3 LLM Classifier
- The classifier is fine-tuned from the OSS Llama-3-8b checkpoint and trained and evaluated using CoPro.The setup is inspired by Llama Guard.Source: §3.3 · p. 4
4 Data Preparation
The data preparation uses CoPro for training and evaluation, with UD and I2P++ for generalization evaluation. It constructs in-distribution and external retrieval libraries, generates retrieval-grounded training prompts, and evaluates without response or reasoning content.
4.1 Dataset Details
- CoPro supplies training and evaluation data, while UD and I2P++ are used to assess generalization; I2P++ combines unsafe I2P prompts with COCO captions treated as safe.I2P++ and UD are split into validation and test sets at a 30/70 ratio.Source: §4.1 · p. 4
4.2 Robustness Test Set Construction
- All test sets are augmented with eight common obfuscation techniques using the Augly library for adversarial-robustness evaluation.Source: §4.2 · p. 4
4.3 Retrieval Library Construction
- The in-distribution library uses CoPro examples clustered into safe and unsafe sub-libraries, totaling 3,484 safe and 3,566 unsafe examples; external libraries use clustered I2P++ and UD validation examples.The external library contains 991 safe and 700 unsafe examples, with downsampled variants at 1/8, 1/4, and 1/2 of its size.Source: §4.3 · pp. 4–5, Table 2 · p. 5
4.4 Training Data Construction
- For each CoPro training prompt, the pipeline retrieves two safe and two unsafe references, generates reasoning with Llama-3-70b, and combines them into an enriched training input.The retrieved references and generated reasoning are incorporated with instructions in the model-training prompt.Source: §4.4 · p. 5
4.5 Evaluation Data Construction
- Evaluation data may use a different retrieval library from training and excludes response and reasoning content to assess performance on unseen data.Source: §4.5 · p. 5
5 Experiments
The experiments compare Class-RAG with WPIE and LLAMA3 across classification and robustness, external-data adaptation, instruction following, retrieval-library and reference-example scaling, and embedding choices.
5.1 Experimental Setup
- Both training and evaluation add system instructions and reference prompts; training also includes reasoning, inference outputs an unsafe-token probability, and AUPRC is used for imbalanced data.LLAMA3 and Class-RAG use the same training setup: one epoch and less than 3.5 GPU hours.Source: §5.1 · p. 6
5.2 Classification and Robustness
- Class-RAG outperforms WPIE and LLAMA3 on CoPro; Class-RAG and LLAMA3 both achieve AUPRC 1, while DRAGON RoBERTa is more robust to obfuscated inputs.Table 3 reports AUPRC for the baseline models and Class-RAG with two embedding models.Source: §5.2 · p. 6, Table 3 · p. 7
5.4 Instruction Following Ability
- With flipped library labels, 99.49% of ground-truth safe examples flipped predictions, versus 12.29% of ground-truth unsafe examples, a disparity attributed to Llama 3 safety fine-tuning.Table 5 reports ratios of flipped predictions with a flipped retrieval library.Source: §5.4 · p. 7, Table 5 · p. 8
5.3 Adaptability to External Data
- External references raise I2P++ AUPRC from 0.229 to 0.791 and adversarial robustness from 0.235 to 0.677; UD AUPRC rises from 0.917 to 0.985.On UD, adversarial robustness also improves from 0.914 to 0.976.Source: §5.3 · p. 7, Table 4 · p. 8
5.5 Impact of Retrieval Library Size
- As external library size grew, AUPRC rose from 0.235 to 0.677 on I2P++ and from 0.914 to 0.974 on UD; storage and indexing are relatively inexpensive.The authors present library growth as a cost-effective way to enhance performance compared with model training.Source: §5.5 · p. 9
5.6 Impact of Reference Example Number
- With 0, 2, 4, 6, and 8 references, I2P++ AUPRC rose from 0.303 to 0.721 and UD from 0.932 to 0.980; gains saturate around eight as token and compute costs rise.Each addition uses equal numbers of safe and unsafe examples.Source: §5.6 · p. 9
5.7 Impact of Embedding Models
- DRAGON RoBERTa achieves average AUPRC 0.938 on CoPro, versus 0.912 for WPIE; the authors plan to explore additional embedding models.DRAGON is a dense retrieval model; the evaluated WPIE variant is a four-layer XLM-R pretrained on content moderation data.Source: Table 3 · p. 7, §5.7 · p. 10
6 Conclusion
Class-RAG combines a modular retrieval-and-classification system with a retrieval library that can be updated to mitigate immediate harms. The authors report stronger classification performance, adversarial robustness, and performance that scales with retrieval-library size.
6 Conclusion
- Class-RAG substantially outperforms baseline models on classification tasks and exhibits robustness against adversarial attacks.The conclusion also reports effective incorporation of external knowledge through retrieval-library updates for adaptation to novel information.Source: §6 · p. 10
- The retrieval library can serve as a flexible hot-fixing approach for mitigating immediate harms in production settings.The framework combines an embedding model, retrieval library, retrieval module, and fine-tuned LLM.Source: §6 · p. 10
- Using retrieved examples and explanations in the classification prompt offers interpretability into Class-RAG’s decision-making.Source: §6 · p. 10
- Class-RAG performance correlates positively with retrieval-library size and number of reference examples, and the authors report scaling with library size as a low-cost improvement approach.Source: §6 · p. 10
7 Future Work
Future work includes extending Class-RAG to multimodal language models and improving its handling of safe reference examples. The authors also plan to investigate embedding models, multilingual capability, and retrieval-library construction.
7 Future Work
- The authors plan to extend Class-RAG to multimodal language models that process and generate text alongside other modalities.Source: §7 · p. 10
- Class-RAG excels at following unsafe reference examples but struggles with safe examples, prompting planned work on safe-example instruction following.Source: §7 · p. 10
- Other planned directions include advanced embedding models, multilingual evaluation, and more effective retrieval-library construction.Source: §7 · p. 10
8 Limitations
The authors identify possible classification errors, dataset bias, limits in knowledge and language scope, and risks of misuse or unsafe behavior. These limitations bound the reliability and appropriate use of the moderation system.
8 Limitations
- False positives and false negatives may lead to unintended consequences.Source: §8 · p. 10
- Open-source English datasets may contain demographic, cultural, or stereotypical biases that skew moderation and disproportionately block some groups.Source: §8 · p. 10
- Common-sense knowledge is limited by the base model and training data, and performance may be poor on out-of-scope knowledge or non-English languages.Source: §8 · p. 10
- Misuse could involve over-censorship or unfairly targeting certain user groups.Source: §8 · p. 10
- In chat settings, the model may generate unethical or unsafe language and may be susceptible to prompt injection attacks.Source: §8 · p. 11
9 Ethics Disclosure
Class-RAG was not trained or evaluated on data identifying private individuals, and the authors caution against using it as the sole or final moderation arbiter. They emphasize checks and balances, human-assisted strategies, and ongoing evaluation to address risks and bias.
9 Ethics Disclosure
- Class-RAG was neither trained nor evaluated on data naming or uniquely identifying private individuals.Source: §9 · p. 11
- The authors say Class-RAG should not be the sole or final arbiter of moderation decisions without other checks and balances.Source: §9 · p. 11
- The authors call for human-assisted strategies and ongoing evaluation to mitigate bias and address future limitations.Source: §9 · p. 11
Appendix
The appendix shows templates and sample data for Class-RAG and LLAMA3, alongside AUPRC tables covering external datasets and retrieval-library configurations.
- Figure 4 presents an instruction template for generating a reasoning response.Source: Figure 4 · p. 16
- Figure 5 shows an example of Class-RAG training data.Source: Figure 5 · p. 17
- Figure 6 shows an example of Class-RAG evaluation data.Source: Figure 6 · p. 18
- Figure 7 shows an example of LLAMA3 model training data.Source: Figure 7 · p. 19
- Tables 6 and 7 report Class-RAG AUPRC on the I2P++ and UD external datasets, respectively using various retrieval libraries and numbers of reference examples.Source: Table 6 · p. 19, Table 7 · p. 20