Research map / Large language models / Alignment and preferences
LEGEND: Leveraging Representation Engineering to Annotate Safety Margin for Preference Datasets
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes LEGEND, a training-free method that uses a safety direction (Standard Margin Vector) in LLM embedding space, derived from harmful/harmless responses, to automatically annotate binned safety margins for preference datasets.
Problem
Preference datasets for safety reward modeling use binary harmless/harmful labels that do not capture how much safer one response is; human margin annotation is costly and existing automatic margin methods (RewardEnsemble) require training multiple reward models.
Approach
Construct a safety direction (SMV) as the normalized mean difference between last-token embeddings of harmless and harmful responses from an annotator LLM, then project each pair's embedding difference onto this direction and bin the resulting values into discrete margins, requiring only inference.
Key results
- On Harmless, LEGEND improves average reward-model accuracy by 2.05% over Origin across Pythia and Llama2-7B-chat reward models, comparable to RewardEnsemble@3 (1.89%).
- LEGEND annotation takes about 23 minutes versus 25-276 minutes for RewardEnsemble@1-3 on the same hardware (roughly 11-fold reduction vs. the ensemble).
- On Safe-RLHF, LEGEND gives +2.34% average accuracy gain, and reaches 63.48% with Pythia-2.8B and 70.88% with Llama2-7B-chat, outperforming RewardEnsemble@3 on several models.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add LEGEND margin annotations to the Harmless preference dataset | Reward models across Pythia-410M, Pythia-1.4B, Pythia-2.8B, and Llama2-7B-chat · Harmless | accuracy: +3.65, +1.99, +1.53, +1.04 | Helped |
| Add LEGEND margin annotations to the Safe-RLHF preference dataset | Reward models across Pythia-410M, Pythia-1.4B, Pythia-2.8B, and Llama2-7B-chat · Safe-RLHF | accuracy: +0.17, +1.75, +5.39, +2.04 | Helped |
| Use LEGEND instead of RewardEnsemble@3 for margin annotation | Reward models across the Harmless and Safe-RLHF benchmarks · Harmless and Safe-RLHF | accuracy | Mixed |
| Remove safety-direction SMV projection from LEGEND (ablation) | LEGEND with Llama2-7B Base annotator · Safe-RLHF | accuracy gain over Origin: -1.04_{\pm 2.42} | Hurt |
| Add binning to LEGEND rather than using continuous margins (ablation) | LEGEND with Llama2-7B Base annotator · Safe-RLHF | accuracy: at least a 1.42% increase | Helped |
| Vary the number of bins from 10 to 3, 5, or 7 (ablation) | LEGEND reward-model training · Safe-RLHF | accuracy gain over Origin: 1.69_{\pm 1.46}; 1.99_{\pm 0.99}; -0.01_{\pm 0.68}; 2.34_{\pm 2.19} | Mixed |
| Replace Llama2-7B Base with Llama2-13B Base as Annotator LLM (ablation) | LEGEND · Safe-RLHF | accuracy gain over Origin: -0.39_{\pm 0.57} | Mixed |
| Use Llama2-13B Base instead of Llama2-7B Base as Annotator LLM (ablation) | LEGEND · Safe-RLHF | accuracy: 52.32; 56.39; 58.75; 70.49 vs 53.73; 59.63; 63.48; 70.88 | Mixed |
| Use a larger Annotator LLM, Llama2-13B Base (ablation) | LEGEND · Harmless | accuracy: 66.61; 73.30; 72.67; 72.63 vs 72.22; 74.39; 71.05; 73.44 | Mixed |
| Use fewer bins for smaller reward models and more bins for larger reward models (ablation) | LEGEND · Safe-RLHF | accuracy | Mixed |
| Use LEGEND instead of Origin for downstream best-of-n selection | Pythia-6B-static-sft policy models · Safe-RLHF test questions | win rate: 7% to 14% | Helped |
| Increase best-of-n sample size from 32 or 64 to 128 or 256 with LEGEND-based reward models | Pythia-2.8B and Qwen-4B-chat reward models · Safe-RLHF test questions | win rate: 7% to 14% | Helped |
Surprising findings
- Using a larger annotator (Llama2-13B) gave less concentrated gains; its margin distribution was more concentrated, and performance fell on smaller reward models.
- Removing the SMV projection lowered accuracy on several reward models (e.g., -3.99% for Pythia-1.4B on Safe-RLHF), showing the projection matters.
Limitations
- Relies on the linear representation property, whose validity is established mainly for safety, truthfulness, and toxicity; other features need templates and theory.
- Annotation consistency with human experts is modest (about 47% diagonal agreement on the 3-bin confusion matrix), and the method depends on the annotator LLM's ability to identify harmless responses.
- Downstream alignment evaluated on only 100 Safe-RLHF test questions with best-of-n and GPT-4 judging.
Experimental setup
Datasets: AdvBench Harmless Safe-RLHF Alpaca
Benchmarks: Harmless Safe-RLHF
Metrics: reward model accuracy win rate annotation time cost Spearman rank correlation annotation consistency with humans
Baselines: Origin RewardEnsemble@1 RewardEnsemble@2 RewardEnsemble@3
Builds on: RewardEnsemble InferAligner margin-based reward model loss Alpaca fine-tuning
Closest papers by meaning
- Larger or Smaller Reward Margins to Select Preferences for LLM Alignment? (ICML 2025 · similarity 0.89)
- Diverging Preferences: When do Annotators Disagree and do Models Know? (ICML 2025 · similarity 0.89)
- Spread Preference Annotation: Direct Preference Judgment for Efficient LLM Alignment (ICLR 2025 · similarity 0.89)
- Reward-Augmented Data Enhances Direct Preference Alignment of LLMs (ICML 2025 · similarity 0.88)
- Data with High and Consistent Preference Difference Are Better for Reward Model (AAAI 2025 · similarity 0.87)
- Less is More: Improving LLM Alignment via Preference Data Selection (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.