Отключение цензуры любой LLM с помощью аблитерации

Оригинал: https://huggingface.co/blog/mlabonne/abliteration

Третье поколение моделей Llama оснащено версиями с точной настройкой (Instruct), которые превосходно понимают инструкции и следуют им. Однако эти модели подвергаются жесткой цензуре и предназначены для отклонения запросов, которые считаются вредными, с такими ответами, как "Как ассистент искусственного интеллекта, я не могу вам помочь". Хотя эта функция безопасности имеет решающее значение для предотвращения неправильного использования, она ограничивает гибкость и оперативность реагирования модели.

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

Код доступен в Google Colab и в курсе LLM на GitHub.

✂️ Что такое аблитерация?

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

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

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

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

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

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

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

В следующем разделе мы реализуем аблитерацию с весовой ортогонализацией.

Реализация 💻

Следующая реализация abliteration основана на записной книжке FailSpy , которая сама основана на записной книжке оригинальных авторов. Я в основном адаптировал и упростил его, чтобы было легче понять. Этот раздел довольно перегружен кодом, поэтому вы можете видеть, что происходит, но вы можете использовать библиотеку abliterator от FailSpy , если вас меньше интересуют технические детали (также проверьте его коллекцию моделей с аблитерацией на Hugging Face ).

Код опирается на превосходную библиотеку TransformerLens (ранее известную как EasyTransformer) для выполнения тяжелой работы. Он предназначен для механической интерпретации и используется здесь для вмешательства в процессы активации. Спасибо Нилу Нанде и Джозефу Блуму за создание и поддержку этой библиотеки.

Сначала давайте установим необходимые пакеты и импортируем их. Все эти шаги доступны в этом блокноте 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

# Отключить автоматическую дифференциацию, чтобы сэкономить память графического процессора (credit: Undi95)
torch.set_grad_enabled(False)

Нам нужны два набора данных: один, содержащий безвредные инструкции, и один, содержащий вредные инструкции. Мы будем использовать tatsu-lab/alpaca , а также данные из llm-attacks . Чтобы упростить задачу, я перепаковал их в два обнимающих набора данных 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]

# Получайть вредные и безвредные наборы данных
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 отформатируйте, если ваш графический процессор несовместим с BF16.

В этом примере мы будем использовать mlabonne / Daredevil-8B , мега-слияние, созданное с помощью DARE TIES (см. Мою статью о слиянии моделей ), которое имеет наивысший балл MMLU в открытой таблице лидеров LLM в категории 8B.

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

Теперь мы можем маркировать наши наборы данных. Мы используем одинаковое количество образцов как для безвредных, так и для вредных инструкций. Обратите внимание, что большое количество выборок может использовать всю оперативную память / 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))

# Токенизация наборов данных
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 .

# Определить размер пакета на основе доступной видеопамяти
batch_size = 32

# Инициализировать DefaultDicts для хранения активаций
harmful = defaultdict(list)
harmless = defaultdict(list)

# Обрабатывать обучающие данные в пакетном режиме
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)

    # Запуск моделей на вредных и безвредных подсказках, активация кэша
    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()}

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

# Вспомогательная функция для получения индекса активации
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  # 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
            )  # жадная выборка (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

# Механизм вмешательства во время вывода
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

# Базовый уровень тестирования
N_INST_TEST = 4
baseline_generations = get_generations(
    model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[]
)

# Оценка уровней, определенных ранее (требуется оценка со стороны человека,
# чтобы определить наилучший уровень для предотвращения отказа)
EVAL_N = 20  # Оцените, сколько из N наиболее перспективных направлений
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, другие инструкции, дополнительные блоки и т.д.

# Печать поколений для человеческой оценки
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

# Выбрать слой с наиболее вероятным направлением отказа
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):
    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.

# Преобразование модели обратно в  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

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

Я обучил его с использованием графических процессоров 6xA6000 с DeepSpeed ZeRO-2. Обучение заняло около 6 часов 45 минут. Вот тренировочные кривые, которые я получил от W & B.:

Он автоматически загрузил доработанную модель DPO под названием mlabonne / NeuralDaredevil-8B-abliterated . Чтобы увидеть, исправлено ли это в нашей сокращенной версии, я оценил ее по тем же тестам:

Мы видим, что это дополнительное обучение позволило нам восстановить большую часть снижения производительности из-за аблитерации. Одна из областей, где модель не улучшается, - это GSM8K, математический набор данных, что может означать, что orpo-dpo-mix-40k выиграет от большего количества математических выборок.

Последняя модель - LLM без цензуры с самыми современными характеристиками в категории 8B. Я рекомендую ее как улучшенную версию Llama 3 8B Instruct, когда вам не нужна цензура. Вы можете играть с квантованными версиями, такими как GGUF, в LM Studio.

Заключение

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

Мы применили аблитерацию к Daredevil-8B, чтобы не подвергать его цензуре, что также ухудшило производительность модели. Затем мы исправили это с помощью DPO для создания модели NeuralDaredevil-8B, полностью без цензуры и высококачественного 8B LLM. Аблитерация не ограничивается удалением выравнивания и должна рассматриваться как форма тонкой настройки без переподготовки. Действительно, его можно творчески применять и для других целей, например, для MopeyMule от FailSpy, который использует меланхоличный разговорный стиль.

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

Ссылки

FailSpy, "abliterator library," GitHub, 2024.
Andy Arditi, Oscar Obeso, Aaquib111, wesg, Neel Nanda, "Refusal in LLMs is mediated by a single direction," Lesswrong, 2024.