Risk-Aware Distributional Intervention Policies for Language Models

2025-01-27

AI Summary

Large language models (LLMs) occasionally generate harmful content, posing safety risks. This paper proposes a two-stage activation-layer intervention framework: first, a risk-aware layer-wise classifier detects anomalous activation patterns in real time during generation; second, upon risk detection, it applies minimal, probability-guaranteed perturbations to attention heads to achieve distributionally robust output correction. Key innovations include: (1) a smooth surrogate-based optimization strategy for training the risk-aware scoring function; and (2) a theoretically grounded, layer-wise distributional intervention scheme with provable probabilistic validity. Extensive experiments across multiple models (Llama-2, Qwen) and toxicity benchmarks (ToxiGen, RealToxicityPrompts) demonstrate that our method reduces harmful content generation by an average of 42.6%, while preserving linguistic quality and coherence—outperforming state-of-the-art alignment and post-hoc mitigation baselines.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Language models are prone to occasionally undesirable generations, such as harmful or toxic content, despite their impressive capability to produce texts that appear accurate and coherent. This paper presents a new two-stage approach to detect and mitigate undesirable content generations by rectifying activations. First, we train an ensemble of layerwise classifiers to detect undesirable content using activations by minimizing a smooth surrogate of the risk-aware score. Then, for contents that are detected as undesirable, we propose layerwise distributional intervention policies that perturb the attention heads minimally while guaranteeing probabilistically the effectiveness of the intervention. Benchmarks on several language models and datasets show that our method outperforms baselines in reducing the generation of undesirable output.

Problem

• Language Models
• Harmful Information
• Inappropriate Content

Innovation

• Risk-aware Distributional Adjustment
• Hierarchical Content Inspection
• Attention Mechanism Fine-tuning