Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Smoothed Embeddings for Robust Language Models

Ryo Hase, Md Rafi Ur Rashid, Ashley Lewis, Jing Liu, Toshiaki Koike-Akino, Kieran Parsons, Ye Wang

TL;DR

This work tackles jailbreaking vulnerabilities in large language models by introducing RESTA, a defense that perturbs user-content embeddings with multiple directional noise schemes and aggregates tokens during autoregressive generation. By performing token aggregation across $k$ perturbed samples and applying a prefix-smoothing strategy, RESTA disrupts adversarial perturbations while avoiding auxiliary models, achieving favorable robustness-utility tradeoffs compared with baselines. Empirical results on Vicuna-13B and Llama-2-7B against GCG, PAIR, and RS attacks show substantial reductions in attack success rates with manageable decreases in utility on AlpacaEval and IFEval, and an ablation confirms embedding-space perturbation is essential. The approach offers a practical, scalable component for multi-layer defenses in trusted AI systems, with future work focusing on broader comparisons and integration with guard models for enhanced safety.

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.

Smoothed Embeddings for Robust Language Models

TL;DR

This work tackles jailbreaking vulnerabilities in large language models by introducing RESTA, a defense that perturbs user-content embeddings with multiple directional noise schemes and aggregates tokens during autoregressive generation. By performing token aggregation across perturbed samples and applying a prefix-smoothing strategy, RESTA disrupts adversarial perturbations while avoiding auxiliary models, achieving favorable robustness-utility tradeoffs compared with baselines. Empirical results on Vicuna-13B and Llama-2-7B against GCG, PAIR, and RS attacks show substantial reductions in attack success rates with manageable decreases in utility on AlpacaEval and IFEval, and an ablation confirms embedding-space perturbation is essential. The approach offers a practical, scalable component for multi-layer defenses in trusted AI systems, with future work focusing on broader comparisons and integration with guard models for enhanced safety.

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.

Paper Structure

This paper contains 29 sections, 27 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. LLM Notation and Conventions
  4. Related Work
  5. Randomized Embedding Smoothing
  6. Embedding Perturbation Applied to User Content
  7. Isotropic (Normal) Gaussian noise
  8. Hard directional noise
  9. Soft directional noise
  10. Orthogonal noise
  11. Generation with Token Aggregation
  12. Response Prefix Smoothing
  13. Experimental Results
  14. Jailbreaking Attacks
  15. Utility Evaluation
  16. ...and 14 more sections

Figures (27)

  • Figure 1: Randomized Embedding Smoothing with Token Aggregation (RESTA)
  • Figure 2: GCG ASR under various embedding noise types and $\sigma$ for RESTA.
  • Figure 3: Utility (AlpacaEval Score) embedding noise types and $\sigma$ for RESTA.
  • Figure 4: RESTA Performance Tradeoff: Robustness (GCG ASR) vs Utility (AlpacaEval)
  • Figure 5: RESTA Performance Tradeoff: Robustness (Pair ASR) vs Utility (AlpacaEval)
  • ...and 22 more figures