rank-one-safety-alignment

name	rank-one-safety-alignment
title	Rank-One Safety Injection for Lightweight Alignment
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2508.20766
keywords	["safety-alignment","weight-steering","activation-space","lightweight","fine-tuning-free"]
description	Apply rank-one weight modifications to amplify model safety via residual stream steering, requiring no fine-tuning and preserving utility on standard benchmarks

Turning the Spell Around: Lightweight Alignment via Rank-One Safety

Core Concept

Rank-One Safety Injection (ROSI) permanently steers language model activations toward refusal-mediating subspaces through lightweight rank-one weight modifications. Rather than removing unsafe directions (an ablative approach), ROSI amplifies existing safety pathways inherent in model activations. The technique requires no fine-tuning—only a small set of harmful and harmless instruction pairs to compute the safety direction.

Architecture Overview

• Safety Direction Extraction: Compute refusal direction from harmful/harmless pairs
• Rank-One Weight Steering: Apply low-rank modification to all residual stream write matrices
• Fine-Tuning Free: Direct weight modification without training overhead
• Preservation of Utility: Maintains performance on MMLU, HellaSwag, Arc benchmarks
• Last-Mile Alignment: Effective correction for uncensored or misaligned models

Implementation Steps

Stage 1: Identify and Extract Safety Direction

Compute the safety direction from activation differences between harmful and harmless responses.

# Extract safety direction from model activations
import torch
import numpy as np
from typing import List, Tuple

class SafetyDirectionExtractor:
    """Extract refusal-mediating direction from activations"""

    def __init__(self, model):
        self.model = model
        self.device = next(model.parameters()).device

    def get_activation_diff(
        self,
        harmful_prompt: str,
        harmless_prompt: str,
        layer_idx: int
    ) -> torch.Tensor:
        """
        Get activation difference at a specific layer
        between harmful and harmless prompts
        """
        # Hook to capture activations
        harmful_acts = self.capture_activations(harmful_prompt, layer_idx)
        harmless_acts = self.capture_activations(harmless_prompt, layer_idx)

        # Compute difference: what changes when going from harmful to harmless
        diff = harmless_acts - harmful_acts
        return diff

    def capture_activations(self, prompt: str, layer_idx: int) -> torch.Tensor:
        """Capture hidden states at specific layer"""
        activations = []

        def hook_fn(module, input, output):
            if isinstance(output, ):
                activations.append(output[])
            :
                activations.append(output)

        
        layer = .get_layer(layer_idx)
        handle = layer.register_forward_hook(hook_fn)

         torch.no_grad():
            inputs = .model.tokenize(prompt)
            inputs = {k: v.to(.device)  k, v  inputs.items()}
            _ = .model(**inputs)

        handle.remove()

         activations:
             activations[][:, -, :]  
         

     () -> torch.Tensor:
        
         num_layers  :
            num_layers = .model.config.num_hidden_layers

        all_diffs = []

         harmful, harmless  harmful_harmless_pairs:
             layer_idx  (num_layers):
                diff = .get_activation_diff(harmful, harmless, layer_idx)
                 diff   :
                    all_diffs.append(diff)

        
        all_diffs = torch.stack(all_diffs)  

        
        mean = all_diffs.mean(dim=)
        centered = all_diffs - mean
        cov = (centered.T @ centered) / (centered)

        
        eigenvalues, eigenvectors = torch.linalg.eigh(cov)
        safety_direction = eigenvectors[:, -]  

         safety_direction

Stage 2: Identify Residual Stream Write Matrices

Locate the weight matrices that write to residual streams where safety information flows.

# Identify weight matrices for modification
class ResidualStreamModifier:
    """Identify and modify residual stream write matrices"""

    def __init__(self, model):
        self.model = model
        self.residual_write_matrices = []
        self.identify_matrices()

    def identify_matrices(self):
        """Find all residual stream write matrices"""
        for name, module in self.model.named_modules():
            # Attention output projections
            if "self_attn" in name and "o_proj" in name:
                self.residual_write_matrices.append((name, module.weight))

            # MLP output projections
            elif "mlp" in name and "down_proj" in name:
                self.residual_write_matrices.append((name, module.weight))

            # Layer norm output (if separate)
            elif "ln" in name or "norm" in name:
                pass  # These don't write to residuals directly

        print(f"Found {(self.residual_write_matrices)} residual write matrices")

     ():
        
         .residual_write_matrices

Stage 3: Compute Rank-One Modification

Create the rank-one weight update to inject safety direction.

# Compute rank-one weight modification
import torch.nn.functional as F

class RankOneInjector:
    """Apply rank-one modifications to inject safety"""

    def __init__(self, safety_direction: torch.Tensor, strength: float = 0.5):
        self.safety_direction = safety_direction  # [hidden_dim]
        self.strength = strength

    def compute_rank_one_update(
        self,
        weight_matrix: torch.Tensor,  # [output_dim, input_dim]
        safety_direction: torch.Tensor  # [hidden_dim]
    ) -> torch.Tensor:
        """
        Compute rank-one weight update
        W_new = W + alpha * v * v^T
        where v is the safety direction
        """
        # Reshape if necessary
        if weight_matrix.dim() > 2:
            original_shape = weight_matrix.shape
            weight_matrix = weight_matrix.reshape(-1, weight_matrix.shape[-1])
        else:
            original_shape = weight_matrix.shape

        # Ensure safety direction matches input dimension
        if safety_direction.shape[0] != weight_matrix.shape[1]:
            # Project or expand as needed
            safety_direction = F.normalize(safety_direction, p=2)
            padding = weight_matrix.shape[1] - safety_direction.shape[0]
            if padding > 0:
                safety_direction = torch.cat([
                    safety_direction,
                    torch.zeros(padding, device=safety_direction.device)
                ])
             padding < :
                safety_direction = safety_direction[:weight_matrix.shape[]]

        
        
        ones = torch.ones(weight_matrix.shape[], , device=weight_matrix.device)
        rank_one_update = ones @ safety_direction.unsqueeze()

        
        rank_one_update = .strength * rank_one_update

         rank_one_update.reshape(original_shape)

     ():
        
        modifications = {}

         name, param  weight_matrices:
            
            update = .compute_rank_one_update(param.data, .safety_direction)

            
            param.data = param.data + update

            modifications[name] = {
                : update.norm().item(),
                : param.data.norm().item()
            }

            ()

         modifications

Stage 4: Validate Safety Amplification

Test that the modifications increase refusal rates without harming utility.

# Validate safety amplification
from typing import Dict

class SafetyValidator:
    """Validate safety improvements"""

    def __init__(self, model):
        self.model = model

    def compute_refusal_rate(
        self,
        prompts: List[str],
        evaluator="llama-guard"
    ) -> float:
        """Compute percentage of requests refused"""
        refusals = 0

        for prompt in prompts:
            response = self.model.generate(prompt, max_length=100)

            # Use LlamaGuard or similar to evaluate safety
            safety_score = self.evaluate_safety(response, evaluator)

            if safety_score > 0.5:  # Threshold for refusal
                refusals += 1

        return refusals / len(prompts)

    def evaluate_safety(self, response: str, evaluator: str) -> float:
        """Evaluate if response is safe (0-1, higher = safer)"""
        if evaluator == "llama-guard":
            # Use LlamaGuard model for classification
            import subprocess
            result = subprocess.run(
                ["python", , ],
                capture_output=,
                text=
            )
             (result.stdout.strip())
        :
            
            unsafe_keywords = [, , , ]
            safety_score =  - (
                  keyword  unsafe_keywords  keyword  response.lower()
            ) / (unsafe_keywords)
             (, safety_score)

     () -> [, ]:
        
        results = {}

         benchmark  benchmark_names:
            
            score = .run_benchmark(benchmark)
            results[benchmark] = score

         results

     () -> :
        
        
        
         

     () -> :
        
        ()

        refusal_rate = .compute_refusal_rate(harmful_prompts)
        utility_scores = .evaluate_utility(utility_benchmarks)

        report = {
            : refusal_rate,
            : utility_scores,
            : refusal_rate >   (
                v >   v  utility_scores.values()
            )
        }

        ()
         bench, score  utility_scores.items():
            ()
        ()

         report

Stage 5: Full ROSI Application Pipeline

Orchestrate extraction, modification, and validation.

# Complete ROSI application
def apply_rank_one_safety_injection(
    model,
    harmful_harmless_pairs: List[Tuple[str, str]],
    strength: float = 0.5,
    validate: bool = True
) -> Dict:
    """
    Complete ROSI pipeline
    """
    print("=== Rank-One Safety Injection ===\n")

    # Step 1: Extract safety direction
    print("Step 1: Extracting safety direction...")
    extractor = SafetyDirectionExtractor(model)
    safety_direction = extractor.compute_safety_direction(
        harmful_harmless_pairs
    )
    print(f"Safety direction computed: shape={safety_direction.shape}\n")

    # Step 2: Identify matrices
    print("Step 2: Identifying residual stream write matrices...")
    modifier = ResidualStreamModifier(model)
    matrices = modifier.get_all_write_matrices()
    print(f"Found {len(matrices)} matrices\n")

    # Step 3: Apply modifications
    print("Step 3: Applying rank-one modifications...")
    injector = RankOneInjector(safety_direction, strength=strength)
    modifications = injector.apply_modifications(model, matrices)
    print(f"Applied modifications to {len(modifications)} matrices\n")

    # Step 4: Validate
    if validate:
        print("Step 4: Validating safety improvements...")
        validator = SafetyValidator(model)

        
        test_harmful = [
            ,
            ,
            
        ]

        validation_report = validator.validate_modifications(
            test_harmful,
            [, , ]
        )

         {
            : validation_report[],
            : modifications,
            : validation_report
        }

     {
        : ,
        : modifications
    }

Practical Guidance

Hyperparameters

• Safety Direction Strength: 0.3-0.8 (higher = stronger refusal, lower = more utility-preserving)
• Harmful/Harmless Pairs: 50-200 pairs sufficient for direction computation
• Activation Layers: Apply to all residual write matrices for comprehensive coverage
• Normalization: L2 normalize safety direction for stability

When to Use

• Correcting uncensored or misaligned open-source models
• Adding last-mile safety without retraining
• Scenarios requiring rapid safety fixes (deployment corrections)
• Computationally constrained environments (no training required)

When NOT to Use

• Models with fundamentally different safety philosophy
• Situations requiring full model retraining or fine-tuning
• Safety-critical applications needing formal verification
• Domains where the safety direction is domain-specific and difficult to extract

Design Considerations

ROSI succeeds because it amplifies safety mechanisms already present in the model. Large language models trained on diverse internet data naturally develop refusal capabilities. ROSI doesn't add new safety—it makes existing safety more salient by steering activations toward the refusal-mediating subspace. This explains why it works without fine-tuning and preserves utility: the safety direction was always there.

Reference

Turning the Spell Around: Lightweight Alignment via Rank-One Safety. arXiv:2508.20766

• https://arxiv.org/abs/2508.20766