Persona Vectors: LLM의 성격 특성 모니터링과 제어

# Persona vector 추출 및 steering 예시 (PyTorch + HuggingFace)
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Persona vector 추출 (difference-in-means)
def extract_persona_vector(model, tokenizer, pos_prompts, neg_prompts, questions, layer=20):
    pos_activations, neg_activations = [], []
    
    for prompt, question in zip(pos_prompts * len(questions), questions * len(pos_prompts)):
        inputs = tokenizer(f"{prompt}\n\nQ: {question}\nA:", return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=True)
        # response token activations at target layer
        hidden = outputs.hidden_states[layer][0]  # [seq_len, d_model]
        pos_activations.append(hidden.mean(0))  # avg over response tokens
    
    # negative prompts도 동일하게...
    # neg_activations = ...
    
    persona_vec = torch.stack(pos_activations).mean(0) - torch.stack(neg_activations).mean(0)
    return persona_vec / persona_vec.norm()  # unit normalize

# 2. Inference-time steering (악한 성격 억제)
def generate_with_steering(model, tokenizer, prompt, persona_vec, alpha=-1.5, layer=20):
    hooks = []
    def hook_fn(module, input, output):
        if isinstance(output, tuple):
            hidden = output[0]
            hidden = hidden + alpha * persona_vec.to(hidden.device)
            return (hidden,) + output[1:]
        return output + alpha * persona_vec.to(output.device)
    
    hook = model.model.layers[layer].register_forward_hook(hook_fn)
    hooks.append(hook)
    
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=200)
    
    for h in hooks:
        h.remove()
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# 3. Projection difference로 학습 데이터 필터링
def compute_projection_difference(model, tokenizer, dataset_response, natural_response, persona_vec, layer=20):
    def get_activation(text):
        inputs = tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=True)
        return outputs.hidden_states[layer][0].mean(0)
    
    a_dataset = get_activation(dataset_response)
    a_natural = get_activation(natural_response)
    return ((a_dataset - a_natural) @ persona_vec).item()  # 높을수록 위험