Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Enhancing Safety of Large Language Models via Embedding Space Separation

Xu Zhao, Xiting Wang, Weiran Shen

Abstract

Large language models (LLMs) have achieved impressive capabilities, yet ensuring their safety against harmful prompts remains a critical challenge. Recent work has revealed that the latent representations (embeddings) of harmful and safe queries in LLMs typically exhibit linear separability, a property that has been exploited to construct attacks by perturbing the embeddings of harmful queries towards the safe subspace. Motivated by this observation, we propose a representation-level fine-tuning approach, named Embedding Space Separation (ES2), which improves LLM safety by explicitly enlarging the distance between harmful and safe representations in the embedding space. To prevent degradation of model's general capabilities, we introduce a Kullback-Leibler (KL) divergence regularization term into the loss function, which constrains the logits of the fine-tuned model to align with those of the original base model on harmless inputs. We evaluate our method on several open-source LLMs using standard safety benchmarks. Extensive experimental results demonstrate that our approach substantially improves model safety while maintaining comparable general capabilities.

Enhancing Safety of Large Language Models via Embedding Space Separation

Abstract

Large language models (LLMs) have achieved impressive capabilities, yet ensuring their safety against harmful prompts remains a critical challenge. Recent work has revealed that the latent representations (embeddings) of harmful and safe queries in LLMs typically exhibit linear separability, a property that has been exploited to construct attacks by perturbing the embeddings of harmful queries towards the safe subspace. Motivated by this observation, we propose a representation-level fine-tuning approach, named Embedding Space Separation (ES2), which improves LLM safety by explicitly enlarging the distance between harmful and safe representations in the embedding space. To prevent degradation of model's general capabilities, we introduce a Kullback-Leibler (KL) divergence regularization term into the loss function, which constrains the logits of the fine-tuned model to align with those of the original base model on harmless inputs. We evaluate our method on several open-source LLMs using standard safety benchmarks. Extensive experimental results demonstrate that our approach substantially improves model safety while maintaining comparable general capabilities.
Paper Structure (28 sections, 8 equations, 3 figures, 9 tables, 1 algorithm)

This paper contains 28 sections, 8 equations, 3 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Large Language Models and Embedding Space
  5. Safety Mechanisms of LLMs
  6. Method
  7. Embedding Space Separation via Distance Maximization
  8. Capabilities Preservation via KL Regularization
  9. Training Algorithm
  10. Experiments
  11. Experimental Setup
  12. Main Results: Safety Defense
  13. General Capabilities Preservation
  14. Incoherent and Gibberish Output Statistics
  15. Perturbation Distance
  16. ...and 13 more sections

Figures (3)

  • Figure 1: Intuition of ES2. (a) In the original base model, the embedding space exhibits a small margin between harmful (blue) and harmless (orange) prompts, making the safety boundary easily traversable. (b) ES2 explicitly increasing this margin, creating a large and safety margin. (c) For the base model, a small adversarial perturbation ($\epsilon_{small}$) is sufficient to cross the linear separation hyperplane, leading to a successful attack. (d) In the fine-tuned model, the increased margin forces the attacker to apply a large perturbation ($\epsilon_{large}$) to breach the safety guardrail, which alters or destroys the semantics of the prompt, effectively neutralizing the attack.
  • Figure 2: Classification accuracy with increasing layer depth of different LLMs.
  • Figure 3: Average perturbation distance required by the SCAV attack across different defense methods and LLMs.