Conceptor를 이용한 Semantic Steering: LLM 내부 표현의 다차원 개념 제어

import torch
import numpy as np

def compute_conceptor(activations: torch.Tensor, aperture: float = 10.0) -> torch.Tensor:
    """
    activations: (N, d) - 양극 개념 문장들의 hidden state (긍정+부정 모두 포함)
    aperture: alpha 값, 클수록 더 많은 방향 포함 (보통 2~50)
    returns: (d, d) conceptor 행렬
    """
    N, d = activations.shape
    # 상관 행렬 계산
    R = (activations.T @ activations) / N  # (d, d)
    # C = R * (R + alpha^{-2} * I)^{-1}
    I = torch.eye(d, device=activations.device)
    C = R @ torch.linalg.inv(R + (1.0 / aperture**2) * I)
    return C

def conceptor_quota(C: torch.Tensor) -> float:
    """레이어 선택 진단값. 높을수록 개념이 해당 레이어에 잘 인코딩됨."""
    d = C.shape[0]
    return (torch.trace(C) / d).item()

def conceptor_not(C: torch.Tensor) -> torch.Tensor:
    """NOT 연산: 해당 개념을 억제"""
    return torch.eye(C.shape[0], device=C.device) - C

def conceptor_and(C_a: torch.Tensor, C_b: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    """AND 연산: 두 개념의 교집합 subspace. overlap이 0.4+ 일 때 의미있음."""
    d = C_a.shape[0]
    I = torch.eye(d, device=C_a.device)
    inv_a = torch.linalg.inv(C_a + eps * I)
    inv_b = torch.linalg.inv(C_b + eps * I)
    return torch.linalg.inv(inv_a + inv_b - I)

# 사용 예시: Gemma/Qwen 등 transformer 모델에 hook 등록
def make_steering_hook(C: torch.Tensor, beta: float = 1.0, mode: str = 'replace'):
    """
    mode='replace': z' = beta * C @ z  (강한 조종, 불안정 가능)
    mode='interpolate': z' = (1-beta)*z + beta * C @ z  (안전)
    """
    def hook(module, input, output):
        # output이 tuple인 경우 처리
        hidden = output[0] if isinstance(output, tuple) else output
        if mode == 'replace':
            steered = beta * (C @ hidden.transpose(-1, -2)).transpose(-1, -2)
        else:
            steered = (1 - beta) * hidden + beta * (C @ hidden.transpose(-1, -2)).transpose(-1, -2)
        if isinstance(output, tuple):
            return (steered,) + output[1:]
        return steered
    return hook

# 레이어 선택: quota가 가장 높은 레이어 찾기
# layer_quotas = [conceptor_quota(C_per_layer[l]) for l in range(num_layers)]
# best_layer = np.argmax(layer_quotas)