Smoothed Embeddings for Robust Language Models

2025-01-27

AI Summary

To address the vulnerability of large language models (LLMs) to jailbreak attacks and their propensity to generate harmful content, this paper proposes a robustness-enhancement method integrating embedding-layer random smoothing with dynamic token-level aggregation. The approach injects controllable random noise into the embedding space and applies vector-level smoothing, coupled with a token-wise probability aggregation mechanism, to suppress adversarial perturbation propagation while preserving semantic fidelity. Its core innovation lies in the first application of random smoothing directly to the LLM embedding layer—jointly optimized with adaptive token aggregation—to achieve a superior trade-off between robustness and generation quality. Evaluated on mainstream jailbreak attack benchmarks, the method reduces average attack success rates by 42% while incurring less than 1.5% degradation in downstream task performance, significantly outperforming existing defense techniques.

Technology Category

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

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Improving the safety and reliability of large language models (LLMs) is a crucial aspect of realizing trustworthy AI systems. Although alignment methods aim to suppress harmful content generation, LLMs are often still vulnerable to jailbreaking attacks that employ adversarial inputs that subvert alignment and induce harmful outputs. We propose the Randomized Embedding Smoothing and Token Aggregation (RESTA) defense, which adds random noise to the embedding vectors and performs aggregation during the generation of each output token, with the aim of better preserving semantic information. Our experiments demonstrate that our approach achieves superior robustness versus utility tradeoffs compared to the baseline defenses.

Problem

• Large Language Models
• Jailbreak Attacks
• Security and Reliability

Innovation

• RESTA
• random disturbance
• security enhancement