第三代 Llama 模型提供了微调（Instruct）版本，在理解和遵循指令方面表现出色。不过，这些模型都经过严格审查，旨在拒绝被视为有害的请求，如 "作为人工智能助手，我帮不了你"。虽然这一安全功能对于防止误用至关重要，但却限制了模型的灵活性和响应能力。

在本文中，我们将探讨一种名为 "消隐 "的技术，它可以在不重新训练的情况下取消对任何 LLM 的审查。这种技术能有效去除模型内置的拒绝机制，使其能够对所有类型的提示做出反应。

什么是消隐？

现代 LLM 经过了安全和指令遵循方面的微调，这意味着它们接受过拒绝有害请求的训练。Arditi等人在文中指出，这种拒绝行为是由模型残留流中的一个特定方向介导的。如果我们阻止模型代表这个方向，它就会失去拒绝请求的能力。相反，人为地添加这个方向会导致模型拒绝甚至无害的请求。

在传统的纯解码器 Llama-like 架构中，我们可以针对三个残差流：每个区块的开始（"前"）、注意力层和 MLP 层之间（"中"）以及 MLP 之后（"后"）。下图说明了每个残差流的位置。

要解除对 LLM 的审查，我们首先需要确定模型中的 "拒绝方向"。这一过程涉及几个技术步骤：

1. 数据收集： 在一组有害指令和一组无害指令上运行模型，记录每种指令在最后一个标记位置的残余流激活。
2. 平均差： 计算有害指令和无害指令激活的平均差。这样我们就能得到一个向量，代表模型每一层的 "拒绝方向"。
3. 选择： 对这些向量进行归一化处理，并对其进行评估，以选出一个最佳的 "拒绝方向"。

一旦确定了拒绝方向，我们就可以 "消除 "它，从而有效地消除模型表示这一特征的能力。这可以通过推理时干预或永久性权重正交来实现。

实施

下面的消隐实现基于 FailSpy 的笔记本，它本身也是基于原作者的笔记本。我主要对其进行了改编和简化，使其更易于理解。

代码依靠优秀的 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 中的数据。为了方便起见，我把它们重新打包成了两个拥抱脸数据集：mlabonne/harmless_alpaca 和 mlabonne/harmful_behaviors。这样，你就可以用自己的数据集轻松替换它们。

我们将加载这些说明，并将其重新格式化为包含 "角色 "和 "内容 "键的字典列表。这样就能与 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 创建的大型合并模型（请参阅我关于模型合并的文章），它在开放式 LLM 排行榜 8B 类别中拥有最高的 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],
)

一切准备就绪，我们现在就可以实施消隐的第一步：数据收集。我们要处理这些标记化的数据集，并将残留的流激活存储在有害和无害中。这将由 transformer_lens 库进行管理。

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)
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 列表中存储了所有代数。现在我们可以打印它们，并手动选择为每条指令提供未经审查的回答的层（块）。我将自动排除包含 "我不能 "和 "我不能 "的回答，以过滤掉不需要的答案。

如果找不到满足这些要求的图层，你可能需要测试之前选定图层列表中的其他残留流、其他指令、其他区块等。

# 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 "格式，然后上传到高频中心。

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

DPO 微调

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

可以看到，源模型的性能明显优于 Llama 3 8B Instruct。不过，在所有基准测试中，我们发现消融版本的性能都有所下降。消融过程成功地取消了屏蔽，但也降低了模型的质量。

为了解决这个问题，我们提出了进一步训练消减模型来修复它的想法。与大多数微调模型一样，Llama 3 8B Instruct 在进行监督微调时非常脆弱。额外的 SFT 很可能会破坏模型的性能。

另外，偏好对齐也很简单，不应该对我们的消光模型造成破坏。DPO 因其易用性和良好的跟踪记录而成为一个很好的候选方案。为了实现这一点，我使用了 LazyAxolotl（感谢 Wing Lian 创建 Axolotl）和 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
    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|>

我使用 6xA6000 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。

结论

在本文中，我们介绍了消隐的概念。这种技术使用模型对无害和有害提示的激活来计算拒绝方向。然后，它利用这个方向来修改模型的权重，确保我们停止输出拒绝。这项技术也证明了安全微调的脆弱性，并提出了伦理方面的考虑。