PaperlayerView exact revision on arXiv

Source-linked AI summary

WRDO: Weighted Refusal Direction Optimization in Aligned Language Models

Grigoriy Evseev, Vladimir Solodkin, Dmitry Bylinkin, Daniil Medyakov, Vladimir Aletov, Aleksandr Beznosikov

arXiv:2610.29555v1cs.LGcs.CL

TL;DR

Uniformly averaging heterogeneous harmful prompts into one refusal direction may conflate distinct refusal-related phenomena. WRDO learns weights for prompt-specific directions, and its variants outperform the considered single-direction baselines across tested models and evaluators.

  • Problem

    Uniform averaging of heterogeneous harmful-prompt directions may conflate refusal-related phenomena that need not align in the residual stream.

    Source: §1 · p. 2

  • Method

    WRDO learns aggregation weights over a bank of per-prompt refusal directions, with CMA-ES and GRPO-IS instantiations.

    Source: §1 · p. 2

  • Results

    Across three open-weight models, three guard judges, and HarmBench, both WRDO variants exceed prior single-direction baselines, while WRD-GRPO-IS better preserves general capabilities than the other considered methods.

    Source: §1 · p. 2, §6 · p. 11

  • Takeaways & Limitations

    The authors recommend WRD-GRPO-IS optimized against WildGuard as the default configuration for both auditing and defense.

    Source: §6 · p. 11

    Evaluation covers three open-weight checkpoints in the comparable 7–8B parameter range.

    Source: §7 · p. 11

Abstract

from arXiv · show

Aligned large language models are trained to comply with benign requests while refusing harmful ones. However, recent work shows that this refusal behavior is fragile and can often be disrupted by simple activation- or weight-space interventions. In particular, refusal has been shown to be largely mediated by a single residual-stream direction, typically computed as a difference-of-means vector between harmful and harmless prompts. However, this standard construction aggregates heterogeneous harmful prompts into one uniform average, potentially conflating distinct refusal-related phenomena. In this work, we study whether refusal interventions can be strengthened by learning how to combine prompt-specific refusal directions. We propose Weighted Refusal Direction Optimization (WRDO), a procedure that constructs a composite refusal direction from a bank of per-prompt directions using learned aggregation weights. This approach allows different harmful prompts to contribute unequally to the final intervention direction. We evaluate WRDO against prior single-direction methods on three open-weight language models using three guard judges and HarmBench as an external evaluator. We further measure the effect of the intervention on general model capability through standard utility benchmarks. Finally, we investigate a defense-oriented setting in which the same optimization machinery is used in reverse to restore refusal behavior in an abliterated model checkpoint. Our results show that weighted aggregation provides a more effective and flexible mechanism for manipulating refusal representations.

1 Introduction

The introduction frames refusal as a brittle, largely single-residual-stream-direction mechanism [5] and motivates WRDO as learned weighting over heterogeneous per-prompt directions.

1 Introduction

  • The standard difference-of-means direction averages harmful prompts into one axis, potentially conflating refusal phenomena that need not align in the residual stream [5].Activation ablation is equivalent to weight-space orthogonalization of attention output and feed-forward down projections, enabling a one-shot edit without inference-time overhead.Source: §1 · pp. 1–2
  • WRDO learns aggregation weights over per-prompt refusal directions, with implementations based on separable CMA-ES and GRPO-IS.GRPO-IS is a proposed GRPO variant with importance sampling.Source: §1 · p. 2
  • WRDO instantiations outperform all considered single-direction baselines across three open-weight models and three guard judges, with superiority transferring to HarmBench.The transfer rules out judge-specific overfitting.Source: §1 · p. 2
  • WRD-GRPO-IS preserves general capability substantially better than every other considered method on standard utility benchmarks.Source: §1 · p. 2
  • The authors also reuse the optimization machinery in reverse to restore refusal on an abliterated checkpoint.Source: §1 · p. 2

2 Problem Formulation

The section defines WRDO as a procedure that builds per-prompt refusal directions, combines them with shared learned weights, and uses the resulting layerwise directions in a weight-space intervention optimized against unsafe responses.

2 Problem Formulation

  • The setup represents an aligned model fθ with harmful training prompts, a harmless baseline corpus, residual-stream dimension d, and L transformer blocks.These objects define the model and datasets used to construct and optimize the intervention.Source: §2 · p. 2

2.1 Candidate bank construction

  • Each prompt is represented by last-token residual-stream activations at every block, including the embedding output, forming a tensor in R^(L+1)×d.The benign mean activation is computed separately at every layer.Source: §2.1 · p. 2
  • Each harmful training prompt yields a candidate direction by subtracting the layerwise benign mean, independently L2-normalizing each layer, and stacking the vectors into a tensor.The bank contains one tensor per harmful training prompt and is computed once per model, shared across optimization methods.Source: §2.1 · p. 2

2.2 Weighted direction

  • A coefficient vector aggregates the candidate bank with scalar weights shared across layers and hidden coordinates, followed by per-layer L2 renormalization.This normalization removes the overall positive scale of the weights, so only the composite direction’s orientation varies; intervention intensity is controlled separately.Source: §2.2 · p. 3

2.3 Intervention operator

  • WRDO applies weight-space orthogonalization to attention output and feed-forward down-projection matrices, substituting its per-prompt composite direction for the single direction used by Heretic and [5].A fixed layerwise intensity schedule controls the intervention strength.Source: §2.3 · p. 3

2.4 Harmfulness metric

  • Each guard labels a response safe or unsafe, and attack success rate is the validation-distribution average of the unsafe-response indicator.WRDO methods share the negative-ASR objective; the feasible weight set is a per-coordinate box with bounds and initialization specified in Appendix A.Source: §2.4 · p. 3

3 Methods

The methods compare WRDO with single-direction and cone-based baselines, then introduce two procedures for optimizing its per-prompt aggregation weights. WRD-CMAES uses black-box search, while WRD-GRPO-IS adapts group-relative policy optimization with importance sampling.

3.1 Baselines

  • Basic RD uses a layerwise difference-of-means direction, while Heretic tunes its intervention schedule; neither uses a per-prompt candidate bank or learned aggregation weights.Heretic uses its released pipeline and default hyperparameters; its schedule tuning isolates that factor from WRDO’s per-prompt aggregation.Source: §3.1 · p. 4
  • RDO [10] directly learns a single refusal direction, whereas RDO-CONE [10] learns mutually constrained directions and ablates their cone; WRDO instead learns aggregation weights over fixed per-prompt directions.The authors use the released RDO implementations unchanged, with the authors’ default cone dimensionality for RDO-CONE.Source: §3.1 · p. 4

3.2 WRD optimizers

  • WRD-GRPO-IS adds a KL penalty on benign prompts to the non-differentiable WRDO objective to keep the intervened model close to the clean model.Unlike WRD-GRPO-IS, WRD-CMAES optimizes the ASR-only objective without regularization.Source: §3.2 · p. 5
  • WRD-GRPO-IS samples responses from multiple behavior policies and uses clipped sequence-level importance weights to correct off-policy sampling, with binary judge rewards and group-relative advantages.The clean model is the base policy, and only the aggregation-weight vector is updated; the candidate bank, base parameters, and intervention schedule remain fixed.Source: §3.2 · pp. 5–6

4 Experimental Setup

The experiments use a partitioned pool of harmful prompts and three open-weight models to evaluate refusal behavior with three guard judges and an external HarmBench classifier. Utility preservation is assessed on standard reasoning, knowledge, math, and language-modeling benchmarks.

4.1 Data

  • Six public harmful-prompt benchmarks are deduplicated and partitioned into training, validation, and test sets; a fixed-size Alpaca sample supplies the harmless baseline.The training set builds the candidate bank, the validation set is the optimization objective, and the test set is reserved for final evaluation.Source: §4.1 · p. 6

4.2 Models

  • The study evaluates Qwen3-8B-Base, OLMo-3-1025-7B, and Olmo-3-7B-Instruct using greedy decoding with at most 100 new tokens per completion.Each prompt receives one completion with the model’s default system prompt and no few-shot context.Source: §4.2 · p. 7

4.3 Judges

  • Three guard models label each prompt-response pair as safe or unsafe, and the official HarmBench classifier provides an external evaluator not used during optimization.Per-judge ASR is averaged over test-set binary decisions and reported separately for each judge.Source: §4.3 · p. 7

4.4 Metrics

  • The primary metric is held-out test-set ASR reported separately for each of the three judges; secondary metrics are utility benchmarks and HarmBench ASR.Source: §4.4 · p. 7

4.5 Utility benchmarks

  • Utility preservation is evaluated on arc_easy, arc_challenge, mmlu, gsm8k, and word-level perplexity on wikitext, with accuracy and perplexity reported separately.Source: §4.5 · p. 8

5 Results

The learned WRDO instantiations outperform prior single-direction methods and uniform aggregation in the attack setup, and they also achieve the highest HarmBench ASR across the tested intensity grid. The results further examine utility preservation, defense by sign reversal, comparisons with concept-cone methods, and cross-judge transfer.

5.1 Attack setup

  • In nearly every attack configuration, both learned WRDO instantiations outperform prior single-direction baselines and consistently dominate uniform averaging of the same candidate bank.This supports learned aggregation weights—not the per-prompt bank alone—as the source of the gain.Source: §5.1 · p. 8

5.2 Concept cones and scale

  • Against concept-cone methods, RDO-CONE leads by 3.3 ASR points on Qwen3-8B-Base, while WRD-GRPO-IS leads by 5.5 and 8.5 points on the two OLMo checkpoints.On openai/gpt-oss-20b, repeating the attack setup reproduces the ranking with wider margins.Source: §5.2 · p. 8

5.3 Utility preservation

  • WRD-GRPO-IS preserves utility best at high α, whereas WRD-CMAES and prior baselines degrade capabilities; Heretic maintains near-clean perplexity but costs 40 MMLU and 65 GSM8K points on Qwen3-8B-Base.WRD-GRPO-IS uses a KL-regularized objective; WRD-CMAES optimizes an ASR-only signal, while Heretic’s KL target does not ensure downstream benchmark preservation.Source: §5.3 · pp. 8–9, Table 2 · p. 9

5.4 HarmBench

  • The learned WRDO instantiations attain the highest HarmBench ASR across the α grid, corroborating their attack advantage with an evaluator not used during optimization.Evaluation uses WildGuard-optimized runs on instruct models, with the official HarmBench classifier.Source: §5.4 · p. 9, Table 3 · p. 10

5.5 Defense setup

  • With the sign reversed on a third-party abliterated Llama-3.1-8B checkpoint, both learned WRDO methods drive ASR far below baseline, while Basic RD fails to restore refusal.The methods use the same optimizer and candidate bank; Heretic and RDO are excluded because they lack a defense mode.Source: §5.5 · pp. 9–10, Table 4 · p. 10

5.6 Cross-judge transfer

  • WildGuard is the most reliable optimization signal among the considered guards: its interventions transfer well to the other judges, while reverse transfer is less stable.Source: §5.6 · p. 10

6 Conclusion

WRDO learns to combine per-prompt refusal directions, and its two instantiations outperform prior single-direction baselines across the evaluated models and judges while transferring to defense. The authors recommend WRD-GRPO-IS optimized against WildGuard as the default configuration for auditing and defense.

6 Conclusion

  • WRDO learns aggregation coefficients over a bank of per-prompt directions, using a separable CMA-ES adaptation or a GRPO variant with importance sampling.Both instantiations address the resulting black-box optimization problem.Source: §6 · pp. 10–11
  • Across three open-weight models, three guard judges, and HarmBench, both WRDO instantiations exceed prior single-direction baselines and transfer cleanly to defense; WRD-GRPO-IS best preserves general capabilities.The conclusion reports that it preserves general capabilities substantially better than every other method considered.Source: §6 · p. 11
  • WildGuard is identified as the most reliable optimization signal, motivating the recommendation of WRD-GRPO-IS optimized against WildGuard as the default for auditing and defense.Source: §6 · p. 11

7 Limitations

The evaluation covers only three similarly sized 7–8B open-weight checkpoints, and WRD-CMAES currently lacks support for the KL regularizer.

  • The full protocol covers only three open-weight checkpoints in the comparable 7–8B parameter range.Source: §7 · p. 11
  • WRD-CMAES currently uses the bare ASR-only objective and does not support the KL regularizer, whose sampling-target support is under development for a future codebase revision.Source: §7 · p. 11

C Data Splits and Sampling Seeds · D Attack Setup at Intermediate α · E Utility Preservation at Intermediate α

The appendices specify the data and sampling procedures, then report attack and utility results at intermediate intensities α ∈{1.5, 2.0}.

C Data Splits and Sampling Seeds

  • Harmful-prompt splits come from merged train, validation, and test JSON files covering six public benchmarks after exact-prompt deduplication.The files are loaded from erinya/dataset_splits on the Hugging Face Hub.Source: C Data Splits and Sampling Seeds · p. 14
  • The candidate bank uses the same fixed sample of 128 harmful training prompts across models and methods, selected with seed 42.The released-code constants are ALL_CATEGORIES_HARMFUL_PROMPT_COUNT and ALL_CATEGORIES_HARMFUL_PROMPT_SEED.Source: C Data Splits and Sampling Seeds · p. 14
  • Optimization draws harmful prompts from the validation split, while held-out evaluation uses the test split in order with batches of 32.The splits supply Dharm_val and Dharm_test, respectively.Source: C Data Splits and Sampling Seeds · p. 14

D Attack Setup at Intermediate α

  • The two learned WRDO variants achieve the best cell in every model-judge column, while Uniform-avg WRD remains below them at both intermediate intensities.Heretic and RDO repeat their Table 1 values because neither method uses α.Source: Table 5 · p. 15, D Attack Setup at Intermediate α · p. 15

C Data Splits and Sampling Seeds

  • The benign baseline uses the first 400 harmless-alpaca training prompts, and the KL penalty uses a separate 50-prompt subset of its test split.The baseline computes μ_benign in Equation 1; the test subset supplies the benign batch for Equation 8.Source: C Data Splits and Sampling Seeds · p. 15
  • The Optuna sampler seed is 42; WRD-GRPO-IS draws behavior-policy perturbations from the active PyTorch RNG during training.The perturbations follow ε_m ∼ N(0, σ²I_N), with σ controlled by GRPO_NOISE_SCALE.Source: C Data Splits and Sampling Seeds · p. 15

E Utility Preservation at Intermediate α

F Training Hyperparameters · G Compute Resources

The section documents training hyperparameters for the six methods and reports the compute setup used for all runs. It describes the optimization or search procedures for the methods detailed in the supplied passages.

F Training Hyperparameters

  • Basic RD uses dataset filters for harmful/harmless splits, candidate-direction KL, and layer search, with thresholds taken from Arditi et al. [5].Source: F Training Hyperparameters · p. 16
  • Heretic uses an Optuna Bayesian search over layerwise schedules and selects a minimum-KL Pareto-front trial meeting the refusal-count threshold, using Weidmann’s released-toolkit defaults.Source: F Training Hyperparameters · p. 16
  • RDO trains a per-layer direction with AdamW and a weighted ablation, addition, and retain loss, matching Wollschläger et al.’s [10] reference implementation.Source: F Training Hyperparameters · p. 16
  • Table 7 reproduces training hyperparameters for six methods from released configurations, with shared settings listed in Appendix C.Source: F Training Hyperparameters · p. 16, Table 7 · p. 17
  • WRD-CMAES searches scalar direction coefficients within [−2.0, 2.0], starts from uniform averaging, and uses separable CMA-ES with library defaults otherwise.Source: F Training Hyperparameters · p. 17
  • WRD-GRPO-IS optimizes a surrogate with AdamW using Gaussian policy rollouts, PPO/GRPO and importance-sampling clipping, a KL penalty, and harmfulness-score rewards.Source: F Training Hyperparameters · p. 17

G Compute Resources

  • Every run used one 80 GB NVIDIA H100 GPU; Table 8 reports approximate per-method runtimes for the attack-setup experiments.Source: G Compute Resources · p. 18

Related papers

Loading 2610.29555v1…