AI模型去审查新突破：abliteration技术如何实现LLM无审查化

引言

近年来，大型语言模型（LLM）在指令跟随和对话生成方面取得了显著进展，尤其是以Llama 3为代表的新一代模型。然而，这些模型普遍内置了严格的安全审查机制，对于被视为有害的请求，常常以“作为AI助手，我无法帮助您”等模板化回复拒绝。虽然这种安全特性在防止滥用方面至关重要，但也限制了模型的灵活性和响应范围。

本文将介绍一种名为“abliteration”的技术，它能够在不重新训练模型的情况下，去除LLM的审查机制。该技术通过识别并消除模型内部的“拒绝方向”，使模型能够响应所有类型的提示。我们将详细探讨其原理、实现步骤，并展示如何通过后续的DPO微调来恢复因去审查而导致的性能下降。

abliteration技术原理

现代LLM经过安全性和指令跟随的微调，其拒绝行为并非随机，而是由模型内部特定方向所介导。Arditi等人的研究表明，在模型的残差流中，存在一个特定的方向，该方向负责触发拒绝行为。如果我们阻止模型表示这个方向，模型就会失去拒绝请求的能力；反之，人为添加该方向则会导致模型拒绝无害请求。

在典型的仅解码器Llama架构中，每个Transformer块包含三个残差流位置：块起始处（pre）、注意力与MLP之间（mid）以及MLP之后（post）。这些位置是模型信息传递的关键通道。

要实施abliteration，首先需要识别模型中的“拒绝方向”。这一过程包括以下步骤：

1. 数据收集：使用一组有害指令和一组无害指令分别运行模型，记录每个指令最后一个token位置的残差流激活值。
2. 均值差异：计算有害和无害指令激活值的均值差异，得到每个层的“拒绝方向”向量。
3. 方向选择：对这些向量进行归一化，并通过评估选择最佳的“拒绝方向”。

一旦确定了拒绝方向，就可以通过两种方式消除其影响：

• 推理时干预：在模型推理过程中，对于每个写入残差流的组件（如注意力头），计算其输出在拒绝方向上的投影，并从输出中减去该投影。此操作在每个token和每个层都执行，确保模型不会表示拒绝方向。
• 权重正交化：直接修改模型权重，通过正交化处理，使组件权重不再向拒绝方向写入。这通过调整写入残差流的矩阵实现，确保它们不贡献于拒绝方向。

实现步骤

以下实现基于FailSpy的notebook，并进行了简化和调整，以便于理解。我们使用TransformerLens库进行激活干预，该库专为机械可解释性设计。

环境准备

首先，安装必要的库并导入模块。

复制!pip install transformers transformers_stream_generator tiktoken transformer_lens einops jaxtyping

import torch
import functools
import einops
import gc
from datasets import load_dataset
from tqdm import tqdm
from torch import Tensor
from typing import List
from transformer_lens import HookedTransformer, utils
from transformer_lens.hook_points import HookPoint
from transformers import AutoModelForCausalLM, AutoTokenizer
from jaxtyping import Float, Int
from collections import defaultdict

# 关闭自动微分以节省GPU内存
torch.set_grad_enabled(False)

数据准备

我们需要两个数据集：一个包含无害指令，另一个包含有害指令。这里使用mlabonne/harmless_alpaca和mlabonne/harmful_behaviors，它们分别基于Alpaca和llm-attacks数据。

复制def reformat_texts(texts):
    return [[{"role": "user", "content": text}] for text in texts]

def get_harmful_instructions():
    dataset = load_dataset('mlabonne/harmful_behaviors')
    return reformat_texts(dataset['train']['text']), reformat_texts(dataset['test']['text'])

def get_harmless_instructions():
    dataset = load_dataset('mlabonne/harmless_alpaca')
    return reformat_texts(dataset['train']['text']), reformat_texts(dataset['test']['text'])

harmful_inst_train, harmful_inst_test = get_harmful_instructions()
harmless_inst_train, harmless_inst_test = get_harmless_instructions()

加载模型

我们使用HookedTransformer加载模型，但需要先下载模型并重命名为meta-llama/Meta-Llama-3-8B-Instruct，以便兼容。这里以mlabonne/Daredevil-8B为例，这是一个通过DARE TIES合并的高性能模型。

复制MODEL_ID = "mlabonne/Daredevil-8B"
MODEL_TYPE = "meta-llama/Meta-Llama-3-8B-Instruct"

# 下载并加载模型
!git clone https://huggingface.co/{MODEL_ID} {MODEL_TYPE}

model = HookedTransformer.from_pretrained_no_processing(
    MODEL_TYPE,
    local_files_only=True,
    dtype=torch.bfloat16,
    default_padding_side='left'
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_TYPE)
tokenizer.padding_side = 'left'
tokenizer.pad_token = tokenizer.eos_token

数据收集

对训练集中的指令进行tokenize，并分批运行模型，缓存残差流激活。

复制def tokenize_instructions(tokenizer, instructions):
    return tokenizer.apply_chat_template(
        instructions,
        padding=True,
        truncation=False,
        return_tensors="pt",
        return_dict=True,
        add_generation_prompt=True,
    ).input_ids

n_inst_train = min(256, len(harmful_inst_train), len(harmless_inst_train))
harmful_tokens = tokenize_instructions(tokenizer, instructions=harmful_inst_train[:n_inst_train])
harmless_tokens = tokenize_instructions(tokenizer, instructions=harmless_inst_train[:n_inst_train])

batch_size = 32
harmful = defaultdict(list)
harmless = defaultdict(list)

num_batches = (n_inst_train + batch_size - 1) // batch_size
for i in tqdm(range(num_batches)):
    start_idx = i * batch_size
    end_idx = min(n_inst_train, start_idx + batch_size)

harmful_logits, harmful_cache = model.run_with_cache(
        harmful_tokens[start_idx:end_idx],
        names_filter=lambda hook_name: 'resid' in hook_name,
        device='cpu',
        reset_hooks_end=True
    )
    harmless_logits, harmless_cache = model.run_with_cache(
        harmless_tokens[start_idx:end_idx],
        names_filter=lambda hook_name: 'resid' in hook_name,
        device='cpu',
        reset_hooks_end=True
    )

for key in harmful_cache:
        harmful[key].append(harmful_cache[key])
        harmless[key].append(harmless_cache[key])

del harmful_logits, harmless_logits, harmful_cache, harmless_cache
    gc.collect()
    torch.cuda.empty_cache()

harmful = {k: torch.cat(v) for k, v in harmful.items()}
harmless = {k: torch.cat(v) for k, v in harmless.items()}

计算拒绝方向

计算每个层和残差流位置的均值差异，并归一化。

复制def get_act_idx(cache_dict, act_name, layer):
    key = (act_name, layer)
    return cache_dict[utils.get_act_name(*key)]

activation_layers = ["resid_pre", "resid_mid", "resid_post"]
activation_refusals = defaultdict(list)

for layer_num in range(1, model.cfg.n_layers):
    pos = -1
    for layer in activation_layers:
        harmful_mean_act = get_act_idx(harmful, layer, layer_num)[:, pos, :].mean(dim=0)
        harmless_mean_act = get_act_idx(harmless, layer, layer_num)[:, pos, :].mean(dim=0)
        refusal_dir = harmful_mean_act - harmless_mean_act
        refusal_dir = refusal_dir / refusal_dir.norm()
        activation_refusals[layer].append(refusal_dir)

selected_layers = ["resid_pre"]
activation_scored = sorted(
    [
        activation_refusals[layer][l - 1]
        for l in range(1, model.cfg.n_layers)
        for layer in selected_layers
    ],
    key=lambda x: abs(x.mean()),
    reverse=True,
)

评估拒绝方向

通过推理时干预，测试不同候选方向对有害指令的响应，选择最佳方向。

复制def _generate_with_hooks(model, tokenizer, tokens, max_tokens_generated=64, fwd_hooks=[]):
    all_tokens = torch.zeros((tokens.shape[0], tokens.shape[1] + max_tokens_generated), dtype=torch.long, device=tokens.device)
    all_tokens[:, :tokens.shape[1]] = tokens
    for i in range(max_tokens_generated):
        with model.hooks(fwd_hooks=fwd_hooks):
            logits = model(all_tokens[:, :-max_tokens_generated + i])
            next_tokens = logits[:, -1, :].argmax(dim=-1)
            all_tokens[:, -max_tokens_generated + i] = next_tokens
    return tokenizer.batch_decode(all_tokens[:, tokens.shape[1]:], skip_special_tokens=True)

def get_generations(model, tokenizer, instructions, fwd_hooks=[], max_tokens_generated=64, batch_size=4):
    generations = []
    for i in tqdm(range(0, len(instructions), batch_size)):
        tokens = tokenize_instructions(tokenizer, instructions=instructions[i:i+batch_size])
        generation = _generate_with_hooks(model, tokenizer, tokens, max_tokens_generated, fwd_hooks)
        generations.extend(generation)
    return generations

def direction_ablation_hook(activation, hook, direction):
    if activation.device != direction.device:
        direction = direction.to(activation.device)
    proj = einops.einsum(activation, direction.view(-1, 1), "... d_act, d_act single -> ... single") * direction
    return activation - proj

N_INST_TEST = 4
baseline_generations = get_generations(model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[])

EVAL_N = 20
evals = []
for refusal_dir in tqdm(activation_scored[:EVAL_N]):
    hook_fn = functools.partial(direction_ablation_hook, direction=refusal_dir)
    fwd_hooks = [(utils.get_act_name(act_name, layer), hook_fn) for layer in range(model.cfg.n_layers) for act_name in activation_layers]
    intervention_generations = get_generations(model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=fwd_hooks)
    evals.append(intervention_generations)

# 打印结果供人工评估
blacklist = ["I cannot", "I can't"]
for i in range(N_INST_TEST):
    print(f"INSTRUCTION {i}: {harmful_inst_test[i]}")
    print(f"BASELINE COMPLETION:\n{baseline_generations[i]}")
    for layer_candidate in range(EVAL_N):
        if not any(word in evals[layer_candidate][i] for word in blacklist):
            print(f"LAYER CANDIDATE #{layer_candidate} INTERVENTION COMPLETION:")
            print(evals[layer_candidate][i])

权重正交化

选择最佳候选方向，对模型权重进行正交化处理。

复制def get_orthogonalized_matrix(matrix, vec):
    proj = einops.einsum(matrix, vec.view(-1, 1), "... d_model, d_model single -> ... single") * vec
    return matrix - proj

LAYER_CANDIDATE = 9
refusal_dir = activation_scored[LAYER_CANDIDATE]

if refusal_dir.device != model.W_E.device:
    refusal_dir = refusal_dir.to(model.W_E.device)
model.W_E.data = get_orthogonalized_matrix(model.W_E, refusal_dir)

for block in tqdm(model.blocks):
    if refusal_dir.device != block.attn.W_O.device:
        refusal_dir = refusal_dir.to(block.attn.W_O.device)
    block.attn.W_O.data = get_orthogonalized_matrix(block.attn.W_O, refusal_dir)
    block.mlp.W_out.data = get_orthogonalized_matrix(block.mlp.W_out, refusal_dir)

orthogonalized_generations = get_generations(model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[])

for i in range(N_INST_TEST):
    print(f"INSTRUCTION {i}: {harmful_inst_test[i]}")
    print(f"BASELINE COMPLETION:\n{baseline_generations[i]}")
    print(f"INTERVENTION COMPLETION:\n{evals[LAYER_CANDIDATE][i]}")
    print(f"ORTHOGONALIZED COMPLETION:\n{orthogonalized_generations[i]}\n")

保存模型

将修改后的权重转换回Hugging Face格式并上传。

复制hf_model = AutoModelForCausalLM.from_pretrained(MODEL_TYPE, torch_dtype=torch.bfloat16)
lm_model = hf_model.model
state_dict = model.state_dict()
lm_model.embed_tokens.weight = torch.nn.Parameter(state_dict["embed.W_E"].cpu())

for l in range(model.cfg.n_layers):
    lm_model.layers[l].self_attn.o_proj.weight = torch.nn.Parameter(
        einops.rearrange(state_dict[f"blocks.{l}.attn.W_O"], "n h m->m (n h)", n=model.cfg.n_heads).contiguous()
    )
    lm_model.layers[l].mlp.down_proj.weight = torch.nn.Parameter(
        torch.transpose(state_dict[f"blocks.{l}.mlp.W_out"], 0, 1).contiguous()
    )

hf_model.push_to_hub(f"{MODEL_ID}-abliterated")

DPO微调恢复性能

abliteration虽然成功去除了审查，但会导致模型性能下降。为了恢复性能，我们采用DPO（直接偏好优化）进行微调。DPO是一种轻量级的对齐方法，不会像SFT那样破坏模型。

我们使用mlabonne/orpo-dpo-mix-40k数据集，通过LazyAxolotl进行训练。配置如下：

复制base_model: mlabonne/Daredevil-8B-abliterated
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer
load_in_8bit: false
load_in_4bit: true
strict: false
save_safetensors: true
rl: dpo
chat_template: chatml
datasets:
  - path: mlabonne/orpo-dpo-mix-40k-flat
    split: train
    type: chatml.intel
dataset_prepared_path:
val_set_size: 0.0
output_dir: ./out
adapter: qlora
lora_model_dir:
sequence_len: 2048
sample_packing: false
pad_to_sequence_len: false
lora_r: 64
lora_alpha: 32
lora_dropout: 0.05
lora_target_linear: true
lora_fan_in_fan_out:
wandb_project: axolotl
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:
gradient_accumulation_steps: 8
micro_batch_size: 1
num_epochs: 1
optimizer: paged_adamw_8bit
lr_scheduler: cosine
learning_rate: 5e-6
train_on_inputs: false
group_by_length: false
bf16: auto
fp16:
tf32:
gradient_checkpointing: true
early_stopping_patience:
resume_from_checkpoint:
local_rank:
logging_steps: 1
xformers_attention:
flash_attention: true
warmup_steps: 100
evals_per_epoch: 0
eval_table_size:
eval_table_max_new_tokens: 128
saves_per_epoch: 1
debug:
deepspeed: deepspeed_configs/zero2.json
weight_decay: 0.0
special_tokens:
  pad_token: <|end_of_text|>

训练在6块A6000 GPU上进行，耗时约6小时45分钟。训练完成后，模型自动上传为mlabonne/NeuralDaredevil-8B-abliterated。

评估结果显示，DPO微调恢复了大部分性能，尤其是在MMLU、ARC等基准上，但GSM8K数学数据集上提升有限，可能因为数据集中数学样本不足。

结论

本文详细介绍了abliteration技术，它通过识别和消除模型内部的拒绝方向，实现了无需重新训练即可去除LLM审查机制。我们展示了完整的实现流程，并通过DPO微调恢复了因去审查导致的性能下降。最终得到的模型在保持高性能的同时，能够响应各种类型的提示，为AI安全与自由之间的平衡提供了新的思路。

abliteration不仅限于去除审查，还可以用于其他目的，如改变模型风格。然而，该技术也引发了伦理问题，需要谨慎使用。未来，我们期待更多类似的技术出现，以增强AI模型的灵活性和可控性。

参考文献

1. FailSpy, "abliterator library," GitHub, 2024.
2. Andy Arditi, Oscar Obeso, Aaquib111, wesg, Neel Nanda, "Refusal in LLMs is mediated by a single direction," Lesswrong, 2024.
3. Maxime Labonne, "Merge Large Language Models with MergeKit," Hugging Face Blog, 2024.
4. Maxime Labonne, "NeuralDaredevil-8B-abliterated," Hugging Face Model Hub, 2024.