Legend: Leveraging Representation Engineering to Annotate Safety Margin for Preference Datasets

TL;DR

This work addresses the challenge of capturing nuanced safety differences in reward-model training by proposing Legend, a train-free framework that leverages representation engineering to discover a safety direction in LLM embeddings. It constructs a Standard Margin Vector (SMV) and annotates margins by projecting response differences onto this direction, enabling automatic, inference-time margin annotation without additional model training. Empirically, Legend improves reward-model accuracy and downstream harmless alignment on benchmark datasets (Harmless and Safe-RLHF) while significantly reducing annotation time, achieving comparable or better performance than training-based approaches like RewardEnsemble@K. The approach demonstrates robust generalization across annotator LLMs and datasets, and offers practical benefits for scalable safe-conversation alignment, with potential extension to other semantic features.

Abstract

The success of the reward model in distinguishing between responses with subtle safety differences depends critically on the high-quality preference dataset, which should capture the fine-grained nuances of harmful and harmless responses. This motivates the need to develop a dataset involving preference margins, which accurately quantify how harmless one response is compared to another. In this paper, we take the first step to propose an effective and cost-efficient framework to promote the margin-enhanced preference dataset development. Our framework, Legend, Leverages representation engineering to annotate preference datasets. It constructs the specific direction within the LLM's embedding space that represents safety. By leveraging this safety direction, Legend can then leverage the semantic distances of paired responses along this direction to annotate margins automatically. We experimentally demonstrate our effectiveness in both reward modeling and harmless alignment for LLMs. Legend also stands out for its efficiency, requiring only the inference time rather than additional training. This efficiency allows for easier implementation and scalability, making Legend particularly valuable for practical applications in aligning LLMs with safe conversations.

Figures (8)

• Figure 1: Win rate of policy models trained with enhanced reward models on the Safe-RLHF. The y-axis of each figure represents the value of $n$ in the best-of-$n$. Legend promotes the harmless alignment ability of policy models, particularly when utilizing large $n$.
• Figure 2: Results of Legend w/o SMV and Legend w/o bin.
• Figure 3: Results of Legend w/ different binning operations.
• Figure 4: Histogram of continuous margins on Safe-RLHF annotated by different Annonator LLMs. Llama2-13B Base has a more concentrated distribution, making it difficult to distinguish semantic differences.
• Figure 5: The Visualization of the margin vectors of the paired responses on the AdvBench dataset with Llama2-7B Base annotator LLM.

Theorems & Definitions (1)

• Remark 1