Enhancing Safety of Large Language Models via Embedding Space Separation

TL;DR

This paper introduces Embedding Space Separation (ES2), a fine-tuning method that enhances large language model safety by increasing the distance between harmful and safe query representations in the embedding space, while preserving overall performance.

Contribution

The paper proposes a novel embedding space separation technique with KL regularization for improving LLM safety without sacrificing capabilities.

Findings

Significant safety improvements on open-source LLMs

Maintains comparable general capabilities

Effective in standard safety benchmarks

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.