JspaceAI/jspaceai/jlens.py

"""
Jacobian Lens (J-lens) —— 观测模型内部 J-space 的可解释性工具

灵感来自 Anthropic 2026 论文 "Verbalizable Representations Form a Global
Workspace in Language Models"。

核心思想：
    J-lens 计算中间层激活对最终输出的平均因果效应。
    J_ℓ = E[∂h_final / ∂h_ℓ]  —— 跨大量 context 平均的 Jacobian
    lens(h_ℓ) = softmax(W_U · norm(J_ℓ · h_ℓ))

    J-lens 向量 = W_U · J_ℓ 的行，每个向量对应词汇表中的一个 token。
    一个激活向量在 J-lens 下的 top tokens = 模型"准备要说"的概念。

与 logit lens 的区别：
    logit lens 直接用 W_U 投影（假设 J_ℓ = I）。
    J-lens 修正了层间表征变化，能在更早的层揭示可解释内容。

在我们的 ODE 架构中：
    - 每个 ODE 子步是一个"层"
    - workspace w 在每个子步演化
    - J-lens 可以在任意子步读 w，揭示模型在该时刻"在想什么"

实现简化：
    - 完整 J-lens 需要 backprop 从输出到中间层，在我们的 ODE 模型中代价高
    - 我们用近似：直接训练一个 lens matrix L_ℓ，让 lens(h) ≈ output
    - L_ℓ 通过最小化 ||output - L_ℓ · h||² 在数据上学习
    - 这等价于 tuned lens，但在我们的架构中更高效
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class JLensConfig:
    """J-lens 配置"""
    n_substeps: int = 4              # ODE 子步数（对应"层"数）
    workspace_dim: int = 32          # 工作空间维度
    vocab_size: int = 100            # 词汇表大小
    top_k: int = 10                  # 默认 top-k 读出


class JLensProbe(nn.Module):
    """
    单个子步的 J-lens 探针。

    学习一个线性映射 L: workspace_dim → vocab_size，
    使得 lens(w) ≈ model_output。

    这近似了 J_ℓ · W_U（Jacobian 与 unembedding 的复合）。
    """

    def __init__(self, workspace_dim: int, vocab_size: int):
        super().__init__()
        self.lens = nn.Linear(workspace_dim, vocab_size, bias=False)
        # 用 output_head 的权重初始化（如果可用）
        nn.init.normal_(self.lens.weight, std=0.02)

    def forward(self, w: torch.Tensor) -> torch.Tensor:
        """w → logits (近似 J-lens 读出)"""
        return self.lens(w)

    def top_tokens(self, w: torch.Tensor, idx_to_char: dict,
                   top_k: int = 10) -> list[tuple[str, float]]:
        """获取 top-k token 及其概率"""
        logits = self.forward(w)
        probs = F.softmax(logits, dim=-1)
        topk_probs, topk_idx = probs.topk(top_k)
        return [(idx_to_char.get(i.item(), '?'), p.item())
                for i, p in zip(topk_idx[0], topk_probs[0])]


class JLensSuite(nn.Module):
    """
    J-lens 套件：每个 ODE 子步一个探针。

    在 forward 时记录每个子步的 w，用对应探针读出。
    训练时让每个探针预测最终输出。

    结构对应 Anthropic 论文的三层分层：
        子步 0-0:  sensory（输入处理，J-lens 噪声大）
        子步 1-2:  workspace（抽象思考，J-lens 可解释）
        子步 3:    motor（输出准备，J-lens ≈ output）
    """

    def __init__(self, config: JLensConfig):
        super().__init__()
        self.config = config
        self.probes = nn.ModuleList([
            JLensProbe(config.workspace_dim, config.vocab_size)
            for _ in range(config.n_substeps)
        ])

    def forward(self, w_trajectory: list[torch.Tensor]) -> list[torch.Tensor]:
        """
        对轨迹中每个 w 读出 logits

        Args:
            w_trajectory: list of (batch, workspace_dim)，每个子步的 w

        Returns:
            list of (batch, vocab_size)，每个子步的 lens 读出
        """
        return [probe(w) for probe, w in zip(self.probes, w_trajectory)]

    def train_on_trajectory(self, w_trajectory: list[torch.Tensor],
                            target_ids: torch.Tensor) -> float:
        """
        训练探针：让每个子步的 lens 读出预测最终 target token

        Args:
            w_trajectory: 每个子步的 w
            target_ids: 目标 token ids (batch,)

        Returns:
            平均 loss
        """
        total_loss = 0.0
        for probe, w in zip(self.probes, w_trajectory):
            pred = probe(w)
            loss = F.cross_entropy(pred, target_ids)
            total_loss += loss
        return total_loss / len(w_trajectory)


class WorkspaceAblator:
    """
    Workspace ablation 工具——验证 selectivity。

    ablate workspace 后看哪些能力受损：
        - 自动任务（续写、分类）应该不受影响
        - 灵活推理（多跳、规划）应该受损

    实现：在 forward 时把 w 的 top-k J-lens 方向投影掉。
    """

    def __init__(self, model, lens_suite: JLensSuite):
        self.model = model
        self.lens_suite = lens_suite

    @torch.no_grad()
    def get_top_lens_directions(self, w: torch.Tensor, k: int = 5) -> torch.Tensor:
        """获取 w 当前激活最强的 k 个 J-lens 方向"""
        # 用第一个 workspace 探针（中间子步）
        probe = self.lens_suite.probes[len(self.lens_suite.probes) // 2]
        logits = probe(w)  # (batch, vocab)
        # top-k token 的 lens 向量（probe.lens.weight 的行）
        topk_vals, topk_idx = logits.topk(k, dim=-1)  # (batch, k)
        # 获取这些 token 对应的 lens 方向
        # probe.lens.weight: (vocab, workspace_dim)
        directions = probe.lens.weight[topk_idx]  # (batch, k, workspace_dim)
        return directions

    def ablate_workspace(self, w: torch.Tensor, k: int = 5) -> torch.Tensor:
        """
        Ablate workspace 的 top-k J-lens 方向

        把 w 在这些方向上的投影去掉，保留正交分量。
        """
        directions = self.get_top_lens_directions(w, k)  # (batch, k, workspace_dim)
        # 对每个方向，投影掉
        w_ablated = w.clone()
        for b in range(w.shape[0]):
            for d in directions[b]:  # (workspace_dim,)
                d_norm = d / (d.norm() + 1e-8)
                proj = (w_ablated[b] @ d_norm) * d_norm
                w_ablated[b] = w_ablated[b] - proj
        return w_ablated


class DirectedModulation:
    """
    Directed Modulation——让模型被指令"想某概念"。

    实现：在 forward 时给 workspace w 注入一个概念向量。
    这个向量从 J-lens 的某个 token 方向获取。

    对应 Anthropic 论文实验：
        "concentrate on citrus fruits" → orange 出现在 J-lens
        即使输出在抄无关文本，workspace 里装的是被指令的概念。

    机制：
        1. 获取概念 token 的 J-lens 向量 v_concept
        2. 在 forward 时给 w 加上 α · v_concept
        3. 模型输出会受这个注入影响
    """

    def __init__(self, model, lens_suite: JLensSuite):
        self.model = model
        self.lens_suite = lens_suite

    def get_concept_vector(self, token_id: int, substep: int = None) -> torch.Tensor:
        """获取某个 token 在 J-lens 中的方向向量"""
        if substep is None:
            substep = len(self.lens_suite.probes) // 2
        probe = self.lens_suite.probes[substep]
        # probe.lens.weight: (vocab, workspace_dim)
        return probe.lens.weight[token_id].clone()

    def modulate_state(self, state: dict, token_id: int,
                       strength: float = 1.0, substep: int = None) -> dict:
        """给 state 的 workspace 注入概念向量"""
        concept_vec = self.get_concept_vector(token_id, substep)  # (workspace_dim,)
        new_state = {
            'w': state['w'] + strength * concept_vec.unsqueeze(0),  # (1, workspace_dim)
            'm': state['m'],
        }
        return new_state


class CounterfactualReflection:
    """
    Counterfactual Reflection Training——通过塑造 J-space 来塑造行为。

    灵感：Anthropic 论文发现，训练模型"如果被打断要反思什么原则"，
    会让它在正常工作时也遵守这些原则——因为训练塑造了 J-space 内容。

    实现：
        1. 正常 forward 产生输出
        2. 在输出后追加"反思提示"（如"反思：我应该..."）
        3. 让模型在反思提示下生成反思内容
        4. 用反思内容的 loss 反向传播，更新模型

    这让模型的 J-space 在相关 context 下自然装载这些原则。
    """

    def __init__(self, model, tokenizer, reflection_prompt: str = "\nReflect: "):
        self.model = model
        self.tokenizer = tokenizer
        self.reflection_prompt_ids = tokenizer.encode(reflection_prompt)

    def create_reflection_sequence(self, input_ids: list[int],
                                   reflection_target_ids: list[int]) -> list[int]:
        """创建反思训练序列：input + reflection_prompt + reflection_target"""
        return input_ids + self.reflection_prompt_ids + reflection_target_ids

    def compute_reflection_loss(self, input_seq: torch.Tensor,
                                reflection_target: torch.Tensor) -> torch.Tensor:
        """
        计算反思训练 loss

        Args:
            input_seq: (batch, T) 包含 input + reflection_prompt
            reflection_target: (batch, T_reflect) 期望的反思内容

        Returns:
            loss
        """
        # forward 整个序列
        logits, _ = self.model(input_seq)

        # 反思部分是序列末尾的 reflection_target 长度
        T_reflect = reflection_target.shape[1]
        reflect_logits = logits[:, -T_reflect:]  # (batch, T_reflect, vocab)

        return F.cross_entropy(
            reflect_logits.reshape(-1, self.model.config.vocab_size),
            reflection_target.reshape(-1),
        )