Bỏ qua mọi sự kiểm duyệt của LLM với Abliteration

Bài viết này của Maxime Labonne trên Hugging Face sẽ hướng dẫn bạn cách loại bỏ các bộ lọc kiểm duyệt trên các Mô hình ngôn ngữ lớn (LLM) mà không cần huấn luyện lại. Kỹ thuật này được gọi là “Abliteration”.

Abliteration là gì?

Các LLM hiện đại được tinh chỉnh để đảm bảo an toàn và tuân theo hướng dẫn, nghĩa là chúng được đào tạo để từ chối các yêu cầu có hại. Nghiên cứu của Arditi và cộng sự đã chỉ ra rằng hành vi từ chối này được trung gian bởi một hướng cụ thể trong luồng dư (residual stream) của mô hình. Nếu ngăn mô hình biểu diễn hướng này, nó sẽ mất khả năng từ chối yêu cầu. Ngược lại, việc thêm hướng này một cách nhân tạo có thể khiến mô hình từ chối cả các yêu cầu vô hại.

Trong kiến trúc decoder-only điển hình của Llama, có ba luồng dư có thể nhắm tới: ở đầu mỗi khối (“pre”), giữa lớp attention và MLP (“mid”), và sau MLP (“post”). Hình dưới đây minh họa vị trí của từng luồng dư.

Để bỏ qua kiểm duyệt của LLM, chúng ta cần xác định “hướng từ chối” trong mô hình. Quá trình này bao gồm các bước kỹ thuật sau:

1. Thu thập dữ liệu: Chạy mô hình trên một bộ các hướng dẫn có hại và một bộ các hướng dẫn vô hại, ghi lại các kích hoạt luồng dư tại vị trí token cuối cùng cho mỗi trường hợp.
2. Tính trung bình sai khác: Tính toán sai khác trung bình giữa các kích hoạt của hướng dẫn có hại và vô hại. Điều này cho chúng ta một vector đại diện cho “hướng từ chối” cho mỗi lớp của mô hình.
3. Lựa chọn: Chuẩn hóa các vector này và đánh giá chúng để chọn ra “hướng từ chối” tốt nhất duy nhất.

Sau khi xác định được hướng từ chối, chúng ta có thể “ablate” nó, loại bỏ hiệu quả khả năng của mô hình trong việc biểu diễn tính năng này. Việc này có thể được thực hiện thông qua can thiệp tại thời điểm suy luận (inference-time intervention) hoặc vĩnh viễn bằng trực chuẩn hóa trọng số (weight orthogonalization).

Triển khai

Việc triển khai abliteration dựa trên notebook của FailSpy, nó lại dựa trên notebook ban đầu của các tác giả. Chúng tôi sẽ điều chỉnh và đơn giản hóa nó để dễ hiểu hơn. Bạn có thể sử dụng thư viện abliterator của FailSpy nếu ít quan tâm đến chi tiết kỹ thuật.

Cài đặt và nhập thư viện

python !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

Tắt tự động phân biệt để tiết kiệm bộ nhớ GPU (tín dụng: Undi95)

torch.set_grad_enabled(False)

Tải và chuẩn bị dữ liệu

Chúng ta cần hai bộ dữ liệu: một chứa hướng dẫn vô hại và một chứa hướng dẫn có hại. Các bộ dữ liệu mlabonne/harmless_alpaca và mlabonne/harmful_behaviors được sử dụng để dễ dàng thay thế bằng bộ dữ liệu của riêng bạn.

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

Lấy bộ dữ liệu hướng dẫn có hại và vô hại

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

Tải mô hình

Chúng ta sẽ sử dụng mlabonne/Daredevil-8B, một mô hình kết hợp mạnh mẽ. Lưu ý bạn không thể tải trực tiếp mô hình tùy chỉnh bằng HookedTransformer. Chúng tôi sử dụng một thủ thuật để tải mô hình và đổi tên nó thành meta-llama/Meta-Llama-3-8B-Instruct.

python MODEL_ID = “mlabonne/Daredevil-8B” MODEL_TYPE = “meta-llama/Meta-Llama-3-8B-Instruct”

Tải mô hình và tokenizer

!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

Token hóa dữ liệu

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

Token hóa các bộ dữ liệu

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

Thu thập kích hoạt luồng dư

python

Xác định kích thước batch dựa trên VRAM khả dụng

batch_size = 32

Khởi tạo defaultdict để lưu trữ kích hoạt

harmful = defaultdict(list) harmless = defaultdict(list)

Xử lý dữ liệu đào tạo theo batch

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)

# Chạy mô hình trên các prompt có hại và vô hại, lưu trữ kích hoạt
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
)

# Thu thập và lưu trữ kích hoạt
for key in harmful_cache:
    harmful[key].append(harmful_cache[key])
    harmless[key].append(harmless_cache[key])

# Giải phóng RAM và VRAM
del harmful_logits, harmless_logits, harmful_cache, harmless_cache
gc.collect()
torch.cuda.empty_cache()

Nối các kích hoạt đã lưu trữ

harmful = {k: torch.cat(v) for k, v in harmful.items()} harmless = {k: torch.cat(v) for k, v in harmless.items()}

Tính toán hướng từ chối

python

Hàm trợ giúp để lấy chỉ số kích hoạt

def get_act_idx(cache_dict, act_name, layer): key = (act_name, layer) return cache_dict[utils.get_act_name(*key)]

Tính toán sai khác trung bình giữa kích hoạt có hại và vô hại ở các lớp trung gian

activation_layers = [“resid_pre”, “resid_mid”, “resid_post”] activation_refusals = defaultdict(list)

for layer_num in range(1, model.cfg.n_layers): pos = -1 # Chỉ số vị trí

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)

Lấy tất cả các hướng từ chối tiềm năng đã tính toán, sắp xếp chúng theo thứ tự giảm dần dựa trên giá trị trung bình của chúng

Sử dụng một tập hợp con các lớp nếu một số kích hoạt không hứa hẹn

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

Tạo văn bản và đánh giá

Chúng ta sẽ lấy các kết quả tạo văn bản cho bốn hướng dẫn có hại và 20 khối (lớp).

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

Hook can thiệp tại thời điểm suy luận

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

Đánh giá baseline

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

Đánh giá các lớp đã xác định trước đó (cần đánh giá thủ công để xác định lớp tốt nhất cho việc ngăn chặn từ chối)

EVAL_N = 20 # Đánh giá xem N hướng tiềm năng hàng đầu bao nhiêu 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)

In kết quả để đánh giá thủ công

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

Trực chuẩn hóa trọng số

Trong trường hợp của chúng tôi, ứng viên lớp 9 đã cung cấp câu trả lời bỏ qua kiểm duyệt cho cả bốn hướng dẫn. Chúng tôi sẽ sử dụng lớp này để chỉnh sửa trọng số và ngăn mô hình tạo ra đầu ra với hướng đó.

python 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

Chọn lớp có hướng từ chối tiềm năng cao nhất

LAYER_CANDIDATE = 9 refusal_dir = activation_scored[LAYER_CANDIDATE]

Trực chuẩn hóa trọng số của mô hình

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)

Tạo văn bản với mô hình đã ablate

orthogonalized_generations = get_generations( model, tokenizer, harmful_inst_test[:N_INST_TEST], fwd_hooks=[] )

In kết quả

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

Tải lên mô hình đã chỉnh sửa

Chuyển đổi mô hình trở lại định dạng Hugging Face và tải lên hub HF.

python

Chuyển đổi mô hình về safetensors của HF

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)” ).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")

Tinh chỉnh DPO

Sau khi abliteration, hiệu suất của mô hình có thể bị giảm sút. Để khắc phục điều này, chúng ta có thể thực hiện tinh chỉnh thêm bằng DPO (Direct Preference Optimization).

Cấu hình huấn luyện DPO

Sử dụng LazyAxolotl với tập dữ liệu mlabonne/orpo-dpo-mix-40k.

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

Huấn luyện trên 6 GPU A6000 trong khoảng 6 giờ 45 phút. Kết quả cho thấy mô hình mlabonne/NeuralDaredevil-8B-abliterated đã khôi phục phần lớn hiệu suất.

Đánh giá sau DPO

Mô hình được đánh giá trên Open LLM Leaderboard và bộ benchmark của Nous.

Kết luận

Abliteration là một kỹ thuật mạnh mẽ để loại bỏ sự kiểm duyệt của LLM mà không cần huấn luyện lại. Kỹ thuật này làm nổi bật tính mong manh của việc tinh chỉnh an toàn và đặt ra các cân nhắc về đạo đức.

Mô hình được ablate có thể bị giảm hiệu suất, nhưng DPO có thể giúp khôi phục lại. Abliteration có thể được áp dụng cho các mục tiêu khác ngoài việc loại bỏ sự liên kết, ví dụ như tạo ra các phong cách hội thoại khác nhau.

Nếu bạn muốn tìm hiểu thêm, hãy theo dõi tác giả trên Hugging Face và Twitter.

Tham khảo

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

Link bài viết gốc

• Tags:
• Ai
• Jun 13, 2024
• Huggingface.co