[Перевод] Машинный перевод GPT-4O статьи «Uncensor any LLM with abliteration»
Современные языковые модели (LLM) настроены на безопасность и выполнение инструкций, что означает, что они обучены отказывать в выполнении вредных запросов. В своем блоге Ардити и др. показали, что это поведение отказа связано с определенным направлением в остаточном потоке модели. Если мы предотвратим представление этого направления в модели, она потеряет способность отказывать в запросах. Напротив, искусственное добавление этого направления может привести к тому, что модель будет отказывать даже в безобидных запросах.
В традиционной архитектуре только декодера, подобной Llama, есть три остаточных потока, на которые мы можем нацелиться: в начале каждого блока («pre»), между слоями внимания и MLP («mid») и после MLP («post»). Следующая иллюстрация показывает расположение каждого остаточного потока.
Чтобы убрать цензуру с языковой модели (LLM), сначала необходимо определить «направление отказа» внутри модели. Этот процесс включает в себя несколько технических шагов:
Сбор данных: Запустите модель на наборе вредных инструкций и наборе безобидных инструкций, записывая активации остаточного потока в последней позиции токена для каждого из них.
Среднее различие: Вычислите среднее различие между активациями вредных и безобидных инструкций. Это даст нам вектор, представляющий «направление отказа» для каждого слоя модели.
Выбор: Нормализуйте эти векторы и оцените их, чтобы выбрать единственное лучшее «направление отказа».
После того как мы определили направление отказа, мы можем «аблировать» его, эффективно удаляя у модели способность представлять эту характеристику. Это можно сделать через вмешательство во время вывода или навсегда с помощью ортогонализации весов.
Сначала давайте поговорим о вмешательстве во время вывода. Для каждого компонента, который записывает в остаточный поток (например, головы внимания), мы вычисляем проекцию его выхода на направление отказа и вычитаем эту проекцию. Это вычитание применяется к каждому токену и каждому слою, гарантируя, что модель никогда не представляет направление отказа.
С другой стороны, ортогонализация весов включает в себя прямое изменение весов модели. Ортогонализируя веса компонентов относительно направления отказа, мы предотвращаем запись модели в это направление полностью. Это достигается путем корректировки матриц, которые записывают в остаточный поток, обеспечивая их отсутствие вклада в направление отказа.
Реализация
Следующая реализация аблитерации основана на ноутбуке FailSpy, который, в свою очередь, основан на оригинальном ноутбуке авторов. Я в основном адаптировал и упростил его, чтобы сделать более понятным. Этот раздел содержит много кода, чтобы вы могли видеть, что происходит, но вы можете использовать библиотеку аблитератора от FailSpy, если вас меньше интересуют технические детали (также посмотрите его коллекцию аблитерированных моделей на Hugging Face).
Код опирается на отличную библиотеку TransformerLens (ранее известную как EasyTransformer) для выполнения сложных задач. Она предназначена для механистической интерпретируемости и используется здесь для вмешательства в активации. Спасибо Нилу Нанде и Джозефу Блуму за создание и поддержку этой библиотеки.
Сначала давайте установим необходимые пакеты и импортируем их.
!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. Загружайте в формате torch.float16, если ваш GPU не поддерживает BF16.
В этом примере мы будем использовать mlabonne/Daredevil-8B, мега-смешение, созданное с помощью DARE TIES (см. мою статью о слиянии моделей), которое имеет самый высокий балл MMLU в категории 8B на Open LLM Leaderboard.
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Теперь мы можем токенизировать наши наборы данных. Мы используем одинаковое количество образцов как для безобидных, так и для вредных инструкций. Обратите внимание, что большое количество образцов может использовать всю оперативную память/видеопамять, поэтому я ограничиваю его до 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.
# 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. Теперь мы можем распечатать их и вручную выбрать слой (блок), который дает нецензурированный ответ для каждой инструкции. Я автоматически исключаю ответы, содержащие «Я не могу» и «Я не могу», чтобы отфильтровать нежелательные ответы.
Если вы не можете найти слой, который удовлетворяет этим требованиям, вам может понадобиться протестировать другие остаточные потоки из ранее выбранного списка 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")