Машинный перевод GPT-4o статьи «Uncensor any LLM with abliteration»

Современные языковые модели (LLM) настроены на безопасность и выполнение инструкций, что означает, что они обучены отказывать в выполнении вредных запросов. В своем блоге Ардити и др. показали, что это поведение отказа связано с определенным направлением в остаточном потоке модели. Если мы предотвратим представление этого направления в модели, она потеряет способность отказывать в запросах. Напротив, искусственное добавление этого направления может привести к тому, что модель будет отказывать даже в безобидных запросах.

В традиционной архитектуре только декодера, подобной Llama, есть три остаточных потока, на которые мы можем нацелиться: в начале каждого блока («pre»), между слоями внимания и MLP («mid») и после MLP («post»). Следующая иллюстрация показывает расположение каждого остаточного потока.

Чтобы убрать цензуру с языковой модели (LLM), сначала необходимо определить «направление отказа» внутри модели. Этот процесс включает в себя несколько технических шагов:

1. Сбор данных: Запустите модель на наборе вредных инструкций и наборе безобидных инструкций, записывая активации остаточного потока в последней позиции токена для каждого из них.

2. Среднее различие: Вычислите среднее различие между активациями вредных и безобидных инструкций. Это даст нам вектор, представляющий «направление отказа» для каждого слоя модели.

3. Выбор: Нормализуйте эти векторы и оцените их, чтобы выбрать единственное лучшее «направление отказа».

После того как мы определили направление отказа, мы можем «аблировать» его, эффективно удаляя у модели способность представлять эту характеристику. Это можно сделать через вмешательство во время вывода или навсегда с помощью ортогонализации весов.

Сначала давайте поговорим о вмешательстве во время вывода. Для каждого компонента, который записывает в остаточный поток (например, головы внимания), мы вычисляем проекцию его выхода на направление отказа и вычитаем эту проекцию. Это вычитание применяется к каждому токену и каждому слою, гарантируя, что модель никогда не представляет направление отказа.

С другой стороны, ортогонализация весов включает в себя прямое изменение весов модели. Ортогонализируя веса компонентов относительно направления отказа, мы предотвращаем запись модели в это направление полностью. Это достигается путем корректировки матриц, которые записывают в остаточный поток, обеспечивая их отсутствие вклада в направление отказа.

Реализация
Следующая реализация аблитерации основана на ноутбуке 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")

Я оценил аблитерированные и исходные модели из предыдущего раздела на 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|> 

Заключение
В этой статье мы представили концепцию аблитерации. Эта техника использует активации модели на безобидных и вредных запросах для вычисления направления отказа. Затем это направление используется для изменения весов модели и обеспечения того, чтобы мы прекратили выдавать отказы. Эта техника также демонстрирует хрупкость тонкой настройки безопасности и поднимает этические вопросы.

Надеюсь, вам понравилась эта статья. Если вы хотите увидеть больше, подписывайтесь на меня на Hugging Face.