SafeMERGE: Fine-Tuning 후 Safety Alignment를 보존하는 선택적 Layer-Wise Model Merging

# SafeMERGE 핵심 로직 (PyTorch + PEFT 기반 의사코드)
import torch
from peft import PeftModel

def compute_safety_subspace(aligned_model, base_model, layer_name):
    """V_i = W_aligned - W_base (레이어별 safety alignment 방향)"""
    W_aligned = dict(aligned_model.named_parameters())[layer_name]
    W_base = dict(base_model.named_parameters())[layer_name]
    return W_aligned - W_base

def cosine_similarity_to_subspace(delta_W, V):
    """LoRA 업데이트가 safety subspace에서 얼마나 벗어났는지 측정"""
    V_norm = V / torch.norm(V, 'fro')
    projection = V_norm @ V_norm.T @ delta_W  # C * delta_W
    cos_sim = torch.nn.functional.cosine_similarity(
        delta_W.flatten().unsqueeze(0),
        projection.flatten().unsqueeze(0)
    )
    return cos_sim.item()

def safe_merge(
    finetuned_lora,   # 파인튜닝된 LoRA 어댑터
    safe_lora,        # safety 데이터로 학습한 LoRA 어댑터
    aligned_model,    # instruct/chat 모델
    base_model,       # base 모델
    tau=0.7,          # cosine similarity 임계값
    alpha=0.8         # 파인튜닝 모델 가중치 (1-alpha = safety 모델 가중치)
):
    merged_weights = {}
    
    for layer_name, delta_W_f in finetuned_lora.items():
        # safety subspace 계산
        V = compute_safety_subspace(aligned_model, base_model, layer_name)
        
        # 해당 레이어의 cosine similarity 측정
        rho = cosine_similarity_to_subspace(delta_W_f, V)
        
        if rho < tau:
            # 안전성이 저하된 레이어 → safety 모델과 선형 병합
            delta_W_s = safe_lora[layer_name]
            merged_weights[layer_name] = alpha * delta_W_f + (1 - alpha) * delta_W_s
            print(f"[MERGE] {layer_name}: rho={rho:.3f} < tau={tau}")
        else:
            # 안전한 레이어 → 파인튜닝 가중치 그대로 유지
            merged_weights[layer_name] = delta_W_f
    
    return merged_weights

# 사용 예시
# merged = safe_merge(finetuned_lora, safe_lora, aligned_model, base_model, tau=0.7, alpha=0.8)
# → 위험한 레이어만 선택적으로 병합된 최종 LoRA 어댑터 반환