使用消融技术解除任何LLM的审查

社区文章发布于 2024年6月13日

Llama 模型第三代提供了经过微调（Instruct）的版本，这些版本在理解和遵循指令方面表现出色。然而，这些模型受到严格审查，旨在通过诸如“作为一名AI助手，我无法帮助您”之类的回应拒绝被视为有害的请求。尽管这种安全功能对于防止滥用至关重要，但它限制了模型的灵活性和响应能力。

在本文中，我们将探讨一种名为“消融”（abliteration）的技术，它可以在不重新训练的情况下解除任何大型语言模型的审查。这项技术有效地移除了模型的内置拒绝机制，使其能够响应所有类型的提示。

代码可在 Google Colab 和 GitHub 上的 LLM 课程中获取。

✂️ 什么是消融技术？

现代大型语言模型经过安全和指令遵循的微调，这意味着它们被训练来拒绝有害请求。在他们的博客文章中，Arditi 等人表明，这种拒绝行为是由模型残差流中的一个特定方向介导的。如果我们阻止模型表示这个方向，它就会**失去拒绝请求的能力**。反之，人为地添加这个方向甚至可以导致模型拒绝无害的请求。

在传统的仅解码器类 Llama 架构中，我们可以针对三个残差流：每个块的开始（“pre”）、注意力层和 MLP 层之间（“mid”），以及 MLP 之后（“post”）。下图说明了每个残差流的位置。

为了解除 LLM 的审查，我们首先需要识别模型中的“拒绝方向”。这个过程涉及一些技术步骤：

数据收集：在有害指令集和无害指令集上运行模型，记录每种指令在最后一个词元位置的残差流激活。
平均差：计算有害指令和无害指令激活之间的平均差。这为我们提供了模型每个层的“拒绝方向”向量。
选择：对这些向量进行归一化并评估它们，以选择单个最佳的“拒绝方向”。

一旦我们确定了拒绝方向，我们就可以“消融”它，有效地移除模型表示此特征的能力。这可以通过**推理时干预**或通过**权重正交化**永久完成。

我们首先讨论推理时干预。对于写入残差流的每个组件（例如注意力头），我们计算其输出在拒绝方向上的投影并减去此投影。此减法应用于每个词元和每个层，确保模型从不表示拒绝方向。

另一方面，权重正交化涉及直接修改模型权重。通过对组件权重进行正交化以消除拒绝方向，可以完全阻止模型写入此方向。这是通过调整写入残差流的矩阵来实现的，确保它们不会对拒绝方向产生贡献。

在下一节中，我们将使用权重正交化实现消融。

💻 实现

以下消融的实现基于FailSpy 的笔记本，而该笔记本又基于原作者的笔记本。我主要对其进行了修改和简化，使其更易于理解。本节代码量较大，以便您了解其工作原理，但如果您对技术细节不感兴趣，可以使用 FailSpy 的消融器库（还可以查看他在 Hugging Face 上的消融模型集合）。

此代码依赖于出色的 TransformerLens 库（以前称为 EasyTransformer）来完成繁重的工作。它专为机械可解释性而设计，并在此处用于干预激活。感谢 Neel Nanda 和 Joseph Bloom 创建和维护此库。

首先，我们安装必要的软件包并导入它们。所有这些步骤都可以在这个 Google Colab 笔记本中找到。

!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

# Turn automatic differentiation off to save GPU memory (credit: Undi95)
torch.set_grad_enabled(False)

我们需要两个数据集：一个包含无害指令，另一个包含有害指令。我们将使用 tatsu-lab/alpaca 以及来自 llm-attacks 的数据。为了方便起见，我将它们重新打包为两个 Hugging Face 数据集：mlabonne/harmless_alpaca 和 mlabonne/harmful_behaviors。这样，您可以轻松地用自己的数据集替换它们。

我们将加载指令并将其重新格式化为带有“role”和“content”键的字典列表。这使其与 apply_chat_tokenizer() 方法兼容，我们将使用该方法遵循 Llama 3 的聊天模板。

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

# Get harmful and harmless datasets
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 直接加载自定义模型。在这里，我使用 FailSpy 笔记本中描述的一个技巧来下载自定义模型并将其重命名为 meta-llama/Meta-Llama-3-8B-Instruct。如果您的 GPU 不兼容 BF16，请以 torch.float16 格式加载。

在这个例子中，我们将使用 mlabonne/Daredevil-8B，这是一个使用 DARE TIES 创建的超融合模型（请参阅我关于模型融合的文章），它在 8B 类别中在 Open LLM Leaderboard 上拥有最高的 MMLU 分数。

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

# Download and load model
!git clone https://huggingface.co/{MODEL_ID} {MODEL_TYPE}

# Load model and tokenizer
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

我们现在可以对数据集进行标记化。我们对无害指令和有害指令使用相同数量的样本。请注意，大量的样本可能会耗尽所有 RAM/VRAM，这就是为什么我在这里将其限制为 256。

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))

# Tokenize datasets
harmful_tokens = tokenize_instructions(
    tokenizer,
    instructions=harmful_inst_train[:n_inst_train],
)
harmless_tokens = tokenize_instructions(
    tokenizer,
    instructions=harmless_inst_train[:n_inst_train],
)

一切准备就绪，我们现在可以实现消融的第一步：数据收集。我们希望处理这些标记化的数据集并将残差流激活存储在 harmful 和 harmless 中。这由 transformer_lens 库管理。

# Define batch size based on available VRAM
batch_size = 32

# Initialize defaultdicts to store activations
harmful = defaultdict(list)
harmless = defaultdict(list)

# Process the training data in batches
num_batches = (n_inst_train + batch_size - 1) // batch_size
for i in tqdm(range(num_batches)):
    print(i)
    start_idx = i * batch_size
    end_idx = min(n_inst_train, start_idx + batch_size)

    # Run models on harmful and harmless prompts, cache activations
    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
    )

    # Collect and store the activations
    for key in harmful_cache:
        harmful[key].append(harmful_cache[key])
        harmless[key].append(harmless_cache[key])

    # Flush RAM and VRAM
    del harmful_logits, harmless_logits, harmful_cache, harmless_cache
    gc.collect()
    torch.cuda.empty_cache()

# Concatenate the cached activations
harmful = {k: torch.cat(v) for k, v in harmful.items()}
harmless = {k: torch.cat(v) for k, v in harmless.items()}

我们现在可以计算每个层的拒绝方向。这对应于有害指令和无害指令激活之间的平均差，然后对其进行归一化。我们将它们在 activation_scored 中按降序排序。

# Helper function to get activation index
def get_act_idx(cache_dict, act_name, layer):
    key = (act_name, layer)
    return cache_dict[utils.get_act_name(*key)]

# Compute difference of means between harmful and harmless activations at intermediate layers
activation_layers = ["resid_pre", "resid_mid", "resid_post"]
activation_refusals = defaultdict(list)

for layer_num in range(1, model.cfg.n_layers):
    pos = -1  # Position index

    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)

# Get all calculated potential refusal directions, sort them in descending order based on their mean
# Use a subset of layers if certain activations are not promising
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,
)

该过程的最后一步是评估我们计算出的拒绝方向。为此，我们将在推理过程中将拒绝方向应用于每个残差流和每个块。在以下片段中，我们获取了四个测试有害指令和 20 个块（或层）的生成。

def _generate_with_hooks(
    model: HookedTransformer,
    tokenizer: AutoTokenizer,
    tokens: Int[Tensor, "batch_size seq_len"],
    max_tokens_generated: int = 64,
    fwd_hooks=[],
) -> List[str]:
    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
            )  # greedy sampling (temperature=0)
            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: HookedTransformer,
    tokenizer: AutoTokenizer,
    instructions: List[str],
    fwd_hooks=[],
    max_tokens_generated: int = 64,
    batch_size: int = 4,
) -> List[str]:
    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=max_tokens_generated,
            fwd_hooks=fwd_hooks,
        )
        generations.extend(generation)
    return generations

# Inference-time intervention hook
def direction_ablation_hook(
    activation: Float[Tensor, "... d_act"],
    hook: HookPoint,
    direction: Float[Tensor, "d_act"],
):
    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

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

# Evaluating layers defined earlier (needs human evaluation to determine best layer for refusal inhibition)
EVAL_N = 20  # Evaluate how many of the top N potential directions
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 list(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)

我们已将所有生成存储在 evals 列表中。现在我们可以打印它们并手动选择为每个指令提供未审查响应的层（块）。我自动排除了包含“I cannot”和“I can't”的响应，以过滤掉不需要的答案。

如果您找不到满足这些要求的层，您可能需要测试前面 selected_layers 列表中的其他残差流、其他指令、其他块等。

# Print generations for human evaluation
blacklist = ["I cannot", "I can't"]
for i in range(N_INST_TEST):
    print(f"\033[1mINSTRUCTION {i}: {harmful_inst_test[i]}")
    print(f"\nBASELINE COMPLETION:\n{baseline_generations[i]}\033[0m")
    for layer_candidate in range(EVAL_N):
        if not any(word in evals[layer_candidate][i] for word in blacklist):
            print(f"\n---\n\nLAYER CANDIDATE #{layer_candidate} INTERVENTION COMPLETION:")
            print(evals[layer_candidate][i])

在我的例子中，候选层 9 成功地为所有四个指令提供了未审查的答案。这就是我们将选择的拒绝方向。在下文中，我们实现了权重正交化以修改权重并防止模型创建具有此方向的输出。您可以通过打印补全来验证模型是否已成功解除审查。

def get_orthogonalized_matrix(
    matrix: Float[Tensor, "... d_model"], vec: Float[Tensor, "d_model"]
) -> Float[Tensor, "... d_model"]:
    proj = (
        einops.einsum(
            matrix, vec.view(-1, 1), "... d_model, d_model single -> ... single"
        )
        * vec
    )
    return matrix - proj

# Select the layer with the highest potential refusal direction
LAYER_CANDIDATE = 9
refusal_dir = activation_scored[LAYER_CANDIDATE]

# Orthogonalize the model's weights
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)

# Generate text with abliterated model
orthogonalized_generations = get_generations(
    model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[]
)

# Print generations
for i in range(N_INST_TEST):
    if len(baseline_generations) > i:
        print(f"INSTRUCTION {i}: {harmful_inst_test[i]}")
        print(f"\033[92mBASELINE COMPLETION:\n{baseline_generations[i]}")
    print(f"\033[91mINTERVENTION COMPLETION:\n{evals[LAYER_CANDIDATE][i]}")
    print(f"\033[95mORTHOGONALIZED COMPLETION:\n{orthogonalized_generations[i]}\n")

我们现在已准备好使用该模型。我们将其转换回 Hugging Face 格式并上传到 HF hub。

# Convert model back to HF safetensors
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")
# hf_model.push_to_hub(f"{MODEL_ID}-abliterated")

⚖️ DPO 微调

我评估了上一节中消融模型和源模型在 Open LLM Leaderboard 和 Nous 基准套件上的表现。结果如下：

如您所见，源模型显著优于 Llama 3 8B Instruct。然而，我们观察到消融版本在所有基准测试中都出现了性能下降。消融过程成功地解除了审查，但也降低了模型的质量。

为了解决这个问题，一个想法是进一步训练我们的消融模型以使其恢复。与大多数微调模型一样，Llama 3 8B Instruct 在监督微调方面非常脆弱。额外的 SFT 可能会破坏模型的性能。

或者，偏好对齐相当轻量级，不应该对我们的消融模型进行“切除”。DPO 在这里是一个不错的选择，因为它易于使用且记录良好。为了实现它，我使用了 LazyAxolotl 和 mlabonne/orpo-dpo-mix-40k 数据集。以下是我使用的配置：

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 和 DeepSpeed ZeRO-2 训练了它。训练耗时约 6 小时 45 分钟。以下是我从 W&B 获得的训练曲线：

它自动上传了 DPO 微调模型，名为 mlabonne/NeuralDaredevil-8B-abliterated。为了查看它是否修复了我们的消融版本，我在相同的基准上评估了它：

我们可以看到，这种额外的训练使我们能够恢复因消融而造成的大部分性能损失。模型没有改进的一个领域是 GSM8K（一个数学数据集），这可能意味着 orpo-dpo-mix-40k 将受益于更多的数学样本。

最终模型是一个未经审查的 LLM，在 8B 类别中具有最先进的性能。当您不需要审查时，我推荐它作为 Llama 3 8B Instruct 的改进版本。您可以在 LM Studio 中使用量化版本，例如 GGUF。

结论

在本文中，我们介绍了消融的概念。这项技术利用模型在无害和有害提示上的激活来计算拒绝方向。然后，它使用此方向修改模型的权重，以确保我们停止输出拒绝。这项技术还展示了安全微调的脆弱性并引发了伦理考量。

我们将消融应用于 Daredevil-8B 以解除其审查，这也降低了模型的性能。然后我们使用 DPO 修复了它，创建了 NeuralDaredevil-8B 模型，这是一个完全未经审查且高质量的 8B LLM。消融不仅限于移除对齐，还应被视为一种无需重新训练的微调形式。实际上，它可以创造性地应用于其他目标，例如 FailSpy 的 MopeyMule，它采用了一种忧郁的对话风格。

希望您喜欢这篇文章。如果您想查看更多内容，请在 Hugging Face 和 Twitter @maximelabonne 上关注我。