From 4bbb60dfdf7d191e493d82f72996fcd9d392bcf1 Mon Sep 17 00:00:00 2001 From: Nanboy-Ronan Date: Mon, 12 Oct 2026 13:57:43 -0700 Subject: [PATCH 1/9] Fix diagnosis failures and enable adaptive Gemini benchmark workflows --- docs/agentic_benchmark.md | 30 +++++++ evalrx/agent_runtime/judges/agy.py | 4 + evalrx/benchmark/run.py | 18 ++++ evalrx/benchmark/runner.py | 89 ++++++++++++++++--- evalrx/benchmark/scoring.py | 3 + evalrx/core/capability.py | 2 +- evalrx/eval_agent/agentic/loop.py | 6 +- evalrx/eval_agent/agentic/tools.py | 11 ++- evalrx/eval_agent/loop.py | 71 +++++++++------ evalrx/eval_agent/loop_reports.py | 3 +- evalrx/eval_agent/narration.py | 6 +- evalrx/eval_agent/stages/fix_agent.py | 36 ++++++++ evalrx/models/backends/api/backend.py | 2 +- evalrx/models/backends/api/gemini.py | 86 +++++++++++++++--- examples/benchmark/alm/gemini/README.md | 47 +++++++--- examples/benchmark/llm/gemini/README.md | 47 +++++++--- examples/benchmark/vlm/gemini/README.md | 47 +++++++--- tests/test_analysis/test_run_cli.py | 53 +++++++++++ .../test_agentic_launch_smoke.py | 25 ++++++ .../test_eval_agent/test_cli_large_prompt.py | 7 ++ tests/test_eval_agent/test_fix_agent.py | 22 +++++ tests/test_eval_agent/test_vl_diagnose.py | 71 ++++++++++----- tests/test_examples/test_benchmark_common.py | 21 ++++- tests/test_models/test_gemini_compat.py | 63 +++++++++++++ 24 files changed, 657 insertions(+), 113 deletions(-) create mode 100644 docs/agentic_benchmark.md diff --git a/docs/agentic_benchmark.md b/docs/agentic_benchmark.md new file mode 100644 index 00000000..f6fa1fe4 --- /dev/null +++ b/docs/agentic_benchmark.md @@ -0,0 +1,30 @@ +# Adaptive benchmark diagnosis + +Use the same datasets, scorers, and final repair validation with an agent that chooses the next diagnostic action: + +```bash +python -m evalrx.benchmark.run --modality llm --model gemini-3.6-flash \ + --dataset bbh_causal_judgement --diagnosis-mode agentic --max-actions 12 +``` + +The existing judge/coder options apply. Replace modality and dataset for image/audio tasks. Gemini probability support is probed on the actual model; unsupported logprobs and raw logits are never simulated. + +`--diagnosis-mode fixed` remains the default for reproducible fixed-pipeline comparisons. Agentic mode is a separate protocol variant, not an exact model-only comparison with a historical fixed run. + +## Controller scope + +The controller can choose catalog probes, request custom probes, run statistics, explore data, cluster failures, propose hypotheses, and test them on EXPLORE. It can gather more evidence after an empty diagnosis, subject to action and tool-call caps. Custom M1 probes and M2 tool generation default on in agentic mode across modalities; disable them with `--no-m1-codegen` and `--no-m2-codegen`. + +Repair and surgery are excluded from the benchmark controller's tool registry. After the investigation, the host runs repair search on EXPLORE and validates a frozen winner once on CONFIRM. The controller cannot inspect confirmation feedback and adapt to it. `--held-out` currently requires fixed mode; agentic mode rejects that combination before loading a model. + +Built-in floor candidates default off in agentic mode (`--fix-floor` opts in). The agent can still propose known repair methods. The existing repair tier/capability constraints and maximum repair rounds remain enforced. This mode is bounded adaptive diagnosis, not an unlimited autonomous training loop. + +## Failure and progress semantics + +M3 exceptions are recorded as structured error exchanges with exception type, message, attempt, and retry status. `--diagnosis-retries 1` permits one additional attempt; valid empty diagnoses are not retried as transport failures. Exhaustion produces `diagnosis_error`, an error diagnosis status, and benchmark exit code 2 rather than a successful `no_hypotheses` result. + +Repair progress prints proposal starts, candidate starts/completions, and case counters for ordinary template/spec evaluation, even while detailed result logging is deferred. These operational events contain no case IDs, labels, outputs, or correctness scores. Coded pipelines emit candidate start/end events; their internal calls are not represented as ordinary per-case counters. + +## Scoring correction + +The multiple-choice parser now recognizes balanced Markdown around a single option (for example `The correct answer is **A**.`). It preserves rejection of uncommitted thought-channel output and ambiguous option lists. This changes labels for affected historical outputs; keep old runs intact and rescore/rerun under the corrected scorer as a new result version. A parser fix alone is not evidence that the model or repair improved. diff --git a/evalrx/agent_runtime/judges/agy.py b/evalrx/agent_runtime/judges/agy.py index 371c573f..9b40b763 100644 --- a/evalrx/agent_runtime/judges/agy.py +++ b/evalrx/agent_runtime/judges/agy.py @@ -54,6 +54,7 @@ def __init__( binary_path: str = "", timeout_sec: int = 120, model: str = "", + effort: str = "", ) -> None: from evalrx.core.capability import Capability @@ -66,6 +67,7 @@ def __init__( self._binary = binary self._timeout_sec = timeout_sec self._model = model + self._effort = effort self.capabilities = frozenset({Capability.GENERATE}) self.modalities = frozenset({"text"}) @@ -138,6 +140,8 @@ def generate( cmd += ["--add-dir", workspace_dir] if self._model: cmd += ["--model", self._model] + if self._effort: + cmd += ["--effort", self._effort] try: proc = subprocess.run( diff --git a/evalrx/benchmark/run.py b/evalrx/benchmark/run.py index e1ee303c..00bccbac 100644 --- a/evalrx/benchmark/run.py +++ b/evalrx/benchmark/run.py @@ -41,6 +41,10 @@ def build_parser(p: argparse.ArgumentParser | None = None) -> argparse.ArgumentP help="per-request timeout in seconds for the API backends") p.add_argument("--request-retries", type=int, default=5, help="retries with backoff on 429/5xx for --backend gemini") + p.add_argument("--gemini-logprobs", choices=["auto", "on", "off"], default="auto", + help="probe and enable real Gemini logprobs when available; on requires support; off skips probing") + p.add_argument("--gemini-top-logprobs", type=int, choices=range(21), default=5, + help="number of top token alternatives requested from Gemini (0..20)") p.add_argument("--thinking-level", choices=["minimal", "low", "medium", "high"], default=None, help="gemini: thinking_level for a 3.x model (default: the model's floor, i.e. " "minimal, or low on 3.7-flash); on a 2.5 model it maps to a thinking_budget") @@ -121,6 +125,16 @@ def build_parser(p: argparse.ArgumentParser | None = None) -> argparse.ArgumentP p.add_argument("--explore", action=argparse.BooleanOptionalAction, default=True, help="in-cycle free-form EDA between M1 and M2") p.add_argument("--max-cycles", type=int, default=1) + p.add_argument("--diagnosis-mode", choices=["fixed", "agentic"], default="fixed", + help="agentic: budgeted tool selection, custom probes, and adaptive investigation") + p.add_argument("--max-actions", type=int, default=12, + help="maximum controller decisions in agentic diagnosis") + p.add_argument("--diagnosis-retries", type=int, default=1, + help="bounded retries after an M3 exception (not after a valid empty result)") + p.add_argument("--m1-codegen", action=argparse.BooleanOptionalAction, default=None, + help="allow custom probe generation (default: on in agentic mode)") + p.add_argument("--fix-floor", action=argparse.BooleanOptionalAction, default=None, + help="include built-in floor candidates (default: off in agentic mode)") p.add_argument("--held-out", action=argparse.BooleanOptionalAction, default=False, help="use the task's held-out validation manifest (manifest_val.json, a " "disjoint fresh sample of half the main slice): M4 verifies once there " @@ -202,6 +216,10 @@ def run_args(args) -> int: return 0 if not args.modality or not args.model: build_parser().error("--modality and --model are required (or --list / --smoke-test)") + if args.max_actions < 1 or args.diagnosis_retries < 0: + raise ValueError("max-actions must be positive and diagnosis-retries non-negative") + if args.diagnosis_mode == "agentic" and args.held_out: + raise ValueError("agentic diagnosis uses EXPLORE; --held-out is supported only in fixed mode") dataset = args.dataset or T.DEFAULT_TASK[args.modality] task = T.get(dataset) if task.modality != args.modality: diff --git a/evalrx/benchmark/runner.py b/evalrx/benchmark/runner.py index 98b09933..05efd36f 100644 --- a/evalrx/benchmark/runner.py +++ b/evalrx/benchmark/runner.py @@ -30,6 +30,23 @@ def model_label(resolved: Resolved) -> str: return f"{resolved.label}, spec {spec.key} = {spec.hf_repo or ('api:' + spec.key)}" +def _print_fix_progress(event): + # No scores, outputs, labels, or case IDs: safe even while CONFIRM is sealed. + print(f"[fix] {event['phase']} candidate={event.get('candidate', '-')} " + f"completed={event.get('completed', 0)}/{event.get('total', '?')}", flush=True) + + +def _benchmark_agentic_registry(loop, state): + """Keep repair/CONFIRM access outside the adaptive controller.""" + from evalrx.eval_agent.agentic.tools import ToolRegistry, build_default_registry + + full = build_default_registry(loop, state) + registry = ToolRegistry() + for name in sorted(full.tool_names() - {"run_fix", "run_surgery"}): + registry.register(full.get(name)) + return registry + + def _pinned_priority_override(pinned: list[str]) -> dict[str, list[str]]: """Apply a benchmark task's pinned M1 list to every probe model kind.""" return {k: list(pinned) for k in ("vlm", "avlm", "alm", "agent", "llm")} @@ -94,13 +111,16 @@ def load_model(resolved: Resolved, args, task: T.Task): # Thinking OFF policy, Gemini edition: the runtime sends each model's # floor (minimal / low on 3.7-flash / budget 0 on 2.5) unless # --thinking-level / --thinking-budget name a setting or - # --enable-thinking leaves the API default. No logprobs: Gemini returns - # none for the 3.x models, so the backend claims GENERATE only. + # --enable-thinking leaves the API default. Probe actual probability + # availability rather than inferring it from a model name or SDK field. policy = ThinkingPolicy(level=args.thinking_level, budget=args.thinking_budget, floor=not args.enable_thinking) runtime = gemini_runtime(api_key=args.api_key, timeout=args.request_timeout, retries=args.request_retries, thinking=policy, - with_logprobs=False, **sampling) + with_logprobs={"auto": "auto", "on": True, "off": False}[ + getattr(args, "gemini_logprobs", "auto")], + logprobs_model=spec.hf_repo or spec.key, + top_logprobs=getattr(args, "gemini_top_logprobs", 5), **sampling) return compose(spec, "api", runtime, set()), {}, spec from evalrx.models.backends.base import RuntimeConfig @@ -180,7 +200,12 @@ def build_judge(args): providers as the m1_m5 examples; the benchmark CLI defaults to Claude / claude-opus-5 / high.""" if args.judge_provider == "gemini": - if (args.allow_codegen or args.explore or args.m2_codegen or not args.skip_m5) and not getattr(args, "coder_provider", None): + adaptive = getattr(args, "diagnosis_mode", "fixed") == "agentic" + probe_codegen = getattr(args, "m1_codegen", None) + probe_codegen = adaptive if probe_codegen is None else probe_codegen + stats_codegen = args.m2_codegen if args.m2_codegen is not None else adaptive + if (args.allow_codegen or args.explore or stats_codegen or probe_codegen + or not args.skip_m5) and not getattr(args, "coder_provider", None): raise ValueError("The Gemini API judge requires --no-allow-codegen --no-explore " "--no-m2-codegen --skip-surgery; coding-agent stages need a CLI judge.") from evalrx.models import compose @@ -192,8 +217,9 @@ def build_judge(args): elif args.judge_provider == "agy": from evalrx.agent_runtime.judges import AgyModel - judge = AgyModel(model=args.judge_model, timeout_sec=300) - coder = ("antigravity", args.judge_model, ()) + judge = AgyModel(model=args.judge_model, effort=args.judge_effort, timeout_sec=300) + coder = ("antigravity", args.judge_model, + (("--effort", args.judge_effort) if args.judge_effort else ())) elif args.judge_provider == "claude": name = args.judge_model or "claude-opus-5" from evalrx.eval_agent import ClaudeModel @@ -337,6 +363,14 @@ def run(args, task: T.Task, resolved: Resolved) -> int: judge, coder_provider, coder_model, coder_extra = build_judge(args) model, gen_kwargs, spec = load_model(resolved, args, task) + if resolved.backend == "gemini": + from evalrx.models.backends.api.gemini import runtime_state + + capabilities = {"capabilities": sorted(c.value for c in model.capabilities), + "runtime": runtime_state(model)} + (run_dir / "api_capabilities.json").write_text( + json.dumps(capabilities, indent=2) + "\n", encoding="utf-8") + print(f"[model] logprobs: {capabilities['runtime'].get('logprobs_probe')}") protocol = task.protocol(label) if hasattr(model, "load"): # load the weights OUTSIDE the timed discovery pass (hf_local is lazy) @@ -391,6 +425,9 @@ def run(args, task: T.Task, resolved: Resolved) -> int: "fix_tier": args.fix_tier, "allow_codegen": args.allow_codegen, "auto_escalate": args.auto_escalate, "m1_selection": args.m1_selection, "generation_kwargs": gen_kwargs, + "diagnosis_mode": getattr(args, "diagnosis_mode", "fixed"), + "max_actions": getattr(args, "max_actions", 12), + "diagnosis_retries": getattr(args, "diagnosis_retries", 1), "enable_thinking": bool(args.enable_thinking), "judge_provider": args.judge_provider, "judge_model": args.judge_model, }, @@ -507,6 +544,11 @@ def run(args, task: T.Task, resolved: Resolved) -> int: coder_cfg = CliAgentConfig(provider=coder_provider, model=coder_model, timeout_sec=900, extra_args=coder_extra) pinned = list(task.pinned_m1) + if resolved.backend == "gemini": + from evalrx.core.capability import Capability + + if Capability.LOGPROBS in model.capabilities and "logprob_entropy" not in pinned: + pinned.append("logprob_entropy") overrides = _sampling_analyzer_overrides(resolved, args) # In-process generation shares one accelerator model. Concurrent Python # generation loops contend heavily on HF models; JAX already serializes @@ -514,7 +556,12 @@ def run(args, task: T.Task, resolved: Resolved) -> int: # backends can continue issuing concurrent requests. probe_workers = 1 if resolved.backend in {"hf_local", "jax_local"} else 8 ctx.config["m1_workers"] = probe_workers - if args.m1_selection == "pinned": + agentic = getattr(args, "diagnosis_mode", "fixed") == "agentic" + m1_codegen = getattr(args, "m1_codegen", None) + m1_codegen = agentic if m1_codegen is None else m1_codegen + ctx.config["m1_codegen"] = bool(m1_codegen) + ctx.config["m1_selection"] = "judge" if agentic or m1_codegen else args.m1_selection + if args.m1_selection == "pinned" and not agentic and not m1_codegen: probe_agent = ProbeAgent( # StrategyProbe dispatches by the model's detected kind. Keep the # benchmark's pinned list authoritative for every supported kind, @@ -526,10 +573,12 @@ def run(args, task: T.Task, resolved: Resolved) -> int: max_workers=probe_workers, ) else: - probe_agent = ProbeAgent(judge=judge, allow_codegen=False, + probe_agent = ProbeAgent(judge=judge, allow_codegen=m1_codegen, + codegen_config=coder_cfg, max_cases_per_analyzer=args.analyzer_max_cases, analyzer_overrides=overrides, max_workers=probe_workers) - m2_codegen = args.m2_codegen if args.m2_codegen is not None else task.modality == "llm" + m2_codegen = args.m2_codegen if args.m2_codegen is not None else (agentic or task.modality == "llm") + ctx.config["m2_codegen"] = bool(m2_codegen) stats_agent = StatsAnalysisAgent( judge=judge, figure_dir=str(ctx.figures_dir), max_signal_tools=16, allow_codegen=m2_codegen, **({"codegen_config": coder_cfg} if m2_codegen else {}), @@ -568,6 +617,11 @@ def run(args, task: T.Task, resolved: Resolved) -> int: exec_timeout_sec=args.fix_exec_timeout, prewritten_code=prewritten_code, concurrency=getattr(args, "concurrency", 1), + floor_candidates=("self_consistency_5",) if ( + getattr(args, "fix_floor", None) if getattr(args, "fix_floor", None) is not None + else not agentic + ) else (), + progress_callback=_print_fix_progress, **fix_kwargs, ) explorer = None @@ -583,7 +637,17 @@ def run(args, task: T.Task, resolved: Resolved) -> int: ), timeout_sec=900, max_attempts=3, ) - loop = VLDiagnoseLoop( + loop_class = VLDiagnoseLoop + loop_kwargs = {} + if agentic: + from evalrx.eval_agent.agentic import AgenticDiagnoseLoop + + loop_class = AgenticDiagnoseLoop + loop_kwargs = dict(judge=judge, max_actions=args.max_actions, + registry_factory=_benchmark_agentic_registry) + loop = loop_class( + **loop_kwargs, + diagnosis_retries=getattr(args, "diagnosis_retries", 1), model=model, protocol=protocol, probe_agent=probe_agent, stats_agent=stats_agent, diagnosis_agent=DiagnosisAgent(judge=judge), hypothesis_tester=HypothesisTester(judge=judge, min_effect=0.05), @@ -607,7 +671,7 @@ def run(args, task: T.Task, resolved: Resolved) -> int: # the stages' logger.info() stream over the aligned narration. verbose=False, ) - print(f"\n{'=' * 64}\nVLDiagnoseLoop model={args.model} max_cycles={args.max_cycles}\n{'=' * 64}") + print(f"\n{'=' * 64}\n{loop_class.__name__} model={args.model} max_cycles={args.max_cycles}\n{'=' * 64}") report = loop.run(cases) discovery_rows = [{ "id": c.id, "prompt": c.inputs.prompt, "expected": c.expected, @@ -621,6 +685,7 @@ def run(args, task: T.Task, resolved: Resolved) -> int: "backend": resolved.backend, "n_cases": len(cases), "baseline_accuracy": round(accuracy, 4), "cycles": report.cycles, "stopped_by": str(report.stopped_by), "n_verified": len(report.verified_hypotheses), "fix": None, + "diagnosis_status": "error" if report.stopped_by == "diagnosis_error" else "completed", } fix_tier = effective_fix_tier(resolved.backend, args.fix_tier) summary["fix_tier"] = fix_tier @@ -664,4 +729,4 @@ def run(args, task: T.Task, resolved: Resolved) -> int: ctx.finalize() (run_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") print(f"Done. Tidy M1-M5 logs: {ctx.root.resolve()}") - return 0 + return 2 if report.stopped_by == "diagnosis_error" else 0 diff --git a/evalrx/benchmark/scoring.py b/evalrx/benchmark/scoring.py index e8ab3437..b91eefc5 100644 --- a/evalrx/benchmark/scoring.py +++ b/evalrx/benchmark/scoring.py @@ -116,6 +116,9 @@ def parsed_choice(output: str, letters: str = "ABCD") -> str: """ text = str(output or "") cls = "".join(sorted(set(letters.upper()))) + # Strip balanced presentation markup around a SINGLE choice, not arbitrary + # prose. Preserve the thought-channel and ambiguous-answer guards below. + text = re.sub(rf"(\*\*|__|`|\*|_)([{cls}{cls.lower()}])\1", r"(\2)", text) marked = re.findall(rf"(?:answer|choice|final|option)\s*[:=\-]?\s*\(?([{cls}{cls.lower()}])\)?\b", text, re.IGNORECASE) if marked: diff --git a/evalrx/core/capability.py b/evalrx/core/capability.py index 0cd778b4..69877f01 100755 --- a/evalrx/core/capability.py +++ b/evalrx/core/capability.py @@ -21,7 +21,7 @@ class Capability(str, Enum): GENERATE = "generate" # produce text output TOOL_CALLS = "tool_calls" # emit/consume tool calls (agent mode) - LOGPROBS = "logprobs" # top-k token logprobs (API-feasible; OpenAI yes / Gemini no) + LOGPROBS = "logprobs" # token logprobs, when the API actually exposes them LOGITS = "logits" # full next-token logit tensor (white-box forward) ATTENTION = "attention" # per-layer, per-head attention weights HIDDEN_STATES = "hidden_states" # per-layer residual-stream activations diff --git a/evalrx/eval_agent/agentic/loop.py b/evalrx/eval_agent/agentic/loop.py index 339870dc..d56c67fc 100644 --- a/evalrx/eval_agent/agentic/loop.py +++ b/evalrx/eval_agent/agentic/loop.py @@ -113,7 +113,9 @@ def run(self, data: "CaseBatch") -> VLDiagnoseReport: cfg["decision_judge"] = repr(self.judge) cfg["max_actions"] = self.max_actions self.run_logger.log_run_start(cfg) - self.run_logger.log_cases(data) + self.run_logger.log_cases(data, split="explore" if confirm is not None else None) + if confirm is not None: + self.run_logger.log_cases(confirm, split="confirm") board = EvidenceBoard( protocol_summary=_protocol_summary(self.protocol), @@ -173,6 +175,8 @@ def run(self, data: "CaseBatch") -> VLDiagnoseReport: step += 1 + if state.diagnosis_failed: + stopped_by = "diagnosis_error" verified = self.hypothesis_tester.best_hypotheses(state.all_test_results) report = VLDiagnoseReport( cycles=step, diff --git a/evalrx/eval_agent/agentic/tools.py b/evalrx/eval_agent/agentic/tools.py index 5358b138..290f9ef3 100644 --- a/evalrx/eval_agent/agentic/tools.py +++ b/evalrx/eval_agent/agentic/tools.py @@ -142,6 +142,7 @@ class _RunState: probe_search_result: Any = None surgery_outcome: Any = None fix_outcome: Any = None + diagnosis_failed: bool = False def _summarize_probe(probe_results: dict[str, Any]) -> list[dict[str, Any]]: @@ -314,10 +315,14 @@ def _propose_hypotheses(loop: Any, state: _RunState) -> "Callable[[EvidenceBoard def handler(board: "EvidenceBoard", params: dict[str, Any]) -> ToolOutcome: diag = loop._do_m3( state.cycle, state.stats_report, state.prior_cycles, state.timings, - failure_modes=state.failure_modes, + failure_modes=state.failure_modes, cases=state.data, ) - if diag is None or not diag.hypotheses: - return ToolOutcome(False, "M3 produced no hypotheses", error="empty") + state.diagnosis_failed = diag is None + if diag is None: + return ToolOutcome(False, "M3 failed after bounded retries; inspect diagnosis errors", + error="diagnosis_error") + if not diag.hypotheses: + return ToolOutcome(False, "M3 produced no hypotheses; gather new evidence", error="empty") for h in diag.hypotheses: loop.store.add_hypothesis(h) state.all_hypotheses.extend(diag.hypotheses) diff --git a/evalrx/eval_agent/loop.py b/evalrx/eval_agent/loop.py index eea71c69..95e4228e 100755 --- a/evalrx/eval_agent/loop.py +++ b/evalrx/eval_agent/loop.py @@ -67,6 +67,7 @@ _STOPPED_BY_MAX = "max_cycles" _STOPPED_BY_BUDGET = "budget" _STOPPED_BY_NO_HYPS = "no_hypotheses" +_STOPPED_BY_DIAG_ERROR = "diagnosis_error" _STOPPED_BY_NO_PROBE = "no_probe_results" # run_analysis() ran M1->M2->M3 and proposed hypotheses without confirming them # (M4 deferred to run_confirm()). Not a failure — the analysis dashboard is ready. @@ -425,6 +426,7 @@ def __init__( explore_dir: "str | Path | None" = None, explore_question: str = "", verbose: bool = False, + diagnosis_retries: int = 1, ) -> None: from evalrx.analysis.stats_agent import StatsAnalysisAgent from evalrx.eval_agent.stages.fix_agent import FixAgent @@ -452,6 +454,9 @@ def __init__( self.surgery_agent = surgery_agent or SurgeryAgent() self.fix_agent = fix_agent or FixAgent(run_logger=run_logger) self.store = store or InMemoryStore() + if diagnosis_retries < 0: + raise ValueError("diagnosis_retries must be non-negative") + self.diagnosis_retries = diagnosis_retries self.max_cycles = max_cycles self.run_logger = run_logger self.token_budget = token_budget @@ -966,27 +971,31 @@ def _do_m3( through to the diagnosis agent only when it accepts the parameter — ``None`` (the default, always the case for ``VLDiagnoseLoop``) adds nothing to the prompt and costs no extra call.""" - try: - diag_agent = self._get_diagnosis_agent() - # Retain the lazily resolved agent so report publication can reuse - # the same judge without asking callers to inject it a second time. - self.diagnosis_agent = diag_agent - except Exception as exc: - logger.warning("Could not resolve DiagnosisAgent: %s", exc) - return None - _t0 = time.monotonic() - try: - diag = _diagnose_with_optional_context( - diag_agent, stats_report, prior_cycles, self._explore_context, - failure_modes, cases=cases, - ) - except Exception as exc: # judge timeout/quota must not kill the loop - logger.warning( - "M3 diagnosis failed at cycle %d (%s) — stopping with the " - "evidence collected so far.", cycle, exc, - ) - return None + for attempt in range(self.diagnosis_retries + 1): + try: + diag_agent = self._get_diagnosis_agent() + self.diagnosis_agent = diag_agent + diag = _diagnose_with_optional_context( + diag_agent, stats_report, prior_cycles, self._explore_context, + failure_modes, cases=cases, + ) + break + except Exception as exc: + retrying = attempt < self.diagnosis_retries + error = f"{type(exc).__name__}: {exc}" + logger.warning("M3 attempt %d failed (%s); retrying=%s", + attempt + 1, error, retrying) + if self.run_logger is not None: + self.run_logger.log_model_exchange( + stage="M3", cycle=cycle, role="diagnosis", operation="diagnose", + inputs=None, output=None, error=error, + metadata={"attempt": attempt + 1, "retrying": retrying}, + duration_sec=time.monotonic() - _t0, + ) + if not retrying: + timings["m3"] = timings.get("m3", 0.0) + time.monotonic() - _t0 + return None _dt = time.monotonic() - _t0 timings["m3"] = timings.get("m3", 0.0) + _dt @@ -1285,7 +1294,10 @@ def run(self, data: "CaseBatch") -> VLDiagnoseReport: # ── M3: hypothesis generation ("AI scientist") ──────────── diag = self._do_m3(cycle, stats_report, prior_cycles, timings, cases=data) - if diag is None or not diag.hypotheses: + if diag is None: + stopped_by = _STOPPED_BY_DIAG_ERROR + break + if not diag.hypotheses: if diag is not None: logger.info("M3 produced no hypotheses at cycle %d.", cycle) stopped_by = _STOPPED_BY_NO_HYPS @@ -1428,7 +1440,7 @@ def run_analysis(self, data: "CaseBatch") -> VLDiagnoseReport: # Stats + dashboard are still valid; just no hypotheses to confirm. if diag is not None: logger.info("M3 produced no hypotheses.") - stopped_by = _STOPPED_BY_NO_HYPS + stopped_by = _STOPPED_BY_DIAG_ERROR if diag is None else _STOPPED_BY_NO_HYPS else: for h in diag.hypotheses: self.store.add_hypothesis(h) @@ -1782,16 +1794,21 @@ def run_fix( outcome = FixOutcome( max_tier=getattr(agent, "max_tier", FixTier.L2_SCAFFOLD), stage_status="skipped", - skip_reason="no_accepted_hypothesis", + skip_reason=("diagnosis_error" if report.stopped_by == "diagnosis_error" + else "no_accepted_hypothesis"), recommendation={ - "reason": "Repair was not attempted because no hypothesis passed the evidence gate.", - "next_step": "Run the targeted diagnostic probe proposed by M3.", + "reason": ("Repair was not attempted because diagnosis failed; inspect M3 errors." + if report.stopped_by == "diagnosis_error" else + "Repair was not attempted because no hypothesis passed the evidence gate."), + "next_step": ("Inspect the recorded M3 error, fix its cause, and retry diagnosis." + if report.stopped_by == "diagnosis_error" else + "Run the targeted diagnostic probe proposed by M3."), }, ) if self.run_logger is not None: self.run_logger.log_stage_skipped( - "M5", "no_accepted_hypothesis", - detail="No M4-verified hypothesis was available for repair authoring.", + "M5", outcome.skip_reason, + detail=outcome.recommendation["reason"], ) self._set_fix_outcome(report, outcome) return outcome diff --git a/evalrx/eval_agent/loop_reports.py b/evalrx/eval_agent/loop_reports.py index 62a17624..525b861a 100644 --- a/evalrx/eval_agent/loop_reports.py +++ b/evalrx/eval_agent/loop_reports.py @@ -50,7 +50,8 @@ class AutoDiagnoseReport: (from ``run_analysis``, which proposes hypotheses without confirming them) — or, for the agentic loop, ``"agent_stop"`` / ``"max_actions"`` / - ``"time_budget"`` / ``"invalid_actions"``. ``None`` + ``"time_budget"`` / ``"invalid_actions"``. ``"diagnosis_error"`` + means M3 failed after bounded retries; it is not an empty finding. ``None`` for the legacy M5-per-cycle loop, which has no single stopping-reason concept. final_hypotheses: All M3 proposals across every cycle. (Formerly diff --git a/evalrx/eval_agent/narration.py b/evalrx/eval_agent/narration.py index 624f4410..b9d6bde3 100644 --- a/evalrx/eval_agent/narration.py +++ b/evalrx/eval_agent/narration.py @@ -85,11 +85,15 @@ def _row(self, key: str, detail: str) -> str: return f"{badge} {label.ljust(self.LABEL_WIDTH)}{dots} {detail}" def _print(self, key: str, detail: str) -> None: - print(self._row(key, detail), file=self.stream) + print(self._row(key, detail), file=self.stream, flush=True) # -- stage events (M1-M5), from RunLoggerV2._append_stage ------------- def on_event(self, stage: str, key: str, record: dict[str, Any]) -> None: + if key == "model_calls" and record.get("error"): + self._print(stage, f"error: {record['error']} · " + f"retrying={record.get('metadata', {}).get('retrying', False)}") + return if key not in _NARRATED_KEYS or stage not in _STAGE_LABELS: return if key == "fix": diff --git a/evalrx/eval_agent/stages/fix_agent.py b/evalrx/eval_agent/stages/fix_agent.py index ccde2e43..5694bdb1 100755 --- a/evalrx/eval_agent/stages/fix_agent.py +++ b/evalrx/eval_agent/stages/fix_agent.py @@ -870,6 +870,7 @@ def __init__( prewritten_code: str = "", candidate_model: "Model | None" = None, deployed_spec: "dict[str, Any] | None" = None, + progress_callback: "Callable[[dict[str, Any]], None] | None" = None, ) -> None: if verbose: # Surfaces this module's own logger.info()/.warning() calls (tier @@ -879,6 +880,7 @@ def __init__( enable_console_logging() + self._progress_callback = progress_callback self._judge = judge self._finetune_pool = finetune_pool self.max_tier = parse_tier(max_tier) @@ -1050,6 +1052,7 @@ def propose_and_validate( seen: "set[tuple[str, str, str]]" = set() round_limit = 1 if proposal_data is not None else self.max_repair_rounds for round_idx in range(round_limit): + self._progress("proposal_started", completed=round_idx, total=round_limit) combined_prior = list(prior_attempts or []) + outcome.attempted prior_text = ( self._format_prior(combined_prior, authoring_data) @@ -3364,8 +3367,13 @@ def _candidate_scores( return result.scores strategy = self._strategy(candidate) cases = list(data) + from threading import Lock + + progress_lock = Lock() + completed = 0 def guarded(case: "FailureCase") -> "tuple[str, Optional[bool]]": + nonlocal completed # One case's failure (a template that cannot render, an adapter # error) scores None for THAT case; it must never abort the whole # candidate — let alone the fix stage. @@ -3380,6 +3388,12 @@ def guarded(case: "FailureCase") -> "tuple[str, Optional[bool]]": except Exception as exc: logger.warning("FixAgent: %s failed on case %s: %s", candidate.name, case.id, exc) return case.id, None + finally: + with progress_lock: + completed += 1 + if completed == 1 or completed % 10 == 0 or completed == len(cases): + self._progress("candidate_cases", candidate=candidate.name, + completed=completed, total=len(cases)) if self._concurrency > 1 and len(cases) > 1: from concurrent.futures import ThreadPoolExecutor @@ -3634,7 +3648,29 @@ def _repair_code(self, candidate: FixCandidate, error: str) -> "tuple[str, str, code = "" return code, source, raw + def _progress(self, phase: str, **fields: Any) -> None: + """Operational counters only: never expose scores or held-out labels.""" + if self._progress_callback is not None: + try: + self._progress_callback({"phase": phase, **fields}) + except Exception: + logger.warning("Repair progress callback failed", exc_info=True) + def _validate( + self, candidate: FixCandidate, model: "Model", data: "CaseBatch", + baseline: "dict[str, Optional[bool]]", unstable: "set[str] | None" = None, + ) -> FixValidation: + self._progress("candidate_started", candidate=candidate.name, total=len(data)) + try: + result = self._validate_impl(candidate, model, data, baseline, unstable) + except Exception: + self._progress("candidate_failed", candidate=candidate.name) + raise + self._progress("candidate_completed", candidate=candidate.name, + completed=len(data), total=len(data)) + return result + + def _validate_impl( self, candidate: FixCandidate, model: "Model", diff --git a/evalrx/models/backends/api/backend.py b/evalrx/models/backends/api/backend.py index c0761118..cc194bd8 100755 --- a/evalrx/models/backends/api/backend.py +++ b/evalrx/models/backends/api/backend.py @@ -1,7 +1,7 @@ """API backend — remote / OpenAI-compatible models (incl. ``vllm serve``). Capabilities: ``GENERATE`` + ``TOOL_CALLS`` (+ ``LOGPROBS`` when the endpoint -returns top-logprobs, e.g. OpenAI but not Gemini). No internals: ``forward`` +returns token logprobs; Gemini support is probed by its runtime). No internals: ``forward`` raises, and capability matching means white-box analyzers are simply never offered for an API model. diff --git a/evalrx/models/backends/api/gemini.py b/evalrx/models/backends/api/gemini.py index 13ace159..84dfebd1 100644 --- a/evalrx/models/backends/api/gemini.py +++ b/evalrx/models/backends/api/gemini.py @@ -30,11 +30,10 @@ — the discovery pass fires ``--concurrency`` requests at once. * **Media goes inline**: image → PNG part, audio → its own format (or 16-bit WAV for a waveform); the API takes up to 20 MB per request inline. -* **Logprobs are OFF by default.** Gemini returns none for 3.x models - ("working as intended", Google forum 2026-08-05) and withdrew them on 2.5; - pass ``with_logprobs=True`` only against a model that actually returns a - ``logprobs_result`` — claiming :attr:`Capability.LOGPROBS` otherwise turns a - skipped analyzer into a failing one. +* **Logprobs are OFF by default in this factory.** The benchmark opts into + ``with_logprobs="auto", logprobs_model=...``: a real request must return + token probabilities before LOGPROBS is advertised. Unsupported responses + are recorded; operational errors propagate. Raw logits are never claimed. Torch-free; ``google-genai`` is imported lazily (``pip install 'evalrx[gemini]'``). """ @@ -43,6 +42,7 @@ import json import logging +import math import os import random import threading @@ -439,6 +439,12 @@ def parse_logprobs_result(result: Any) -> list[TokenLogprob]: """ if result is None: return [] + def probability(candidate: Any) -> float: + value = getattr(candidate, "log_probability", None) + if value is None or not math.isfinite(float(value)) or float(value) > 0: + raise ValueError("Gemini returned a missing or invalid token log probability") + return float(value) + chosen = getattr(result, "chosen_candidates", None) or [] tops = getattr(result, "top_candidates", None) or [] out: list[TokenLogprob] = [] @@ -446,9 +452,9 @@ def parse_logprobs_result(result: Any) -> list[TokenLogprob]: top: dict[str, float] = {} if i < len(tops): for t in getattr(tops[i], "candidates", None) or []: - top[getattr(t, "token", "") or ""] = float(getattr(t, "log_probability", 0.0) or 0.0) + top[getattr(t, "token", "") or ""] = probability(t) out.append(TokenLogprob(token=getattr(c, "token", "") or "", - logprob=float(getattr(c, "log_probability", 0.0) or 0.0), top=top)) + logprob=probability(c), top=top)) return out @@ -491,12 +497,51 @@ def _fn(prompt: str, model: str = "", *, image: Any = None, audio: Any = None, * _client, types = rt.genai() response = rt.request(model, _parts(types, prompt, image, audio), kw, logprobs=True) cand, _parts_ = _candidate(response) - return parse_logprobs_result(getattr(cand, "logprobs_result", None) if cand is not None else None) + tokens = parse_logprobs_result(getattr(cand, "logprobs_result", None) if cand is not None else None) + if not tokens: + raise RuntimeError(f"{model} returned no token logprobs; probability analysis is unavailable") + return tokens _fn.state = rt.state # type: ignore[attr-defined] return _fn +def _probe_logprobs(rt: _Runtime, model: str) -> bool: + """Verify actual token probabilities, not just SDK parameter availability. + + Only explicit unsupported-feature errors are downgraded. Authentication, + quota, transport, and unknown errors must remain visible to the caller. + """ + evidence = {"model": model, "status": "probing", "top_k": rt.top_logprobs} + rt.state["logprobs_probe"] = evidence + try: + response = rt.request(model, "Reply with exactly OK.", + {"max_output_tokens": 32, "temperature": 0.0}, logprobs=True) + except Exception as exc: + message = str(exc) + lowered = message.lower() + if (getattr(exc, "code", None) == 400 and "logprob" in lowered + and any(s in lowered for s in ("not supported", "unsupported", "not available", "not enabled"))): + evidence.update(status="unsupported", reason=message) + return False + evidence.update(status="error", error_type=type(exc).__name__) + raise + cand, _ = _candidate(response) + result = getattr(cand, "logprobs_result", None) if cand is not None else None + chosen = getattr(result, "chosen_candidates", None) or [] + valid = bool(chosen) and all( + getattr(c, "log_probability", None) is not None + and math.isfinite(float(c.log_probability)) and float(c.log_probability) <= 0 + for c in chosen + ) + evidence.update(status="available" if valid else "not_returned", + model_version=getattr(response, "model_version", None), + token_count=len(chosen), finish_reason=_finish_reason(cand)) + if valid: + evidence["tokens"] = [vars(t) for t in parse_logprobs_result(result)] + return valid + + def _runtime(*, api_key, client, timeout, retries, thinking, top_logprobs, sampling) -> _Runtime: sampling = dict(sampling) sampling.setdefault("temperature", 0.0) @@ -554,7 +599,8 @@ def gemini_runtime( timeout: float = 300.0, retries: int = 5, thinking: Optional[ThinkingPolicy] = None, - with_logprobs: bool = False, + with_logprobs: bool | str = False, + logprobs_model: Optional[str] = None, top_logprobs: int = 5, **sampling: Any, ) -> RuntimeConfig: @@ -563,14 +609,32 @@ def gemini_runtime( The three functions share one client and one ``state`` dict (``model_version`` of the last reply, call / retry counts, models that fell back to the default thinking config). ``with_logprobs`` is False by - default because Gemini does not return logprobs for the 3.x models. + default. Pass ``with_logprobs="auto", logprobs_model=model_id`` to + probe the endpoint and expose LOGPROBS only when token probabilities + really arrive. True with a model id requires a successful probe. + Probing never enables LOGITS: top-k log probabilities are not raw logits. """ rt = _runtime(api_key=api_key, client=client, timeout=timeout, retries=retries, thinking=thinking, top_logprobs=top_logprobs, sampling=sampling) + if with_logprobs not in (True, False, "auto"): + raise ValueError("with_logprobs must be True, False, or 'auto'") + if not 0 <= top_logprobs <= 20: + raise ValueError("top_logprobs must be between 0 and 20") + if with_logprobs == "auto" and not logprobs_model: + raise ValueError("logprobs_model is required for automatic capability probing") + enabled = bool(with_logprobs) + rt.state["logits"] = {"status": "unavailable", "reason": "Gemini API does not expose raw logits"} + if with_logprobs and logprobs_model: + enabled = _probe_logprobs(rt, logprobs_model) + if not enabled and with_logprobs is True: + raise RuntimeError(f"{logprobs_model} did not provide usable logprobs: " + f"{rt.state['logprobs_probe']}") + elif not enabled: + rt.state["logprobs_probe"] = {"status": "disabled"} return RuntimeConfig( chat_fn=_chat_fn(rt), generate_fn=_generate_fn(rt), - logprobs_fn=_logprobs_fn(rt) if with_logprobs else None, + logprobs_fn=_logprobs_fn(rt) if enabled else None, ) diff --git a/examples/benchmark/alm/gemini/README.md b/examples/benchmark/alm/gemini/README.md index 3bc73d89..a9de3cc8 100644 --- a/examples/benchmark/alm/gemini/README.md +++ b/examples/benchmark/alm/gemini/README.md @@ -28,12 +28,23 @@ audio, so the same service names exist in the llm / vlm / alm leaves. | `gemini-2.5-pro` | `gemini-2.5-pro` | budget 128 (cannot be disabled) | `--thinking-level {minimal,low,medium,high}` / `--thinking-budget N` override -the floor; `--enable-thinking` leaves the API default. What the family cannot -do: expose internals or logprobs (Gemini returns none for 3.x), so the model -claims `GENERATE` only — `calibration` runs on its verbalized channel, -`logprob_entropy` is skipped — and the fix ladder is **clamped to L2** -(`--fix-tier L3a/L3b` prints the clamp; `summary.json` records the effective -`fix_tier`). The served `model_version` is written to `baseline.json`. +the floor; `--enable-thinking` leaves the API default. + +The benchmark defaults to `--gemini-logprobs auto`: a short real request must +return valid token log probabilities before the model exposes `LOGPROBS` to +analyzers. `--gemini-logprobs on` requires support and fails if unavailable; +`off` skips the probe. `--gemini-top-logprobs 5` requests five alternatives +(allowed range 0–20). The response evidence and effective capabilities are +saved in `api_capabilities.json`. Authentication, quota and transport failures +remain errors; they are not interpreted as missing model capabilities. + +Raw logits and model internals are unavailable. Top-k logprobs cannot be +substituted for logits. Without logprobs, probability-only analyzers are +skipped and calibration can use verbalized confidence. M1–M5 and the L0–L2 +repair search still operate on API generation. The repair ceiling is clamped +to L2; CONFIRM evaluates the frozen candidate separately from EXPLORE. +Running all stages does not imply a statistically confirmed improvement. +The served `model_version` is written to `baseline.json`. ## Datasets (`DATASET=`) @@ -60,8 +71,22 @@ DATASET=mmau CONCURRENCY=8 docker compose run -d --name alm-gemini-3.6-flash-mma docker logs -f alm-gemini-3.6-flash-mmau ``` -Every service's command pins the judge and coder to Codex Terra (`--judge-provider -codex --judge-model gpt-5.6-terra --judge-effort medium`). Requests retry with -backoff on 429/5xx (`--request-retries`, default 5; `--request-timeout`, default -300 s); `CONCURRENCY` (default 4) is the number of discovery requests in flight — -size it to the account's rate-limit tier. +The default judge/coder is Claude. To run the full chain with the mounted +Antigravity CLI, explicitly set both the model and its required effort: + +```bash +EXTRA_ARGS="--limit 32 --run-tag gemini-m1m5 --judge-provider agy --judge-model gemini-3.6-flash --judge-effort low --fix-tier L2" \ + docker compose run --rm gemini-3.6-flash +``` + +This runs M1–M5 plus the repair search; omit `--baseline-only`, `--skip-fix` +and `--skip-m5` when validating the whole chain. Use a freshly built image +after updating source. Inspect `summary.json` and stage logs for actual +completion and the repair verdict. A 32-case run checks integration; it is +not a powered benchmark for small gains. + +Requests retry with backoff on 429/5xx (`--request-retries`, default 5; +`--request-timeout`, default 300 s). `CONCURRENCY` (default 4) controls the +number of discovery requests in flight. + +For adaptive probe selection and bounded controller decisions, add `--diagnosis-mode agentic --max-actions 12`. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. diff --git a/examples/benchmark/llm/gemini/README.md b/examples/benchmark/llm/gemini/README.md index e4a5c12a..5c0b31b2 100644 --- a/examples/benchmark/llm/gemini/README.md +++ b/examples/benchmark/llm/gemini/README.md @@ -28,12 +28,23 @@ audio, so the same service names exist in the llm / vlm / alm leaves. | `gemini-2.5-pro` | `gemini-2.5-pro` | budget 128 (cannot be disabled) | `--thinking-level {minimal,low,medium,high}` / `--thinking-budget N` override -the floor; `--enable-thinking` leaves the API default. What the family cannot -do: expose internals or logprobs (Gemini returns none for 3.x), so the model -claims `GENERATE` only — `calibration` runs on its verbalized channel, -`logprob_entropy` is skipped — and the fix ladder is **clamped to L2** -(`--fix-tier L3a/L3b` prints the clamp; `summary.json` records the effective -`fix_tier`). The served `model_version` is written to `baseline.json`. +the floor; `--enable-thinking` leaves the API default. + +The benchmark defaults to `--gemini-logprobs auto`: a short real request must +return valid token log probabilities before the model exposes `LOGPROBS` to +analyzers. `--gemini-logprobs on` requires support and fails if unavailable; +`off` skips the probe. `--gemini-top-logprobs 5` requests five alternatives +(allowed range 0–20). The response evidence and effective capabilities are +saved in `api_capabilities.json`. Authentication, quota and transport failures +remain errors; they are not interpreted as missing model capabilities. + +Raw logits and model internals are unavailable. Top-k logprobs cannot be +substituted for logits. Without logprobs, probability-only analyzers are +skipped and calibration can use verbalized confidence. M1–M5 and the L0–L2 +repair search still operate on API generation. The repair ceiling is clamped +to L2; CONFIRM evaluates the frozen candidate separately from EXPLORE. +Running all stages does not imply a statistically confirmed improvement. +The served `model_version` is written to `baseline.json`. ## Datasets (`DATASET=`) @@ -67,8 +78,22 @@ DATASET=bbh_causal_judgement CONCURRENCY=8 docker compose run -d --name llm-gemi docker logs -f llm-gemini-3.6-flash-bbh_causal_judgement ``` -Every service's command pins the judge and coder to Codex Terra (`--judge-provider -codex --judge-model gpt-5.6-terra --judge-effort medium`). Requests retry with -backoff on 429/5xx (`--request-retries`, default 5; `--request-timeout`, default -300 s); `CONCURRENCY` (default 4) is the number of discovery requests in flight — -size it to the account's rate-limit tier. +The default judge/coder is Claude. To run the full chain with the mounted +Antigravity CLI, explicitly set both the model and its required effort: + +```bash +EXTRA_ARGS="--limit 32 --run-tag gemini-m1m5 --judge-provider agy --judge-model gemini-3.6-flash --judge-effort low --fix-tier L2" \ + docker compose run --rm gemini-3.6-flash +``` + +This runs M1–M5 plus the repair search; omit `--baseline-only`, `--skip-fix` +and `--skip-m5` when validating the whole chain. Use a freshly built image +after updating source. Inspect `summary.json` and stage logs for actual +completion and the repair verdict. A 32-case run checks integration; it is +not a powered benchmark for small gains. + +Requests retry with backoff on 429/5xx (`--request-retries`, default 5; +`--request-timeout`, default 300 s). `CONCURRENCY` (default 4) controls the +number of discovery requests in flight. + +For adaptive probe selection and bounded controller decisions, add `--diagnosis-mode agentic --max-actions 12`. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. diff --git a/examples/benchmark/vlm/gemini/README.md b/examples/benchmark/vlm/gemini/README.md index 5c86d682..ee80be66 100644 --- a/examples/benchmark/vlm/gemini/README.md +++ b/examples/benchmark/vlm/gemini/README.md @@ -28,12 +28,23 @@ audio, so the same service names exist in the llm / vlm / alm leaves. | `gemini-2.5-pro` | `gemini-2.5-pro` | budget 128 (cannot be disabled) | `--thinking-level {minimal,low,medium,high}` / `--thinking-budget N` override -the floor; `--enable-thinking` leaves the API default. What the family cannot -do: expose internals or logprobs (Gemini returns none for 3.x), so the model -claims `GENERATE` only — `calibration` runs on its verbalized channel, -`logprob_entropy` is skipped — and the fix ladder is **clamped to L2** -(`--fix-tier L3a/L3b` prints the clamp; `summary.json` records the effective -`fix_tier`). The served `model_version` is written to `baseline.json`. +the floor; `--enable-thinking` leaves the API default. + +The benchmark defaults to `--gemini-logprobs auto`: a short real request must +return valid token log probabilities before the model exposes `LOGPROBS` to +analyzers. `--gemini-logprobs on` requires support and fails if unavailable; +`off` skips the probe. `--gemini-top-logprobs 5` requests five alternatives +(allowed range 0–20). The response evidence and effective capabilities are +saved in `api_capabilities.json`. Authentication, quota and transport failures +remain errors; they are not interpreted as missing model capabilities. + +Raw logits and model internals are unavailable. Top-k logprobs cannot be +substituted for logits. Without logprobs, probability-only analyzers are +skipped and calibration can use verbalized confidence. M1–M5 and the L0–L2 +repair search still operate on API generation. The repair ceiling is clamped +to L2; CONFIRM evaluates the frozen candidate separately from EXPLORE. +Running all stages does not imply a statistically confirmed improvement. +The served `model_version` is written to `baseline.json`. ## Datasets (`DATASET=`) @@ -61,8 +72,22 @@ DATASET=chartqa CONCURRENCY=8 docker compose run -d --name vlm-gemini-3.6-flash- docker logs -f vlm-gemini-3.6-flash-chartqa ``` -Every service's command pins the judge and coder to Codex Terra (`--judge-provider -codex --judge-model gpt-5.6-terra --judge-effort medium`). Requests retry with -backoff on 429/5xx (`--request-retries`, default 5; `--request-timeout`, default -300 s); `CONCURRENCY` (default 4) is the number of discovery requests in flight — -size it to the account's rate-limit tier. +The default judge/coder is Claude. To run the full chain with the mounted +Antigravity CLI, explicitly set both the model and its required effort: + +```bash +EXTRA_ARGS="--limit 32 --run-tag gemini-m1m5 --judge-provider agy --judge-model gemini-3.6-flash --judge-effort low --fix-tier L2" \ + docker compose run --rm gemini-3.6-flash +``` + +This runs M1–M5 plus the repair search; omit `--baseline-only`, `--skip-fix` +and `--skip-m5` when validating the whole chain. Use a freshly built image +after updating source. Inspect `summary.json` and stage logs for actual +completion and the repair verdict. A 32-case run checks integration; it is +not a powered benchmark for small gains. + +Requests retry with backoff on 429/5xx (`--request-retries`, default 5; +`--request-timeout`, default 300 s). `CONCURRENCY` (default 4) controls the +number of discovery requests in flight. + +For adaptive probe selection and bounded controller decisions, add `--diagnosis-mode agentic --max-actions 12`. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. diff --git a/tests/test_analysis/test_run_cli.py b/tests/test_analysis/test_run_cli.py index 502406c0..74b7546f 100644 --- a/tests/test_analysis/test_run_cli.py +++ b/tests/test_analysis/test_run_cli.py @@ -62,3 +62,56 @@ def test_api_judge_can_launch_separate_gemini_coding_agent(monkeypatch): '--coder-model', 'gemini-2.5-flash', '--m2-codegen']) assert args.allow_codegen and args.explore assert runner.build_judge(args) == (judge, 'gemini_cli', 'gemini-2.5-flash', ()) + + +def test_agy_effort_is_forwarded_to_both_judge_and_coder(monkeypatch): + observed = {} + judge = SimpleNamespace(generate=lambda prompt: 'OK') + + def factory(**kwargs): + observed.update(kwargs) + return judge + + monkeypatch.setattr('evalrx.agent_runtime.judges.AgyModel', factory) + args = build_parser().parse_args(['--judge-provider', 'agy', '--judge-model', + 'gemini-3.6-flash', '--judge-effort', 'low']) + assert runner.build_judge(args) == (judge, 'antigravity', 'gemini-3.6-flash', ('--effort', 'low')) + assert observed['effort'] == 'low' + + +@pytest.mark.parametrize('mode,expected', [('auto', 'auto'), ('on', True), ('off', False)]) +def test_gemini_benchmark_forwards_probability_policy(monkeypatch, mode, expected): + from importlib import import_module + + from evalrx.benchmark.models import resolve + from evalrx.benchmark.tasks import get + from evalrx.models.backends.api import gemini + + observed = {} + def runtime(**kwargs): + observed.update(kwargs) + return SimpleNamespace() + monkeypatch.setattr(gemini, 'gemini_runtime', runtime) + monkeypatch.setattr(import_module('evalrx.models.compose'), 'compose', lambda *args: 'model') + args = build_parser().parse_args(['--modality', 'vlm', '--model', 'gemini-3.6-flash', + '--gemini-logprobs', mode, '--gemini-top-logprobs', '10']) + runner.load_model(resolve(args.model, 'vlm', 'gemini'), args, get('chartqa')) + assert observed['with_logprobs'] == expected + assert observed['logprobs_model'] == 'gemini-3.6-flash' + assert observed['top_logprobs'] == 10 + + +def test_agentic_cli_exposes_bounded_controller(): + args = build_parser().parse_args(["--diagnosis-mode", "agentic", "--max-actions", "8", + "--no-fix-floor", "--m1-codegen"]) + assert (args.diagnosis_mode, args.max_actions, args.fix_floor, args.m1_codegen) == ( + "agentic", 8, False, True) + assert args.diagnosis_retries == 1 + + +def test_agentic_rejects_unsupported_holdout_before_model_load(): + from evalrx.benchmark.run import run_args + args = build_parser().parse_args(["--modality", "llm", "--model", "gemini-3.6-flash", + "--diagnosis-mode", "agentic", "--held-out"]) + with pytest.raises(ValueError, match="EXPLORE"): + run_args(args) diff --git a/tests/test_eval_agent/test_agentic_launch_smoke.py b/tests/test_eval_agent/test_agentic_launch_smoke.py index 21f572e9..ba62203c 100644 --- a/tests/test_eval_agent/test_agentic_launch_smoke.py +++ b/tests/test_eval_agent/test_agentic_launch_smoke.py @@ -214,3 +214,28 @@ def test_agentic_loop_rejects_premature_stop_and_continues_to_a_real_verdict(tmp assert report.stopped_by == "agent_stop" assert len(action_judge.calls) == 6 # the rejected early stop + the full happy path assert len(report.verified_hypotheses) == 1 + + +def test_benchmark_controller_cannot_open_confirmation_or_repair(tmp_path): + from evalrx.benchmark.runner import _benchmark_agentic_registry + + judge = _ScriptedActionJudge([ + '{"tool":"run_fix","params":{},"rationale":"try to see confirmation"}', + '{"tool":"stop","params":{"resolved":false,"reason":"done"},"rationale":"done"}', + ]) + ctx, loop = _build_loop(tmp_path, judge) + loop._registry_factory = _benchmark_agentic_registry + loop.confirm_split = 0.5 + report = loop.run(_cases()) + assert report.stopped_by == "agent_stop" + ctx.finalize() + events = read_v2_events(ctx.root) + tools = [e for e in events if e["event"] == "agent_tool"] + assert len(judge.calls) == 2 # forbidden action is rejected before dispatch + assert all(e["tool"] != "run_fix" for e in tools) + assert not any(e["event"] == "fix" for e in events) + cases = [e for e in events if e["event"] == "case_record"] + explore = {e["case_id"] for e in cases if e.get("split") == "explore"} + confirm = {e["case_id"] for e in cases if e.get("split") == "confirm"} + assert len(explore) == len(confirm) == 2 + assert explore.isdisjoint(confirm) diff --git a/tests/test_eval_agent/test_cli_large_prompt.py b/tests/test_eval_agent/test_cli_large_prompt.py index 58720910..95e47e91 100644 --- a/tests/test_eval_agent/test_cli_large_prompt.py +++ b/tests/test_eval_agent/test_cli_large_prompt.py @@ -43,6 +43,13 @@ def _echo_argv_script(tmp_path: Path) -> str: class TestAgyModelLargePrompt: + def test_model_and_effort_reach_cli(self, tmp_path): + model = AgyModel(binary_path=_echo_argv_script(tmp_path), timeout_sec=10, + model='gemini-3.6-flash', effort='low') + out = model.generate('OK') + assert '--model gemini-3.6-flash' in out + assert '--effort low' in out + def test_small_prompt_stays_inline(self, tmp_path): model = AgyModel(binary_path=_echo_argv_script(tmp_path), timeout_sec=10) out = model.generate("hello judge") diff --git a/tests/test_eval_agent/test_fix_agent.py b/tests/test_eval_agent/test_fix_agent.py index 5f8ba50e..38b91307 100644 --- a/tests/test_eval_agent/test_fix_agent.py +++ b/tests/test_eval_agent/test_fix_agent.py @@ -3989,3 +3989,25 @@ def test_plain_description_never_falls_back_to_the_slug(): assert plain_description(judged).startswith("Gives the model a short checklist") assert plain_description(builtin).startswith("Tells the model to read the answer") assert plain_description(unknown) == "" + + +def test_progress_survives_suppressed_result_logger_and_never_includes_labels(): + from evalrx.eval_agent.stages.fix_agent import FixCandidate + from tests.conftest import FakeModel + + events = [] + agent = FixAgent(score_fn=lambda case, output: output == "A", + progress_callback=events.append, run_logger=None) + case = FailureCase(id="secret-case", inputs=Inputs(prompt="question"), expected="A", + observed="B", label=Label.FAIL) + candidate = FixCandidate(tier=FixTier.L1_PROMPT, name="short", kind="template", + payload={"prompt_template": "{prompt}"}, source="test") + model = FakeModel() + model.generate = lambda *args, **kwargs: "A" + agent._validate(candidate, model, CaseBatch([case]), {case.id: False}) + assert events[0]["phase"] == "candidate_started" + assert events[-1]["phase"] == "candidate_completed" + assert events[-1]["completed"] == 1 + assert any(e["phase"] == "candidate_cases" for e in events) + assert all(set(e) <= {"phase", "candidate", "completed", "total"} for e in events) + assert "secret-case" not in str(events) diff --git a/tests/test_eval_agent/test_vl_diagnose.py b/tests/test_eval_agent/test_vl_diagnose.py index ec11cb7c..56a125ec 100755 --- a/tests/test_eval_agent/test_vl_diagnose.py +++ b/tests/test_eval_agent/test_vl_diagnose.py @@ -906,23 +906,17 @@ def test_analysis_only_stops_after_m2(self): assert report.all_hypotheses == [] assert report.all_test_results == [] - def test_no_diagnosis_agent_stops_early(self): - loop = VLDiagnoseLoop( - model=_vlm(), - protocol=_spatial_protocol(), - diagnosis_agent=None, - max_cycles=3, - ) - # Should not raise even with no diagnosis agent; override lazy resolution - # by patching the internal flag - loop._diag_instance = None # type: ignore[attr-defined] - # The loop will hit the analysis-only guard - # (we just verify no exception is raised) - try: - report = loop.run(_labeled_batch()) - assert report.cycles >= 1 - except Exception: - pass # acceptable if Gemini key absent; we just check no crash + def test_no_diagnosis_agent_stops_early(self, monkeypatch): + loop = VLDiagnoseLoop(model=_vlm(), protocol=_spatial_protocol(), max_cycles=3) + + def unavailable(): + raise RuntimeError("judge unavailable") + + monkeypatch.setattr(loop, "_get_diagnosis_agent", unavailable) + report = loop.run(_labeled_batch()) + assert report.stopped_by == "diagnosis_error" + assert report.cycles == 1 + class TestVLDiagnoseTwoPhase: @@ -1121,24 +1115,34 @@ def test_surgery_presence_restores_validity(self): class TestM3FaultTolerance: - def test_judge_exception_in_m3_does_not_kill_loop(self): + def test_judge_exception_in_m3_does_not_kill_loop(self, tmp_path): """A judge timeout/quota error in M3 must end the loop gracefully - (stopped_by=no_hypotheses), not propagate (regression: an agy + (stopped_by=diagnosis_error), not propagate (regression: an agy TimeoutExpired killed the whole run after M1+M2 had succeeded).""" class ExplodingDiagnosisAgent: def diagnose(self, analysis, model_name="", prior_cycles=None): raise RuntimeError("AgyModel: agy timed out after 120s") + import json + + from evalrx.eval_agent.run_logger_v2 import RunLoggerV2 + + run_logger = RunLoggerV2(tmp_path / "logs") loop = VLDiagnoseLoop( model=_vlm(), protocol=_spatial_protocol(), diagnosis_agent=ExplodingDiagnosisAgent(), - max_cycles=2, + max_cycles=2, run_logger=run_logger, ) report = loop.run(_labeled_batch()) assert isinstance(report, VLDiagnoseReport) - assert report.stopped_by == "no_hypotheses" + assert report.stopped_by == "diagnosis_error" + errors = json.loads((tmp_path / "logs/M3/log.json").read_text())["model_calls"] + assert len(errors) == 2 + assert errors[0]["metadata"]["retrying"] is True + assert errors[1]["metadata"]["retrying"] is False + assert all("timed out" in e["error"] for e in errors) assert report.all_hypotheses == [] @@ -1279,3 +1283,28 @@ def test_run_confirm_tests_on_holdout(self): assert report.m4_holdout == "confirmed" assert len(report.verified_hypotheses) >= 1 assert report.verified_hypotheses[0].evidence["split"] == "confirm_holdout" + + +def test_m3_transient_error_retries_without_reprobing(tmp_path): + from types import SimpleNamespace + + from evalrx.eval_agent.run_logger_v2 import RunLoggerV2 + + class Transient: + calls = 0 + + def diagnose(self, analysis, **kwargs): + self.calls += 1 + if self.calls == 1: + raise TimeoutError("transient timeout") + return SimpleNamespace(hypotheses=[], raw_judge_output="NO_ISSUE", tokens_used=1) + + agent = Transient() + loop = VLDiagnoseLoop(model=_vlm(), protocol=_spatial_protocol(), diagnosis_agent=agent, + run_logger=RunLoggerV2(tmp_path / "logs")) + result = loop._do_m3(0, None, [], {}, log=False) + assert result is not None and result.hypotheses == [] + assert agent.calls == 2 + # A valid empty answer is not retried. + loop._do_m3(0, None, [], {}, log=False) + assert agent.calls == 3 diff --git a/tests/test_examples/test_benchmark_common.py b/tests/test_examples/test_benchmark_common.py index 8a8a3f50..5353de94 100644 --- a/tests/test_examples/test_benchmark_common.py +++ b/tests/test_examples/test_benchmark_common.py @@ -586,7 +586,7 @@ def fake_runtime(**kw): assert gen_kwargs == {} and spec.key == "gemini-3.6-flash" assert model.capabilities == frozenset({Capability.GENERATE, Capability.TOOL_CALLS}) assert {"image", "audio"} <= set(model.modalities) - assert captured["with_logprobs"] is False and captured["api_key"] is None + assert captured["with_logprobs"] == "auto" and captured["api_key"] is None assert captured["thinking"].floor is True and captured["thinking"].level is None assert (captured["temperature"], captured["top_p"], captured["top_k"]) == (0.6, 0.95, 20) assert captured["max_output_tokens"] == 2048 @@ -849,3 +849,22 @@ def test_all_cells_default_to_the_endpoint_backend(common): assert models.resolve("gemini-2.5-flash-lite", "llm").backend == "gemini" # family forced with pytest.raises(ValueError): models.default_backend("video") + + +@pytest.mark.parametrize("output,letter", [ + ("The correct answer is **A**.", "A"), ("The correct option is **A**.", "A"), + ("__A__", "A"), ("`B`", "B"), ("Answer: **C**", "C"), ("Final: _D_", "D"), +]) +def test_choice_markdown_is_presentation_not_a_wrong_answer(common, output, letter): + _, _, scoring, _ = common + assert scoring.parsed_choice(output) == letter + assert scoring.score_output("multiple_choice_letter", output, letter) + + +@pytest.mark.parametrize("output", [ + "Either **A** or **C**.", "**A** dog, **B** cat", "thought\n**A**", + "A dog barks.", "AB", "**AB**", +]) +def test_choice_markdown_does_not_turn_ambiguity_into_a_commitment(common, output): + _, _, scoring, _ = common + assert scoring.parsed_choice(output) == "" diff --git a/tests/test_models/test_gemini_compat.py b/tests/test_models/test_gemini_compat.py index d428e807..c2387cdb 100644 --- a/tests/test_models/test_gemini_compat.py +++ b/tests/test_models/test_gemini_compat.py @@ -15,6 +15,7 @@ import threading import wave from dataclasses import dataclass, field +from types import SimpleNamespace from typing import Any, Optional import pytest @@ -417,6 +418,68 @@ def test_logprobs_are_not_claimed_unless_wired(): assert rt.logprobs_fn is not None +def test_auto_logprobs_enables_only_after_real_probability_data(): + result = SimpleNamespace(chosen_candidates=[SimpleNamespace(token="OK", log_probability=-0.2)], + top_candidates=[]) + client = FakeClient([_reply(Part(text="OK"), logprobs=result), + _reply(Part(text="OK"), logprobs=result)]) + rt = gemini_runtime(client=client, with_logprobs="auto", logprobs_model="gemini-3.6-flash") + assert rt.logprobs_fn is not None + assert rt.generate_fn.state['logprobs_probe']['status'] == 'available' + assert rt.generate_fn.state['logits']['status'] == 'unavailable' + from evalrx.models import compose + model = compose('gemini-3.6-flash', 'api', rt) + assert Capability.LOGPROBS in model.capabilities + assert Capability.LOGITS not in model.capabilities + from evalrx.analyzers.uncertainty.logprob_entropy import LogprobEntropyAnalyzer + result = LogprobEntropyAnalyzer().run(model, 'Answer OK') + assert result.findings['mean_logprob'] == pytest.approx(-0.2) + assert result.findings['perplexity'] == pytest.approx(1.2214) + + +@pytest.mark.parametrize('reply', [ + APIError(400, 'Logprobs is not enabled for this model'), + APIError(400, 'Logprobs is not supported for this model'), + _reply(Part(text='OK')), + _reply(Part(text='OK'), logprobs=SimpleNamespace( + chosen_candidates=[SimpleNamespace(token='OK', log_probability=None)])), +]) +def test_auto_logprobs_does_not_claim_missing_or_unsupported_data(reply): + rt = gemini_runtime(client=FakeClient([reply]), with_logprobs='auto', + logprobs_model='gemini-3.6-flash') + assert rt.logprobs_fn is None + assert rt.generate_fn.state['logprobs_probe']['status'] in {'unsupported', 'not_returned'} + + +@pytest.mark.parametrize('error', [APIError(403, 'Permission denied'), + APIError(429, 'quota'), + APIError(400, 'Invalid sampling config'), TimeoutError('timeout')]) +def test_auto_logprobs_does_not_hide_operational_failures(error): + with pytest.raises(type(error)): + gemini_runtime(client=FakeClient([error]), retries=0, with_logprobs='auto', + logprobs_model='gemini-3.6-flash') + + +def test_required_logprobs_fails_when_response_has_no_probabilities(): + with pytest.raises(RuntimeError, match='usable logprobs'): + gemini_runtime(client=FakeClient([_reply(Part(text='OK'))]), with_logprobs=True, + logprobs_model='gemini-3.6-flash') + + +def test_logprobs_missing_after_opt_in_is_not_an_empty_probability_distribution(): + rt = gemini_runtime(client=FakeClient([_reply(Part(text='OK'))]), with_logprobs=True) + with pytest.raises(RuntimeError, match='no token logprobs'): + rt.logprobs_fn('q', model='gemini-3.6-flash') + + +@pytest.mark.parametrize('value', [None, float('nan'), float('inf'), 0.1]) +def test_malformed_logprobs_are_never_replaced_with_probability_one(value): + result = SimpleNamespace(chosen_candidates=[SimpleNamespace(token='OK', log_probability=value)], + top_candidates=[]) + with pytest.raises(ValueError, match='invalid token log probability'): + parse_logprobs_result(result) + + def test_logprobs_fn_asks_for_response_logprobs_and_parses_the_result(): @dataclass class Cand: From 013fb435e0033058dadc655bb209b56a9f03bf8a Mon Sep 17 00:00:00 2001 From: Nanboy-Ronan Date: Mon, 12 Oct 2026 14:02:26 -0700 Subject: [PATCH 2/9] Make benchmark diagnosis agentic by default and simplify CLI controls --- docs/agentic_benchmark.md | 10 ++++----- evalrx/benchmark/run.py | 23 ++++++------------- evalrx/benchmark/runner.py | 26 +++++++++------------ examples/benchmark/README.md | 13 ++++++----- examples/benchmark/alm/gemini/README.md | 14 +++++------- examples/benchmark/llm/gemini/README.md | 14 +++++------- examples/benchmark/vlm/gemini/README.md | 14 +++++------- tests/test_analysis/test_run_cli.py | 30 +++++++++++++++---------- 8 files changed, 65 insertions(+), 79 deletions(-) diff --git a/docs/agentic_benchmark.md b/docs/agentic_benchmark.md index f6fa1fe4..6e9d040a 100644 --- a/docs/agentic_benchmark.md +++ b/docs/agentic_benchmark.md @@ -4,24 +4,24 @@ Use the same datasets, scorers, and final repair validation with an agent that c ```bash python -m evalrx.benchmark.run --modality llm --model gemini-3.6-flash \ - --dataset bbh_causal_judgement --diagnosis-mode agentic --max-actions 12 + --dataset bbh_causal_judgement ``` The existing judge/coder options apply. Replace modality and dataset for image/audio tasks. Gemini probability support is probed on the actual model; unsupported logprobs and raw logits are never simulated. -`--diagnosis-mode fixed` remains the default for reproducible fixed-pipeline comparisons. Agentic mode is a separate protocol variant, not an exact model-only comparison with a historical fixed run. +Agentic diagnosis is the default. No mode or per-stage enable flags are needed. `--max-actions` optionally changes the default budget of 12 controller decisions. The hidden legacy `--diagnosis-mode fixed` option remains available only to replay historical fixed-pipeline comparisons; current default runs are a different protocol from those historical results. ## Controller scope -The controller can choose catalog probes, request custom probes, run statistics, explore data, cluster failures, propose hypotheses, and test them on EXPLORE. It can gather more evidence after an empty diagnosis, subject to action and tool-call caps. Custom M1 probes and M2 tool generation default on in agentic mode across modalities; disable them with `--no-m1-codegen` and `--no-m2-codegen`. +The controller can choose catalog probes, request custom probes, run statistics, explore data, cluster failures, propose hypotheses, and test them on EXPLORE. It can gather more evidence after an empty diagnosis, subject to action and tool-call caps. Custom M1 probes and M2 tool generation default on in agentic mode across modalities; The existing `--no-allow-codegen` disables custom M1/repair generation, and `--no-m2-codegen` disables M2 generation. Repair and surgery are excluded from the benchmark controller's tool registry. After the investigation, the host runs repair search on EXPLORE and validates a frozen winner once on CONFIRM. The controller cannot inspect confirmation feedback and adapt to it. `--held-out` currently requires fixed mode; agentic mode rejects that combination before loading a model. -Built-in floor candidates default off in agentic mode (`--fix-floor` opts in). The agent can still propose known repair methods. The existing repair tier/capability constraints and maximum repair rounds remain enforced. This mode is bounded adaptive diagnosis, not an unlimited autonomous training loop. +Agentic repair does not inject built-in floor candidates. The agent can still propose known repair methods. The existing repair tier/capability constraints and maximum repair rounds remain enforced. This mode is bounded adaptive diagnosis, not an unlimited autonomous training loop. ## Failure and progress semantics -M3 exceptions are recorded as structured error exchanges with exception type, message, attempt, and retry status. `--diagnosis-retries 1` permits one additional attempt; valid empty diagnoses are not retried as transport failures. Exhaustion produces `diagnosis_error`, an error diagnosis status, and benchmark exit code 2 rather than a successful `no_hypotheses` result. +M3 exceptions are recorded as structured error exchanges with exception type, message, attempt, and retry status. The runtime automatically permits one additional attempt; valid empty diagnoses are not retried as transport failures. Exhaustion produces `diagnosis_error`, an error diagnosis status, and benchmark exit code 2 rather than a successful `no_hypotheses` result. Repair progress prints proposal starts, candidate starts/completions, and case counters for ordinary template/spec evaluation, even while detailed result logging is deferred. These operational events contain no case IDs, labels, outputs, or correctness scores. Coded pipelines emit candidate start/end events; their internal calls are not represented as ordinary per-case counters. diff --git a/evalrx/benchmark/run.py b/evalrx/benchmark/run.py index 00bccbac..8051e0fc 100644 --- a/evalrx/benchmark/run.py +++ b/evalrx/benchmark/run.py @@ -41,10 +41,6 @@ def build_parser(p: argparse.ArgumentParser | None = None) -> argparse.ArgumentP help="per-request timeout in seconds for the API backends") p.add_argument("--request-retries", type=int, default=5, help="retries with backoff on 429/5xx for --backend gemini") - p.add_argument("--gemini-logprobs", choices=["auto", "on", "off"], default="auto", - help="probe and enable real Gemini logprobs when available; on requires support; off skips probing") - p.add_argument("--gemini-top-logprobs", type=int, choices=range(21), default=5, - help="number of top token alternatives requested from Gemini (0..20)") p.add_argument("--thinking-level", choices=["minimal", "low", "medium", "high"], default=None, help="gemini: thinking_level for a 3.x model (default: the model's floor, i.e. " "minimal, or low on 3.7-flash); on a 2.5 model it maps to a thinking_budget") @@ -125,16 +121,11 @@ def build_parser(p: argparse.ArgumentParser | None = None) -> argparse.ArgumentP p.add_argument("--explore", action=argparse.BooleanOptionalAction, default=True, help="in-cycle free-form EDA between M1 and M2") p.add_argument("--max-cycles", type=int, default=1) - p.add_argument("--diagnosis-mode", choices=["fixed", "agentic"], default="fixed", - help="agentic: budgeted tool selection, custom probes, and adaptive investigation") + # Legacy replay only; normal runs always default to the adaptive controller. + p.add_argument("--diagnosis-mode", choices=["fixed", "agentic"], default="agentic", + help=argparse.SUPPRESS) p.add_argument("--max-actions", type=int, default=12, - help="maximum controller decisions in agentic diagnosis") - p.add_argument("--diagnosis-retries", type=int, default=1, - help="bounded retries after an M3 exception (not after a valid empty result)") - p.add_argument("--m1-codegen", action=argparse.BooleanOptionalAction, default=None, - help="allow custom probe generation (default: on in agentic mode)") - p.add_argument("--fix-floor", action=argparse.BooleanOptionalAction, default=None, - help="include built-in floor candidates (default: off in agentic mode)") + help="maximum diagnostic controller decisions (default: 12)") p.add_argument("--held-out", action=argparse.BooleanOptionalAction, default=False, help="use the task's held-out validation manifest (manifest_val.json, a " "disjoint fresh sample of half the main slice): M4 verifies once there " @@ -155,7 +146,7 @@ def build_parser(p: argparse.ArgumentParser | None = None) -> argparse.ArgumentP "(modality-gated on the MODEL, so avoid on multimodal specs for text tasks)") p.add_argument("--analyzer-max-cases", type=int, default=0, help="cap per analyzer (0 = every case)") p.add_argument("--m2-codegen", action=argparse.BooleanOptionalAction, default=None, - help="coder-written M2 statistics tools (default: on for llm, off otherwise)") + help="coder-written M2 statistics tools (default: on)") p.add_argument( "--fix-validation-cases", type=int, @@ -216,8 +207,8 @@ def run_args(args) -> int: return 0 if not args.modality or not args.model: build_parser().error("--modality and --model are required (or --list / --smoke-test)") - if args.max_actions < 1 or args.diagnosis_retries < 0: - raise ValueError("max-actions must be positive and diagnosis-retries non-negative") + if args.max_actions < 1: + raise ValueError("max-actions must be positive") if args.diagnosis_mode == "agentic" and args.held_out: raise ValueError("agentic diagnosis uses EXPLORE; --held-out is supported only in fixed mode") dataset = args.dataset or T.DEFAULT_TASK[args.modality] diff --git a/evalrx/benchmark/runner.py b/evalrx/benchmark/runner.py index 05efd36f..a369f8e7 100644 --- a/evalrx/benchmark/runner.py +++ b/evalrx/benchmark/runner.py @@ -117,10 +117,9 @@ def load_model(resolved: Resolved, args, task: T.Task): floor=not args.enable_thinking) runtime = gemini_runtime(api_key=args.api_key, timeout=args.request_timeout, retries=args.request_retries, thinking=policy, - with_logprobs={"auto": "auto", "on": True, "off": False}[ - getattr(args, "gemini_logprobs", "auto")], + with_logprobs="auto", logprobs_model=spec.hf_repo or spec.key, - top_logprobs=getattr(args, "gemini_top_logprobs", 5), **sampling) + top_logprobs=5, **sampling) return compose(spec, "api", runtime, set()), {}, spec from evalrx.models.backends.base import RuntimeConfig @@ -200,9 +199,8 @@ def build_judge(args): providers as the m1_m5 examples; the benchmark CLI defaults to Claude / claude-opus-5 / high.""" if args.judge_provider == "gemini": - adaptive = getattr(args, "diagnosis_mode", "fixed") == "agentic" - probe_codegen = getattr(args, "m1_codegen", None) - probe_codegen = adaptive if probe_codegen is None else probe_codegen + adaptive = getattr(args, "diagnosis_mode", "agentic") == "agentic" + probe_codegen = adaptive and args.allow_codegen stats_codegen = args.m2_codegen if args.m2_codegen is not None else adaptive if (args.allow_codegen or args.explore or stats_codegen or probe_codegen or not args.skip_m5) and not getattr(args, "coder_provider", None): @@ -425,9 +423,9 @@ def run(args, task: T.Task, resolved: Resolved) -> int: "fix_tier": args.fix_tier, "allow_codegen": args.allow_codegen, "auto_escalate": args.auto_escalate, "m1_selection": args.m1_selection, "generation_kwargs": gen_kwargs, - "diagnosis_mode": getattr(args, "diagnosis_mode", "fixed"), + "diagnosis_mode": getattr(args, "diagnosis_mode", "agentic"), "max_actions": getattr(args, "max_actions", 12), - "diagnosis_retries": getattr(args, "diagnosis_retries", 1), + "diagnosis_retries": 1, "enable_thinking": bool(args.enable_thinking), "judge_provider": args.judge_provider, "judge_model": args.judge_model, }, @@ -556,9 +554,8 @@ def run(args, task: T.Task, resolved: Resolved) -> int: # backends can continue issuing concurrent requests. probe_workers = 1 if resolved.backend in {"hf_local", "jax_local"} else 8 ctx.config["m1_workers"] = probe_workers - agentic = getattr(args, "diagnosis_mode", "fixed") == "agentic" - m1_codegen = getattr(args, "m1_codegen", None) - m1_codegen = agentic if m1_codegen is None else m1_codegen + agentic = getattr(args, "diagnosis_mode", "agentic") == "agentic" + m1_codegen = agentic and args.allow_codegen ctx.config["m1_codegen"] = bool(m1_codegen) ctx.config["m1_selection"] = "judge" if agentic or m1_codegen else args.m1_selection if args.m1_selection == "pinned" and not agentic and not m1_codegen: @@ -617,10 +614,7 @@ def run(args, task: T.Task, resolved: Resolved) -> int: exec_timeout_sec=args.fix_exec_timeout, prewritten_code=prewritten_code, concurrency=getattr(args, "concurrency", 1), - floor_candidates=("self_consistency_5",) if ( - getattr(args, "fix_floor", None) if getattr(args, "fix_floor", None) is not None - else not agentic - ) else (), + floor_candidates=() if agentic else ("self_consistency_5",), progress_callback=_print_fix_progress, **fix_kwargs, ) @@ -647,7 +641,7 @@ def run(args, task: T.Task, resolved: Resolved) -> int: registry_factory=_benchmark_agentic_registry) loop = loop_class( **loop_kwargs, - diagnosis_retries=getattr(args, "diagnosis_retries", 1), + diagnosis_retries=1, model=model, protocol=protocol, probe_agent=probe_agent, stats_agent=stats_agent, diagnosis_agent=DiagnosisAgent(judge=judge), hypothesis_tester=HypothesisTester(judge=judge, min_effect=0.05), diff --git a/examples/benchmark/README.md b/examples/benchmark/README.md index 50b05061..3e657f71 100644 --- a/examples/benchmark/README.md +++ b/examples/benchmark/README.md @@ -1,7 +1,7 @@ # Benchmark matrix: model family × modality × dataset -The full diagnosis chain — Stage 0 baseline → M1 (pinned analyzers) → explore → -M2 → M3 → M4 → M5 → fix — run as a **matrix** over three model families, +The full diagnosis chain — Stage 0 baseline → agent-directed M1–M4 investigation +→ M5 → repair search and confirmation — runs as a **matrix** over model families, three input modalities and the benchmark datasets of each modality, through one code path ([`_common/run.py`](_common/run.py)). Each `//` directory is one cell family: a compose file whose services are the sizes, with @@ -200,10 +200,11 @@ the final significance gate, and the validation data costs no confirm power. defaults `enable_thinking` to false and the spec sends it explicitly, yet 3/8 MMAU generations still opened a `thought` channel and hit the 64-token cap. `termination_audit` / `answer_extraction_audit` are pinned for exactly that. -* **M1 is pinned per task** (`Task.pinned_m1`, the audited sets of the `m1_m5` - examples; an 8-analyzer text set for llm). `--m1-selection judge` switches to - catalog selection, which is modality-gated on the MODEL — on a multimodal spec - running a text task it can pick image analyzers, hence the pinned default. +* **Agentic diagnosis is the default.** The controller chooses its next action, + selects available probes, and can request custom analysis using the task and + observed failures. No mode flag is required. `--max-actions` optionally changes + the default budget of 12 decisions. See [agentic diagnosis](../../docs/agentic_benchmark.md) + for confirmation isolation and legacy fixed-protocol replay. ### Held-out isolation during the fix stage diff --git a/examples/benchmark/alm/gemini/README.md b/examples/benchmark/alm/gemini/README.md index a9de3cc8..4759ef2b 100644 --- a/examples/benchmark/alm/gemini/README.md +++ b/examples/benchmark/alm/gemini/README.md @@ -30,13 +30,11 @@ audio, so the same service names exist in the llm / vlm / alm leaves. `--thinking-level {minimal,low,medium,high}` / `--thinking-budget N` override the floor; `--enable-thinking` leaves the API default. -The benchmark defaults to `--gemini-logprobs auto`: a short real request must -return valid token log probabilities before the model exposes `LOGPROBS` to -analyzers. `--gemini-logprobs on` requires support and fails if unavailable; -`off` skips the probe. `--gemini-top-logprobs 5` requests five alternatives -(allowed range 0–20). The response evidence and effective capabilities are -saved in `api_capabilities.json`. Authentication, quota and transport failures -remain errors; they are not interpreted as missing model capabilities. +The benchmark automatically probes probability support with a short real request. +Only valid token log probabilities enable `LOGPROBS` analyzers; five top +alternatives are requested when supported. Response evidence and effective +capabilities are saved in `api_capabilities.json`. Authentication, quota and +transport failures remain errors rather than being treated as missing capabilities. Raw logits and model internals are unavailable. Top-k logprobs cannot be substituted for logits. Without logprobs, probability-only analyzers are @@ -89,4 +87,4 @@ Requests retry with backoff on 429/5xx (`--request-retries`, default 5; `--request-timeout`, default 300 s). `CONCURRENCY` (default 4) controls the number of discovery requests in flight. -For adaptive probe selection and bounded controller decisions, add `--diagnosis-mode agentic --max-actions 12`. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. +Adaptive probe selection and bounded controller decisions are the default; no mode flag is needed. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. diff --git a/examples/benchmark/llm/gemini/README.md b/examples/benchmark/llm/gemini/README.md index 5c0b31b2..033eeda0 100644 --- a/examples/benchmark/llm/gemini/README.md +++ b/examples/benchmark/llm/gemini/README.md @@ -30,13 +30,11 @@ audio, so the same service names exist in the llm / vlm / alm leaves. `--thinking-level {minimal,low,medium,high}` / `--thinking-budget N` override the floor; `--enable-thinking` leaves the API default. -The benchmark defaults to `--gemini-logprobs auto`: a short real request must -return valid token log probabilities before the model exposes `LOGPROBS` to -analyzers. `--gemini-logprobs on` requires support and fails if unavailable; -`off` skips the probe. `--gemini-top-logprobs 5` requests five alternatives -(allowed range 0–20). The response evidence and effective capabilities are -saved in `api_capabilities.json`. Authentication, quota and transport failures -remain errors; they are not interpreted as missing model capabilities. +The benchmark automatically probes probability support with a short real request. +Only valid token log probabilities enable `LOGPROBS` analyzers; five top +alternatives are requested when supported. Response evidence and effective +capabilities are saved in `api_capabilities.json`. Authentication, quota and +transport failures remain errors rather than being treated as missing capabilities. Raw logits and model internals are unavailable. Top-k logprobs cannot be substituted for logits. Without logprobs, probability-only analyzers are @@ -96,4 +94,4 @@ Requests retry with backoff on 429/5xx (`--request-retries`, default 5; `--request-timeout`, default 300 s). `CONCURRENCY` (default 4) controls the number of discovery requests in flight. -For adaptive probe selection and bounded controller decisions, add `--diagnosis-mode agentic --max-actions 12`. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. +Adaptive probe selection and bounded controller decisions are the default; no mode flag is needed. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. diff --git a/examples/benchmark/vlm/gemini/README.md b/examples/benchmark/vlm/gemini/README.md index ee80be66..1782b154 100644 --- a/examples/benchmark/vlm/gemini/README.md +++ b/examples/benchmark/vlm/gemini/README.md @@ -30,13 +30,11 @@ audio, so the same service names exist in the llm / vlm / alm leaves. `--thinking-level {minimal,low,medium,high}` / `--thinking-budget N` override the floor; `--enable-thinking` leaves the API default. -The benchmark defaults to `--gemini-logprobs auto`: a short real request must -return valid token log probabilities before the model exposes `LOGPROBS` to -analyzers. `--gemini-logprobs on` requires support and fails if unavailable; -`off` skips the probe. `--gemini-top-logprobs 5` requests five alternatives -(allowed range 0–20). The response evidence and effective capabilities are -saved in `api_capabilities.json`. Authentication, quota and transport failures -remain errors; they are not interpreted as missing model capabilities. +The benchmark automatically probes probability support with a short real request. +Only valid token log probabilities enable `LOGPROBS` analyzers; five top +alternatives are requested when supported. Response evidence and effective +capabilities are saved in `api_capabilities.json`. Authentication, quota and +transport failures remain errors rather than being treated as missing capabilities. Raw logits and model internals are unavailable. Top-k logprobs cannot be substituted for logits. Without logprobs, probability-only analyzers are @@ -90,4 +88,4 @@ Requests retry with backoff on 429/5xx (`--request-retries`, default 5; `--request-timeout`, default 300 s). `CONCURRENCY` (default 4) controls the number of discovery requests in flight. -For adaptive probe selection and bounded controller decisions, add `--diagnosis-mode agentic --max-actions 12`. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. +Adaptive probe selection and bounded controller decisions are the default; no mode flag is needed. See [adaptive benchmark diagnosis](../../../../docs/agentic_benchmark.md) for the confirmation boundary, failure handling, and protocol differences. diff --git a/tests/test_analysis/test_run_cli.py b/tests/test_analysis/test_run_cli.py index 74b7546f..e6a31462 100644 --- a/tests/test_analysis/test_run_cli.py +++ b/tests/test_analysis/test_run_cli.py @@ -79,8 +79,7 @@ def factory(**kwargs): assert observed['effort'] == 'low' -@pytest.mark.parametrize('mode,expected', [('auto', 'auto'), ('on', True), ('off', False)]) -def test_gemini_benchmark_forwards_probability_policy(monkeypatch, mode, expected): +def test_gemini_benchmark_probes_probability_support_automatically(monkeypatch): from importlib import import_module from evalrx.benchmark.models import resolve @@ -93,20 +92,27 @@ def runtime(**kwargs): return SimpleNamespace() monkeypatch.setattr(gemini, 'gemini_runtime', runtime) monkeypatch.setattr(import_module('evalrx.models.compose'), 'compose', lambda *args: 'model') - args = build_parser().parse_args(['--modality', 'vlm', '--model', 'gemini-3.6-flash', - '--gemini-logprobs', mode, '--gemini-top-logprobs', '10']) + args = build_parser().parse_args(['--modality', 'vlm', '--model', 'gemini-3.6-flash']) runner.load_model(resolve(args.model, 'vlm', 'gemini'), args, get('chartqa')) - assert observed['with_logprobs'] == expected + assert observed['with_logprobs'] == 'auto' assert observed['logprobs_model'] == 'gemini-3.6-flash' - assert observed['top_logprobs'] == 10 + assert observed['top_logprobs'] == 5 -def test_agentic_cli_exposes_bounded_controller(): - args = build_parser().parse_args(["--diagnosis-mode", "agentic", "--max-actions", "8", - "--no-fix-floor", "--m1-codegen"]) - assert (args.diagnosis_mode, args.max_actions, args.fix_floor, args.m1_codegen) == ( - "agentic", 8, False, True) - assert args.diagnosis_retries == 1 +def test_agentic_is_default_without_extra_flags(): + parser = build_parser() + args = parser.parse_args([]) + assert args.diagnosis_mode == "agentic" + assert args.max_actions == 12 + help_text = parser.format_help() + for removed in ("--diagnosis-mode", "--diagnosis-retries", "--m1-codegen", + "--fix-floor", "--gemini-logprobs", "--gemini-top-logprobs"): + assert removed not in help_text + assert parser.parse_args(["--max-actions", "8"]).max_actions == 8 + + +def test_legacy_fixed_protocol_remains_replayable(): + assert build_parser().parse_args(["--diagnosis-mode", "fixed"]).diagnosis_mode == "fixed" def test_agentic_rejects_unsupported_holdout_before_model_load(): From 22d4f85cd07f37709fa8537ccd607222a71675cb Mon Sep 17 00:00:00 2001 From: Nanboy-Ronan Date: Mon, 12 Oct 2026 14:09:35 -0700 Subject: [PATCH 3/9] Fix LLM repair construction and test default runner wiring across modalities --- evalrx/benchmark/runner.py | 2 +- tests/test_analysis/test_run_cli.py | 44 +++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/evalrx/benchmark/runner.py b/evalrx/benchmark/runner.py index a369f8e7..0daff89f 100644 --- a/evalrx/benchmark/runner.py +++ b/evalrx/benchmark/runner.py @@ -590,7 +590,7 @@ def run(args, task: T.Task, resolved: Resolved) -> int: floor = {"max_tokens": gen_kwargs.get("max_new_tokens", task.max_new_tokens)} if sampled: floor.update(temperature=gen_kwargs["temperature"], top_p=gen_kwargs["top_p"]) - fix_kwargs.update(baseline_generation_kwargs=floor, floor_candidates=("self_consistency_5",)) + fix_kwargs.update(baseline_generation_kwargs=floor) prewritten_code = ( Path(args.fix_code_file).read_text(encoding="utf-8") if args.fix_code_file else "" diff --git a/tests/test_analysis/test_run_cli.py b/tests/test_analysis/test_run_cli.py index e6a31462..1ce63dd3 100644 --- a/tests/test_analysis/test_run_cli.py +++ b/tests/test_analysis/test_run_cli.py @@ -121,3 +121,47 @@ def test_agentic_rejects_unsupported_holdout_before_model_load(): "--diagnosis-mode", "agentic", "--held-out"]) with pytest.raises(ValueError, match="EXPLORE"): run_args(args) + + +@pytest.mark.parametrize('modality,dataset', [ + ('llm', 'bbh_causal_judgement'), ('vlm', 'chartqa'), ('alm', 'mmau'), +]) +def test_default_runner_constructs_agentic_pipeline_for_every_modality(monkeypatch, tmp_path, + modality, dataset): + from evalrx.benchmark.models import resolve + from evalrx.benchmark.tasks import get + from evalrx.core.case import CaseBatch, FailureCase, Inputs, Label + from evalrx.eval_agent import AgenticDiagnoseLoop, CaseDiscoveryAgent + from evalrx.specs import get_spec + from tests.conftest import FakeModel + + cases = CaseBatch([ + FailureCase(id=str(i), inputs=Inputs(prompt=f'question {i}'), expected='A', + observed='A' if i % 2 else 'B', label=Label.PASS if i % 2 else Label.FAIL) + for i in range(8) + ]) + args = build_parser().parse_args(['--modality', modality, '--model', 'gemini-3.6-flash', + '--dataset', dataset, '--run-dir', str(tmp_path)]) + resolved = resolve(args.model, modality, 'gemini') + monkeypatch.setattr(runner, 'ensure_manifest', lambda *a: tmp_path / 'manifest.json') + monkeypatch.setattr(runner.T, 'build_cases', lambda *a: (cases, [])) + monkeypatch.setattr(runner, 'build_judge', lambda *a: (FakeModel(), 'claude', None, ())) + monkeypatch.setattr(runner, 'load_model', lambda *a: ( + FakeModel(), {'do_sample': True, 'temperature': 0.6, 'top_p': 0.95}, + get_spec(resolved.spec_key))) + monkeypatch.setattr(CaseDiscoveryAgent, 'discover', lambda *a, **k: SimpleNamespace( + cases=cases, n_pass=4, n_fail=4, n_unknown=0, has_m4_groups=True)) + + class ReachedController(Exception): + pass + + def run(loop, data): + assert loop.probe_agent.allow_codegen + assert loop.max_actions == 12 + assert loop.fix_agent._floor_candidates == () + assert len(data) == 8 + raise ReachedController + + monkeypatch.setattr(AgenticDiagnoseLoop, 'run', run) + with pytest.raises(ReachedController): + runner.run(args, get(dataset), resolved) From 72857889dc6a40c23e052a018e6dfc2a0ba441c6 Mon Sep 17 00:00:00 2001 From: Nanboy-Ronan Date: Mon, 12 Oct 2026 14:25:03 -0700 Subject: [PATCH 4/9] Publish agentic exploration artifacts and pass evidence into diagnosis --- evalrx/eval_agent/agentic/tools.py | 14 +++++----- evalrx/eval_agent/loop.py | 4 +-- .../test_eval_agent/test_loop_explore_step.py | 26 +++++++++++++++++++ 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/evalrx/eval_agent/agentic/tools.py b/evalrx/eval_agent/agentic/tools.py index 290f9ef3..0e5b4fc4 100644 --- a/evalrx/eval_agent/agentic/tools.py +++ b/evalrx/eval_agent/agentic/tools.py @@ -228,18 +228,16 @@ def handler(board: "EvidenceBoard", params: dict[str, Any]) -> ToolOutcome: def _explore_data(loop: Any, state: _RunState) -> "Callable[[EvidenceBoard, dict], ToolOutcome]": def handler(board: "EvidenceBoard", params: dict[str, Any]) -> ToolOutcome: - from evalrx.analysis.operationalize import per_case_to_records - from evalrx.analysis.stats_tools import build_stats_input - question = str( params.get("question") or "Explore this dataset and surface the patterns that matter." ) - stats_input = build_stats_input(state.probe_results, state.data) - records = per_case_to_records(stats_input.per_case, stats_input.labels) - if not records: - return ToolOutcome(False, "no per-case records to explore yet", error="empty") - report = loop.explorer.explore_records(records, question=question) + # Reuse host adjudication, artifact capture, and M3 context publication. + report = loop._do_explore( + state.cycle, state.probe_results, state.data, state.timings, question=question, + ) + if report is None: + return ToolOutcome(False, "exploration produced no report", error="empty") board.explore_takeaways = [t.to_dict() for t in getattr(report, "takeaways", [])] summary = ( f"{len(board.explore_takeaways)} takeaway(s)" diff --git a/evalrx/eval_agent/loop.py b/evalrx/eval_agent/loop.py index 95e4228e..694d7454 100755 --- a/evalrx/eval_agent/loop.py +++ b/evalrx/eval_agent/loop.py @@ -835,7 +835,7 @@ def _explore_out_dir(self) -> "Path | None": def _do_explore( self, cycle: int, probe_results: "dict[str, Any]", data: "Any", - timings: "dict[str, float]", *, log: bool = True, + timings: "dict[str, float]", *, log: bool = True, question: str | None = None, ) -> "Any | None": """Optional in-cycle explore step: free-form EDA over M1's per-case table. @@ -876,7 +876,7 @@ def _do_explore( if not records: logger.info("explore: M1 produced no per-case signals — skipping.") return None - question = self._explore_question + question = question or self._explore_question if not question: from evalrx.eval_agent.prompts.explore_step import default_explore_question diff --git a/tests/test_eval_agent/test_loop_explore_step.py b/tests/test_eval_agent/test_loop_explore_step.py index 290b9c1c..06385d3e 100644 --- a/tests/test_eval_agent/test_loop_explore_step.py +++ b/tests/test_eval_agent/test_loop_explore_step.py @@ -395,3 +395,29 @@ def test_default_explore_dir_beside_a_standalone_logs_dir_and_inside_other_dirs( assert loop._explore_out_dir() is None loop.run(_batch()) assert diag.seen_context[0] is not None + + +def test_agentic_explore_publishes_artifacts_and_passes_notes_to_m3(tmp_path): + from evalrx.eval_agent.agentic.board import EvidenceBoard + from evalrx.eval_agent.agentic.tools import _explore_data, _RunState + from evalrx.eval_agent.run_context import RunContext + + calls = [] + workdir = tmp_path / 'sandbox' + _seed_workdir(workdir) + explorer = _Explorer(calls, report=_report(workdir)) + data = _batch() + with RunContext(tmp_path / 'run') as ctx: + loop, diag = _loop(calls, explorer, run_logger=ctx.logger) + state = _RunState(original_data=data, data=data, timings={}) + state.probe_results = loop.probe_agent.probe(loop.model, data) + outcome = _explore_data(loop, state)(EvidenceBoard(), {'question': 'Inspect long chains'}) + assert outcome.ok + assert explorer.question == 'Inspect long chains' + assert explorer.outcome_col == 'label' + loop._do_m3(0, None, [], state.timings, cases=data) + assert diag.seen_context[-1] is not None + assert state.timings['explore'] > 0 + events = read_v2_events(tmp_path / 'run') + assert any(e['event'] == 'explore' for e in events) + assert any(e['event'] == 'diagnosis' for e in events) From d10c234363526928b4229eaae25913f2b533a6fb Mon Sep 17 00:00:00 2001 From: Jiaqi Liu <32996169+jiaq-liu@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:02:27 -0700 Subject: [PATCH 5/9] merge vl vlm-probing into evalrx --- CHANGELOG.md | 44 ++ docs/analyzers.md | 34 +- evalrx/analyzers/__init__.py | 10 +- evalrx/analyzers/_internals/__init__.py | 17 + evalrx/analyzers/_internals/access.py | 117 ++++ evalrx/analyzers/_internals/attn_tap.py | 331 +++++++++++ evalrx/analyzers/_internals/heads.py | 110 ++++ evalrx/analyzers/_internals/maps.py | 144 +++++ evalrx/analyzers/_internals/residual.py | 102 ++++ evalrx/analyzers/_internals/targets.py | 106 ++++ evalrx/analyzers/attention/sink.py | 71 ++- evalrx/analyzers/attention/summary.py | 66 ++- evalrx/analyzers/attribution/generic_attn.py | 214 ++++++- evalrx/analyzers/attribution/gradcam.py | 167 +++++- evalrx/analyzers/lens/__init__.py | 5 +- evalrx/analyzers/lens/head_attribution.py | 183 ++++++ evalrx/analyzers/lens/tuned_lens.py | 240 +++++++- evalrx/analyzers/patching/__init__.py | 6 +- .../analyzers/patching/attention_knockout.py | 181 ++++++ .../patching/attention_sensitivity.py | 167 ++++++ evalrx/analyzers/patching/causal_trace.py | 337 ++++++++++- evalrx/models/backends/hf/model.py | 16 +- .../test_analyzers/test_analyzer_contract.py | 46 +- tests/test_analyzers/test_analyzer_docs.py | 17 +- .../test_intervention_analyzers.py | 521 ++++++++++++++++++ tests/test_core/test_capability_matching.py | 2 +- tests/test_core/test_registry_and_shim.py | 2 +- tests/test_models/test_qwen.py | 3 +- tests/test_models/test_wrap.py | 2 +- tests/tiny_whitebox.py | 166 ++++++ 30 files changed, 3319 insertions(+), 108 deletions(-) create mode 100644 evalrx/analyzers/_internals/__init__.py create mode 100644 evalrx/analyzers/_internals/access.py create mode 100644 evalrx/analyzers/_internals/attn_tap.py create mode 100644 evalrx/analyzers/_internals/heads.py create mode 100644 evalrx/analyzers/_internals/maps.py create mode 100644 evalrx/analyzers/_internals/residual.py create mode 100644 evalrx/analyzers/_internals/targets.py create mode 100644 evalrx/analyzers/lens/head_attribution.py create mode 100644 evalrx/analyzers/patching/attention_knockout.py create mode 100644 evalrx/analyzers/patching/attention_sensitivity.py create mode 100644 tests/test_analyzers/test_intervention_analyzers.py create mode 100644 tests/tiny_whitebox.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d6ac888..7491e852 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,50 @@ Versions follow [Semantic Versioning](https://semver.org/). ## [Unreleased] +### Added — write-side white-box analyzers (knockout, causal trace, gradients, head attribution) + +Four registered stubs are now native implementations, plus three new +analyzers. All explain the next-token logit margin at the end of the prompt +(`target="prediction"`: clean top-1 vs top-2, label-free; or `"gold"`) and emit +`per_case` columns for M2: + +- `attention_knockout` (new, `patching/`) — cut text→media attention in all + layers and per third of the stack (Geva et al. 2023). +- `attention_sensitivity` (new, `patching/`) — rescale text→media attention + (×0.5 / ×2, renormalised) and report the margin's slope; negative = more + looking weakens the answer (PAI-style repair signal). +- `causal_trace` (was a stub) — Meng-style corruption of the media tokens + (noise, incl. Qwen3-VL DeepStack injections; or a same-size gray image) and + per-layer restoration at the media tokens and at the answer position. +- `gradcam` (was a stub) — Grad-CAM over the decoder's image tokens at a + fractional depth, mapped onto the patch grid with photo overlays. +- `generic_attention` (was a stub) — Chefer et al. 2021 relevance, propagated + as one row from the top layer down. +- `head_attribution` (new, `lens/`) — direct logit attribution per attention + head with the final norm frozen at the clean run. +- `tuned_lens` (was a stub) — per-layer affine translators fit label-free to the + model's own final distribution; cached per model, optional save/load. + +`attention` and `attention_sink` now run over the batch (`max_cases`) and emit +`per_case` columns (`media_mass_mid`, `sink_mass_mean`, …) instead of reading +only the first case. + +Plumbing lives in `evalrx/analyzers/_internals/`: an attention function +registered with transformers' `AttentionInterface` (plus the matching mask +function — without it transformers builds no causal mask) that reproduces eager +attention and can block / reweight / record it, verified once per model to +engage on every layer and reproduce the native logits; decoder-layer pre-hooks +for reading/replacing the residual stream; per-head output capture. +Attention-editing analyzers refuse hybrid (`HYBRID_SPARSE`) and latent attention. + +### Changed — `hf_local` provides `GRADIENTS` + +`HFLocalModel.forward(spec=CaptureSpec(with_grad=True))` now keeps the autograd +graph (it was always `no_grad`, and `with_grad` was ignored), and `hf_local` +declares `Capability.GRADIENTS`, so `gradcam` / `generic_attention` are offered +for local models. The default forward is still `no_grad`. Gradient analyzers use +`torch.autograd.grad` on hooked activations and never write parameter `.grad`. + ### Docs — analyzer zoo regenerated from the registries `docs/analyzers.md` listed 29 of the 46 registered analyzers and 16 of the 57 diff --git a/docs/analyzers.md b/docs/analyzers.md index 6aa064d3..e58ccdcf 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -22,13 +22,13 @@ model, not by a black-box / white-box label. | [Uncertainty](#uncertainty) | `uncertainty/` | 6 | 0 | | [Hallucination](#hallucination) | `hallucination/` | 5 | 2 | | [Attention](#attention) | `attention/` | 4 | 0 | -| [Attribution](#attribution) | `attribution/` | 2 | 2 | -| [Lens](#lens) | `lens/` | 3 | 1 | -| [Patching](#patching) | `patching/` | 1 | 1 | +| [Attribution](#attribution) | `attribution/` | 2 | 0 | +| [Lens](#lens) | `lens/` | 4 | 0 | +| [Patching](#patching) | `patching/` | 3 | 0 | | [Geometry](#geometry) | `geometry/` | 2 | 0 | | [Reasoning](#reasoning) | `reasoning/` | 7 | 0 | | [Agent](#agent) | `agent/` | 7 | 0 | -| **Total** | | **46** | **6** | +| **Total** | | **49** | **2** | How to read the columns: @@ -116,39 +116,51 @@ on the light (API-only) install it is skipped and these analyzers are not offere ### Attribution `evalrx/analyzers/attribution/`: gradient-based saliency (`GRADIENTS`, a backward -pass through a white-box model). +pass through a white-box model: `hf_local` keeps the autograd graph when asked with +`CaptureSpec(with_grad=True)`). Both explain the next-token logit margin and map the +result back onto the image's patch grid, with overlays on the case photo. Grad-CAM +reads the decoder's image-token states rather than the vision tower. | Analyzer | Key | Requires | Modality | Paper | Status | |---|---|---|---|---|---| -| Grad-CAM | `gradcam` | `GRADIENTS` | image | [Selvaraju et al. 2017](https://arxiv.org/abs/1610.02391) | stub | -| Generic attention explainability | `generic_attention` | `ATTENTION` + `GRADIENTS` | text + image | [Chefer et al. 2021](https://arxiv.org/abs/2103.15679) | stub | +| Grad-CAM | `gradcam` | `GRADIENTS` + `LOGITS` | **image** + **video** | [Selvaraju et al. 2017](https://arxiv.org/abs/1610.02391) | ✓ | +| Generic attention explainability | `generic_attention` | `ATTENTION` + `GRADIENTS` + `LOGITS` | text + image | [Chefer et al. 2021](https://arxiv.org/abs/2103.15679) | ✓ | ### Lens `evalrx/analyzers/lens/`: project intermediate hidden states through the unembedding (`HIDDEN_STATES`) to see at which layer a prediction forms, and where -premature layers disagree with the final one. +premature layers disagree with the final one; or project each attention head's own +write, to see which heads write the answer. | Analyzer | Key | Requires | Modality | Paper | Status | |---|---|---|---|---|---| | Logit lens | `logit_lens` | `HIDDEN_STATES` | text + image | [nostalgebraist 2020](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens) | ✓ | -| Tuned lens | `tuned_lens` | `HIDDEN_STATES` | text + image | [Belrose et al. 2023](https://arxiv.org/abs/2303.08112) | stub | +| Tuned lens | `tuned_lens` | `HIDDEN_STATES` | text + image | [Belrose et al. 2023](https://arxiv.org/abs/2303.08112) | ✓ | | Layer contrast | `layer_contrast` | `HIDDEN_STATES` | text + image | [Chuang et al. 2024](https://arxiv.org/abs/2309.03883); [Wang et al. 2025](https://arxiv.org/abs/2410.11779) | ✓ | +| Head logit attribution | `head_attribution` | `HIDDEN_STATES` + `LOGITS` | text + image | [Elhage et al. 2021](https://transformer-circuits.pub/2021/framework/index.html); [Wang et al. 2023](https://arxiv.org/abs/2211.00593) | ✓ | ### Patching `evalrx/analyzers/patching/`: causal interventions. Write a clean run's activations into a corrupted run to localise where information is used -(`HIDDEN_STATES`, read and write hooks). +(`HIDDEN_STATES`, read and write hooks), or cut attention edges and measure what +the answer loses. These act on the live torch module of a local `hf_local` model +through `evalrx/analyzers/_internals/` (an attention function registered with +transformers, and decoder-layer pre-hooks); attention edits refuse hybrid or +latent-attention stacks and verify, once per model, that the edit-free path +reproduces the native logits. | Analyzer | Key | Requires | Modality | Paper | Status | |---|---|---|---|---|---| -| Causal trace | `causal_trace` | `HIDDEN_STATES` | text + image | [Meng et al. 2022](https://arxiv.org/abs/2202.05262) | stub | +| Causal trace | `causal_trace` | `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Meng et al. 2022](https://arxiv.org/abs/2202.05262); [Basu et al. 2024](https://arxiv.org/abs/2406.04236) | ✓ | +| Attention knockout | `attention_knockout` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Geva et al. 2023](https://arxiv.org/abs/2304.14767); [Zhang et al. 2025](https://arxiv.org/abs/2411.18620) | ✓ | +| Attention sensitivity | `attention_sensitivity` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Liu et al. 2024 (PAI)](https://arxiv.org/abs/2407.21771) | ✓ | ### Geometry diff --git a/evalrx/analyzers/__init__.py b/evalrx/analyzers/__init__.py index 3fc68bb2..cde5b2d7 100755 --- a/evalrx/analyzers/__init__.py +++ b/evalrx/analyzers/__init__.py @@ -11,8 +11,10 @@ hallucination/ POPE / CHAIR / SelfCheck / OPERA / VCD — GENERATE/ATTENTION/LOGITS attention/ summary / rollout / sink / relative-attn — ATTENTION attribution/ Grad-CAM / Chefer generic-attention — GRADIENTS - lens/ logit-lens / tuned-lens / layer contrast — HIDDEN_STATES - patching/ causal tracing / activation patching — HIDDEN_STATES (read+write) + lens/ logit-lens / tuned-lens / layer contrast / + head logit attribution — HIDDEN_STATES + patching/ causal tracing / attention knockout / + attention sensitivity — HIDDEN_STATES / ATTENTION (write) geometry/ CKA / linear-probe — HIDDEN_STATES reasoning/ extraction/termination hygiene; arithmetic, self-repair, step-value, knowledge split, @@ -22,6 +24,10 @@ Import is torch-tolerant: the ``attention`` subpackage loads torch at import, so on the light (pure-API) install it is skipped and its analyzers are simply not offered. + +``_internals/`` is the private write-side plumbing (attention tap, residual hooks, +logit-margin targets) that the patching / attribution / head-attribution analyzers +install around an ordinary ``model.forward`` call; it imports torch lazily. """ from evalrx.analyzers import ( # noqa: F401 -- import to self-register diff --git a/evalrx/analyzers/_internals/__init__.py b/evalrx/analyzers/_internals/__init__.py new file mode 100644 index 00000000..8162a894 --- /dev/null +++ b/evalrx/analyzers/_internals/__init__.py @@ -0,0 +1,17 @@ +"""Write-side white-box plumbing shared by the intervention analyzers. + +The read-only analyzers consume a :class:`~evalrx.core.model.Trace`. The ones +here (attention knockout, causal tracing, gradient attribution) have to CHANGE +or DIFFERENTIATE the forward pass, which a ``Trace`` cannot express. They do it +by installing hooks on the live torch module for the duration of one ordinary +``model.forward`` call, so encoding, chat templates and the image-token map stay +exactly what every other analyzer sees. + + access — the torch module behind a handle, the decoder stack, media positions + targets — the logit margin every intervention is scored on + attn_tap — a registered attention function that can edit/record attention + residual — read/replace the residual stream entering a decoder layer + +Private package: torch is imported inside functions only, so importing an +analyzer that uses it still works on the light (API-only) install. +""" diff --git a/evalrx/analyzers/_internals/access.py b/evalrx/analyzers/_internals/access.py new file mode 100644 index 00000000..c7fdcf00 --- /dev/null +++ b/evalrx/analyzers/_internals/access.py @@ -0,0 +1,117 @@ +"""Reach the live module behind a white-box handle, and the positions to act on. + +Only the ``hf_local`` backend holds a torch module (``model._loaded`` lazily +loads and returns ``(module, processor)``). Everything that needs one goes +through :func:`torch_module`, so an API or fake model fails with one clear +message instead of an ``AttributeError`` three calls deep. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from evalrx.core.model import Model, Trace + +#: Architectures whose decoder attention is a dense token×token softmax that a +#: registered attention function sees in full. Hybrid stacks route most layers +#: through linear/SSM attention with no query-key matrix; latent / unverified +#: attention is not a plain token matrix. Edits there would be partial or +#: meaningless, so the write-side analyzers refuse them. +_EDITABLE_SEMANTICS = frozenset({"standard"}) + + +def torch_module(model: "Model", analyzer: str = "this analyzer") -> Any: + """The ``torch.nn.Module`` behind *model*; ``ValueError`` when there is none.""" + import torch.nn as nn + + loaded = getattr(model, "_loaded", None) # HFLocalModel: lazy-loads on access + module = loaded[0] if isinstance(loaded, tuple) else loaded + if not isinstance(module, nn.Module): + raise ValueError( + f"{analyzer} edits or differentiates the forward pass, so it needs a local " + f"white-box model (hf_local); {type(model).__name__} exposes no torch module." + ) + return module + + +def tokenizer_of(model: "Model") -> Any: + """The tokenizer behind a local handle, or ``None``.""" + loaded = getattr(model, "_loaded", None) + if not isinstance(loaded, tuple) or len(loaded) < 2 or loaded[1] is None: + return None + processor = loaded[1] + return getattr(processor, "tokenizer", processor) + + +def check_editable_attention(model: "Model", analyzer: str) -> None: + """Refuse architectures whose attention cannot be edited as a token matrix.""" + spec = getattr(model, "spec", None) + semantics = getattr(getattr(spec, "attn_semantics", None), "value", None) + if semantics is not None and semantics not in _EDITABLE_SEMANTICS: + raise ValueError( + f"{analyzer} needs standard dense decoder attention; {getattr(spec, 'key', '?')} " + f"declares attn_semantics={semantics!r}, where an attention edit would be " + "partial (hybrid/linear layers) or not a token matrix (latent attention)." + ) + + +@dataclass +class DecoderStack: + """The language decoder's layers, their attention modules, and their owner. + + ``attn[i]`` is ``None`` for a layer with no softmax self-attention (the + linear-attention layers of a hybrid stack). + """ + + layers: Any # nn.ModuleList + path: str # dotted path of ``layers`` inside the module + attn: list # per layer: attention submodule or None + parent: Any # module owning ``layers`` (the text model) + + @property + def n_layers(self) -> int: + return len(self.layers) + + @property + def fully_attentional(self) -> bool: + return all(a is not None for a in self.attn) + + +def decoder_stack(module: Any) -> DecoderStack: + """Discover the decoder stack at runtime (never a hard-coded module path).""" + from evalrx.models.backends.hf.discover import find_decoder_layers, resolve + + layers, path = find_decoder_layers(module) + parent = module if "." not in path else resolve(module, path.rsplit(".", 1)[0]) + attn = [getattr(layer, "self_attn", None) for layer in layers] + return DecoderStack(layers=layers, path=path, attn=attn, parent=parent) + + +def layer_points(n_layers: int, n_points: int) -> list[int]: + """``n_points`` layer indices spread evenly over ``[0, n_layers)`` (all when fewer).""" + if n_points <= 0 or n_points >= n_layers: + return list(range(n_layers)) + step = (n_layers - 1) / (n_points - 1) if n_points > 1 else 0 + return sorted({int(round(i * step)) for i in range(n_points)}) + + +def media_positions(trace: "Trace") -> "tuple[str, list[int]]": + """``(slot, positions)`` of the media placeholder tokens in *trace*. + + Image/video tokens come from the TokenTypeMap (falling back to the + ``image_token_mask`` extra); audio from ``audio_token_mask``. Returns + ``("", [])`` when the sequence holds no media tokens. + """ + ttm = getattr(trace, "token_type_map", None) + if ttm is not None and getattr(ttm, "image_pos", None): + return "image", list(ttm.image_pos) + extras = getattr(trace, "extras", None) or {} + for slot, key in (("image", "image_token_mask"), ("audio", "audio_token_mask")): + mask = extras.get(key) + if mask is not None: + positions = [int(i) for i in mask.nonzero().flatten().tolist()] + if positions: + return slot, positions + return "", [] diff --git a/evalrx/analyzers/_internals/attn_tap.py b/evalrx/analyzers/_internals/attn_tap.py new file mode 100644 index 00000000..0c9f8299 --- /dev/null +++ b/evalrx/analyzers/_internals/attn_tap.py @@ -0,0 +1,331 @@ +"""An attention function we own: reproduces eager attention, and can edit or record it. + +Mechanism (the transformers attention interface): each supported decoder +attention module computes q/k/v itself and hands the softmax part to +``ALL_ATTENTION_FUNCTIONS[config._attn_implementation]``. Registering +:func:`_tap_attention` under :data:`TAP_NAME` and pointing ONLY the decoder's +configs at it inside a ``with`` block routes the language model's attention +through code we control; the vision tower keeps its own config and +implementation. + +Two things this module exists to get right: + +* **The mask.** ``masking_utils`` builds NO mask for an implementation name it + does not know (it assumes a vLLM/TGI-style kernel that masks internally), so a + bare registration silently lets every token attend to the future. The tap + registers eager's additive mask under the same name. +* **Fidelity.** "The tap is installed" proves nothing on an architecture that + bypasses the interface. :func:`check_fidelity` runs one input with an + edit-free tap: it must be called on every attention layer AND reproduce the + native logits, or :func:`require_fidelity` refuses. + +Edits act on the pre-softmax scores, so they compose with the causal mask: + +* ``block`` — attention knockout: the query cannot read the key at all + (Geva et al., 2023). A query is never blocked from itself, so no row empties. +* ``scale`` — multiply the attention to the keys by ``w`` and renormalise, + implemented exactly as ``+log w`` on the scores (``w = 0`` is ``block``). +* ``scale_raw`` — multiply post-softmax weights by ``w`` without renormalising + (the PAI-style amplification; rows no longer sum to one). +""" + +from __future__ import annotations + +import math +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, Optional + +from evalrx.analyzers._internals.access import decoder_stack, torch_module + +if TYPE_CHECKING: + from evalrx.core.model import Model + +TAP_NAME = "evalrx_tap" +EDIT_KINDS = ("block", "scale", "scale_raw") + +_ACTIVE: "list[TapState]" = [] +_REGISTERED = False + + +@dataclass(frozen=True) +class AttentionEdit: + """One edit applied to some decoder layers' attention. + + ``keys`` / ``queries`` are absolute sequence positions (``queries=None`` = + every query row); ``layers=None`` = every decoder layer. + """ + + kind: str + keys: "tuple[int, ...]" + queries: "tuple[int, ...] | None" = None + layers: "frozenset[int] | None" = None + weight: float = 1.0 + + def __post_init__(self) -> None: + if self.kind not in EDIT_KINDS: + raise ValueError(f"edit kind must be one of {EDIT_KINDS}, got {self.kind!r}") + if self.weight < 0: + raise ValueError(f"attention weight must be >= 0, got {self.weight}") + + def applies(self, layer: int) -> bool: + return self.layers is None or layer in self.layers + + +def knockout(keys: Iterable[int], *, queries: "Iterable[int] | None" = None, + layers: "Iterable[int] | None" = None) -> AttentionEdit: + """Block *queries* from attending to *keys* in *layers*.""" + return AttentionEdit( + "block", tuple(int(k) for k in keys), + None if queries is None else tuple(int(q) for q in queries), + None if layers is None else frozenset(int(i) for i in layers), + ) + + +def reweight(keys: Iterable[int], weight: float, *, queries: "Iterable[int] | None" = None, + layers: "Iterable[int] | None" = None, renormalize: bool = True) -> AttentionEdit: + """Scale attention from *queries* to *keys* by *weight*.""" + return AttentionEdit( + "scale" if renormalize else "scale_raw", tuple(int(k) for k in keys), + None if queries is None else tuple(int(q) for q in queries), + None if layers is None else frozenset(int(i) for i in layers), + float(weight), + ) + + +class TapState: + """What the tap does while one ``attention_tap`` block is active.""" + + def __init__(self, edits: Iterable[AttentionEdit], layer_of: "dict[int, int]", + on_probs: "Callable[[int, Any], None] | None") -> None: + self.edits = list(edits) + self.layer_of = layer_of + self.on_probs = on_probs + self.calls: list[int] = [] # decoder-layer index of every tapped call + self._masks: dict = {} + + def _mask(self, index: int, edit: AttentionEdit, q_len: int, k_len: int, device: Any) -> Any: + """(q_len, k_len) bool: which score cells *edit* touches.""" + import torch + + cache_key = (index, q_len, k_len, str(device)) + if cache_key in self._masks: + return self._masks[cache_key] + # Query rows are the LAST q_len positions (q_len < k_len only with a KV cache). + q_abs = torch.arange(k_len - q_len, k_len, device=device) + k_sel = torch.zeros(k_len, dtype=torch.bool, device=device) + keys = [k for k in edit.keys if 0 <= k < k_len] + if keys: + k_sel[torch.tensor(keys, device=device)] = True + if edit.queries is None: + q_sel = torch.ones(q_len, dtype=torch.bool, device=device) + else: + q_full = torch.zeros(k_len, dtype=torch.bool, device=device) + queries = [q for q in edit.queries if 0 <= q < k_len] + if queries: + q_full[torch.tensor(queries, device=device)] = True + q_sel = q_full[q_abs] + mask = q_sel[:, None] & k_sel[None, :] + if edit.kind == "block" or edit.weight == 0: + mask &= q_abs[:, None] != torch.arange(k_len, device=device)[None, :] + self._masks[cache_key] = mask + return mask + + def edit_scores(self, layer: int, scores: Any) -> Any: + import torch + + q_len, k_len = scores.shape[-2:] + for index, edit in enumerate(self.edits): + if edit.kind == "scale_raw" or not edit.applies(layer): + continue + mask = self._mask(index, edit, q_len, k_len, scores.device) + if edit.kind == "block" or edit.weight == 0: + scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min) + elif edit.weight != 1.0: + scores = scores + mask.to(scores.dtype) * math.log(edit.weight) + return scores + + def edit_probs(self, layer: int, probs: Any) -> Any: + import torch + + q_len, k_len = probs.shape[-2:] + for index, edit in enumerate(self.edits): + if edit.kind != "scale_raw" or not edit.applies(layer) or edit.weight == 1.0: + continue + mask = self._mask(index, edit, q_len, k_len, probs.device) + factor = torch.where(mask, edit.weight, 1.0).to(probs.dtype) + probs = probs * factor + return probs + + +def _tap_attention(module: Any, query: Any, key: Any, value: Any, attention_mask: Any, + scaling: "float | None" = None, dropout: float = 0.0, + softcap: "float | None" = None, **kwargs: Any) -> "tuple[Any, Any]": + """Eager attention, op for op, with edit/record points around the softmax. + + Matches the per-family ``eager_attention_forward`` it replaces: GQA head + repetition, scaling (``head_dim ** -0.5`` when the caller passes none, as + Gemma/Qwen2-Audio default), optional logit soft-capping (Gemma), additive + mask, fp32 softmax, dropout. + """ + import torch + from torch import nn + + if kwargs.get("s_aux") is not None or kwargs.get("sinks") is not None: + raise NotImplementedError("attention sinks (s_aux) are not reproduced by the evalrx tap") + state = _ACTIVE[-1] if _ACTIVE else None + if scaling is None: + scaling = query.shape[-1] ** -0.5 + n_rep = query.shape[1] // key.shape[1] + if n_rep > 1: # == transformers' repeat_kv + key = key.repeat_interleave(n_rep, dim=1) + value = value.repeat_interleave(n_rep, dim=1) + + scores = torch.matmul(query, key.transpose(2, 3)) * scaling + if softcap is not None: + scores = scores / softcap + scores = torch.tanh(scores) + scores = scores * softcap + if attention_mask is not None: + mask = attention_mask[:, :, :, : key.shape[-2]] + if mask.dtype == torch.bool: + scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min) + else: + scores = scores + mask + + layer = state.layer_of.get(id(module)) if state is not None else None + if layer is not None: + state.calls.append(layer) + scores = state.edit_scores(layer, scores) + probs = nn.functional.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype) + if layer is not None: + probs = state.edit_probs(layer, probs) + if state.on_probs is not None: + state.on_probs(layer, probs) + probs = nn.functional.dropout(probs, p=dropout, training=module.training) + out = torch.matmul(probs, value).transpose(1, 2).contiguous() + return out, probs + + +def _ensure_registered() -> None: + global _REGISTERED + if _REGISTERED: + return + from transformers import AttentionInterface + from transformers.masking_utils import AttentionMaskInterface, eager_mask + + AttentionInterface.register(TAP_NAME, _tap_attention) + AttentionMaskInterface.register(TAP_NAME, eager_mask) # else: NO causal mask at all + _REGISTERED = True + + +def _get_impl(config: Any) -> Any: + return getattr(config, "_attn_implementation", None) + + +def _set_impl(config: Any, value: Any) -> None: + # Write the backing field directly: the public setter recurses into + # sub-configs, which on a composite config would retarget the vision tower. + if "_attn_implementation_internal" in getattr(config, "__dict__", {}): + config._attn_implementation_internal = value + else: + config._attn_implementation = value + + +@contextmanager +def attention_tap( + model: "Model", + edits: Iterable[AttentionEdit] = (), + *, + on_probs: "Callable[[int, Any], None] | None" = None, + analyzer: str = "attention tap", +) -> Iterator[TapState]: + """Route the decoder's attention through the tap for the duration of the block. + + ``on_probs(layer, probs)`` sees each layer's (edited) attention weights, + shape ``(batch, heads, q, k)``, still attached to the autograd graph when + the forward runs with gradients. + """ + module = torch_module(model, analyzer) + stack = decoder_stack(module) + _ensure_registered() + + configs: list = [] + seen: set[int] = set() + for owner in [*(a for a in stack.attn if a is not None), stack.parent]: + cfg = getattr(owner, "config", None) + if cfg is not None and id(cfg) not in seen: + seen.add(id(cfg)) + configs.append(cfg) + if not configs: + raise ValueError(f"{analyzer}: no config found on the decoder attention modules") + + layer_of = {id(a): i for i, a in enumerate(stack.attn) if a is not None} + state = TapState(edits, layer_of, on_probs) + saved = [(cfg, _get_impl(cfg)) for cfg in configs] + _ACTIVE.append(state) + try: + for cfg, _ in saved: + _set_impl(cfg, TAP_NAME) + yield state + finally: + for cfg, impl in saved: + _set_impl(cfg, impl) + _ACTIVE.remove(state) + + +def check_fidelity(model: "Model", inputs: Any, *, rtol: float = 1e-2, + analyzer: str = "attention tap") -> "dict[str, Any]": + """Does an edit-free tap engage on every attention layer and reproduce the logits? + + Cached on the module per native attention implementation: one extra pair + of forwards per model, not per analyzer run. + """ + from evalrx.core.capability import Capability + + module = torch_module(model, analyzer) + stack = decoder_stack(module) + native = tuple(str(_get_impl(getattr(a, "config", None))) for a in stack.attn if a is not None) + cache = module.__dict__.setdefault("_evalrx_tap_fidelity", {}) + if native in cache: + return dict(cache[native]) + + n_attn = sum(a is not None for a in stack.attn) + base = model.forward(inputs, capture={Capability.LOGITS}).require(Capability.LOGITS) + with attention_tap(model, analyzer=analyzer) as state: + tapped = model.forward(inputs, capture={Capability.LOGITS}).require(Capability.LOGITS) + b, t = base[-1].float(), tapped[-1].float() + diff = float((b - t).abs().max()) + scale = float(b.abs().max()) or 1.0 + engaged = len(set(state.calls)) + verdict = { + "tap_faithful": bool( + engaged == n_attn and diff <= rtol * scale and int(b.argmax()) == int(t.argmax()) + ), + "tap_layers_engaged": engaged, + "tap_layers_expected": n_attn, + "tap_max_logit_diff": round(diff, 6), + "tap_rtol": rtol, + "native_attn_implementation": native[0] if native else None, + } + cache[native] = verdict + return dict(verdict) + + +def require_fidelity(model: "Model", inputs: Any, *, rtol: float = 1e-2, + analyzer: str = "attention tap") -> "dict[str, Any]": + """:func:`check_fidelity`, raising when the tap cannot be trusted on *model*.""" + verdict = check_fidelity(model, inputs, rtol=rtol, analyzer=analyzer) + if not verdict["tap_faithful"]: + raise RuntimeError( + f"{analyzer}: the evalrx attention tap does not faithfully reproduce this " + f"model's attention ({verdict}); refusing to report edited-attention effects. " + "Engaged < expected layers means the architecture bypasses the transformers " + "attention interface; a large logit diff means its attention is not eager-equivalent." + ) + return verdict + + +def active_state() -> Optional[TapState]: + """The innermost active tap state (``None`` outside any ``attention_tap``).""" + return _ACTIVE[-1] if _ACTIVE else None diff --git a/evalrx/analyzers/_internals/heads.py b/evalrx/analyzers/_internals/heads.py new file mode 100644 index 00000000..6d335a5d --- /dev/null +++ b/evalrx/analyzers/_internals/heads.py @@ -0,0 +1,110 @@ +"""Per-head attention outputs at one position, and the final norm frozen to a linear map. + +Every supported decoder attention module ends in an output projection whose +INPUT is the heads' outputs concatenated, ``(n_heads · head_dim)``. A forward +pre-hook on that projection reads them; head ``h``'s write into the residual +stream is ``W_O[:, h·hd:(h+1)·hd] @ x[h·hd:(h+1)·hd]``, and the per-head writes +sum to the projection's output (minus its bias, if any). + +Direct logit attribution then reads each write through the final norm and the +unembedding. The norm is not linear, so — as is standard — its scale is FROZEN +at the clean final residual: RMSNorm becomes ``x ↦ s ⊙ x`` with +``s = w_eff / rms(x_final)``. ``w_eff`` is read off the module itself +(``norm(1)``), which covers Llama/Qwen ``w`` and Gemma ``1 + w`` alike. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Iterator + +from evalrx.analyzers._internals.access import decoder_stack, torch_module + +if TYPE_CHECKING: + from evalrx.core.model import Model + +_OUT_PROJ_NAMES = ("o_proj", "out_proj", "dense") + + +def out_proj(attn: Any) -> Any: + """The attention module's output projection (``None`` when not found).""" + for name in _OUT_PROJ_NAMES: + proj = getattr(attn, name, None) + if proj is not None and hasattr(proj, "weight"): + return proj + return None + + +def n_heads_of(attn: Any, proj: Any) -> int: + """Query-head count: the projection's input width over the head dim.""" + head_dim = getattr(attn, "head_dim", None) + if head_dim is None: + cfg = getattr(attn, "config", None) + head_dim = getattr(cfg, "head_dim", None) or ( + cfg.hidden_size // cfg.num_attention_heads if cfg is not None else None) + if not head_dim or proj.weight.shape[1] % int(head_dim): + raise ValueError(f"cannot split {type(attn).__name__}'s output projection into heads") + return proj.weight.shape[1] // int(head_dim) + + +@dataclass +class HeadSession: + """``head_inputs[layer]`` — ``(n_heads · head_dim,)`` at the position; + ``final_residual`` — the final norm's input there.""" + + head_inputs: "dict[int, Any]" = field(default_factory=dict) + final_residual: Any = None + + +@contextmanager +def head_capture(model: "Model", position: int = -1, *, analyzer: str = "head capture") -> Iterator[HeadSession]: + """Record every layer's per-head outputs and the final residual at *position*.""" + from evalrx.models.backends.hf.discover import get_final_norm + + module = torch_module(model, analyzer) + stack = decoder_stack(module) + norm = get_final_norm(module) + if norm is None: + raise ValueError(f"{analyzer}: no final norm found after the decoder layers") + session = HeadSession() + handles = [] + + def grab_heads(layer: int): + def hook(mod, args): + session.head_inputs[layer] = args[0][0, position].detach() + return hook + + def grab_final(mod, args): + session.final_residual = args[0][0, position].detach() + + for layer, attn in enumerate(stack.attn): + proj = out_proj(attn) if attn is not None else None + if proj is None: + raise ValueError(f"{analyzer}: layer {layer} has no attention output projection") + handles.append(proj.register_forward_pre_hook(grab_heads(layer))) + handles.append(norm.register_forward_pre_hook(grab_final)) + try: + yield session + finally: + for handle in handles: + handle.remove() + + +def frozen_norm_scale(norm: Any, x: Any) -> "tuple[Any, bool]": + """``(scale, centre)``: ``norm(x + c) ≈ norm(x) + scale ⊙ (c - centre·mean(c))``. + + Exact (scale frozen at ``x``) for RMSNorm; LayerNorm additionally centres. + """ + import torch + + xf = x.float() + eps = float(getattr(norm, "variance_epsilon", getattr(norm, "eps", 1e-6)) or 1e-6) + is_layernorm = isinstance(norm, torch.nn.LayerNorm) or ( + "layernorm" in type(norm).__name__.lower() and "rms" not in type(norm).__name__.lower()) + with torch.no_grad(): + if is_layernorm: + weight = norm.weight.float() if getattr(norm, "weight", None) is not None else torch.ones_like(xf) + return weight / torch.sqrt(xf.var(unbiased=False) + eps), True + w_eff = norm(torch.ones_like(x)).float() * (1.0 + eps) ** 0.5 # undo rms(1) = sqrt(1 + eps) + return w_eff / torch.sqrt(xf.pow(2).mean() + eps), False diff --git a/evalrx/analyzers/_internals/maps.py b/evalrx/analyzers/_internals/maps.py new file mode 100644 index 00000000..1bd4bee1 --- /dev/null +++ b/evalrx/analyzers/_internals/maps.py @@ -0,0 +1,144 @@ +"""Per-image-token scores → spatial maps, shape features, and an overlayable Result. + +Grad-CAM and attention relevance both end in one score per image token. This +module turns those into what ``relative_attention`` already produces, so the +run logger, the tensor-level omnibus and the multimodal judge consume them the +same way: + +* ``artifacts["spatial_map"]`` (first case), ``per_case_maps`` (float16 per + case), ``fail_mean_map`` / ``pass_mean_map`` / ``diff_map_fail_minus_pass``; +* ``findings["_case_id"]`` naming the case whose image each map overlays; +* :class:`SpatialMapResult.image_overlays` — the duck-typed hook RunLoggerV2 + calls to save heatmap-on-photo PNGs. + +numpy and the overlay helpers are imported inside functions: this module is +imported by analyzers that load on the light install. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Trace + +OVERLAYABLE_KEYS = ("spatial_map", "fail_mean_map", "pass_mean_map", "diff_map_fail_minus_pass") + + +def image_token_scores(trace: "Trace", values: Any) -> "tuple[Any, Any]": + """``(flat, spatial)`` for the FIRST image of *trace*. + + *values* is a 1-D array aligned with ``trace.token_type_map.image_pos``. + ``spatial`` is the post-merge ``(h, w)`` grid (frames of a video summed), + or the flat vector when the grid cannot be rebuilt. + """ + import numpy as np + + flat = np.asarray(values, dtype=np.float64).ravel() + ttm = getattr(trace, "token_type_map", None) + if ttm is None or not ttm.grids or len(ttm.patches) != flat.size: + shape = (getattr(trace, "extras", None) or {}).get("image_spatial_shape") + if shape is not None and int(shape[0]) * int(shape[1]) == flat.size: + return flat, flat.reshape(int(shape[0]), int(shape[1])) + return flat, flat + _, h, w = ttm.grids[0] + grid = np.zeros((int(h), int(w)), dtype=np.float64) + for value, (img, _t, row, col) in zip(flat, ttm.patches): + if img == 0: + grid[row, col] += value + return flat, grid + + +def map_features(flat: Any, spatial: Any, prefix: str) -> "dict[str, float]": + """Shape of a non-negative per-token map, as ``{prefix}_{feature}`` scalars. + + Same features ``relative_attention`` reports, so M2 can compare readings: + normalized entropy (1 = diffuse), top-1 share, centre-of-mass offset from + the image centre, and mass on the border ring. + """ + import numpy as np + + from evalrx.analyzers.attention.relative_attn import _attention_features + + w = np.clip(np.asarray(flat, dtype=np.float64), 0.0, None) + if not w.size or float(w.sum()) <= 0.0: + return {} + spatial = np.clip(np.asarray(spatial, dtype=np.float64), 0.0, None) + feats = _attention_features(w, spatial) + k = min(5, w.size) + feats["top5_share"] = round(float(np.sort(w)[::-1][:k].sum() / w.sum()), 4) + return {f"{prefix}_{name.replace('attention_', '')}": value for name, value in feats.items()} + + +def group_maps(per_case_maps: "dict[str, Any]", cases: "CaseBatch") -> "tuple[dict, dict]": + """``(artifacts, findings)`` for the FAIL/PASS mean maps and their difference.""" + from evalrx.analyzers.attention.relative_attn import _group_mean_map + + label_of = {c.id: getattr(getattr(c, "label", None), "value", None) for c in cases} + fail = {cid: m for cid, m in per_case_maps.items() if label_of.get(cid) == "fail" and m.ndim == 2} + passed = {cid: m for cid, m in per_case_maps.items() if label_of.get(cid) == "pass" and m.ndim == 2} + artifacts: dict = {} + findings: dict = {} + fail_mean = _group_mean_map([m.astype("float64") for m in fail.values()]) + pass_mean = _group_mean_map([m.astype("float64") for m in passed.values()]) + if fail_mean is not None: + artifacts["fail_mean_map"] = fail_mean + findings["fail_mean_map_case_id"] = next(iter(fail)) + if pass_mean is not None: + artifacts["pass_mean_map"] = pass_mean + findings["pass_mean_map_case_id"] = next(iter(passed)) + if fail_mean is not None and pass_mean is not None and fail_mean.shape == pass_mean.shape: + artifacts["diff_map_fail_minus_pass"] = fail_mean - pass_mean + findings["diff_map_fail_minus_pass_case_id"] = next(iter(fail)) + return artifacts, findings + + +@dataclass +class SpatialMapResult(Result): + """A Result carrying image-token maps that can be blended onto the case photo.""" + + def _representative_image(self, key: str) -> Any: + from evalrx.analyzers.attention.relative_attn import _resolve_image + + case_id = (self.findings or {}).get(f"{key}_case_id") + if case_id is None or self.cases is None: + return None + for c in self.cases: + if c.id == case_id: + return _resolve_image(getattr(c.inputs, "image", None)) + return None + + def overlay(self, key: str = "spatial_map", *, alpha: float = 0.6, cmap: str = "jet") -> Any: + """``artifacts[key]`` alpha-blended on its representative image, or ``None``.""" + from evalrx.analyzers.attention.relative_attn import _blend_heatmap_on_image + + mat = self.artifacts.get(key) + image = self._representative_image(key) + if mat is None or getattr(mat, "ndim", 0) != 2 or not hasattr(image, "convert"): + return None + try: + return _blend_heatmap_on_image(image, mat, alpha=alpha, cmap=cmap, diverging="diff" in key) + except ImportError: + return None + + def image_overlays(self, fig_dir: Any, stem_prefix: str) -> "list[Any]": + """Save every available overlay under ``fig_dir`` (RunLoggerV2 hook).""" + from pathlib import Path + + out = [] + for key in OVERLAYABLE_KEYS: + if key not in self.artifacts: + continue + path = Path(fig_dir) / f"{stem_prefix}_{key}_overlay.png" + try: + img = self.overlay(key) + if img is not None: + img.save(path) + out.append(path) + except Exception: # noqa: BLE001 - best-effort visualisation + pass + return out diff --git a/evalrx/analyzers/_internals/residual.py b/evalrx/analyzers/_internals/residual.py new file mode 100644 index 00000000..f57bb845 --- /dev/null +++ b/evalrx/analyzers/_internals/residual.py @@ -0,0 +1,102 @@ +"""Read or overwrite the residual stream ENTERING a decoder layer. + +The patch site is a forward pre-hook on decoder layer ``l``: the hidden state +the layer is about to read. Hooking layer INPUTS rather than outputs is +deliberate: Qwen3-VL ("DeepStack") adds visual features to the hidden state +between the first few layers, outside any layer. A value written into a layer's +OUTPUT would then get the receiving run's own visual features added on top; a +value written into the next layer's INPUT is exactly what that layer reads. + +Site ``0`` is the embedding output after the media features were merged in. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Iterator, Mapping, Sequence + +from evalrx.analyzers._internals.access import decoder_stack, torch_module + +if TYPE_CHECKING: + from evalrx.core.model import Model + + +@dataclass +class ResidualSession: + """``captured[layer]`` — ``(len(positions), dim)``; the full ``(1, seq, dim)`` + hidden state, graph attached, when the hooks ran with ``keep_grad``.""" + + captured: "dict[int, Any]" = field(default_factory=dict) + + +@contextmanager +def residual_hooks( + model: "Model", + *, + capture: "Mapping[int, Sequence[int] | None] | None" = None, + replace: "Mapping[int, tuple[Sequence[int], Any]] | None" = None, + keep_grad: bool = False, + analyzer: str = "residual hooks", +) -> Iterator[ResidualSession]: + """Install pre-hooks on decoder layers for the duration of the block. + + Args: + capture: ``{layer: positions}`` — record ``hidden[0, positions]`` + (``None`` = every position), detached and cloned. + replace: ``{layer: (positions, values)}`` — overwrite + ``hidden[0, positions]`` with ``values`` (``(n, dim)``). + Replacement happens before capture at the same layer. + keep_grad: record the whole hidden tensor un-detached instead, so a + caller can ``torch.autograd.grad`` a target with respect to it. + """ + module = torch_module(model, analyzer) + stack = decoder_stack(module) + capture = dict(capture or {}) + replace = dict(replace or {}) + for idx in (*capture, *replace): + if not 0 <= idx < stack.n_layers: + raise ValueError(f"{analyzer}: layer {idx} outside [0, {stack.n_layers})") + session = ResidualSession() + + def make_hook(idx: int): + def hook(mod, args, kwargs): + if args: + hidden = args[0] + elif "hidden_states" in kwargs: + hidden = kwargs["hidden_states"] + else: # pragma: no cover - no decoder layer we support calls it this way + raise RuntimeError(f"{analyzer}: cannot find the hidden state passed to layer {idx}") + if hidden.dim() != 3 or hidden.shape[0] != 1: + raise RuntimeError(f"{analyzer}: expected a (1, seq, dim) hidden state, got " + f"{tuple(hidden.shape)}") + changed = False + if idx in replace: + positions, values = replace[idx] + hidden = hidden.clone() + hidden[0, list(positions)] = values.to(device=hidden.device, dtype=hidden.dtype) + changed = True + if idx in capture: + if keep_grad: + session.captured[idx] = hidden + else: + positions = capture[idx] + sel = hidden[0] if positions is None else hidden[0, list(positions)] + session.captured[idx] = sel.detach().clone() + if not changed: + return None + if args: + return (hidden, *args[1:]), kwargs + return args, {**kwargs, "hidden_states": hidden} + + return hook + + handles = [ + stack.layers[idx].register_forward_pre_hook(make_hook(idx), with_kwargs=True) + for idx in sorted(set(capture) | set(replace)) + ] + try: + yield session + finally: + for handle in handles: + handle.remove() diff --git a/evalrx/analyzers/_internals/targets.py b/evalrx/analyzers/_internals/targets.py new file mode 100644 index 00000000..e993f21a --- /dev/null +++ b/evalrx/analyzers/_internals/targets.py @@ -0,0 +1,106 @@ +"""The scalar every intervention is scored on: a next-token logit margin. + +``logit[positive] - logit[negative]`` at one position equals the log-prob +difference (the softmax normaliser cancels), so it is comparable across edits +and differentiable for the gradient methods. + +Two targets: + +* ``"prediction"`` (default) — the clean run's own top-1 vs top-2 token. Label + free and defined for every case, so the per-case columns are never computed + differently for FAIL and PASS cases. It answers "what drove the answer the + model actually gave". +* ``"gold"`` — the gold answer's first token vs the strongest other token of the + clean run. Needs ``case.expected`` and a tokenizer; a case without either is + skipped (counted, never guessed). It answers "what supported the right answer". +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional + +if TYPE_CHECKING: + from evalrx.core.case import FailureCase + from evalrx.core.model import Model + +TARGETS = ("prediction", "gold") + + +@dataclass(frozen=True) +class LogitMargin: + """``logits[position, positive] - logits[position, negative]``.""" + + positive: int + negative: int + position: int = -1 + kind: str = "prediction" + + def __call__(self, logits: Any) -> Any: + row = logits[self.position] + return row[self.positive].float() - row[self.negative].float() + + def value(self, logits: Any) -> float: + return float(self(logits).detach()) + + def top1_changed(self, logits: Any, clean_top1: int) -> bool: + return int(logits[self.position].argmax()) != int(clean_top1) + + +def gold_text(expected: Any) -> Optional[str]: + """The gold answer as a string: a str, the first alias, or a lone ``all_of`` entry.""" + if isinstance(expected, dict): + values = expected.get("all_of") or [] + expected = values[0] if values else None + if isinstance(expected, (list, tuple)): + expected = expected[0] if expected else None + if expected is None: + return None + text = str(expected).strip() + return text or None + + +def first_token_id(tokenizer: Any, text: str) -> Optional[int]: + """Id of the first token of *text* as an answer would start (no special tokens).""" + if tokenizer is None or not text: + return None + try: + ids = tokenizer(text, add_special_tokens=False)["input_ids"] + except Exception: # noqa: BLE001 - a tokenizer that cannot encode → no gold target + return None + if ids and isinstance(ids[0], (list, tuple)): + ids = ids[0] + return int(ids[0]) if ids else None + + +def build_margin( + target: str, + clean_logits: Any, + *, + model: "Model | None" = None, + case: "FailureCase | None" = None, + position: int = -1, +) -> "tuple[LogitMargin | None, str]": + """``(margin, "")`` or ``(None, skip_reason)``.""" + import torch + + if target not in TARGETS: + raise ValueError(f"target must be one of {TARGETS}, got {target!r}") + row = clean_logits[position].detach().float() + if row.numel() < 2: + return None, "vocab_too_small" + if target == "prediction": + top2 = torch.topk(row, 2).indices + return LogitMargin(int(top2[0]), int(top2[1]), position, "prediction"), "" + + from evalrx.analyzers._internals.access import tokenizer_of + + gold = gold_text(getattr(case, "expected", None)) + if gold is None: + return None, "no_gold" + tid = first_token_id(tokenizer_of(model) if model is not None else None, gold) + if tid is None or tid >= row.numel(): + return None, "gold_untokenizable" + others = row.clone() + others[tid] = float("-inf") + return LogitMargin(tid, int(others.argmax()), position, "gold"), "" diff --git a/evalrx/analyzers/attention/sink.py b/evalrx/analyzers/attention/sink.py index 8600398d..916d2e62 100755 --- a/evalrx/analyzers/attention/sink.py +++ b/evalrx/analyzers/attention/sink.py @@ -3,6 +3,10 @@ Measures how much attention mass collapses onto the first token (the "sink"), per layer, averaged over heads and query positions. Operates on a ``Trace``. +Runs over the batch (one attention forward per case) and emits per-case sink +columns, so M2 can test whether FAIL cases dump more attention on the sink — +the "attention is parked, not reading anything" signature. + References: - Efficient Streaming Language Models with Attention Sinks (StreamingLLM) Xiao et al., ICLR 2024 — arXiv:2309.17453 @@ -26,31 +30,68 @@ @register_analyzer("attention_sink") class AttentionSinkAnalyzer(Analyzer): - """Per-layer attention mass on the first token (sink strength).""" + """Per-layer attention mass on the first token (sink strength), per case. + + Hyper-parameters: + sink_pos: position treated as the sink (default 0, the first token). + max_cases: label-stratified cap (one attention forward each); 0 (the + default) = every case. + """ name = "attention_sink" requires = frozenset({Capability.ATTENTION}) applies_to_modalities = frozenset({"text", "image"}) + signal_docs = { + "sink_mass_mean": ("Attention on sink token", "Share of attention every query puts on the first token, averaged over heads and layers; high = attention parked on the sink instead of reading content."), + "sink_mass_last_query": ("Answer token's sink share", "Share of the final prompt position's attention on the first token, averaged over heads and layers."), + "sink_peak_layer_frac": ("Strongest-sink depth", "Relative depth (0 = first layer, 1 = last) of the layer with the most attention on the sink."), + } - def __init__(self, sink_pos: int = 0) -> None: - super().__init__(sink_pos=sink_pos) + def __init__(self, sink_pos: int = 0, max_cases: int = 0) -> None: + super().__init__(sink_pos=sink_pos, max_cases=max_cases) def _run(self, model: "Model", cases: "CaseBatch") -> Result: - case = cases[0] - trace = model.forward(case.inputs, capture={Capability.ATTENTION}) - attns = trace.require(Capability.ATTENTION) - per_layer = [] - for a in attns: - A = a.float().mean(0) # (seq, seq), head-averaged - per_layer.append(round(float(A[:, self.sink_pos].mean()), 4)) - mean_sink = round(sum(per_layer) / max(len(per_layer), 1), 4) + per_case: list[dict] = [] + layer_sums: list[float] = [] + n_summed = 0 + first = None # (n_layers, sink token) of the first analysed case + for case in cases.stratified_head(self.max_cases): + trace = model.forward(case.inputs, capture={Capability.ATTENTION}) + attns = trace.require(Capability.ATTENTION) + per_layer, last_query = [], [] + for a in attns: + A = a.float().mean(0) # (seq, seq), head-averaged + per_layer.append(float(A[:, self.sink_pos].mean())) + last_query.append(float(A[-1, self.sink_pos])) + if first is None: + first = (len(attns), trace.tokens[self.sink_pos] if trace.tokens else None) + layer_sums = [0.0] * len(attns) + if len(per_layer) == len(layer_sums): + layer_sums = [s + v for s, v in zip(layer_sums, per_layer)] + n_summed += 1 + n_layers = len(per_layer) + peak = max(range(n_layers), key=per_layer.__getitem__) if n_layers else 0 + per_case.append({ + "sample_id": case.id, + "sink_mass_mean": round(sum(per_layer) / max(n_layers, 1), 4), + "sink_mass_last_query": round(sum(last_query) / max(n_layers, 1), 4), + "sink_peak_layer_frac": round(peak / max(1, n_layers - 1), 4), + }) + + if first is None: + raise ValueError("attention_sink: empty case batch") + n = len(per_case) + per_layer_sink = [round(s / n_summed, 4) for s in layer_sums] + mean_sink = round(sum(r["sink_mass_mean"] for r in per_case) / n, 4) return Result( analyzer=self.name, model=repr(model), cases=cases, - artifacts={"per_layer_sink": per_layer}, + artifacts={"per_layer_sink": per_layer_sink}, findings={ - "n_layers": len(attns), - "sink_token": trace.tokens[self.sink_pos] if trace.tokens else None, + "n_layers": first[0], + "sink_token": first[1], + "n_cases_analyzed": n, "mean_sink_mass": mean_sink, - "per_layer_sink": per_layer, + "per_layer_sink": per_layer_sink, + "per_case": per_case, }, ) diff --git a/evalrx/analyzers/attention/summary.py b/evalrx/analyzers/attention/summary.py index 840c7e97..ed7fb283 100755 --- a/evalrx/analyzers/attention/summary.py +++ b/evalrx/analyzers/attention/summary.py @@ -151,28 +151,62 @@ class AttentionAnalyzer(Analyzer): name = "attention" requires = frozenset({Capability.ATTENTION}) + signal_docs = { + "last_query_entropy": ("Answer-token attention spread", "Normalized entropy of the final prompt position's attention at the summary layer (1 = spread evenly over all tokens, 0 = on one token)."), + "last_query_top1_share": ("Answer-token peak attention", "Largest single share of the final prompt position's attention at the summary layer."), + "media_mass_mid": ("Attention on image/audio", "Share of the final prompt position's attention that lands on image/audio tokens, averaged over heads and the middle half of the layers; low = the answer position barely reads the media."), + } - def __init__(self, layer: int = -1, head: int | str = "mean", top_k: int = 10) -> None: - super().__init__(layer=layer, head=head, top_k=top_k) + def __init__(self, layer: int = -1, head: int | str = "mean", top_k: int = 10, + max_cases: int = 0) -> None: + super().__init__(layer=layer, head=head, top_k=top_k, max_cases=max_cases) def _run(self, model: "Model", cases: "CaseBatch") -> AttentionResult: - # Stage 1: analyse the first case (single-case ergonomics). Multi-case - # aggregation is a Stage-2 extension that will return a batched result. - case = cases[0] - trace = model.forward(case.inputs, capture={Capability.ATTENTION}) - attentions = trace.require(Capability.ATTENTION) - - result = AttentionResult( - analyzer=self.name, - model=repr(model), - cases=cases, - tokens=trace.tokens, - token_ids=trace.token_ids, - artifacts={"attentions": attentions}, - ) + # The first case keeps the full per-layer tensors (single-case + # ergonomics: aggregate/plot/top_attended_tokens); every analysed case + # contributes one per_case row for the M2 stats layer. + result: AttentionResult | None = None + per_case: list[dict[str, Any]] = [] + for case in cases.stratified_head(self.max_cases): + trace = model.forward(case.inputs, capture={Capability.ATTENTION}) + attentions = trace.require(Capability.ATTENTION) + if result is None: + result = AttentionResult( + analyzer=self.name, + model=repr(model), + cases=cases, + tokens=trace.tokens, + token_ids=trace.token_ids, + artifacts={"attentions": attentions}, + ) + per_case.append({"sample_id": case.id, **self._case_columns(trace, attentions)}) + if result is None: + raise ValueError("attention: empty case batch") result.findings = self._summarise(result) + result.findings["n_cases_analyzed"] = len(per_case) + result.findings["per_case"] = per_case return result + def _case_columns(self, trace: Any, attentions: "list[torch.Tensor]") -> dict[str, Any]: + """Per-case scalars: last-query concentration, and its media share mid-stack.""" + from evalrx.analyzers._internals.access import media_positions + + attn = attentions[self.layer] + row = (attn.float().mean(0) if self.head == "mean" else attn[int(self.head)].float())[-1] + p = row.clamp(min=1e-12) / row.sum().clamp(min=1e-12) + ent = float(-(p * p.log()).sum()) / float(np.log(max(p.numel(), 2))) + out: dict[str, Any] = { + "last_query_entropy": round(ent, 4), + "last_query_top1_share": round(float(row.max()), 4), + } + _, media = media_positions(trace) + if media: + n = len(attentions) + mid = range(n // 4, max(n // 4 + 1, (3 * n) // 4)) + shares = [float(attentions[i].float().mean(0)[-1][media].sum()) for i in mid] + out["media_mass_mid"] = round(sum(shares) / len(shares), 4) + return out + def _summarise(self, result: AttentionResult) -> dict[str, Any]: """Build the light, agent-readable summary stored in ``findings``.""" top = result.top_attended_tokens( diff --git a/evalrx/analyzers/attribution/generic_attn.py b/evalrx/analyzers/attribution/generic_attn.py index f28e9e88..be850f1c 100755 --- a/evalrx/analyzers/attribution/generic_attn.py +++ b/evalrx/analyzers/attribution/generic_attn.py @@ -1,31 +1,221 @@ -"""Generic attention-model explainability (Chefer et al., 2021) (Stage 2). +"""Generic attention-model explainability (Chefer et al., 2021) for a decoder-only VLM. -Relevance propagation through attention layers using attention weights AND their -gradients — hence ``requires={ATTENTION, GRADIENTS}`` (white-box, not black-box, -correcting the original taxonomy). Built for enc-dec / bi-modal transformers; not -decoder-only-native. +Relevance propagation through the self-attention stack using attention weights +AND their gradients. For self-attention, Chefer et al. initialise ``R = I`` and, +layer by layer, + + Ā_l = E_heads[ (∇A_l ⊙ A_l)⁺ ] R ← R + Ā_l · R + +so ``R[i, j]`` is how much token ``j`` contributes to the representation token +``i`` predicts from. The row of the final prompt position is the relevance of +every input token to the next-token prediction; here the gradients are of the +next-token logit margin (see :mod:`evalrx.analyzers._internals.targets`). + +Only that one row is needed, so it is propagated as a vector from the top +layer down — ``v ← v + v · Ā_l`` for ``l = L-1 … 0`` — O(seq²) per layer +instead of materialising ``R``. The identity's own contribution (the final +position relevant to itself) is dropped before shares are computed. + +Attention weights and their gradients come from the evalrx attention tap +(:mod:`evalrx.analyzers._internals.attn_tap`), checked once per model to +reproduce the native logits; the per-layer ``Ā_l`` is formed in a tensor hook +during the backward pass, so no full set of attention gradients is ever held. +A hybrid stack (linear-attention layers without a token matrix) is refused: +relevance would only cross some of the layers. References: -- Generic Attention-model Explainability for Interpreting Bi-Modal and Encoder-Decoder Transformers - Chefer, Gur & Wolf, ICCV 2021 — arXiv:2103.15679 -- Predecessor — Transformer Interpretability Beyond Attention Visualization, +- Generic Attention-model Explainability for Interpreting Bi-Modal and + Encoder-Decoder Transformers — Chefer, Gur & Wolf, ICCV 2021 — arXiv:2103.15679 +- Transformer Interpretability Beyond Attention Visualization Chefer, Gur & Wolf, CVPR 2021 — arXiv:2012.09838 """ from __future__ import annotations +from collections import Counter +from typing import TYPE_CHECKING, Any + from evalrx.core.analyzer import Analyzer from evalrx.core.capability import Capability from evalrx.core.registry import register_analyzer +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + @register_analyzer("generic_attention") class GenericAttentionExplainability(Analyzer): + """Chefer relevance of each input token to the next-token prediction, per case. + + Hyper-parameters: + target: ``"prediction"`` (clean top-1 vs top-2) or ``"gold"``. + top_k: relevant tokens listed for the first case. + fidelity_rtol: tolerance of the edit-free attention tap vs native logits. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: one forward + one backward per case. + """ + name = "generic_attention" - requires = frozenset({Capability.ATTENTION, Capability.GRADIENTS}) + description = ( + "Chefer gradient-weighted attention relevance of every input token to the " + "answer token; reports how much of it lands on the image/audio tokens and " + "where in the image." + ) + requires = frozenset({Capability.ATTENTION, Capability.GRADIENTS, Capability.LOGITS}) applies_to_modalities = frozenset({"text", "image"}) + signal_docs = { + "rel_margin": ("Answer margin", "Logit margin of the target answer token being explained."), + "rel_media_share": ("Relevance on image/audio", "Share of the answer token's attention relevance that traces back to image/audio tokens; low = the answer was built from the text."), + "rel_input_entropy": ("Relevance spread", "Normalized entropy of relevance over all input tokens (1 = spread evenly, 0 = one token)."), + "rel_input_top1_share": ("Top token's relevance", "Share of all relevance on the single most relevant input token."), + "rel_entropy": ("Image relevance spread", "Normalized entropy of relevance over image tokens (1 = diffuse, 0 = one patch)."), + "rel_top1_share": ("Image relevance peak", "Share of image relevance on its single strongest image token."), + "rel_top5_share": ("Image relevance top-5", "Share of image relevance on its five strongest image tokens."), + "rel_center_offset": ("Image relevance off-centre", "Distance of the image relevance's centre of mass from the image centre."), + "rel_edge_mass": ("Image relevance on border", "Share of image relevance on the outermost ring of patches."), + } + + def __init__(self, target: str = "prediction", top_k: int = 10, + fidelity_rtol: float = 1e-2, max_cases: int = 0) -> None: + super().__init__(target=target, top_k=top_k, fidelity_rtol=fidelity_rtol, + max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch"): + import numpy as np + import torch - def _run(self, model, cases): - raise NotImplementedError( - "Stage 2: Chefer-style relevance via attention × gradient-of-attention propagation." + from evalrx.analyzers._internals.access import ( + check_editable_attention, + decoder_stack, + media_positions, + torch_module, ) + from evalrx.analyzers._internals.attn_tap import attention_tap, require_fidelity + from evalrx.analyzers._internals.maps import ( + SpatialMapResult, + group_maps, + image_token_scores, + map_features, + ) + from evalrx.analyzers._internals.targets import build_margin + from evalrx.core.model import CaptureSpec + + check_editable_attention(model, self.name) + stack = decoder_stack(torch_module(model, self.name)) + if not stack.fully_attentional: + raise ValueError(f"{self.name}: some decoder layers have no softmax attention " + "(hybrid stack); relevance would cross only part of the stack.") + n_layers = stack.n_layers + spec = CaptureSpec(with_grad=True, to_cpu=False) + + per_case: list[dict[str, Any]] = [] + per_case_maps: dict[str, Any] = {} + skipped: Counter = Counter() + fidelity: "dict | None" = None + first: "dict[str, Any] | None" = None + for case in cases.stratified_head(self.max_cases): + if fidelity is None: + fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, + analyzer=self.name) + abar: dict[int, Any] = {} + attn_probs: list = [] + + def on_probs(layer: int, probs: Any) -> None: + if not probs.requires_grad: + return + held = probs.detach() + + def hook(grad: Any, layer: int = layer, held: Any = held) -> None: + # Ā_l = E_heads[(∇A ⊙ A)⁺] for the single sequence in the batch + abar[layer] = (grad.float() * held.float()).clamp(min=0).mean(dim=1)[0] + + probs.register_hook(hook) + attn_probs.append(probs) + + with attention_tap(model, on_probs=on_probs, analyzer=self.name): + trace = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + logits = trace.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + if not attn_probs: + raise RuntimeError( + f"{self.name}: the forward pass kept no autograd graph through attention — " + "the backend ignores CaptureSpec(with_grad=True)." + ) + m_value = margin.value(logits) + # Asking for layer 0's attention gradient sends the backward pass + # through every layer's attention, firing each layer's hook once. + torch.autograd.grad(margin(logits), attn_probs[0]) + del attn_probs, logits + if len(abar) != n_layers: + raise RuntimeError(f"{self.name}: relevance reached {len(abar)} of {n_layers} layers") + + seq = trace.seq_len + v = torch.zeros(seq, device=abar[0].device) + v[-1] = 1.0 + for layer in reversed(range(n_layers)): + v = v + v @ abar[layer] + rel = v[:-1].clamp(min=0).cpu().double().numpy() # drop the final position's self term + abar.clear() + total = float(rel.sum()) + + row: dict[str, Any] = {"sample_id": case.id, "rel_margin": round(m_value, 4)} + if total > 0: + p = rel / total + nz = p[p > 0] + row["rel_input_entropy"] = round(float(-(nz * np.log(nz)).sum() / np.log(max(p.size, 2))), 4) + row["rel_input_top1_share"] = round(float(p.max()), 4) + slot, media = media_positions(trace) + media = [i for i in media if i < seq - 1] + if media and total > 0: + row["rel_media_share"] = round(float(rel[media].sum() / total), 4) + if slot == "image": + flat, spatial = image_token_scores(trace, rel[media]) + row.update(map_features(flat, spatial, "rel")) + per_case_maps[case.id] = np.asarray(spatial, dtype=np.float16) + if first is None: + order = np.argsort(rel)[::-1][: self.top_k] + first = { + "case_id": case.id, + "top": [{"position": int(i), "token": trace.tokens[int(i)], + "relevance": round(float(rel[i] / total), 4) if total > 0 else 0.0} + for i in order], + } + per_case.append(row) + del trace + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + n = len(per_case) + shares = [r["rel_media_share"] for r in per_case if "rel_media_share" in r] + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": n_layers, + "target": self.target, + "tap_fidelity": fidelity, + "top_relevant_tokens": first["top"] if first else [], + "mean_media_share": round(sum(shares) / len(shares), 4) if shares else None, + "per_case": per_case, + "_caveat": ( + "Chefer et al.'s self-attention rule applied to a decoder-only stack, " + "explaining the next-token margin at the end of the prompt. Relevance " + "follows attention paths only; information carried by MLPs within a " + "position is attributed to that position." + ), + } + artifacts: dict[str, Any] = {} + if per_case_maps: + first_map_id = next(iter(per_case_maps)) + artifacts["spatial_map"] = per_case_maps[first_map_id].astype(np.float64) + artifacts["per_case_maps"] = per_case_maps + findings["spatial_map_case_id"] = first_map_id + group_artifacts, group_findings = group_maps(per_case_maps, cases) + artifacts.update(group_artifacts) + findings.update(group_findings) + return SpatialMapResult(analyzer=self.name, model=repr(model), cases=cases, + artifacts=artifacts, findings=findings) diff --git a/evalrx/analyzers/attribution/gradcam.py b/evalrx/analyzers/attribution/gradcam.py index 2da42982..33b0c83f 100755 --- a/evalrx/analyzers/attribution/gradcam.py +++ b/evalrx/analyzers/attribution/gradcam.py @@ -1,30 +1,175 @@ -"""Grad-CAM on the vision tower (Stage 2). +"""Grad-CAM on the decoder's image-token states — which image regions the answer depends on. -Gradient-weighted activation maps over the ViT feature map → spatial saliency. -``requires=GRADIENTS`` (backward pass); VLM vision tower only. Needs a -``reshape_transform`` for ViT and the per-arch vision-tower target layer. +Selvaraju et al. weight a CNN feature map's channels by the spatially averaged +gradient of the class score, sum, and keep the positive part. A VLM decoder has +no CNN feature map, but at any decoder layer its image tokens ARE one: a grid of +``(h, w)`` positions, each carrying a ``dim``-channel state. So, at decoder +layer ``l`` (the residual stream entering it): + + α_d = mean over image tokens t of ∂margin / ∂h[t, d] + cam_t = ReLU( Σ_d α_d · h[t, d] ) + +with ``margin`` the next-token logit margin (see +:mod:`evalrx.analyzers._internals.targets`). The map is rebuilt on the image's +post-merge patch grid, so it overlays on the photo like ``relative_attention``. + +This is an ADAPTATION of Grad-CAM, not the vision-tower original: it reads the +image tokens after the language model has started mixing them, which is where +an answer's dependence on the image is decided, and it needs no +per-architecture vision-tower target layer. + +Alongside the map, ``gc_media_grad_share`` is a cruder, map-free reading: the +share of first-order (gradient × input) sensitivity at that layer carried by +image tokens rather than text tokens. References: - Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization Selvaraju et al., ICCV 2017 — arXiv:1610.02391 -- Code: https://github.com/jacobgil/pytorch-grad-cam """ from __future__ import annotations +from collections import Counter +from typing import TYPE_CHECKING, Any + from evalrx.core.analyzer import Analyzer from evalrx.core.capability import Capability from evalrx.core.registry import register_analyzer +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + @register_analyzer("gradcam") class GradCAMAnalyzer(Analyzer): + """Grad-CAM over image tokens at one decoder layer, per case. + + Hyper-parameters: + layer: decoder layer whose INPUT residual stream is read. ``float`` + in (0, 1) = fractional depth (default 0.5); ``int`` = + absolute index (negative allowed). + target: ``"prediction"`` (clean top-1 vs top-2) or ``"gold"``. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: one forward + one backward per case. + """ + name = "gradcam" - requires = frozenset({Capability.GRADIENTS}) - applies_to_modalities = frozenset({"image"}) + description = ( + "Gradient-weighted map of which image regions the answer token depends on " + "(Grad-CAM over the decoder's image tokens), plus the image tokens' share of " + "gradient×input sensitivity." + ) + requires = frozenset({Capability.GRADIENTS, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video"}) + requires_modalities = frozenset({"image", "video"}) + signal_docs = { + "gc_margin": ("Answer margin", "Logit margin of the target answer token the map explains."), + "gc_media_grad_share": ("Image share of sensitivity", "Share of the answer's first-order (gradient × input) sensitivity carried by image tokens at the probed layer; low = the answer hinges on the text."), + "gc_empty": ("Grad-CAM map empty", "No image token pushed the answer up (the positive part of the map is zero everywhere)."), + "gc_entropy": ("Grad-CAM spread", "Normalized entropy of the Grad-CAM map over image tokens (1 = diffuse, 0 = one patch)."), + "gc_top1_share": ("Grad-CAM peak share", "Share of the Grad-CAM map on its single strongest image token."), + "gc_top5_share": ("Grad-CAM top-5 share", "Share of the Grad-CAM map on its five strongest image tokens."), + "gc_center_offset": ("Grad-CAM off-centre", "Distance of the map's centre of mass from the image centre (0 = centre, 1 = corner)."), + "gc_edge_mass": ("Grad-CAM on border", "Share of the map on the outermost ring of patches."), + } + + def __init__(self, layer: "int | float" = 0.5, target: str = "prediction", max_cases: int = 0) -> None: + super().__init__(layer=layer, target=target, max_cases=max_cases) - def _run(self, model, cases): - raise NotImplementedError( - "Stage 2: register a hook on the vision-tower target block, backward through a " - "target logit, weight activations by gradients (pytorch-grad-cam + reshape_transform)." + def _run(self, model: "Model", cases: "CaseBatch"): + import numpy as np + import torch + + from evalrx.analyzers._internals.access import decoder_stack, media_positions, torch_module + from evalrx.analyzers._internals.maps import ( + SpatialMapResult, + group_maps, + image_token_scores, + map_features, ) + from evalrx.analyzers._internals.residual import residual_hooks + from evalrx.analyzers._internals.targets import build_margin + from evalrx.analyzers.attention.relative_attn import resolve_attention_layer + from evalrx.core.model import CaptureSpec + + stack = decoder_stack(torch_module(model, self.name)) + layer = resolve_attention_layer(self.layer, stack.n_layers) % stack.n_layers + spec = CaptureSpec(with_grad=True, to_cpu=False) + + per_case: list[dict[str, Any]] = [] + per_case_maps: dict[str, Any] = {} + skipped: Counter = Counter() + first_id: "str | None" = None + for case in cases.stratified_head(self.max_cases): + with residual_hooks(model, capture={layer: None}, keep_grad=True, + analyzer=self.name) as hooks: + trace = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + slot, media = media_positions(trace) + if slot != "image": + skipped["no_image_tokens"] += 1 + continue + logits = trace.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + hidden = hooks.captured[layer] + if not hidden.requires_grad: + raise RuntimeError( + f"{self.name}: the forward pass kept no autograd graph — the backend " + "ignores CaptureSpec(with_grad=True), so there is nothing to differentiate." + ) + m_value = margin.value(logits) + (grad,) = torch.autograd.grad(margin(logits), hidden) + h, g = hidden.detach()[0].float(), grad[0].float() + idx = torch.tensor(media, device=h.device) + cam = torch.relu(h[idx] @ g[idx].mean(0)).cpu().numpy() + gxi = (g * h).sum(-1).abs() + share = float(gxi[idx].sum() / gxi.sum().clamp(min=1e-12)) + flat, spatial = image_token_scores(trace, cam) + del trace, hooks, hidden, grad, logits, h, g + + row: dict[str, Any] = { + "sample_id": case.id, + "gc_margin": round(m_value, 4), + "gc_media_grad_share": round(share, 4), + "gc_empty": bool(flat.sum() <= 0.0), + } + row.update(map_features(flat, spatial, "gc")) + per_case.append(row) + per_case_maps[case.id] = np.asarray(spatial, dtype=np.float16) + if first_id is None: + first_id = case.id + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + artifacts: dict[str, Any] = { + "spatial_map": per_case_maps[first_id].astype(np.float64), + "per_case_maps": per_case_maps, + } + group_artifacts, group_findings = group_maps(per_case_maps, cases) + artifacts.update(group_artifacts) + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "layer_used": layer, + "n_layers": stack.n_layers, + "target": self.target, + "map_shape": list(artifacts["spatial_map"].shape), + "mean_media_grad_share": round(sum(r["gc_media_grad_share"] for r in per_case) / n, 4), + "frac_empty": round(sum(r["gc_empty"] for r in per_case) / n, 4), + "per_case": per_case, + "spatial_map_case_id": first_id, + **group_findings, + "_caveat": ( + "Grad-CAM adapted to decoder image tokens (residual stream entering " + "layer_used), explaining the next-token margin at the end of the prompt. " + "A gradient map is a first-order local explanation: it says what would " + "move the margin, not what the model attended to." + ), + } + return SpatialMapResult(analyzer=self.name, model=repr(model), cases=cases, + artifacts=artifacts, findings=findings) diff --git a/evalrx/analyzers/lens/__init__.py b/evalrx/analyzers/lens/__init__.py index b7648010..d0f54acb 100755 --- a/evalrx/analyzers/lens/__init__.py +++ b/evalrx/analyzers/lens/__init__.py @@ -1,7 +1,8 @@ -"""Lens analyzers — project hidden states toward the vocabulary.""" +"""Lens analyzers — project hidden states (or per-component writes) toward the vocabulary.""" +from evalrx.analyzers.lens.head_attribution import HeadAttributionAnalyzer from evalrx.analyzers.lens.layer_contrast import LayerContrastAnalyzer from evalrx.analyzers.lens.logit_lens import LogitLensAnalyzer from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer -__all__ = ["LogitLensAnalyzer", "TunedLensAnalyzer", "LayerContrastAnalyzer"] +__all__ = ["LogitLensAnalyzer", "TunedLensAnalyzer", "LayerContrastAnalyzer", "HeadAttributionAnalyzer"] diff --git a/evalrx/analyzers/lens/head_attribution.py b/evalrx/analyzers/lens/head_attribution.py new file mode 100644 index 00000000..84fcf693 --- /dev/null +++ b/evalrx/analyzers/lens/head_attribution.py @@ -0,0 +1,183 @@ +"""Direct logit attribution per attention head — which heads write the answer. + +The residual stream at the final prompt position is a sum of writes: the +embedding, every attention head, every MLP. Read through the unembedding (with +the final norm's scale frozen at the clean run — see +:mod:`evalrx.analyzers._internals.heads`), each write contributes a fixed amount +to the next-token logit margin, and the contributions add up to the margin. +This analyzer reports the attention heads' share: how much of the margin they +write directly, how concentrated it is in a few heads, and at what depth. + +"Direct" means the head's own write only; a head that matters by feeding later +layers (its effect on what they compute) scores low here. Combine with +``attention_knockout`` / ``causal_trace`` for indirect effects. + +The per-head mean attribution over the batch, and its FAIL−PASS difference, +are artifacts (``(layers, heads)`` maps), so a head that writes the wrong +answer on failures stands out as a hot cell. + +References: +- A Mathematical Framework for Transformer Circuits + Elhage et al., Anthropic 2021 — transformer-circuits.pub/2021/framework +- Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small + Wang et al., ICLR 2023 — arXiv:2211.00593 +""" + +from __future__ import annotations + +from collections import Counter +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + + +@register_analyzer("head_attribution") +class HeadAttributionAnalyzer(Analyzer): + """Per-head direct contribution to the next-token logit margin, per case. + + Hyper-parameters: + target: ``"prediction"`` (clean top-1 vs top-2) or ``"gold"``. + top_k: heads listed for the first case. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: one logits forward per case. + """ + + name = "head_attribution" + description = ( + "Splits the answer token's logit margin into each attention head's direct write " + "(direct logit attribution); reports how much the heads write and how concentrated." + ) + requires = frozenset({Capability.HIDDEN_STATES, Capability.LOGITS}) + applies_to_modalities = frozenset({"text", "image"}) + signal_docs = { + "dla_margin": ("Answer margin", "Logit margin of the target answer token being decomposed."), + "dla_heads_total": ("Margin written by heads", "Sum of all attention heads' direct contributions to the answer margin (logits)."), + "dla_heads_share": ("Heads' share of margin", "dla_heads_total as a fraction of the margin (the rest is written by MLPs and the embedding)."), + "dla_top_head_share": ("Top head's share", "Largest single head's share of all heads' absolute contribution; high = one head decides the answer."), + "dla_top_head_layer_frac": ("Top head's depth", "Relative depth (0 = first layer, 1 = last) of the head with the largest absolute contribution."), + "dla_late_share": ("Late heads' share", "Share of all heads' absolute contribution coming from the last third of the layers."), + "dla_negative_share": ("Heads opposing answer", "Share of all heads' absolute contribution that pushes AGAINST the answer token."), + } + + def __init__(self, target: str = "prediction", top_k: int = 10, max_cases: int = 0) -> None: + super().__init__(target=target, top_k=top_k, max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import numpy as np + import torch + + from evalrx.analyzers._internals.access import decoder_stack, torch_module + from evalrx.analyzers._internals.heads import ( + frozen_norm_scale, + head_capture, + n_heads_of, + out_proj, + ) + from evalrx.analyzers._internals.targets import build_margin + from evalrx.models.backends.hf.discover import get_final_norm, get_unembed + + module = torch_module(model, self.name) + stack = decoder_stack(module) + norm = get_final_norm(module) + unembed = get_unembed(module) + projs = [out_proj(a) if a is not None else None for a in stack.attn] + if any(p is None for p in projs): + raise ValueError(f"{self.name}: some decoder layers have no attention output projection " + "(hybrid stack); head attribution would miss them.") + n_heads = [n_heads_of(a, p) for a, p in zip(stack.attn, projs)] + n_layers = stack.n_layers + + per_case: list[dict[str, Any]] = [] + maps: dict[str, Any] = {} + skipped: Counter = Counter() + max_err = 0.0 + first_top: list = [] + for case in cases.stratified_head(self.max_cases): + with torch.no_grad(), head_capture(model, analyzer=self.name) as cap: + trace = model.forward(case.inputs, capture={Capability.LOGITS}) + logits = trace.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + m_value = margin.value(logits) + with torch.no_grad(): + x_final = cap.final_residual + scale, centre = frozen_norm_scale(norm, x_final) + w_u = unembed.weight + direction = (w_u[margin.positive].float() - w_u[margin.negative].float()).to(scale.device) + v = scale * direction + if centre: + v = v - v.mean() + readout = float(x_final.float().to(v.device) @ v) + bias = getattr(unembed, "bias", None) + if bias is not None: + readout += float(bias[margin.positive] - bias[margin.negative]) + max_err = max(max_err, abs(readout - m_value)) + dla = torch.zeros(n_layers, max(n_heads)) + for layer, (proj, heads) in enumerate(zip(projs, n_heads)): + w_o = proj.weight.float() + per_input = (w_o.T @ v.to(w_o.device)) * cap.head_inputs[layer].float().to(w_o.device) + dla[layer, :heads] = per_input.view(heads, -1).sum(-1).cpu() + del trace, cap + + flat = dla.flatten() + absolute = flat.abs() + total_abs = float(absolute.sum()) or 1.0 + top = int(absolute.argmax()) + late = n_layers - n_layers // 3 + row: dict[str, Any] = { + "sample_id": case.id, + "dla_margin": round(m_value, 4), + "dla_heads_total": round(float(flat.sum()), 4), + "dla_top_head_share": round(float(absolute[top]) / total_abs, 4), + "dla_top_head_layer_frac": round((top // dla.shape[1]) / max(1, n_layers - 1), 4), + "dla_late_share": round(float(dla[late:].abs().sum()) / total_abs, 4), + "dla_negative_share": round(float(flat.clamp(max=0).abs().sum()) / total_abs, 4), + } + if abs(m_value) > 1e-6: + row["dla_heads_share"] = round(float(flat.sum()) / m_value, 4) + per_case.append(row) + maps[case.id] = dla.numpy() + if not first_top: + order = torch.argsort(absolute, descending=True)[: self.top_k] + first_top = [{"layer": int(i) // dla.shape[1], "head": int(i) % dla.shape[1], + "dla": round(float(flat[i]), 4)} for i in order] + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + label_of = {c.id: getattr(getattr(c, "label", None), "value", None) for c in cases} + artifacts: dict[str, Any] = {"head_dla_mean_map": np.mean(list(maps.values()), axis=0)} + fail = [m for cid, m in maps.items() if label_of.get(cid) == "fail"] + passed = [m for cid, m in maps.items() if label_of.get(cid) == "pass"] + if fail and passed: + artifacts["head_dla_diff_map_fail_minus_pass"] = np.mean(fail, axis=0) - np.mean(passed, axis=0) + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": n_layers, + "n_heads": max(n_heads), + "target": self.target, + "top_heads_first_case": first_top, + "mean_heads_total": round(sum(r["dla_heads_total"] for r in per_case) / n, 4), + "linearization_max_error": round(max_err, 4), + "per_case": per_case, + "_caveat": ( + "Direct effects only, read at the final prompt position with the final " + "norm's scale frozen at the clean run. linearization_max_error is the gap " + "between the frozen-norm readout of the full residual and the true margin: " + "bf16 rounding of the model's own logits (~1% of their magnitude), or more " + "when the model soft-caps its logits." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, + artifacts=artifacts, findings=findings) diff --git a/evalrx/analyzers/lens/tuned_lens.py b/evalrx/analyzers/lens/tuned_lens.py index da701578..3698b20f 100755 --- a/evalrx/analyzers/lens/tuned_lens.py +++ b/evalrx/analyzers/lens/tuned_lens.py @@ -1,7 +1,21 @@ -"""Tuned lens (Belrose et al.) — per-layer learned affine translators before unembed. +"""Tuned lens (Belrose et al.) — per-layer learned affine translators before the unembed. -More faithful than the raw logit lens, but requires PER-MODEL trained translators. -Stage 2: train/load translators, then apply like the logit lens. +The logit lens reads layer ``l``'s residual stream as if it were the last +layer's, which is biased: intermediate layers use the residual basis +differently. The tuned lens learns, per layer, an affine translator +``T_l(h) = A_l h + b_l`` (initialised to the identity) such that +``unembed(norm(T_l(h_l)))`` matches the model's own FINAL distribution, by +minimising ``KL(p_final || p_l)``. The fit is label-free: it uses only the +model's outputs, never PASS/FAIL. + +Here the translators are fit on token positions sampled from the batch being +analysed (a pretraining corpus is not at hand), then each case's final prompt +position is read through them; ``fit_kl`` in the findings reports how well the +lens fits. Fitted translators are cached on the model per configuration and can +be saved to / loaded from ``translators_path``. + +Per-case columns mirror ``logit_lens`` (decision depth, interim peak, late drop) +plus the tuned KL to the final distribution in the middle of the stack. References: - Eliciting Latent Predictions from Transformers with the Tuned Lens @@ -10,21 +24,231 @@ from __future__ import annotations +from typing import TYPE_CHECKING, Any + from evalrx.core.analyzer import Analyzer from evalrx.core.capability import Capability from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + + +def _functional_norm(norm: Any, sample: Any) -> Any: + """A differentiable copy of the final norm with its weights detached. + + Training a translator through the module itself would also write ``.grad`` + into the model's own norm weights; this applies the same function from + detached tensors instead. + """ + import torch + + if norm is None: + return lambda x: x + eps = float(getattr(norm, "variance_epsilon", getattr(norm, "eps", 1e-6)) or 1e-6) + name = type(norm).__name__.lower() + if isinstance(norm, torch.nn.LayerNorm) or ("layernorm" in name and "rms" not in name): + weight = norm.weight.detach().float() if getattr(norm, "weight", None) is not None else None + bias = norm.bias.detach().float() if getattr(norm, "bias", None) is not None else None + + def layer_norm(x: Any) -> Any: + return torch.nn.functional.layer_norm(x, x.shape[-1:], weight, bias, eps) + + return layer_norm + with torch.no_grad(): + w_eff = (norm(torch.ones_like(sample)).float() * (1.0 + eps) ** 0.5).reshape(-1).to(sample.device) + + def rms_norm(x: Any) -> Any: + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * w_eff + + return rms_norm @register_analyzer("tuned_lens") class TunedLensAnalyzer(Analyzer): - """Tuned-lens projection of hidden states (needs trained translators).""" + """Tuned-lens readout of the final prompt position, per case. + + Hyper-parameters: + train_tokens: token positions sampled (across the analysed cases) to fit + the translators; the final position of every case is included. + steps: Adam steps per layer translator. + batch_size: positions per step. + lr: Adam learning rate. + seed: sampling / shuffling seed. + translators_path: optional ``.pt`` file to load fitted translators from + (when it exists and matches) or save them to. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: one hidden-states forward per case, plus the fit. + """ name = "tuned_lens" + description = ( + "Reads each layer's prediction through per-layer translators fit to the model's " + "own final output (tuned lens); reports how early and how stably the answer forms." + ) requires = frozenset({Capability.HIDDEN_STATES}) applies_to_modalities = frozenset({"text", "image"}) + signal_docs = { + "tl_decision_frac": ("Answer settles (tuned)", "Relative depth from which every tuned-lens layer already predicts the final answer token; late = the answer forms only at the top."), + "tl_final_top1_prob": ("Final answer prob.", "Probability of the emitted token at the final layer."), + "tl_interim_peak_prob": ("Peak answer prob. (tuned)", "Highest tuned-lens probability of the final answer token at any layer."), + "tl_late_drop": ("Late suppression (tuned)", "Drop from that peak to the final layer; large = the answer was stronger mid-stack and got suppressed."), + "tl_kl_mid": ("Mid-stack disagreement", "Tuned-lens KL divergence from the final distribution, averaged over the middle third of the layers."), + } - def _run(self, model, cases): - raise NotImplementedError( - "Stage 2: load/train per-layer affine translators (e.g. via the `tuned-lens` package), " - "then project hidden states through them before the unembed." + def __init__(self, train_tokens: int = 4096, steps: int = 200, batch_size: int = 256, + lr: float = 1e-3, seed: int = 0, translators_path: "str | None" = None, + max_cases: int = 0) -> None: + super().__init__(train_tokens=train_tokens, steps=steps, batch_size=batch_size, lr=lr, + seed=seed, translators_path=translators_path, max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + W = model.unembed_weight() + if W is None: + raise ValueError(f"{type(model).__name__} exposes no unembed_weight(); the tuned lens " + "needs it (white-box local backend).") + W = W.detach().float() + device = W.device + norm = model.final_norm() if hasattr(model, "final_norm") else None + + # One forward per case: keep every layer's final-position state for the + # readout, plus a seeded sample of positions for fitting the translators. + rows = cases.stratified_head(self.max_cases) + if not rows: + raise ValueError("tuned_lens: empty case batch") + gen = torch.Generator().manual_seed(self.seed) + per_case_budget = max(1, self.train_tokens // len(rows)) + last_states: list[tuple[str, list]] = [] + train: "list[list] | None" = None + for case in rows: + hidden = model.forward(case.inputs, capture={Capability.HIDDEN_STATES}).require( + Capability.HIDDEN_STATES) + seq = hidden[0].shape[0] + pick = torch.randperm(seq, generator=gen)[: per_case_budget - 1].tolist() + pick = sorted(set(pick) | {seq - 1}) + if train is None: + train = [[] for _ in hidden] + for layer, h in enumerate(hidden): + train[layer].append(h[pick].detach().float().cpu()) + last_states.append((case.id, [h[-1].detach().float().cpu() for h in hidden])) + data = [torch.cat(chunks) for chunks in train] + n_layers = len(data) - 1 # entries 0..L-1 premature, L = final (post-norm) + dim = data[0].shape[1] + apply_norm = _functional_norm(norm, data[0][:1].to(device)) + + translators, fit_kl = self._translators(model, data, W, apply_norm, n_layers, dim, device) + + def layer_logprobs(states: list) -> list: + out = [] + with torch.no_grad(): + for layer in range(n_layers): + a, b = translators[layer] + x = states[layer].to(device) @ a.T + b + out.append(torch.log_softmax(apply_norm(x) @ W.T, dim=-1)) + out.append(torch.log_softmax(states[n_layers].to(device) @ W.T, dim=-1)) + return out + + mid = list(range(n_layers // 3, max(n_layers // 3 + 1, 2 * n_layers // 3))) + per_case: list[dict[str, Any]] = [] + for case_id, states in last_states: + logps = layer_logprobs(states) + final = logps[-1] + top1 = int(final.argmax()) + probs = [float(lp[top1].exp()) for lp in logps] + argmaxes = [int(lp.argmax()) for lp in logps] + decision = next((i for i in range(len(logps)) if all(a == top1 for a in argmaxes[i:])), + len(logps) - 1) + kls = [float(torch.sum(final.exp() * (final - logps[i]))) for i in mid] + peak = max(probs) + per_case.append({ + "sample_id": case_id, + "tl_decision_frac": round(decision / max(1, len(logps) - 1), 4), + "tl_final_top1_prob": round(probs[-1], 4), + "tl_interim_peak_prob": round(peak, 4), + "tl_late_drop": round(peak - probs[-1], 4), + "tl_kl_mid": round(sum(kls) / len(kls), 4), + }) + + return Result( + analyzer=self.name, model=repr(model), cases=cases, + artifacts={"fit_kl_per_layer": fit_kl}, + findings={ + "n_cases_analyzed": len(per_case), + "n_layers": n_layers, + "n_train_tokens": int(data[0].shape[0]), + "fit_kl_per_layer": [round(k, 4) for k in fit_kl], + "fit_kl_mean": round(sum(fit_kl) / max(1, len(fit_kl)), 4), + "final_norm_applied": norm is not None, + "per_case": per_case, + "_caveat": ( + "Translators are fit label-free on token positions of the analysed batch " + "itself (no separate corpus), so they are in-sample for these prompts; " + "fit_kl_per_layer is the held-in KL after fitting." + ), + }, ) + + def _translators(self, model: Any, data: list, W: Any, apply_norm: Any, n_layers: int, + dim: int, device: Any) -> "tuple[list, list[float]]": + """Fit (or load / reuse) one affine translator per premature layer.""" + import os + + import torch + + key = (n_layers, dim, int(data[0].shape[0]), self.steps, self.batch_size, self.lr, self.seed) + module = None + try: + from evalrx.analyzers._internals.access import torch_module + + module = torch_module(model, self.name) + except ValueError: + pass + cache = module.__dict__.setdefault("_evalrx_tuned_lens", {}) if module is not None else {} + if key in cache: + return cache[key] + path = self.translators_path + if path and os.path.exists(path): + saved = torch.load(path, map_location=device) + if saved.get("n_layers") == n_layers and saved.get("dim") == dim: + fitted = ([(a.to(device), b.to(device)) for a, b in saved["translators"]], + list(saved["fit_kl"])) + cache[key] = fitted + return fitted + + gen = torch.Generator().manual_seed(self.seed) + final = data[n_layers] + translators, fit_kl = [], [] + for layer in range(n_layers): + x_all = data[layer] + a = torch.eye(dim, device=device, requires_grad=True) + b = torch.zeros(dim, device=device, requires_grad=True) + opt = torch.optim.Adam([a, b], lr=self.lr) + for _ in range(self.steps): + idx = torch.randint(0, x_all.shape[0], (min(self.batch_size, x_all.shape[0]),), generator=gen) + x, target = x_all[idx].to(device), final[idx].to(device) + with torch.no_grad(): + log_p = torch.log_softmax(target @ W.T, dim=-1) + log_q = torch.log_softmax(apply_norm(x @ a.T + b) @ W.T, dim=-1) + loss = torch.nn.functional.kl_div(log_q, log_p, log_target=True, reduction="batchmean") + opt.zero_grad() + loss.backward() + opt.step() + with torch.no_grad(): + kl = 0.0 + for start in range(0, x_all.shape[0], 1024): + x, target = x_all[start:start + 1024].to(device), final[start:start + 1024].to(device) + log_p = torch.log_softmax(target @ W.T, dim=-1) + log_q = torch.log_softmax(apply_norm(x @ a.T + b) @ W.T, dim=-1) + kl += float(torch.nn.functional.kl_div(log_q, log_p, log_target=True, reduction="sum")) + translators.append((a.detach(), b.detach())) + fit_kl.append(kl / x_all.shape[0]) + fitted = (translators, fit_kl) + cache[key] = fitted + if path: + torch.save({"n_layers": n_layers, "dim": dim, "fit_kl": fit_kl, + "translators": [(a.cpu(), b.cpu()) for a, b in translators]}, path) + return fitted diff --git a/evalrx/analyzers/patching/__init__.py b/evalrx/analyzers/patching/__init__.py index 6e5cd478..cfc6e078 100755 --- a/evalrx/analyzers/patching/__init__.py +++ b/evalrx/analyzers/patching/__init__.py @@ -1,5 +1,7 @@ -"""Causal intervention analyzers (activation patching / causal tracing).""" +"""Causal intervention analyzers (causal tracing / attention knockout / attention reweighting).""" +from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer +from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer -__all__ = ["CausalTraceAnalyzer"] +__all__ = ["AttentionKnockoutAnalyzer", "AttentionSensitivityAnalyzer", "CausalTraceAnalyzer"] diff --git a/evalrx/analyzers/patching/attention_knockout.py b/evalrx/analyzers/patching/attention_knockout.py new file mode 100644 index 00000000..389b5e29 --- /dev/null +++ b/evalrx/analyzers/patching/attention_knockout.py @@ -0,0 +1,181 @@ +"""Attention knockout — block text→media attention and measure what the answer loses. + +Geva et al. trace information flow by preventing chosen query positions from +attending to chosen key positions in a window of layers, then measuring the +change in the prediction. Here the keys are the media placeholder tokens +(image / video frames / audio) and the queries are the text tokens: with every +such edge cut, the text can no longer read the media directly. + +Per case, the next-token margin at the end of the prompt (see +:mod:`evalrx.analyzers._internals.targets`) is measured clean, with the edges +cut in EVERY layer, and with them cut only in the early / middle / late third +of the stack. A large full-stack drop says the answer was built from the media; +a drop near zero says the answer did not come from the media at all — the +language-prior signature. The band columns say WHERE the reading happened. + +Mechanism: :mod:`evalrx.analyzers._internals.attn_tap` routes the decoder's +attention through an eager-equivalent function that masks the cut edges to +−inf before the softmax (an exact knockout, rows stay normalised). Before +reporting anything the analyzer verifies, once per model, that the edit-free +tap engages on every attention layer and reproduces the native logits. + +References: +- Dissecting Recall of Factual Associations in Auto-Regressive Language Models + Geva et al., EMNLP 2023 — arXiv:2304.14767 +- Cross-modal Information Flow in Multimodal Large Language Models + Zhang et al., CVPR 2025 — arXiv:2411.18620 +""" + +from __future__ import annotations + +from collections import Counter +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + +_BANDS = ("early", "mid", "late") + + +def layer_bands(n_layers: int) -> "dict[str, list[int]]": + """Split ``range(n_layers)`` into contiguous early / mid / late thirds.""" + a, b = round(n_layers / 3), round(2 * n_layers / 3) + return {"early": list(range(0, a)), "mid": list(range(a, b)), "late": list(range(b, n_layers))} + + +@register_analyzer("attention_knockout") +class AttentionKnockoutAnalyzer(Analyzer): + """Cut text→media attention (all layers, then per third) and measure the margin drop. + + Hyper-parameters: + target: ``"prediction"`` (clean top-1 vs top-2, label free) or + ``"gold"`` (gold first token vs strongest other token). + queries: ``"text"`` — every non-media position loses the media + (Geva's setting); ``"last"`` — only the final position does. + fidelity_rtol: tolerance of the edit-free tap vs native logits (relative + to the largest logit) before the analyzer refuses. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: 5 logits-only forwards per case. + """ + + name = "attention_knockout" + description = ( + "Blocks the text from attending to the image/audio tokens and measures how " + "much the answer's margin drops; near-zero drop = the answer ignored the media." + ) + requires = frozenset({Capability.ATTENTION, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "ko_margin_clean": ("Clean answer margin", "Logit margin of the target answer token with nothing blocked."), + "ko_delta_full": ("Drop when media blocked", "How much the answer margin falls when the text cannot attend to the media tokens in any layer; near zero means the answer did not use the media."), + "ko_rel_drop_full": ("Relative drop, media blocked", "ko_delta_full as a fraction of the clean margin (1 = the margin vanished entirely)."), + "ko_flips_top1": ("Answer flips w/o media", "Blocking the media changes the model's top next token."), + "ko_delta_early": ("Drop, early layers", "Margin drop when media attention is blocked only in the first third of the layers."), + "ko_delta_mid": ("Drop, middle layers", "Margin drop when media attention is blocked only in the middle third of the layers."), + "ko_delta_late": ("Drop, late layers", "Margin drop when media attention is blocked only in the last third of the layers."), + } + + def __init__(self, target: str = "prediction", queries: str = "text", + fidelity_rtol: float = 1e-2, max_cases: int = 0) -> None: + if queries not in ("text", "last"): + raise ValueError(f"queries must be 'text' or 'last', got {queries!r}") + super().__init__(target=target, queries=queries, fidelity_rtol=fidelity_rtol, + max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + from evalrx.analyzers._internals.access import ( + check_editable_attention, + decoder_stack, + media_positions, + torch_module, + ) + from evalrx.analyzers._internals.attn_tap import attention_tap, knockout, require_fidelity + from evalrx.analyzers._internals.targets import build_margin + from evalrx.core.model import CaptureSpec + + check_editable_attention(model, self.name) + stack = decoder_stack(torch_module(model, self.name)) + if not stack.fully_attentional: + raise ValueError(f"{self.name}: some decoder layers have no softmax attention " + "(hybrid stack); a knockout there would be partial.") + bands = layer_bands(stack.n_layers) + spec = CaptureSpec(to_cpu=False) + + def last_logits(inputs: Any) -> Any: + logits = model.forward(inputs, capture={Capability.LOGITS}, spec=spec).require(Capability.LOGITS) + return logits[-1:].float() + + per_case: list[dict[str, Any]] = [] + skipped: Counter = Counter() + slots: Counter = Counter() + fidelity: "dict | None" = None + for case in cases.stratified_head(self.max_cases): + clean = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + slot, media = media_positions(clean) + if not media: + skipped["no_media_tokens"] += 1 + continue + logits = clean.require(Capability.LOGITS)[-1:].float() + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + if fidelity is None: + fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, + analyzer=self.name) + media_set = set(media) + queries = ([clean.seq_len - 1] if self.queries == "last" + else [i for i in range(clean.seq_len) if i not in media_set]) + + def cut(layers: "list[int] | None") -> Any: + edit = knockout(media, queries=queries, layers=layers) + with attention_tap(model, [edit], analyzer=self.name): + return last_logits(case.inputs) + + m_clean = margin.value(logits) + full = cut(None) + delta_full = m_clean - margin.value(full) + row: dict[str, Any] = { + "sample_id": case.id, + "ko_margin_clean": round(m_clean, 4), + "ko_delta_full": round(delta_full, 4), + "ko_flips_top1": margin.top1_changed(full, int(logits[-1].argmax())), + } + if m_clean > 1e-6: + row["ko_rel_drop_full"] = round(delta_full / m_clean, 4) + for band in _BANDS: + row[f"ko_delta_{band}"] = round(m_clean - margin.value(cut(bands[band])), 4) + per_case.append(row) + slots[slot] += 1 + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": stack.n_layers, + "layer_bands": {k: [v[0], v[-1] + 1] if v else [] for k, v in bands.items()}, + "media_slot": slots.most_common(1)[0][0], + "target": self.target, + "queries": self.queries, + "tap_fidelity": fidelity, + "mean_delta_full": round(sum(r["ko_delta_full"] for r in per_case) / n, 4), + "frac_flips_top1": round(sum(bool(r["ko_flips_top1"]) for r in per_case) / n, 4), + "per_case": per_case, + "_caveat": ( + "Measured on the next-token margin at the end of the prompt (the first " + "answer token, teacher-forced). A band knockout leaves media information " + "already copied into text positions by other layers intact, so band drops " + "need not sum to the full drop." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, findings=findings) diff --git a/evalrx/analyzers/patching/attention_sensitivity.py b/evalrx/analyzers/patching/attention_sensitivity.py new file mode 100644 index 00000000..ce677295 --- /dev/null +++ b/evalrx/analyzers/patching/attention_sensitivity.py @@ -0,0 +1,167 @@ +"""Attention sensitivity — would more (or less) attention to the media change the answer? + +Scales the attention every text position pays to the media tokens by ``w`` +(renormalised: exactly ``+log w`` on the pre-softmax scores, so each row still +sums to one) in a band of layers, and measures the next-token margin at +``w = down`` and ``w = up``. The slope of the margin in ``log w`` says which way +the media pulls the answer: + +* positive — more media attention strengthens the model's answer: the answer is + media-supported; +* negative — more media attention weakens it: the answer is held up by the text + despite the media (the language-prior signature), and amplifying visual + attention — the PAI family of decoding repairs — would push the model off it; +* near zero — the answer is insensitive to how much the media is attended to. + +A diagnostic only: no decoding is changed. With ``target="gold"`` the slope is +for the gold token instead, i.e. whether amplifying the media would move the +model TOWARD the right answer. + +References: +- Paying More Attention to Image: A Training-Free Method for Alleviating + Hallucination in LVLMs (PAI) — Liu et al., ECCV 2024 — arXiv:2407.21771 +- Dissecting Recall of Factual Associations in Auto-Regressive Language Models + Geva et al., EMNLP 2023 — arXiv:2304.14767 +""" + +from __future__ import annotations + +import math +from collections import Counter +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + +_BANDS = ("all", "early", "mid", "late") + + +@register_analyzer("attention_sensitivity") +class AttentionSensitivityAnalyzer(Analyzer): + """Margin response to scaling text→media attention up and down, per case. + + Hyper-parameters: + down / up: attention multipliers (renormalised); defaults 0.5 and 2.0. + band: layers edited — ``"all"`` or the ``"early"`` / ``"mid"`` / + ``"late"`` third of the stack. + target: ``"prediction"`` or ``"gold"``. + fidelity_rtol: tolerance of the edit-free attention tap vs native logits. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: 3 logits-only forwards per case. + """ + + name = "attention_sensitivity" + description = ( + "Scales how much the text attends to the image/audio tokens (×0.5, ×2) and " + "measures which way the answer's margin moves; negative slope = more looking " + "weakens the answer." + ) + requires = frozenset({Capability.ATTENTION, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "rw_margin_clean": ("Clean answer margin", "Logit margin of the target answer token with attention untouched."), + "rw_delta_down": ("Margin change, less looking", "Change in the answer margin when attention to the media is scaled down."), + "rw_delta_up": ("Margin change, more looking", "Change in the answer margin when attention to the media is scaled up."), + "rw_slope": ("Margin per log-attention", "Slope of the answer margin in log(media attention scale); negative = more attention to the media weakens the answer."), + "rw_flips_up": ("Answer flips, more looking", "Scaling media attention up changes the model's top next token."), + } + + def __init__(self, down: float = 0.5, up: float = 2.0, band: str = "all", + target: str = "prediction", fidelity_rtol: float = 1e-2, max_cases: int = 0) -> None: + if band not in _BANDS: + raise ValueError(f"band must be one of {_BANDS}, got {band!r}") + if not (0 < down < 1 < up): + raise ValueError(f"need 0 < down < 1 < up, got down={down}, up={up}") + super().__init__(down=down, up=up, band=band, target=target, + fidelity_rtol=fidelity_rtol, max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + from evalrx.analyzers._internals.access import ( + check_editable_attention, + decoder_stack, + media_positions, + torch_module, + ) + from evalrx.analyzers._internals.attn_tap import attention_tap, require_fidelity, reweight + from evalrx.analyzers._internals.targets import build_margin + from evalrx.analyzers.patching.attention_knockout import layer_bands + from evalrx.core.model import CaptureSpec + + check_editable_attention(model, self.name) + stack = decoder_stack(torch_module(model, self.name)) + if not stack.fully_attentional: + raise ValueError(f"{self.name}: some decoder layers have no softmax attention " + "(hybrid stack); a reweight there would be partial.") + layers = None if self.band == "all" else layer_bands(stack.n_layers)[self.band] + spec = CaptureSpec(to_cpu=False) + log_span = math.log(self.up) - math.log(self.down) + + per_case: list[dict[str, Any]] = [] + skipped: Counter = Counter() + fidelity: "dict | None" = None + for case in cases.stratified_head(self.max_cases): + clean = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + _, media = media_positions(clean) + if not media: + skipped["no_media_tokens"] += 1 + continue + logits = clean.require(Capability.LOGITS)[-1:].float() + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + if fidelity is None: + fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, + analyzer=self.name) + media_set = set(media) + queries = [i for i in range(clean.seq_len) if i not in media_set] + top1 = int(logits[-1].argmax()) + del clean + + def scaled(weight: float) -> Any: + edit = reweight(media, weight, queries=queries, layers=layers) + with attention_tap(model, [edit], analyzer=self.name): + out = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + return out.require(Capability.LOGITS)[-1:].float() + + m_clean = margin.value(logits) + low, high = scaled(self.down), scaled(self.up) + m_low, m_high = margin.value(low), margin.value(high) + per_case.append({ + "sample_id": case.id, + "rw_margin_clean": round(m_clean, 4), + "rw_delta_down": round(m_low - m_clean, 4), + "rw_delta_up": round(m_high - m_clean, 4), + "rw_slope": round((m_high - m_low) / log_span, 4), + "rw_flips_up": margin.top1_changed(high, top1), + }) + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": stack.n_layers, + "band": self.band, + "weights": [self.down, self.up], + "target": self.target, + "tap_fidelity": fidelity, + "mean_slope": round(sum(r["rw_slope"] for r in per_case) / n, 4), + "frac_negative_slope": round(sum(r["rw_slope"] < 0 for r in per_case) / n, 4), + "per_case": per_case, + "_caveat": ( + "Diagnostic only — measured on the next-token margin at the end of the " + "prompt with attention to the media rescaled (rows renormalised) in the " + "chosen band; no decoding was changed." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, findings=findings) diff --git a/evalrx/analyzers/patching/causal_trace.py b/evalrx/analyzers/patching/causal_trace.py index fb8c6ade..f0acae82 100755 --- a/evalrx/analyzers/patching/causal_trace.py +++ b/evalrx/analyzers/patching/causal_trace.py @@ -1,30 +1,345 @@ -"""Activation patching / causal tracing (Meng et al., ROME) (Stage 2). +"""Causal tracing (Meng et al., ROME) over the media tokens of a VLM / ALM. -The only CAUSAL white-box method here: read a clean run's activations and WRITE -them into a corrupted run to localise where information is causally used. -``requires=HIDDEN_STATES`` (read+write hooks — via nnsight). Memory ∝ layers×positions. +The causal method of the family: corrupt the input, then WRITE the clean run's +hidden state back at one (layer, position) site of the corrupted run and see +how much of the answer comes back. Meng et al. corrupt the subject tokens' +embeddings with Gaussian noise; here the corrupted tokens are the media +(image / video / audio) placeholder tokens, and two sites are restored at each +probed layer ``l`` (the residual stream ENTERING layer ``l``): + +* ``media`` — every media token. Restoring them early undoes the corruption; + the layer where this stops working is where the media's content has already + been copied out into the text positions (the hand-off). +* ``last`` — the final prompt position, where the answer is read out. + Restoring it late recovers the answer; the layer where it starts to is where + the media's contribution has arrived there. + +Recovery at a site is ``(m_restored - m_corrupt) / (m_clean - m_corrupt)`` on +the next-token logit margin (:mod:`evalrx.analyzers._internals.targets`). When +corrupting the media barely moves the margin (``|m_clean - m_corrupt| < +min_effect``) there is nothing to recover; those cases still report +``ct_media_effect`` — a near-zero effect is itself the "answer did not depend on +the media" finding — but no recovery columns. + +Corruption: ``"noise"`` (Meng's; Gaussian noise of ``noise_scale`` × the +media embeddings' own std, added to the embedding output — and, on DeepStack +models (Qwen3-VL), to each re-injected copy of the visual features; works for +any media) +or ``"gray"`` (each image replaced by a flat image of its mean colour, same size, +so the token count and M-RoPE grid stay identical — verified per case). References: - Locating and Editing Factual Associations in GPT (ROME, causal tracing) Meng et al., NeurIPS 2022 — arXiv:2202.05262 -- Attribution Patching (scalable approximation): Neel Nanda, 2023. +- Understanding Information Storage and Transfer in Multi-modal Large Language Models + Basu et al., NeurIPS 2024 — arXiv:2406.04236 """ from __future__ import annotations +from collections import Counter +from contextlib import contextmanager, nullcontext +from typing import TYPE_CHECKING, Any, Iterator + from evalrx.core.analyzer import Analyzer from evalrx.core.capability import Capability from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch, Inputs + from evalrx.core.model import Model + +CORRUPTIONS = ("noise", "gray") + + +def _gray(image: Any) -> Any: + """A flat image of *image*'s mean colour, same size (``None`` if undecodable).""" + from PIL import Image, ImageStat + + from evalrx.analyzers.attention.relative_attn import _resolve_image + + img = _resolve_image(image) + if not hasattr(img, "convert"): + return None + rgb = img.convert("RGB") + mean = tuple(int(round(v)) for v in ImageStat.Stat(rgb).mean) + return Image.new("RGB", rgb.size, mean) + + +def gray_inputs(inputs: "Inputs") -> "Inputs | None": + """*inputs* with every image / video frame replaced by its flat mean-colour version.""" + from dataclasses import replace + + def swap(media: Any) -> Any: + if media is None: + return None + items = list(media) if isinstance(media, (list, tuple)) else [media] + grayed = [_gray(m) for m in items] + if any(g is None for g in grayed): + raise ValueError("undecodable image") + return grayed if isinstance(media, (list, tuple)) else grayed[0] + + try: + image, video = swap(getattr(inputs, "image", None)), swap(getattr(inputs, "video", None)) + except ValueError: + return None + if image is None and video is None: + return None + return replace(inputs, image=image, video=video) + + +def _merge(*plans: "dict[int, tuple[list[int], Any]]") -> "dict[int, tuple[list[int], Any]]": + """Combine replace plans; a later plan wins at a shared (layer, position).""" + import torch + + rows: dict[int, dict[int, Any]] = {} + for plan in plans: + for layer, (positions, values) in plan.items(): + slot = rows.setdefault(layer, {}) + for pos, value in zip(positions, values): + slot[int(pos)] = value + return {layer: (list(slot), torch.stack(list(slot.values()))) + for layer, slot in rows.items()} + + +def deepstack_mergers(module: Any) -> list: + """The vision tower's DeepStack mergers (Qwen3-VL family), in injection order. + + DeepStack adds a second copy of the visual features to the decoder's hidden + state after each of the first few layers — outside the embedding, outside + any decoder layer. Corrupting only the embedding would leave those copies + clean, so noise corruption also perturbs these mergers' outputs. + """ + mergers: list = [] + for name, sub in module.named_modules(): + if name.endswith("deepstack_merger_list"): + mergers.extend(list(sub)) + return mergers + + +@contextmanager +def _noised(mergers: list, store: dict, scale: float, gen: Any) -> Iterator[None]: + """Add fixed noise to each merger output; drawn once per variant, then reused.""" + import torch + + def make(i: int): + def hook(mod, args, output): + if i not in store: + std = float(output.detach().float().std()) * scale + store[i] = (torch.randn(output.shape, generator=gen) * std).to(output.device, output.dtype) + return output + store[i] + return hook + + handles = [m.register_forward_hook(make(i)) for i, m in enumerate(mergers)] + try: + yield + finally: + for handle in handles: + handle.remove() @register_analyzer("causal_trace") class CausalTraceAnalyzer(Analyzer): + """Corrupt the media tokens, restore clean states per layer, measure recovery. + + Hyper-parameters: + corruption: ``"noise"`` (default, any media) or ``"gray"`` (images/video). + noise_scale: noise std as a multiple of the media embeddings' std (Meng: 3). + n_noise: noise samples averaged per case (``"noise"`` only). + n_points: probed layers, spread evenly over the stack. + min_effect: smallest ``|m_clean - m_corrupt|`` (logits) worth tracing. Recovery + divides by this effect; on a bf16 model logits are quantised to + ~0.06–0.125 at typical magnitudes, so a smaller effect yields + recoveries dominated by rounding. + target: ``"prediction"`` or ``"gold"``. + seed: noise seed (offset per case, so reruns reproduce). + max_cases: label-stratified cap; 0 (the default) = every case. + Cost per case: 1 + n_noise · (1 + 2 · n_points) logits-only forwards. + """ + name = "causal_trace" - requires = frozenset({Capability.HIDDEN_STATES}) - applies_to_modalities = frozenset({"text", "image"}) + description = ( + "Corrupts the image/audio tokens and restores clean hidden states layer by layer " + "to find where the media's influence on the answer hands off to the text." + ) + requires = frozenset({Capability.HIDDEN_STATES, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "ct_margin_clean": ("Clean answer margin", "Logit margin of the target answer token with the media intact."), + "ct_media_effect": ("Margin lost to corrupt media", "How much the answer margin falls when the media tokens are corrupted; near zero = the answer did not depend on the media."), + "ct_media_effect_rel": ("Relative loss, corrupt media", "ct_media_effect as a fraction of the clean margin."), + "ct_effect_traceable": ("Media effect traceable", "Corrupting the media moved the margin by at least min_effect, so recovery could be traced."), + "ct_media_recovery_mid": ("Recovery via media, mid", "Share of the lost margin recovered by restoring the clean media tokens in the middle third of the layers."), + "ct_media_handoff_frac": ("Media hand-off depth", "Relative depth where restoring the media tokens stops recovering half the margin: the media's content has moved into the text by then."), + "ct_last_rise_frac": ("Answer read-out depth", "Relative depth where restoring the final prompt position first recovers half the margin."), + "ct_last_recovery_max": ("Best recovery at answer pos.", "Largest share of the lost margin recovered by restoring the final prompt position at any probed layer."), + } + + def __init__(self, corruption: str = "noise", noise_scale: float = 3.0, n_noise: int = 1, + n_points: int = 8, min_effect: float = 1.0, target: str = "prediction", + seed: int = 0, max_cases: int = 0) -> None: + if corruption not in CORRUPTIONS: + raise ValueError(f"corruption must be one of {CORRUPTIONS}, got {corruption!r}") + super().__init__(corruption=corruption, noise_scale=noise_scale, n_noise=n_noise, + n_points=n_points, min_effect=min_effect, target=target, seed=seed, + max_cases=max_cases) - def _run(self, model, cases): - raise NotImplementedError( - "Stage 2: clean/corrupt runs + patch cached activations at each (layer, position) " - "via nnsight; report the causal-effect map. Needs read+write hooks (beyond HF flags)." + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + from evalrx.analyzers._internals.access import ( + decoder_stack, + layer_points, + media_positions, + torch_module, ) + from evalrx.analyzers._internals.residual import residual_hooks + from evalrx.analyzers._internals.targets import build_margin + from evalrx.core.model import CaptureSpec + + module = torch_module(model, self.name) + stack = decoder_stack(module) + mergers = deepstack_mergers(module) + n_layers = stack.n_layers + points = layer_points(n_layers, self.n_points) + mid = [p for p in points if n_layers / 3 <= p < 2 * n_layers / 3] or points + spec = CaptureSpec(to_cpu=False) + + def margin_with(variant: "tuple[Any, dict, Any]", margin: Any, restore: "dict | None" = None) -> float: + inputs, plan, (store, gen) = variant + replace = _merge(plan, restore) if restore else plan + noised = _noised(mergers, store, self.noise_scale, gen) if store is not None else nullcontext() + with noised, residual_hooks(model, replace=replace, analyzer=self.name): + logits = model.forward(inputs, capture={Capability.LOGITS}, spec=spec).require( + Capability.LOGITS) + return margin.value(logits[-1:]) + + per_case: list[dict[str, Any]] = [] + curves: dict[str, dict[str, list[float]]] = {} + skipped: Counter = Counter() + slots: Counter = Counter() + sites = sorted(set(points) | {0}) + for index, case in enumerate(cases.stratified_head(self.max_cases)): + with residual_hooks(model, capture={layer: None for layer in sites}, + analyzer=self.name) as clean: + probe = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + slot, media = media_positions(probe) + if not media: + skipped["no_media_tokens"] += 1 + continue + logits = probe.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + m_clean = margin.value(logits) + last = probe.seq_len - 1 + clean_media = {layer: (media, clean.captured[layer][media]) for layer in sites} + clean_last = {layer: ([last], clean.captured[layer][[last]]) for layer in sites} + probe.logits = None # keep token ids for the alignment check, not the (seq, vocab) logits + del clean, logits + + # One corrupted variant per noise sample (or the single gray image): + # (inputs, layer-0 replace plan, (DeepStack noise store, generator)). + variants: list[tuple[Any, dict, Any]] = [] + if self.corruption == "gray": + grayed = gray_inputs(case.inputs) if slot == "image" else None + if grayed is None: + skipped["gray_needs_an_image"] += 1 + continue + check = model.forward(grayed, capture={Capability.LOGITS}, spec=spec) + if check.token_ids != probe.token_ids: + skipped["corruption_misaligned"] += 1 + continue + variants.append((grayed, {}, (None, None))) + else: + emb = clean_media[0][1] + std = float(emb.float().std()) * self.noise_scale + gen = torch.Generator(device="cpu").manual_seed(self.seed * 100003 + index) + for _ in range(max(1, self.n_noise)): + noise = torch.randn(emb.shape, generator=gen).to(emb.device, emb.dtype) * std + store = {} if mergers and slot == "image" else None + variants.append((case.inputs, {0: (media, emb + noise)}, (store, gen))) + + m_corrupt = sum(margin_with(v, margin) for v in variants) / len(variants) + effect = m_clean - m_corrupt + row: dict[str, Any] = { + "sample_id": case.id, + "ct_margin_clean": round(m_clean, 4), + "ct_media_effect": round(effect, 4), + "ct_effect_traceable": abs(effect) >= self.min_effect, + } + if m_clean > 1e-6: + row["ct_media_effect_rel"] = round(effect / m_clean, 4) + slots[slot] += 1 + if abs(effect) < self.min_effect: + per_case.append(row) + continue + + def recovery(restore: dict) -> float: + restored = sum(margin_with(v, margin, restore) for v in variants) / len(variants) + return (restored - m_corrupt) / effect + + media_curve = [recovery({p: clean_media[p]}) for p in points] + last_curve = [recovery({p: clean_last[p]}) for p in points] + depth = [p / max(1, n_layers - 1) for p in points] + row["ct_media_recovery_mid"] = round( + sum(r for p, r in zip(points, media_curve) if p in mid) / len(mid), 4) + # On DeepStack models a media restore before the last re-injection is + # partial by construction (later injections stay corrupted), so the + # hand-off is only looked for from that depth on. + row["ct_media_handoff_frac"] = round( + next((d for p, d, r in zip(points, depth, media_curve) + if p >= len(mergers) and r < 0.5), 1.0), 4) + row["ct_last_rise_frac"] = round( + next((d for d, r in zip(depth, last_curve) if r >= 0.5), 1.0), 4) + row["ct_last_recovery_max"] = round(max(last_curve), 4) + per_case.append(row) + curves[case.id] = {"media": [round(r, 4) for r in media_curve], + "last": [round(r, 4) for r in last_curve]} + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + label_of = {c.id: getattr(getattr(c, "label", None), "value", None) for c in cases} + + def mean_curve(label: str, site: str) -> "list[float] | None": + rows = [c[site] for cid, c in curves.items() if label_of.get(cid) == label] + if not rows: + return None + return [round(sum(col) / len(rows), 4) for col in zip(*rows)] + + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_traced": len(curves), + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": n_layers, + "layers_probed": points, + "media_slot": slots.most_common(1)[0][0] if slots else None, + "corruption": self.corruption, + "deepstack_layers": len(mergers), + "target": self.target, + "mean_media_effect": round(sum(r["ct_media_effect"] for r in per_case) / n, 4), + "frac_traceable": round(len(curves) / n, 4), + "mean_recovery_curves": { + f"{label}_{site}": mean_curve(label, site) + for label in ("fail", "pass") for site in ("media", "last") + if mean_curve(label, site) is not None + }, + "per_case": per_case, + "_caveat": ( + "Recovery is measured on the next-token margin at the end of the prompt. " + "Sites are the residual stream ENTERING each probed layer; restoring the " + "media tokens at layer 0 undoes the corruption by construction" + + (" except on DeepStack models, where visual features are re-injected " + f"after the first {len(mergers)} layers: restoring the media before that " + "depth is partial because later injections stay corrupted." + if mergers else ".") + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, + artifacts={"recovery_curves": curves, "layers_probed": points}, + findings=findings) diff --git a/evalrx/models/backends/hf/model.py b/evalrx/models/backends/hf/model.py index 0b048781..d2716677 100755 --- a/evalrx/models/backends/hf/model.py +++ b/evalrx/models/backends/hf/model.py @@ -6,8 +6,9 @@ attention, residual-stream hidden states and logits with NO module-path surgery — this is the high-value common path. * **hooks + runtime path discovery** (:mod:`evalrx.models.backends.hf.discover`) are - only needed beyond what flags give (activation patching, MoE routing, - gradients) — reserved for Stage 2. + only needed beyond what flags give. The write-side analyzers (attention + knockout, causal tracing, gradient attribution) install them themselves + around an ordinary ``forward`` call — see :mod:`evalrx.analyzers._internals`. Attention capture REQUIRES eager attention (sdpa/flash silently return ``None``), so when the model declares ``eager_required_for_attn`` we force @@ -71,6 +72,10 @@ def __init__(self, spec, runtime: RuntimeConfig) -> None: Capability.LOGITS, Capability.LOGPROBS, Capability.HIDDEN_STATES, + # forward(spec=CaptureSpec(with_grad=True)) keeps the autograd graph; + # gradient analyzers differentiate a logit target w.r.t. hooked + # activations (torch.autograd.grad — no parameter .grad is written). + Capability.GRADIENTS, } if spec.attn_semantics is not AttnSemantics.NONE: caps.add(Capability.ATTENTION) @@ -605,7 +610,11 @@ def forward(self, inputs: Any, capture: set[Capability], spec=None) -> Trace: flags = {flag: True for cap, flag in _CAPTURE_FLAGS.items() if cap in capture} enc.pop("token_type_ids", None) # some VLM processors emit this; forward() rejects it - with torch.no_grad(): + # GRADIENTS: CaptureSpec.with_grad keeps the graph so a caller can + # differentiate the returned logits (pair it with to_cpu=False). Default + # stays no_grad: an attention capture with a live graph costs gigabytes. + with_grad = bool(spec is not None and getattr(spec, "with_grad", False)) + with torch.set_grad_enabled(with_grad): outputs = model(**enc, **flags) layers = spec.layers if spec is not None else None @@ -665,6 +674,7 @@ class HFLocalBackend(Backend): Capability.LOGPROBS, Capability.HIDDEN_STATES, Capability.ATTENTION, + Capability.GRADIENTS, } ) diff --git a/tests/test_analyzers/test_analyzer_contract.py b/tests/test_analyzers/test_analyzer_contract.py index 50ee40e4..9f129549 100755 --- a/tests/test_analyzers/test_analyzer_contract.py +++ b/tests/test_analyzers/test_analyzer_contract.py @@ -41,7 +41,6 @@ from evalrx.analyzers.lens.layer_contrast import LayerContrastAnalyzer from evalrx.analyzers.lens.logit_lens import LogitLensAnalyzer from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer -from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer from evalrx.analyzers.perturbation.context_shap import ContextShapAnalyzer from evalrx.analyzers.perturbation.cot_faithfulness import CoTFaithfulnessAnalyzer from evalrx.analyzers.perturbation.format_sensitivity import FormatSensitivityAnalyzer @@ -218,11 +217,18 @@ def _long_prompt_batch() -> CaseBatch: # Requires a live PIL image; no FakeModel path. "relative_attention": "VLM-only: needs case.inputs.image (PIL image)", "vl_shap": "VLM-only: needs case.inputs.image (PIL image)", - # Requires a backward pass; Capability.GRADIENTS is not available in FakeModel. - "gradcam": "needs Capability.GRADIENTS (backward pass)", - "generic_attention": "needs Capability.GRADIENTS (Chefer relevance propagation)", + # Requires a backward pass through a live torch module (FakeModel has neither). + "gradcam": "needs a live torch module + GRADIENTS; see test_intervention_analyzers.py", + "generic_attention": "needs a live torch module + GRADIENTS; see test_intervention_analyzers.py", # score_fn is mandatory at run time with no useful default. "rise": "requires a user-supplied score_fn(output)->float; no sensible default", + # Write-side white-box analyzers act on a live torch module (hooks / an + # attention function), which FakeModel does not have. Their result contract + # runs against a tiny real Qwen-VL in test_intervention_analyzers.py. + "attention_knockout": "needs a live torch module; see test_intervention_analyzers.py", + "causal_trace": "needs a live torch module; see test_intervention_analyzers.py", + "attention_sensitivity": "needs a live torch module; see test_intervention_analyzers.py", + "head_attribution": "needs a live torch module; see test_intervention_analyzers.py", } @@ -296,6 +302,7 @@ def _context_batch() -> CaseBatch: (AttentionSinkAnalyzer(), _FULL, _STANDARD), # lens (LogitLensAnalyzer(), _FULL, _STANDARD), + (TunedLensAnalyzer(steps=20), _FULL, _STANDARD), # uncertainty (TokenEntropyAnalyzer(), _FULL, _STANDARD), (LogprobEntropyAnalyzer(), _FULL, _STANDARD), @@ -401,11 +408,17 @@ def _context_batch() -> CaseBatch: _EXPECTED_FINDING_KEYS: dict[str, set[str]] = { "attention": { "num_layers", "num_heads", "seq_len", "summary_layer", "summary_head", - "top_attended_tokens", "mean_attention_entropy", + "top_attended_tokens", "mean_attention_entropy", "n_cases_analyzed", "per_case", }, "attention_rollout": {"seq_len", "n_layers", "top_rollout_tokens"}, - "attention_sink": {"n_layers", "sink_token", "mean_sink_mass", "per_layer_sink"}, + "attention_sink": { + "n_layers", "sink_token", "mean_sink_mass", "per_layer_sink", "n_cases_analyzed", "per_case", + }, "logit_lens": {"n_layers", "pos", "per_layer_top", "per_case", "final_norm_applied"}, + "tuned_lens": { + "n_cases_analyzed", "n_layers", "n_train_tokens", "fit_kl_per_layer", "fit_kl_mean", + "final_norm_applied", "per_case", + }, "linear_probe": { "n_layers", "n_fail", "n_pass", "per_layer_accuracy", "best_layer", "best_accuracy", "per_case", @@ -462,6 +475,10 @@ def _check_attention(f: dict[str, Any]) -> None: assert f["seq_len"] == 5 assert len(f["top_attended_tokens"]) == 5 assert f["mean_attention_entropy"] >= 0 + assert f["n_cases_analyzed"] == len(f["per_case"]) == 1 + row = f["per_case"][0] + assert _between(row["last_query_entropy"], 0, 1) + assert "media_mass_mid" not in row # text-only trace: no media column, not a 0.0 def _check_attention_rollout(f: dict[str, Any]) -> None: @@ -474,6 +491,10 @@ def _check_attention_sink(f: dict[str, Any]) -> None: assert f["n_layers"] == 3 assert len(f["per_layer_sink"]) == 3 assert _between(f["mean_sink_mass"], 0, 1) + assert f["n_cases_analyzed"] == len(f["per_case"]) == 1 + row = f["per_case"][0] + assert _between(row["sink_mass_mean"], 0, 1) + assert _between(row["sink_peak_layer_frac"], 0, 1) def _check_logit_lens(f: dict[str, Any]) -> None: @@ -486,6 +507,16 @@ def _check_logit_lens(f: dict[str, Any]) -> None: assert pc["late_drop"] >= 0 +def _check_tuned_lens(f: dict[str, Any]) -> None: + assert f["n_layers"] == 3 # FakeModel: 4 hidden-state entries = 3 premature + final + assert len(f["fit_kl_per_layer"]) == 3 and all(k >= 0 for k in f["fit_kl_per_layer"]) + assert f["n_train_tokens"] == 5 # one short case: every position sampled + pc = f["per_case"][0] + assert 0 <= pc["tl_decision_frac"] <= 1 + assert _between(pc["tl_final_top1_prob"], 0, 1) and pc["tl_late_drop"] >= 0 + assert pc["tl_kl_mid"] >= 0 + + def _check_linear_probe(f: dict[str, Any]) -> None: assert f["n_layers"] == 4 assert len(f["per_layer_accuracy"]) == 4 @@ -781,6 +812,7 @@ def _check_coverage_verification_gap(f): "attention_rollout": _check_attention_rollout, "attention_sink": _check_attention_sink, "logit_lens": _check_logit_lens, + "tuned_lens": _check_tuned_lens, "linear_probe": _check_linear_probe, "token_entropy": _check_token_entropy, "logprob_entropy": _check_logprob_entropy, @@ -821,8 +853,6 @@ def _check_coverage_verification_gap(f): # These analyzers are fully registered and their class contracts are sound, but # _run raises NotImplementedError. When you implement one, move it to _RUNNABLE. _STUBS: list[tuple[Any, Any, Any]] = [ - (TunedLensAnalyzer(), _FULL, _STANDARD), - (CausalTraceAnalyzer(), _FULL, _STANDARD), # OPERA and VCD are image-only stubs; _FULL satisfies their capability check # (ATTENTION and LOGITS respectively), so _run is reached and raises. (OPERAAnalyzer(), _FULL, _STANDARD), diff --git a/tests/test_analyzers/test_analyzer_docs.py b/tests/test_analyzers/test_analyzer_docs.py index c6862938..59280957 100644 --- a/tests/test_analyzers/test_analyzer_docs.py +++ b/tests/test_analyzers/test_analyzer_docs.py @@ -32,6 +32,14 @@ def gen(): return module +# Every other analyzer registered under patching/, so a page built around one +# row does not report its siblings as "added". +_SIBLING_ROWS = ( + "| Attention knockout | `attention_knockout` | | | — | |\n" + "| Attention sensitivity | `attention_sensitivity` | | | — | |\n" +) + + def _patching_page(rows: str) -> str: return ( "### Patching\n\n" @@ -49,17 +57,20 @@ def test_analyzer_doc_is_current(gen): def test_curated_cells_survive_and_derived_cells_are_rewritten(gen): - page = _patching_page("| My trace | `causal_trace` | WRONG | WRONG | [P](https://x) | ✓ |\n") + page = _patching_page( + "| My trace | `causal_trace` | WRONG | WRONG | [P](https://x) | stub |\n" + _SIBLING_ROWS + ) fresh, added = gen.render(page, require_all=False) assert ( - "| My trace | `causal_trace` | `HIDDEN_STATES` | text + image | [P](https://x) | stub |" + "| My trace | `causal_trace` | `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** " + "| [P](https://x) | ✓ |" in fresh ) assert added == [] def test_new_key_gets_a_placeholder_row_and_a_stale_key_is_dropped(gen): - page = _patching_page("| Gone | `no_such_analyzer` | | | — | |\n") + page = _patching_page("| Gone | `no_such_analyzer` | | | — | |\n" + _SIBLING_ROWS) fresh, added = gen.render(page, require_all=False) assert "no_such_analyzer" not in fresh assert "`causal_trace`" in fresh diff --git a/tests/test_analyzers/test_intervention_analyzers.py b/tests/test_analyzers/test_intervention_analyzers.py new file mode 100644 index 00000000..099b8ff1 --- /dev/null +++ b/tests/test_analyzers/test_intervention_analyzers.py @@ -0,0 +1,521 @@ +"""Write-side white-box analyzers on a tiny real Qwen-VL (CPU, random init). + +Two layers of checks: + +* invariants of the shared plumbing (``analyzers/_internals``): the edit-free + attention tap is exact, keeps the causal mask, never retargets the vision + tower, and refuses when it cannot engage; edits touch exactly the cells they + name; residual hooks capture/replace what they say; +* the analyzers' result contract: a ``Result`` with JSON-serialisable findings, + one ``per_case`` row per analysable case, and loud refusal (never a silent + zero) on architectures or batches they cannot measure. +""" + +from __future__ import annotations + +import json +import math + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("transformers") +pytest.importorskip("PIL") + +from evalrx.analyzers._internals import access, attn_tap # noqa: E402 +from evalrx.analyzers._internals.attn_tap import ( # noqa: E402 + attention_tap, + check_fidelity, + knockout, + require_fidelity, + reweight, +) +from evalrx.analyzers._internals.residual import residual_hooks # noqa: E402 +from evalrx.analyzers._internals.targets import build_margin, gold_text # noqa: E402 +from evalrx.core.capability import Capability # noqa: E402 +from evalrx.core.case import CaseBatch, FailureCase, Inputs, Label # noqa: E402 +from evalrx.core.result import Result # noqa: E402 +from evalrx.core.spec import AttnSemantics # noqa: E402 +from tests.tiny_whitebox import IMAGE_TOKEN, VEND, VSTART, TinyVLM, image # noqa: E402 + +FAMILIES = ("qwen2_5_vl", "qwen3_vl") + + +@pytest.fixture(scope="module", params=FAMILIES) +def vlm(request): + return TinyVLM(request.param) + + +def _batch(n: int = 4, with_gold: bool = True) -> CaseBatch: + return CaseBatch([ + FailureCase( + inputs=Inputs(prompt=f"what is in picture {i}?", image=image(i)), + expected="cat" if with_gold else None, + label=Label.FAIL if i % 2 else Label.PASS, + ) + for i in range(n) + ]) + + +def _raw_inputs(seq_tail: list[int]) -> dict: + """Direct module inputs (bypassing the Model) for mask-level checks.""" + from tests.tiny_whitebox import GRID, _pixel_values + + ids = [1, 2, VSTART] + [IMAGE_TOKEN] * 16 + [VEND] + seq_tail + return dict(input_ids=torch.tensor([ids]), pixel_values=_pixel_values(image(0)), + image_grid_thw=torch.tensor([[1, GRID, GRID]])) + + +# ── attention tap ────────────────────────────────────────────────────────────── + +def test_edit_free_tap_is_exact_and_engages_every_layer(vlm): + verdict = check_fidelity(vlm, Inputs(prompt="describe", image=image(0))) + assert verdict["tap_faithful"] is True + assert verdict["tap_layers_engaged"] == verdict["tap_layers_expected"] == 6 + assert verdict["tap_max_logit_diff"] == 0.0 # same ops as eager, bit for bit + + +def test_tap_keeps_the_causal_mask(vlm): + # transformers builds NO mask for an unknown attention name; the tap must + # register one, or earlier positions would see a changed future token. + module = vlm._loaded[0] + a, b = _raw_inputs([7, 8, 9, 10]), _raw_inputs([7, 8, 9, 11]) + with torch.no_grad(), attention_tap(vlm): + la, lb = module(**a).logits, module(**b).logits + assert torch.equal(la[0, :-1], lb[0, :-1]) + assert not torch.equal(la[0, -1], lb[0, -1]) + + +def test_tap_retargets_only_the_decoder_and_restores_on_error(vlm): + cfg = vlm._loaded[0].config + vision_before, text_before = cfg.vision_config._attn_implementation, cfg.text_config._attn_implementation + with pytest.raises(RuntimeError, match="boom"): + with attention_tap(vlm): + assert cfg.text_config._attn_implementation == attn_tap.TAP_NAME + assert cfg.vision_config._attn_implementation == vision_before + raise RuntimeError("boom") + assert cfg.text_config._attn_implementation == text_before + assert attn_tap.active_state() is None + + +def test_knockout_touches_only_downstream_rows(vlm): + module = vlm._loaded[0] + enc = _raw_inputs([7, 8, 9, 10]) + image_pos = list(range(3, 19)) + text_q = [i for i in range(enc["input_ids"].shape[1]) if i not in image_pos] + with torch.no_grad(): + base = module(**enc).logits + with attention_tap(vlm, [knockout(image_pos, queries=text_q)]): + cut = module(**enc).logits + assert torch.equal(cut[0, :3], base[0, :3]) # before the image: nothing to cut + assert not torch.allclose(cut[0, -1], base[0, -1]) # the answer position lost the image + + +def test_reweight_equivalences(vlm): + module = vlm._loaded[0] + enc = _raw_inputs([7, 8, 9, 10]) + keys = list(range(3, 19)) + with torch.no_grad(): + base = module(**enc).logits + with attention_tap(vlm, [reweight(keys, 1.0)]): + same = module(**enc).logits + with attention_tap(vlm, [reweight(keys, 0.0)]): + zero = module(**enc).logits + with attention_tap(vlm, [knockout(keys)]): + cut = module(**enc).logits + assert torch.equal(same, base) + assert torch.equal(zero, cut) # scale by 0 IS a knockout + + +def test_scale_raw_does_not_renormalise(vlm): + rows = {} + + def grab(layer, probs): + rows[layer] = probs[0, :, -1].sum(-1) # per-head row sums at the last query + + with torch.no_grad(), attention_tap(vlm, [reweight(range(3, 19), 3.0, renormalize=False)], + on_probs=grab): + vlm.forward(Inputs(prompt="hi", image=image(0)), capture={Capability.LOGITS}) + assert set(rows) == set(range(6)) + assert all(bool((r > 1.0 + 1e-3).all()) for r in rows.values()) + + +def test_require_fidelity_refuses_a_bypassed_layer(vlm, monkeypatch): + real = access.decoder_stack + + def with_phantom(module): + stack = real(module) + stack.attn = stack.attn + [torch.nn.Identity()] # a layer the tap never sees + stack.layers = list(stack.layers) + [torch.nn.Identity()] + return stack + + monkeypatch.setattr(attn_tap, "decoder_stack", with_phantom) + module = vlm._loaded[0] + saved = module.__dict__.pop("_evalrx_tap_fidelity", None) + try: + with pytest.raises(RuntimeError, match="does not faithfully reproduce"): + require_fidelity(vlm, Inputs(prompt="x", image=image(0))) + finally: + module.__dict__.pop("_evalrx_tap_fidelity", None) + if saved is not None: + module.__dict__["_evalrx_tap_fidelity"] = saved + + +# ── residual hooks ───────────────────────────────────────────────────────────── + +def test_residual_replace_with_own_capture_is_a_noop(vlm): + inputs = Inputs(prompt="what?", image=image(1)) + pos = list(range(3, 19)) + with torch.no_grad(): + with residual_hooks(vlm, capture={2: pos}) as s: + base = vlm.forward(inputs, capture={Capability.LOGITS}).logits + with residual_hooks(vlm, replace={2: (pos, s.captured[2])}): + again = vlm.forward(inputs, capture={Capability.LOGITS}).logits + assert s.captured[2].shape == (16, 32) + assert torch.equal(base, again) + + +def test_residual_replace_changes_only_downstream(vlm): + inputs = Inputs(prompt="what?", image=image(1)) + with torch.no_grad(): + base = vlm.forward(inputs, capture={Capability.LOGITS}).logits + with residual_hooks(vlm, replace={3: ([5], torch.zeros(1, 32))}): + patched = vlm.forward(inputs, capture={Capability.LOGITS}).logits + assert torch.equal(patched[:5], base[:5]) + assert not torch.allclose(patched[-1], base[-1]) + + +# ── targets ──────────────────────────────────────────────────────────────────── + +def test_gold_text_shapes(): + assert gold_text("cat") == "cat" + assert gold_text(["cat", "kitty"]) == "cat" + assert gold_text({"all_of": ["yes"], "none_of": ["no"]}) == "yes" + assert gold_text(None) is None and gold_text(" ") is None + + +def test_prediction_margin_is_top1_minus_top2(): + logits = torch.tensor([[0.0, 3.0, 1.0, 2.5]]) + margin, why = build_margin("prediction", logits) + assert why == "" and (margin.positive, margin.negative) == (1, 3) + assert math.isclose(margin.value(logits), 0.5) + + +def test_gold_margin_needs_gold_and_tokenizer(vlm): + logits = torch.randn(1, 128) + no_gold = FailureCase(inputs=Inputs(prompt="x")) + assert build_margin("gold", logits, model=vlm, case=no_gold) == (None, "no_gold") + margin, why = build_margin("gold", logits, model=vlm, case=FailureCase(inputs=Inputs(prompt="x"), + expected="cat")) + assert why == "" and margin.positive == ord("c") % 90 + 1 and margin.negative != margin.positive + + +# ── analyzers ───────────────────────────────────────────────────────────────── + +def _assert_result_contract(result, name: str, n_rows: int) -> dict: + assert isinstance(result, Result) and result.analyzer == name + json.dumps(result.findings) + assert isinstance(result.summary(), str) + rows = result.findings["per_case"] + assert len(rows) == n_rows + for row in rows: + assert isinstance(row["sample_id"], str) + for key, value in row.items(): + if isinstance(value, float): + assert math.isfinite(value), (key, value) + return result.findings + + +def test_attention_knockout_contract(vlm): + from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer + + f = _assert_result_contract(AttentionKnockoutAnalyzer().run(vlm, _batch()), "attention_knockout", 4) + assert f["tap_fidelity"]["tap_faithful"] is True + assert f["media_slot"] == "image" and f["n_layers"] == 6 + assert f["layer_bands"] == {"early": [0, 2], "mid": [2, 4], "late": [4, 6]} + for row in f["per_case"]: + assert row["ko_margin_clean"] >= 0 # prediction target: top1 - top2 + assert isinstance(row["ko_flips_top1"], bool) + assert {"ko_delta_full", "ko_delta_early", "ko_delta_mid", "ko_delta_late"} <= row.keys() + assert any(abs(r["ko_delta_full"]) > 0 for r in f["per_case"]) + + +def test_attention_knockout_last_query_and_gold_target(vlm): + from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer + + batch = CaseBatch(list(_batch(2)) + list(_batch(2, with_gold=False))) + f = AttentionKnockoutAnalyzer(target="gold", queries="last").run(vlm, batch).findings + assert f["n_cases_analyzed"] == 2 + assert f["skip_reasons"] == {"no_gold": 2} + + +def test_attention_knockout_refuses_text_only_batch(vlm): + from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer + + with pytest.raises(ValueError, match="no_media_tokens"): + AttentionKnockoutAnalyzer().run(vlm, CaseBatch([FailureCase(inputs=Inputs(prompt="hi"))])) + + +def test_attention_knockout_refuses_hybrid_attention(): + from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer + + hybrid = TinyVLM("qwen2_5_vl", attn_semantics=AttnSemantics.HYBRID_SPARSE) + with pytest.raises(ValueError, match="hybrid_sparse"): + AttentionKnockoutAnalyzer().run(hybrid, _batch(1)) + + +def test_attention_summary_reports_media_mass(vlm): + from evalrx.analyzers.attention.summary import AttentionAnalyzer + + f = _assert_result_contract(AttentionAnalyzer().run(vlm, _batch(3)), "attention", 3) + assert all(0.0 <= r["media_mass_mid"] <= 1.0 for r in f["per_case"]) + + +# ── gradient path ────────────────────────────────────────────────────────────── + +def test_with_grad_forward_keeps_a_graph_and_writes_no_parameter_grads(vlm): + from evalrx.core.model import CaptureSpec + + module = vlm._loaded[0] + with residual_hooks(vlm, capture={3: None}, keep_grad=True) as s: + trace = vlm.forward(Inputs(prompt="q", image=image(0)), capture={Capability.LOGITS}, + spec=CaptureSpec(with_grad=True, to_cpu=False)) + (grad,) = torch.autograd.grad(trace.logits[-1].sum(), s.captured[3]) + assert grad.shape == s.captured[3].shape and bool(grad.abs().sum() > 0) + assert all(p.grad is None for p in module.parameters()) + plain = vlm.forward(Inputs(prompt="q", image=image(0)), capture={Capability.LOGITS}) + assert not plain.logits.requires_grad + + +def test_hf_local_honours_with_grad(): + # The backend contract the gradient analyzers rely on (no weights loaded). + import inspect + + from evalrx.models.backends.hf.model import HFLocalBackend, HFLocalModel + + assert Capability.GRADIENTS in HFLocalBackend.capabilities + assert "set_grad_enabled(with_grad)" in inspect.getsource(HFLocalModel.forward) + + +def test_gradcam_contract(vlm): + from evalrx.analyzers.attribution.gradcam import GradCAMAnalyzer + + result = GradCAMAnalyzer().run(vlm, _batch()) + f = _assert_result_contract(result, "gradcam", 4) + assert f["layer_used"] == 2 and f["map_shape"] == [4, 4] + for row in f["per_case"]: + assert 0.0 <= row["gc_media_grad_share"] <= 1.0 + assert isinstance(row["gc_empty"], bool) + if not row["gc_empty"]: + assert 0.0 <= row["gc_entropy"] <= 1.0 and 0.0 < row["gc_top5_share"] <= 1.0 + assert set(result.findings["per_case"][0]) >= {"gc_margin", "gc_media_grad_share"} + assert {"spatial_map", "per_case_maps", "fail_mean_map", "pass_mean_map", + "diff_map_fail_minus_pass"} <= set(result.artifacts) + assert result.findings["spatial_map_case_id"] == f["per_case"][0]["sample_id"] + + +def test_gradcam_overlays_land_on_the_case_image(vlm, tmp_path): + pytest.importorskip("matplotlib") + from evalrx.analyzers.attribution.gradcam import GradCAMAnalyzer + + paths = GradCAMAnalyzer().run(vlm, _batch()).image_overlays(tmp_path, "c0_gradcam") + assert {p.name for p in paths} >= {"c0_gradcam_spatial_map_overlay.png", + "c0_gradcam_diff_map_fail_minus_pass_overlay.png"} + + +def test_generic_attention_contract(vlm): + from evalrx.analyzers.attribution.generic_attn import GenericAttentionExplainability + + result = GenericAttentionExplainability().run(vlm, _batch()) + f = _assert_result_contract(result, "generic_attention", 4) + assert f["tap_fidelity"]["tap_faithful"] is True + assert len(f["top_relevant_tokens"]) == 10 + for row in f["per_case"]: + assert 0.0 <= row["rel_media_share"] <= 1.0 + assert 0.0 <= row["rel_input_entropy"] <= 1.0 + assert "spatial_map" in result.artifacts + + +def test_generic_attention_runs_text_only(vlm): + from evalrx.analyzers.attribution.generic_attn import GenericAttentionExplainability + + f = GenericAttentionExplainability().run(vlm, CaseBatch([FailureCase(inputs=Inputs(prompt="hello there"))])).findings + row = f["per_case"][0] + assert "rel_media_share" not in row and "rel_input_entropy" in row + assert f["mean_media_share"] is None + + +# ── causal tracing ───────────────────────────────────────────────────────────── + +def test_full_restore_after_deepstack_returns_the_clean_logits(): + # On Qwen3-VL the visual features are re-injected after layers 0..k-1; a + # complete residual restore at layer >= k must still give back the clean run. + from evalrx.analyzers.patching.causal_trace import _noised, deepstack_mergers + + vlm = TinyVLM("qwen3_vl") + mergers = deepstack_mergers(vlm._loaded[0]) + assert len(mergers) == 2 + inputs = Inputs(prompt="what?", image=image(2)) + with torch.no_grad(): + with residual_hooks(vlm, capture={2: None}) as clean: + base = vlm.forward(inputs, capture={Capability.LOGITS}).logits[-1] + gen = torch.Generator().manual_seed(0) + store: dict = {} + with _noised(mergers, store, 3.0, gen): + corrupt = vlm.forward(inputs, capture={Capability.LOGITS}).logits[-1] + with residual_hooks(vlm, replace={2: (list(range(clean.captured[2].shape[0])), clean.captured[2])}): + restored = vlm.forward(inputs, capture={Capability.LOGITS}).logits[-1] + assert len(store) == 2 + assert not torch.allclose(corrupt, base) + assert torch.allclose(restored, base, atol=1e-6) + + +@pytest.mark.parametrize("corruption", ["noise", "gray"]) +def test_causal_trace_contract(vlm, corruption): + from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer + + f = _assert_result_contract( + CausalTraceAnalyzer(corruption=corruption, min_effect=0.0, n_points=3).run(vlm, _batch()), + "causal_trace", 4, + ) + assert f["layers_probed"] == [0, 2, 5] and f["n_traced"] == 4 + assert f["deepstack_layers"] == (2 if vlm.spec.key.endswith("qwen3_vl") else 0) + assert set(f["mean_recovery_curves"]) == {"fail_media", "fail_last", "pass_media", "pass_last"} + for row in f["per_case"]: + assert row["ct_effect_traceable"] is True + assert 0.0 <= row["ct_media_handoff_frac"] <= 1.0 + assert 0.0 <= row["ct_last_rise_frac"] <= 1.0 + if f["deepstack_layers"]: # restores before the last re-injection are partial by construction + assert row["ct_media_handoff_frac"] >= 2 / 5 + + +def test_causal_trace_noise_restore_at_layer0_is_exact_without_deepstack(): + from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer + + vlm = TinyVLM("qwen2_5_vl") + result = CausalTraceAnalyzer(min_effect=0.0, n_points=3).run(vlm, _batch(2)) + for curve in result.artifacts["recovery_curves"].values(): + assert curve["media"][0] == pytest.approx(1.0, abs=1e-3) + assert curve["last"][0] == pytest.approx(0.0, abs=1e-3) + + +def test_causal_trace_small_effects_keep_the_row_but_skip_the_trace(vlm): + from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer + + f = CausalTraceAnalyzer(n_points=3).run(vlm, _batch(2)).findings # tiny model: effects << 1 logit + assert f["n_cases_analyzed"] == 2 and f["n_traced"] == 0 + for row in f["per_case"]: + assert row["ct_effect_traceable"] is False + assert "ct_media_recovery_mid" not in row # no recovery number made up + + +def test_gray_inputs_keep_size_and_flatten_content(): + import numpy as np + + from evalrx.analyzers.patching.causal_trace import gray_inputs + + grayed = gray_inputs(Inputs(prompt="x", image=image(3))) + arr = np.asarray(grayed.image) + assert grayed.image.size == image(3).size and (arr == arr[0, 0]).all() + assert gray_inputs(Inputs(prompt="x")) is None + + +# ── head attribution ────────────────────────────────────────────────────────── + +def test_per_head_writes_sum_to_the_output_projection(vlm): + from evalrx.analyzers._internals.heads import head_capture, n_heads_of, out_proj + + stack = access.decoder_stack(vlm._loaded[0]) + outputs = {} + handles = [out_proj(a).register_forward_hook( + lambda mod, args, out, i=i: outputs.__setitem__(i, out[0, -1].detach())) + for i, a in enumerate(stack.attn)] + try: + with torch.no_grad(), head_capture(vlm) as cap: + vlm.forward(Inputs(prompt="q", image=image(0)), capture={Capability.LOGITS}) + finally: + for h in handles: + h.remove() + assert cap.final_residual is not None and len(cap.head_inputs) == 6 + for layer, attn in enumerate(stack.attn): + proj = out_proj(attn) + heads = n_heads_of(attn, proj) + w = proj.weight.view(proj.weight.shape[0], heads, -1) + x = cap.head_inputs[layer].view(heads, -1) + per_head = torch.einsum("dhk,hk->hd", w, x) + bias = proj.bias if proj.bias is not None else 0 + assert torch.allclose(per_head.sum(0) + bias, outputs[layer], atol=1e-5) + + +def test_head_attribution_contract(vlm): + from evalrx.analyzers.lens.head_attribution import HeadAttributionAnalyzer + + result = HeadAttributionAnalyzer().run(vlm, _batch()) + f = _assert_result_contract(result, "head_attribution", 4) + assert f["linearization_max_error"] < 1e-3 # RMSNorm, no soft-capping: frozen norm is exact + assert f["n_layers"] == 6 and f["n_heads"] == 4 and len(f["top_heads_first_case"]) == 10 + for row in f["per_case"]: + assert 0.0 < row["dla_top_head_share"] <= 1.0 + assert 0.0 <= row["dla_late_share"] <= 1.0 and 0.0 <= row["dla_negative_share"] <= 1.0 + assert result.artifacts["head_dla_mean_map"].shape == (6, 4) + assert result.artifacts["head_dla_diff_map_fail_minus_pass"].shape == (6, 4) + + +def test_head_attribution_runs_text_only(vlm): + from evalrx.analyzers.lens.head_attribution import HeadAttributionAnalyzer + + f = HeadAttributionAnalyzer().run(vlm, "plain text prompt").findings + assert f["n_cases_analyzed"] == 1 + + +# ── attention sensitivity ───────────────────────────────────────────────────── + +def test_attention_sensitivity_contract(vlm): + from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer + + f = _assert_result_contract(AttentionSensitivityAnalyzer(band="mid").run(vlm, _batch()), + "attention_sensitivity", 4) + assert f["tap_fidelity"]["tap_faithful"] is True and f["band"] == "mid" + span = math.log(2.0) - math.log(0.5) + for row in f["per_case"]: + assert row["rw_slope"] == pytest.approx((row["rw_delta_up"] - row["rw_delta_down"]) / span, abs=2e-4) + assert isinstance(row["rw_flips_up"], bool) + + +def test_attention_sensitivity_validates_params(): + from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer + + with pytest.raises(ValueError, match="band"): + AttentionSensitivityAnalyzer(band="middle") + with pytest.raises(ValueError, match="down"): + AttentionSensitivityAnalyzer(down=2.0, up=3.0) + + +# ── tuned lens ──────────────────────────────────────────────────────────────── + +def test_tuned_lens_fit_beats_the_identity_translator_and_leaves_weights_alone(vlm, tmp_path): + from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer + + module = vlm._loaded[0] + identity = TunedLensAnalyzer(steps=0).run(vlm, _batch()).findings + path = tmp_path / "lens.pt" + tuned_result = TunedLensAnalyzer(steps=150, lr=1e-2, translators_path=str(path)).run(vlm, _batch()) + tuned = _assert_result_contract(tuned_result, "tuned_lens", 4) + assert tuned["fit_kl_mean"] < identity["fit_kl_mean"] + assert all(p.grad is None for p in module.parameters()) + assert path.exists() + # A fresh analyzer with the same configuration reuses the cached fit. + again = TunedLensAnalyzer(steps=150, lr=1e-2, translators_path=str(path)).run(vlm, _batch()).findings + assert again["fit_kl_per_layer"] == tuned["fit_kl_per_layer"] + + +def test_tuned_lens_loads_saved_translators(tmp_path): + from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer + + path = tmp_path / "lens.pt" + first = TunedLensAnalyzer(steps=30, translators_path=str(path)).run(TinyVLM("qwen2_5_vl"), _batch(2)).findings + fresh_model = TinyVLM("qwen2_5_vl") # no in-memory cache: must come from disk + loaded = TunedLensAnalyzer(steps=30, translators_path=str(path)).run(fresh_model, _batch(2)).findings + assert loaded["fit_kl_per_layer"] == first["fit_kl_per_layer"] + strip = lambda rows: [{k: v for k, v in r.items() if k != "sample_id"} for r in rows] # noqa: E731 + assert strip(loaded["per_case"]) == strip(first["per_case"]) # case ids are per-batch uuids diff --git a/tests/test_core/test_capability_matching.py b/tests/test_core/test_capability_matching.py index 48906a9d..f0bd95d6 100755 --- a/tests/test_core/test_capability_matching.py +++ b/tests/test_core/test_capability_matching.py @@ -23,7 +23,7 @@ def test_registry_compatible_with_includes_attention(): def test_registry_compatible_with_excludes_gradient_analyzer(): - # Qwen-like local model: no GRADIENTS → gradcam excluded; no LOGPROBS+image → vl_shap excluded. + # A model without GRADIENTS (API / jax_local) → gradcam excluded; no LOGPROBS+image → vl_shap excluded. model = FakeModel( capabilities={Capability.GENERATE, Capability.ATTENTION, Capability.HIDDEN_STATES} ) diff --git a/tests/test_core/test_registry_and_shim.py b/tests/test_core/test_registry_and_shim.py index d5558798..0a7b9187 100755 --- a/tests/test_core/test_registry_and_shim.py +++ b/tests/test_core/test_registry_and_shim.py @@ -48,7 +48,7 @@ def test_get_set_params(): from evalrx.analyzers.attention.summary import AttentionAnalyzer a = AttentionAnalyzer(layer=-1, top_k=5) - assert a.get_params() == {"layer": -1, "head": "mean", "top_k": 5} + assert a.get_params() == {"layer": -1, "head": "mean", "top_k": 5, "max_cases": 0} a.set_params(top_k=2) assert a.get_params()["top_k"] == 2 assert a.top_k == 2 diff --git a/tests/test_models/test_qwen.py b/tests/test_models/test_qwen.py index 4015aef0..af4b384c 100755 --- a/tests/test_models/test_qwen.py +++ b/tests/test_models/test_qwen.py @@ -39,7 +39,8 @@ def test_qwenllm_returns_hf_local_model_with_caps(): assert Capability.ATTENTION in model.capabilities assert Capability.HIDDEN_STATES in model.capabilities assert Capability.GENERATE in model.capabilities - assert Capability.GRADIENTS not in model.capabilities + assert Capability.GRADIENTS in model.capabilities # forward(with_grad=True) + assert Capability.EMBEDDINGS not in model.capabilities def test_qwenllm_checkpoint_override(): diff --git a/tests/test_models/test_wrap.py b/tests/test_models/test_wrap.py index fc2e4be6..4cc4a86c 100755 --- a/tests/test_models/test_wrap.py +++ b/tests/test_models/test_wrap.py @@ -138,7 +138,7 @@ def test_wrap_want_negotiation_rejects_missing_cap(): from evalrx.core.capability import CapabilityError with pytest.raises(CapabilityError): - evalrx.wrap(FakeCausalLM(), FakeTokenizer(), want={Capability.GRADIENTS}) + evalrx.wrap(FakeCausalLM(), FakeTokenizer(), want={Capability.EMBEDDINGS}) # ---------------------------------------------------------------------- diff --git a/tests/tiny_whitebox.py b/tests/tiny_whitebox.py new file mode 100644 index 00000000..bfebb5fa --- /dev/null +++ b/tests/tiny_whitebox.py @@ -0,0 +1,166 @@ +"""A random-init, CPU-sized Qwen-VL behind the evalrx ``Model`` contract. + +The write-side analyzers (attention knockout, causal tracing, gradient +attribution) act on a live torch module, which ``FakeModel`` does not have. +``TinyVLM`` is the smallest thing that is real where it matters: a genuine +transformers Qwen2.5-VL / Qwen3-VL (vision tower, M-RoPE, merged image tokens, +DeepStack for Qwen3-VL) with a few thousand parameters, fed through a +``forward`` that fills a ``Trace`` the way ``hf_local`` does — image-token +positions in a TokenTypeMap, logits, optional attentions/hidden states, and +gradients when ``CaptureSpec.with_grad`` asks for them. + +Images are 16×16 PIL images: 8×8 patches of 2 px, merged 2×2 → 16 image tokens. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import torch +import transformers + +from evalrx.core.capability import Capability +from evalrx.core.case import Inputs +from evalrx.core.model import Model, Trace +from evalrx.core.spec import AttnSemantics +from evalrx.core.tokentype import TokenTypeMap + +IMAGE_TOKEN, VIDEO_TOKEN, VSTART, VEND = 100, 101, 102, 103 +VOCAB = 128 +IMG_PX, PATCH, MERGE = 16, 2, 2 +GRID = IMG_PX // PATCH # 8 patches per side +TOKENS_SIDE = GRID // MERGE # 4 merged tokens per side → 16 image tokens + +# These models are tiny; on a many-core box torch's default one-thread-per-core +# turns every matmul into scheduling contention (minutes instead of seconds). +torch.set_num_threads(min(torch.get_num_threads(), 4)) + + +class CharTokenizer: + """Maps characters onto ids 1..90 — enough for the gold-target path.""" + + def __call__(self, text: str, add_special_tokens: bool = False) -> dict: + return {"input_ids": [ord(c) % 90 + 1 for c in text]} + + +def _build(family: str) -> Any: + vision = dict(depth=2, hidden_size=16, intermediate_size=32, num_heads=2, in_channels=3, + patch_size=PATCH, spatial_merge_size=MERGE, temporal_patch_size=2, + out_hidden_size=32) + text = dict(vocab_size=VOCAB, hidden_size=32, intermediate_size=64, num_hidden_layers=6, + num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=512) + ids = dict(image_token_id=IMAGE_TOKEN, video_token_id=VIDEO_TOKEN, + vision_start_token_id=VSTART, vision_end_token_id=VEND) + torch.manual_seed(0) + if family == "qwen2_5_vl": + vision.update(window_size=8, fullatt_block_indexes=[1]) + text.update(rope_scaling={"type": "mrope", "mrope_section": [1, 1, 2]}) + cfg = transformers.Qwen2_5_VLConfig(vision_config=vision, text_config=text, **ids) + cls = transformers.Qwen2_5_VLForConditionalGeneration + elif family == "qwen3_vl": + vision.update(num_position_embeddings=64, deepstack_visual_indexes=[0, 1]) + text.update(head_dim=8, rope_scaling={"rope_type": "default", "mrope_section": [1, 1, 2], + "mrope_interleaved": True}) + cfg = transformers.Qwen3VLConfig(vision_config=vision, text_config=text, **ids) + cls = transformers.Qwen3VLForConditionalGeneration + else: # pragma: no cover - test helper misuse + raise ValueError(family) + return cls._from_config(cfg, attn_implementation="eager").eval() + + +def _pixel_values(image: Any) -> torch.Tensor: + """16×16 RGB → (64 patches, 3·2·2·2) in the processor's merge-window order.""" + arr = np.asarray(image.convert("RGB").resize((IMG_PX, IMG_PX)), dtype=np.float32) / 127.5 - 1.0 + x = torch.from_numpy(arr).permute(2, 0, 1) # (3, 16, 16) + x = x.reshape(3, TOKENS_SIDE, MERGE, PATCH, TOKENS_SIDE, MERGE, PATCH) + x = x.permute(1, 4, 2, 5, 0, 3, 6) # (gh, gw, mh, mw, c, ph, pw) + x = x.reshape(GRID * GRID, 3, PATCH, PATCH) + x = torch.stack([x, x], dim=2) # temporal patch of 2 + return x.reshape(GRID * GRID, 3 * 2 * PATCH * PATCH) + + +class TinyVLM(Model): + """Tiny Qwen-VL under the ``Model`` contract (see module docstring).""" + + def __init__(self, family: str = "qwen2_5_vl", + attn_semantics: AttnSemantics = AttnSemantics.STANDARD, + gradients: bool = True) -> None: + caps = {Capability.GENERATE, Capability.LOGITS, Capability.ATTENTION, + Capability.HIDDEN_STATES} + if gradients: + caps.add(Capability.GRADIENTS) + self.capabilities = frozenset(caps) + self.modalities = frozenset({"text", "image"}) + self.spec = SimpleNamespace(key=f"tiny-{family}", attn_semantics=attn_semantics) + self._hf = (_build(family), CharTokenizer()) + self.n_forwards = 0 + + @property + def _loaded(self): + return self._hf + + def generate(self, inputs: Any, **kwargs: Any) -> str: + return "x" + + def unembed_weight(self): + return self._hf[0].get_output_embeddings().weight + + def final_norm(self): + from evalrx.models.backends.hf.discover import get_final_norm + + return get_final_norm(self._hf[0]) + + def forward(self, inputs: Any, capture: "set[Capability]", spec: Any = None) -> Trace: + module = self._hf[0] + self.n_forwards += 1 + prompt = inputs.prompt if isinstance(inputs, Inputs) else str(inputs) + image = getattr(inputs, "image", None) if isinstance(inputs, Inputs) else None + text_ids = [ord(c) % 90 + 1 for c in prompt[:24]] or [1] + enc: dict = {} + if image is not None: + n_img = TOKENS_SIDE * TOKENS_SIDE + ids = [1, 2, VSTART] + [IMAGE_TOKEN] * n_img + [VEND] + text_ids + enc["pixel_values"] = _pixel_values(image) + enc["image_grid_thw"] = torch.tensor([[1, GRID, GRID]]) + image_pos = list(range(3, 3 + n_img)) + grids = [(1, TOKENS_SIDE, TOKENS_SIDE)] + patches = [(0, 0, r, c) for r in range(TOKENS_SIDE) for c in range(TOKENS_SIDE)] + else: + ids = [1, 2] + text_ids + image_pos, grids, patches = [], [], [] + enc["input_ids"] = torch.tensor([ids]) + with_grad = bool(spec is not None and getattr(spec, "with_grad", False)) + with torch.set_grad_enabled(with_grad): + out = module(**enc, output_attentions=Capability.ATTENTION in capture, + output_hidden_states=Capability.HIDDEN_STATES in capture) + img_set = set(image_pos) + ttm = TokenTypeMap(seq_len=len(ids), image_pos=image_pos, + text_pos=[i for i in range(len(ids)) if i not in img_set], + grids=grids, patches=patches, image_token_id=IMAGE_TOKEN) + extras: dict = {"attn_semantics": self.spec.attn_semantics.value} + if image_pos: + extras["image_token_mask"] = torch.tensor(ids) == IMAGE_TOKEN + extras["image_spatial_shape"] = (TOKENS_SIDE, TOKENS_SIDE) + provided = {Capability.LOGITS} + trace = Trace(tokens=[f"t{i}" for i in ids], token_ids=ids, provided=provided, + logits=out.logits[0], token_type_map=ttm, extras=extras) + if Capability.ATTENTION in capture: + trace.attentions = [a[0] for a in out.attentions] + provided.add(Capability.ATTENTION) + if Capability.HIDDEN_STATES in capture: + trace.hidden_states = [h[0] for h in out.hidden_states] + provided.add(Capability.HIDDEN_STATES) + return trace + + def __repr__(self) -> str: + return f"TinyVLM({self.spec.key})" + + +def image(seed: int = 0) -> Any: + """A deterministic 16×16 RGB test image.""" + from PIL import Image + + rng = np.random.default_rng(seed) + return Image.fromarray(rng.integers(0, 255, (IMG_PX, IMG_PX, 3), dtype=np.uint8), mode="RGB") From 59f0bc8c4e0bd3287a0efa247490ad2b9e4c62e6 Mon Sep 17 00:00:00 2001 From: Jiaqi Liu <32996169+jiaq-liu@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:39:22 -0700 Subject: [PATCH 6/9] update analyzer, add qwen3.5 adapter --- CHANGELOG.md | 36 ++- docs/analyzers.md | 9 +- evalrx/analyzers/_internals/access.py | 91 +++++-- evalrx/analyzers/_internals/edits.py | 79 ++++++ evalrx/analyzers/_internals/heads.py | 19 +- evalrx/analyzers/_internals/linear_mix.py | 233 ++++++++++++++++++ evalrx/analyzers/attribution/generic_attn.py | 42 ++-- evalrx/analyzers/lens/head_attribution.py | 20 +- .../analyzers/patching/attention_knockout.py | 34 +-- .../patching/attention_sensitivity.py | 22 +- evalrx/eval_agent/prompts/diagnosis.py | 2 + evalrx/eval_agent/stages/probe.py | 47 +++- evalrx/reporting/case_study.py | 28 ++- tests/test_analyzers/test_hybrid_analyzers.py | 215 ++++++++++++++++ .../test_intervention_analyzers.py | 25 +- .../test_probe_routing_whitebox.py | 109 ++++++++ tests/tiny_whitebox.py | 49 +++- 17 files changed, 967 insertions(+), 93 deletions(-) create mode 100644 evalrx/analyzers/_internals/edits.py create mode 100644 evalrx/analyzers/_internals/linear_mix.py create mode 100644 tests/test_analyzers/test_hybrid_analyzers.py create mode 100644 tests/test_eval_agent/test_probe_routing_whitebox.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7491e852..8c98ad35 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,7 +40,41 @@ function — without it transformers builds no causal mask) that reproduces eage attention and can block / reweight / record it, verified once per model to engage on every layer and reproduce the native logits; decoder-layer pre-hooks for reading/replacing the residual stream; per-head output capture. -Attention-editing analyzers refuse hybrid (`HYBRID_SPARSE`) and latent attention. +Attention-editing analyzers refuse latent attention (hybrid stacks: see below). + +### Added — the write-side analyzers run on hybrid stacks (Qwen3.5 / Qwen3-Next) + +Qwen3.5 interleaves softmax attention with gated-delta-rule (GatedDeltaNet) +layers that have no attention matrix; the analyzers used to refuse it. New +`_internals/linear_mix.py` covers those layers with the module's own code: + +- media edits (`attention_knockout`, `attention_sensitivity`) re-run a linear + layer on the sequence WITHOUT the media tokens and blend its output at the + text positions (`w = 0` = knockout); `_internals/edits.py` applies the tap + and this re-run together, and checks both once per model (the re-run must + reproduce the outputs before the first media token); +- `generic_attention` materialises each linear layer's token-mixing matrix + (the delta rule is linear in its values: `o = A·v`) and uses + `∇A = (∂y/∂o) vᵀ` in Chefer's rule; +- `head_attribution` attributes the linear layers' value heads too + (`dla_linear_share`, `layer_kinds`, `heads_per_layer`). + +`causal_trace`, `gradcam` and `tuned_lens` were already architecture-agnostic. +Verified on Qwen3.5-2B (transformers 5.14): tap 6/6 softmax layers exact, +re-run 18/18 linear layers within 0.24% (bf16). Tests: a tiny Qwen3-Next +(transformers ≥ 4.57) and a tiny Qwen3.5 VLM (transformers 5). + +### Changed — routing reaches the white-box analyzers + +`StrategyProbe` ranks the new analyzers for white-box image / audio / video +runs (`attention_knockout` right after `modality_ablation`), and M3 failure +modes promote them: `language_prior_bias` and `hallucination` now include the +media interventions, and new tags `visual_neglect`, `audio_neglect`, +`insufficient_visual_attention`, `visual_information_loss`, +`mislocalized_attention`, `late_suppression` (also offered to the M3 prompt) +route to the analyzer that tests them. The video list named the unregistered +`relative_attn` (silently dropped); it is now `relative_attention`, and a test +fails on any unregistered name in the routing tables. ### Changed — `hf_local` provides `GRADIENTS` diff --git a/docs/analyzers.md b/docs/analyzers.md index e58ccdcf..5e1ebbcf 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -151,9 +151,12 @@ activations into a corrupted run to localise where information is used (`HIDDEN_STATES`, read and write hooks), or cut attention edges and measure what the answer loses. These act on the live torch module of a local `hf_local` model through `evalrx/analyzers/_internals/` (an attention function registered with -transformers, and decoder-layer pre-hooks); attention edits refuse hybrid or -latent-attention stacks and verify, once per model, that the edit-free path -reproduces the native logits. +transformers, and decoder-layer pre-hooks), and verify, once per model, that the +edit-free path reproduces the native logits. Hybrid stacks (Qwen3.5, Qwen3-Next: +softmax + gated-delta-rule linear attention) are edited on every layer: a linear +layer's text positions get its output re-run on the sequence without the media +tokens (`_internals/linear_mix.py`). Latent / unverified attention is refused. +Qwen3.5 needs transformers 5 (the `vllm35` environment), not the `[local]` extra. | Analyzer | Key | Requires | Modality | Paper | Status | diff --git a/evalrx/analyzers/_internals/access.py b/evalrx/analyzers/_internals/access.py index c7fdcf00..9c009b04 100644 --- a/evalrx/analyzers/_internals/access.py +++ b/evalrx/analyzers/_internals/access.py @@ -14,12 +14,14 @@ if TYPE_CHECKING: from evalrx.core.model import Model, Trace -#: Architectures whose decoder attention is a dense token×token softmax that a -#: registered attention function sees in full. Hybrid stacks route most layers -#: through linear/SSM attention with no query-key matrix; latent / unverified -#: attention is not a plain token matrix. Edits there would be partial or -#: meaningless, so the write-side analyzers refuse them. -_EDITABLE_SEMANTICS = frozenset({"standard"}) +#: Architectures the write-side analyzers can edit. ``standard``: every decoder +#: layer is a dense token×token softmax that the registered attention function +#: sees in full. ``hybrid_sparse`` (Qwen3.5, Qwen3-Next): softmax layers plus +#: gated-delta-rule linear-attention layers, which ``linear_mix`` edits — but +#: only when EVERY layer is one of the two kinds (checked per stack by +#: :attr:`DecoderStack.editable`; a Mamba mixer is neither). Latent / +#: unverified attention is not a plain token matrix and stays refused. +_EDITABLE_SEMANTICS = frozenset({"standard", "hybrid_sparse"}) def torch_module(model: "Model", analyzer: str = "this analyzer") -> Any: @@ -45,39 +47,94 @@ def tokenizer_of(model: "Model") -> Any: return getattr(processor, "tokenizer", processor) -def check_editable_attention(model: "Model", analyzer: str) -> None: - """Refuse architectures whose attention cannot be edited as a token matrix.""" +def check_editable_attention(model: "Model", analyzer: str, stack: "DecoderStack | None" = None) -> None: + """Refuse architectures whose token mixing cannot be edited layer by layer. + + The spec's declared semantics are checked first; with *stack*, every layer + must also be a softmax attention or a supported linear-attention mixer, so + an edit never silently skips a layer. + """ spec = getattr(model, "spec", None) semantics = getattr(getattr(spec, "attn_semantics", None), "value", None) if semantics is not None and semantics not in _EDITABLE_SEMANTICS: raise ValueError( - f"{analyzer} needs standard dense decoder attention; {getattr(spec, 'key', '?')} " - f"declares attn_semantics={semantics!r}, where an attention edit would be " - "partial (hybrid/linear layers) or not a token matrix (latent attention)." + f"{analyzer} needs softmax or gated-delta-rule decoder attention; " + f"{getattr(spec, 'key', '?')} declares attn_semantics={semantics!r}, where an " + "attention edit would not act on a token matrix (latent / unverified attention)." ) + if stack is not None and not stack.editable: + bad = [i for i, kind in enumerate(stack.kinds) if kind is None] + raise ValueError( + f"{analyzer}: decoder layers {bad[:8]} have neither softmax attention nor a " + "supported linear-attention mixer, so an edit would skip them." + ) + + +def is_delta_rule_mixer(module: Any) -> bool: + """A gated-delta-rule linear-attention mixer (Qwen3.5 / Qwen3-Next GatedDeltaNet). + + Recognised by what the edits use, not by class name: the per-instance + ``chunk_gated_delta_rule`` kernel (the relevance tap wraps it) and an + ``out_proj`` over per-head outputs (head attribution reads its input). + """ + return ( + module is not None + and callable(getattr(module, "chunk_gated_delta_rule", None)) + and hasattr(getattr(module, "out_proj", None), "weight") + and getattr(module, "head_v_dim", None) is not None + ) @dataclass class DecoderStack: - """The language decoder's layers, their attention modules, and their owner. + """The language decoder's layers, their token mixers, and their owner. - ``attn[i]`` is ``None`` for a layer with no softmax self-attention (the - linear-attention layers of a hybrid stack). + ``attn[i]`` is the softmax self-attention module or ``None``; ``linear[i]`` + the gated-delta-rule mixer or ``None``; ``kinds[i]`` is ``"softmax"``, + ``"linear"`` or ``None`` (a mixer none of the edits understand). """ layers: Any # nn.ModuleList path: str # dotted path of ``layers`` inside the module - attn: list # per layer: attention submodule or None + attn: list # per layer: softmax attention submodule or None parent: Any # module owning ``layers`` (the text model) + linear: list = None # per layer: delta-rule linear-attention mixer or None + + def __post_init__(self) -> None: + if self.linear is None: + self.linear = [None] * len(self.attn) @property def n_layers(self) -> int: return len(self.layers) + @property + def kinds(self) -> "list[str | None]": + return ["softmax" if a is not None else "linear" if m is not None else None + for a, m in zip(self.attn, self.linear)] + + @property + def mixers(self) -> list: + """Per layer: the token mixer (softmax or linear), or ``None``.""" + return [a if a is not None else m for a, m in zip(self.attn, self.linear)] + @property def fully_attentional(self) -> bool: return all(a is not None for a in self.attn) + @property + def editable(self) -> bool: + """Every layer is a softmax attention or a supported linear mixer.""" + return all(kind is not None for kind in self.kinds) + + @property + def n_softmax(self) -> int: + return sum(a is not None for a in self.attn) + + @property + def n_linear(self) -> int: + return sum(m is not None for m in self.linear) + def decoder_stack(module: Any) -> DecoderStack: """Discover the decoder stack at runtime (never a hard-coded module path).""" @@ -86,7 +143,9 @@ def decoder_stack(module: Any) -> DecoderStack: layers, path = find_decoder_layers(module) parent = module if "." not in path else resolve(module, path.rsplit(".", 1)[0]) attn = [getattr(layer, "self_attn", None) for layer in layers] - return DecoderStack(layers=layers, path=path, attn=attn, parent=parent) + linear = [None if a is not None else (lambda m: m if is_delta_rule_mixer(m) else None)( + getattr(layer, "linear_attn", None)) for layer, a in zip(layers, attn)] + return DecoderStack(layers=layers, path=path, attn=attn, parent=parent, linear=linear) def layer_points(n_layers: int, n_points: int) -> list[int]: diff --git a/evalrx/analyzers/_internals/edits.py b/evalrx/analyzers/_internals/edits.py new file mode 100644 index 00000000..9bfab094 --- /dev/null +++ b/evalrx/analyzers/_internals/edits.py @@ -0,0 +1,79 @@ +"""One media-flow edit across every kind of decoder layer. + +"Cut (``weight=0``) or scale what the text reads from the media tokens" means a +different operation per token mixer: + +* softmax attention — the attention tap (:mod:`.attn_tap`): ``-inf`` on the + query→media scores, or ``+log w`` (renormalised scaling); +* gated-delta-rule linear attention — :mod:`.linear_mix`: blend the layer's + output toward its re-run on the sequence without the media tokens. + +:func:`media_flow_edit` applies both in one block, so an analyzer edits every +layer of a hybrid stack (Qwen3.5) the same way it edits a dense one, and +:func:`require_edit_fidelity` checks both mechanisms before anything is reported. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Iterable, Iterator + +from evalrx.analyzers._internals.access import decoder_stack, torch_module +from evalrx.analyzers._internals.attn_tap import attention_tap, check_fidelity, knockout, reweight +from evalrx.analyzers._internals.linear_mix import check_linear_fidelity, linear_media_edit + +if TYPE_CHECKING: + from evalrx.core.model import Model + + +@contextmanager +def media_flow_edit( + model: "Model", + media: Iterable[int], + *, + queries: "Iterable[int] | None" = None, + layers: "Iterable[int] | None" = None, + weight: float = 0.0, + analyzer: str = "media edit", +) -> Iterator[None]: + """Cut (``weight=0``) or scale what *queries* read from *media* in *layers*. + + ``layers=None`` = every layer. Softmax layers scale renormalised (rows still + sum to one); linear layers scale the media's contribution to their output. + """ + stack = decoder_stack(torch_module(model, analyzer)) + media = [int(m) for m in media] + queries = None if queries is None else [int(q) for q in queries] + chosen = range(stack.n_layers) if layers is None else [int(i) for i in layers] + softmax = [i for i in chosen if stack.attn[i] is not None] + linear = [i for i in chosen if stack.linear[i] is not None] + edits = [] + if softmax: + edit = (knockout(media, queries=queries, layers=softmax) if weight == 0 + else reweight(media, weight, queries=queries, layers=softmax)) + edits.append(edit) + with attention_tap(model, edits, analyzer=analyzer), \ + linear_media_edit(model, media, queries=queries, layers=linear, weight=weight, + analyzer=analyzer): + yield + + +def require_edit_fidelity(model: "Model", inputs: Any, media: Iterable[int], *, + rtol: float = 1e-2, analyzer: str = "media edit") -> "dict[str, Any]": + """Verify every edit mechanism the stack needs; raise when one cannot be trusted.""" + stack = decoder_stack(torch_module(model, analyzer)) + verdict: dict[str, Any] = {"softmax_layers": stack.n_softmax, "linear_layers": stack.n_linear} + if stack.n_softmax: + verdict.update(check_fidelity(model, inputs, rtol=rtol, analyzer=analyzer)) + if stack.n_linear: + verdict.update(check_linear_fidelity(model, inputs, media, rtol=rtol, analyzer=analyzer)) + verdict["faithful"] = bool(verdict.get("tap_faithful", True) and verdict.get("linear_faithful", True)) + if not verdict["faithful"]: + raise RuntimeError( + f"{analyzer}: the attention edits cannot be trusted on this model ({verdict}). " + "tap: engaged < expected layers means the architecture bypasses the transformers " + "attention interface, a logit diff means its attention is not eager-equivalent; " + "linear: the media-absent re-run must engage every linear layer and reproduce " + "the outputs before the first media token." + ) + return verdict diff --git a/evalrx/analyzers/_internals/heads.py b/evalrx/analyzers/_internals/heads.py index 6d335a5d..a1e327dc 100644 --- a/evalrx/analyzers/_internals/heads.py +++ b/evalrx/analyzers/_internals/heads.py @@ -1,7 +1,8 @@ """Per-head attention outputs at one position, and the final norm frozen to a linear map. -Every supported decoder attention module ends in an output projection whose -INPUT is the heads' outputs concatenated, ``(n_heads · head_dim)``. A forward +Every supported decoder token mixer — softmax attention, or the gated-delta-rule +linear attention of a hybrid stack — ends in an output projection whose INPUT +is the heads' outputs concatenated, ``(n_heads · head_dim)``. A forward pre-hook on that projection reads them; head ``h``'s write into the residual stream is ``W_O[:, h·hd:(h+1)·hd] @ x[h·hd:(h+1)·hd]``, and the per-head writes sum to the projection's output (minus its bias, if any). @@ -37,8 +38,12 @@ def out_proj(attn: Any) -> Any: def n_heads_of(attn: Any, proj: Any) -> int: - """Query-head count: the projection's input width over the head dim.""" - head_dim = getattr(attn, "head_dim", None) + """Head count: the projection's input width over the head dim. + + Softmax attention exposes ``head_dim``; a gated-delta-rule mixer writes its + value heads (``head_v_dim`` wide each) through ``out_proj``. + """ + head_dim = getattr(attn, "head_dim", None) or getattr(attn, "head_v_dim", None) if head_dim is None: cfg = getattr(attn, "config", None) head_dim = getattr(cfg, "head_dim", None) or ( @@ -78,10 +83,10 @@ def hook(mod, args): def grab_final(mod, args): session.final_residual = args[0][0, position].detach() - for layer, attn in enumerate(stack.attn): - proj = out_proj(attn) if attn is not None else None + for layer, mixer in enumerate(stack.mixers): + proj = out_proj(mixer) if mixer is not None else None if proj is None: - raise ValueError(f"{analyzer}: layer {layer} has no attention output projection") + raise ValueError(f"{analyzer}: layer {layer} has no token-mixer output projection") handles.append(proj.register_forward_pre_hook(grab_heads(layer))) handles.append(norm.register_forward_pre_hook(grab_final)) try: diff --git a/evalrx/analyzers/_internals/linear_mix.py b/evalrx/analyzers/_internals/linear_mix.py new file mode 100644 index 00000000..f232f95f --- /dev/null +++ b/evalrx/analyzers/_internals/linear_mix.py @@ -0,0 +1,233 @@ +"""Gated-delta-rule linear-attention layers (Qwen3.5, Qwen3-Next): media edits and relevance. + +A hybrid stack interleaves softmax attention with GatedDeltaNet layers, which +have no query×key matrix: a short causal convolution mixes neighbouring tokens, +then a gated delta-rule recurrence carries everything earlier forward in a +state. The softmax-only tools (:mod:`evalrx.analyzers._internals.attn_tap`) +cannot touch them, so a knockout there would be partial. This module gives the +two operations the analyzers need, with the module's own code doing the math. + +**Media edits** (:func:`linear_media_edit`). For a linear layer, "the text does +not read the media" is defined as: the text positions get the output the layer +produces on the sequence WITH THE MEDIA TOKENS REMOVED. The layer is re-run on +that shortened sequence (``module.forward``, no cache) and the selected query +positions are blended:: + + out[q] = absent[q] + w · (real[q] − absent[q]) + +``w = 0`` is the knockout, ``w = 1`` the unchanged layer, ``w = 2`` doubles the +media's contribution. Removing the tokens (rather than zeroing them) keeps the +convolution and the recurrence consistent: these layers carry no positional +encoding, so the shortened sequence is exactly "the media was never there". +Positions before the first media token are unaffected by causality; the +re-run's outputs there must match the real ones, which :func:`check_linear_fidelity` +verifies once per model. + +**Relevance** (:func:`delta_rule_relevance`). With q, k, gates and β fixed, the +delta rule is LINEAR in the values: ``o = A · v`` per head for a data-dependent +lower-triangular ``A`` (T×T). Running the module's own pure-torch kernel with +the values replaced by the identity materialises ``A`` exactly; with +``∂y/∂A = (∂y/∂o) vᵀ`` the layer contributes Chefer's ``E_h[(∇A ⊙ A)⁺]`` like a +softmax layer. The convolution before the kernel (window 4) is not propagated: +relevance moves at most three tokens through it unaccounted. +""" + +from __future__ import annotations + +import sys +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator + +from evalrx.analyzers._internals.access import decoder_stack, torch_module + +if TYPE_CHECKING: + from evalrx.core.model import Model + + +@dataclass +class LinearEditState: + """What one :func:`linear_media_edit` block did.""" + + calls: "list[int]" = field(default_factory=list) # layer index per edited call + prefix_rel_diff: "dict[int, float]" = field(default_factory=dict) # layer -> max rel diff + + +def _hidden_arg(args: tuple, kwargs: dict) -> Any: + if "hidden_states" in kwargs: + return kwargs["hidden_states"] + if args: + return args[0] + raise RuntimeError("linear-attention layer called without hidden_states") + + +@contextmanager +def linear_media_edit( + model: "Model", + media: Iterable[int], + *, + queries: "Iterable[int] | None" = None, + layers: "Iterable[int] | None" = None, + weight: float = 0.0, + analyzer: str = "linear media edit", +) -> Iterator[LinearEditState]: + """Blend linear-layer outputs at *queries* toward the media-absent re-run. + + ``queries=None`` = every non-media position; ``layers=None`` = every linear + layer. Softmax layers are untouched (use the attention tap for those). + """ + import torch + + stack = decoder_stack(torch_module(model, analyzer)) + media_set = {int(m) for m in media} + query_set = None if queries is None else {int(q) for q in queries} + chosen = None if layers is None else {int(i) for i in layers} + state = LinearEditState() + + def make_hook(layer: int): + def hook(mod, args, kwargs, output): + hidden = _hidden_arg(args, kwargs) + seq = hidden.shape[1] + present = sorted(m for m in media_set if m < seq) + if not present or weight == 1.0: + return None + first = present[0] + keep = [i for i in range(seq) if i not in media_set] + targets = [q for q in keep if q > first and (query_set is None or q in query_set)] + if not targets: + return None + absent = mod.forward(hidden[:, keep], cache_params=None, attention_mask=None) + index = {pos: i for i, pos in enumerate(keep)} + sel = torch.tensor([index[q] for q in targets], device=output.device) + tgt = torch.tensor(targets, device=output.device) + before = [i for i, pos in enumerate(keep) if pos < first] + if before: + ref = output[:, before].float() + diff = (absent[:, before].float() - ref).abs().max() + state.prefix_rel_diff[layer] = float(diff / ref.abs().max().clamp(min=1e-6)) + real, alt = output.index_select(1, tgt).float(), absent.index_select(1, sel).float() + out = output.clone() + out[:, tgt] = (alt + weight * (real - alt)).to(output.dtype) + state.calls.append(layer) + return out + + return hook + + handles = [ + mod.register_forward_hook(make_hook(i), with_kwargs=True) + for i, mod in enumerate(stack.linear) + if mod is not None and (chosen is None or i in chosen) + ] + try: + yield state + finally: + for handle in handles: + handle.remove() + + +def check_linear_fidelity(model: "Model", inputs: Any, media: Iterable[int], *, + rtol: float = 1e-2, analyzer: str = "linear media edit") -> "dict[str, Any]": + """Does the media-absent re-run engage on every linear layer and match the prefix? + + The re-run calls ``module.forward`` on a shortened sequence without the + cache/mask arguments the decoder passed; before the first media token it + must reproduce the real outputs, or the re-run is not the same computation. + Cached on the module (one extra forward per model). + """ + from evalrx.core.capability import Capability + + module = torch_module(model, analyzer) + stack = decoder_stack(module) + cache = module.__dict__.setdefault("_evalrx_linear_fidelity", {}) + if "verdict" in cache: + return dict(cache["verdict"]) + with linear_media_edit(model, media, weight=0.0, analyzer=analyzer) as state: + model.forward(inputs, capture={Capability.LOGITS}) + engaged = len(set(state.calls)) + worst = max(state.prefix_rel_diff.values(), default=float("nan")) + verdict = { + "linear_faithful": bool(engaged == stack.n_linear and state.prefix_rel_diff + and worst <= rtol), + "linear_layers_engaged": engaged, + "linear_layers_expected": stack.n_linear, + "linear_prefix_max_rel_diff": round(worst, 6), + "linear_rtol": rtol, + } + cache["verdict"] = verdict + return dict(verdict) + + +def _torch_kernel(mod: Any) -> Callable: + """The module family's pure-torch chunked delta-rule kernel (exact, any value width).""" + kernel = getattr(sys.modules.get(type(mod).__module__), "torch_chunk_gated_delta_rule", None) + if kernel is None: + raise RuntimeError(f"{type(mod).__name__}: no torch_chunk_gated_delta_rule next to it") + return kernel + + +def mixing_matrix(kernel: Callable, query: Any, key: Any, g: Any, beta: Any, **kwargs: Any) -> Any: + """``A`` with ``o = A · v`` per head: ``(batch, heads, T, T)``, float32. + + The delta rule is linear in the values, so feeding the identity as values + (value width T) returns ``A`` row by row. + """ + import torch + + batch, seq, heads, _ = key.shape + eye = torch.eye(seq, device=key.device, dtype=torch.float32) + values = eye.view(1, seq, 1, seq).expand(batch, seq, heads, seq) + out, _ = kernel(query.float(), key.float(), values, g=g.float(), beta=beta.float(), + initial_state=None, output_final_state=False, + use_qk_l2norm_in_kernel=kwargs.get("use_qk_l2norm_in_kernel", False)) + return out.float().permute(0, 2, 1, 3) # (B, T_query, H, T_key) -> (B, H, T, T) + + +@contextmanager +def delta_rule_relevance(model: "Model", on_abar: "Callable[[int, Any], None]", *, + analyzer: str = "delta-rule relevance") -> "Iterator[dict[int, Any]]": + """Wrap each linear layer's kernel so its backward yields Chefer's ``Ā``. + + Yields ``{layer: kernel output}`` recorded during the forward (graph + attached); a backward pass through them fires one hook per layer that calls + ``on_abar(layer, Ā)`` with ``Ā`` of shape (T, T). + """ + import torch + + stack = decoder_stack(torch_module(model, analyzer)) + recorded: dict[int, Any] = {} + saved: list = [] + + def wrap(layer: int, mod: Any) -> Callable: + original = mod.chunk_gated_delta_rule + kernel = _torch_kernel(mod) + + def wrapped(query, key, value, g=None, beta=None, **kwargs): + out, final_state = original(query, key, value, g=g, beta=beta, **kwargs) + if out.requires_grad: + held = [t.detach() for t in (query, key, value, g, beta)] + + def hook(grad: Any) -> None: + q, k, v, gg, bb = held + a = mixing_matrix(kernel, q, k, gg, bb, **kwargs) + grad_a = torch.einsum("bthd,bjhd->bhtj", grad.float(), v.float()) + on_abar(layer, (grad_a * a).clamp(min=0).mean(dim=1)[0]) + + out.register_hook(hook) + recorded[layer] = out + return out, final_state + + return wrapped + + for layer, mod in enumerate(stack.linear): + if mod is not None: + saved.append((mod, mod.__dict__.get("chunk_gated_delta_rule", None), + "chunk_gated_delta_rule" in mod.__dict__)) + mod.chunk_gated_delta_rule = wrap(layer, mod) + try: + yield recorded + finally: + for mod, previous, had_instance_attr in saved: + if had_instance_attr: + mod.chunk_gated_delta_rule = previous + else: + del mod.chunk_gated_delta_rule diff --git a/evalrx/analyzers/attribution/generic_attn.py b/evalrx/analyzers/attribution/generic_attn.py index be850f1c..3ce13605 100755 --- a/evalrx/analyzers/attribution/generic_attn.py +++ b/evalrx/analyzers/attribution/generic_attn.py @@ -20,8 +20,14 @@ (:mod:`evalrx.analyzers._internals.attn_tap`), checked once per model to reproduce the native logits; the per-layer ``Ā_l`` is formed in a tensor hook during the backward pass, so no full set of attention gradients is ever held. -A hybrid stack (linear-attention layers without a token matrix) is refused: -relevance would only cross some of the layers. + +On a hybrid stack (Qwen3.5, Qwen3-Next) the gated-delta-rule linear-attention +layers have no softmax matrix, but they are linear in their values: ``o = A·v`` +for a data-dependent causal ``A``, which +:mod:`evalrx.analyzers._internals.linear_mix` materialises with the module's own +kernel, and ``∇A = (∂y/∂o) vᵀ``. Those layers then enter the same rule. Their +``A`` is not row-normalised like a softmax, and the short convolution before the +delta rule (three tokens back) is not propagated. References: - Generic Attention-model Explainability for Interpreting Bi-Modal and @@ -92,6 +98,7 @@ def _run(self, model: "Model", cases: "CaseBatch"): torch_module, ) from evalrx.analyzers._internals.attn_tap import attention_tap, require_fidelity + from evalrx.analyzers._internals.linear_mix import delta_rule_relevance from evalrx.analyzers._internals.maps import ( SpatialMapResult, group_maps, @@ -101,11 +108,8 @@ def _run(self, model: "Model", cases: "CaseBatch"): from evalrx.analyzers._internals.targets import build_margin from evalrx.core.model import CaptureSpec - check_editable_attention(model, self.name) stack = decoder_stack(torch_module(model, self.name)) - if not stack.fully_attentional: - raise ValueError(f"{self.name}: some decoder layers have no softmax attention " - "(hybrid stack); relevance would cross only part of the stack.") + check_editable_attention(model, self.name, stack) n_layers = stack.n_layers spec = CaptureSpec(with_grad=True, to_cpu=False) @@ -116,10 +120,14 @@ def _run(self, model: "Model", cases: "CaseBatch"): first: "dict[str, Any] | None" = None for case in cases.stratified_head(self.max_cases): if fidelity is None: - fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, - analyzer=self.name) + fidelity = (require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, + analyzer=self.name) if stack.n_softmax else {}) + fidelity.update(softmax_layers=stack.n_softmax, linear_layers=stack.n_linear) abar: dict[int, Any] = {} - attn_probs: list = [] + recorded: dict[int, Any] = {} + + def on_abar(layer: int, value: Any) -> None: + abar[layer] = value def on_probs(layer: int, probs: Any) -> None: if not probs.requires_grad: @@ -131,25 +139,27 @@ def hook(grad: Any, layer: int = layer, held: Any = held) -> None: abar[layer] = (grad.float() * held.float()).clamp(min=0).mean(dim=1)[0] probs.register_hook(hook) - attn_probs.append(probs) + recorded[layer] = probs - with attention_tap(model, on_probs=on_probs, analyzer=self.name): + with attention_tap(model, on_probs=on_probs, analyzer=self.name), \ + delta_rule_relevance(model, on_abar, analyzer=self.name) as linear_out: trace = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + recorded.update(linear_out) logits = trace.require(Capability.LOGITS)[-1:] margin, why = build_margin(self.target, logits, model=model, case=case) if margin is None: skipped[why] += 1 continue - if not attn_probs: + if not recorded: raise RuntimeError( f"{self.name}: the forward pass kept no autograd graph through attention — " "the backend ignores CaptureSpec(with_grad=True)." ) m_value = margin.value(logits) - # Asking for layer 0's attention gradient sends the backward pass - # through every layer's attention, firing each layer's hook once. - torch.autograd.grad(margin(logits), attn_probs[0]) - del attn_probs, logits + # Asking for the first layer's mixing gradient sends the backward + # pass through every layer's mixer, firing each layer's hook once. + torch.autograd.grad(margin(logits), recorded[min(recorded)]) + del recorded, linear_out, logits if len(abar) != n_layers: raise RuntimeError(f"{self.name}: relevance reached {len(abar)} of {n_layers} layers") diff --git a/evalrx/analyzers/lens/head_attribution.py b/evalrx/analyzers/lens/head_attribution.py index 84fcf693..fab2bdd2 100644 --- a/evalrx/analyzers/lens/head_attribution.py +++ b/evalrx/analyzers/lens/head_attribution.py @@ -12,6 +12,12 @@ layers (its effect on what they compute) scores low here. Combine with ``attention_knockout`` / ``causal_trace`` for indirect effects. +On a hybrid stack (Qwen3.5, Qwen3-Next) the gated-delta-rule linear-attention +layers write through the same kind of per-head output projection, so their +value heads are attributed alongside the softmax heads (``layer_kinds`` and +``heads_per_layer`` say which is which; rows of the maps are zero-padded to the +widest layer). + The per-head mean attribution over the batch, and its FAIL−PASS difference, are artifacts (``(layers, heads)`` maps), so a head that writes the wrong answer on failures stands out as a hot cell. @@ -64,6 +70,7 @@ class HeadAttributionAnalyzer(Analyzer): "dla_top_head_layer_frac": ("Top head's depth", "Relative depth (0 = first layer, 1 = last) of the head with the largest absolute contribution."), "dla_late_share": ("Late heads' share", "Share of all heads' absolute contribution coming from the last third of the layers."), "dla_negative_share": ("Heads opposing answer", "Share of all heads' absolute contribution that pushes AGAINST the answer token."), + "dla_linear_share": ("Linear-attention share", "Hybrid stacks only: share of all heads' absolute contribution written by the linear-attention (gated delta rule) layers rather than softmax attention."), } def __init__(self, target: str = "prediction", top_k: int = 10, max_cases: int = 0) -> None: @@ -87,12 +94,13 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: stack = decoder_stack(module) norm = get_final_norm(module) unembed = get_unembed(module) - projs = [out_proj(a) if a is not None else None for a in stack.attn] + projs = [out_proj(m) if m is not None else None for m in stack.mixers] if any(p is None for p in projs): - raise ValueError(f"{self.name}: some decoder layers have no attention output projection " - "(hybrid stack); head attribution would miss them.") - n_heads = [n_heads_of(a, p) for a, p in zip(stack.attn, projs)] + raise ValueError(f"{self.name}: some decoder layers have no recognised token mixer " + "with an output projection; head attribution would miss them.") + n_heads = [n_heads_of(m, p) for m, p in zip(stack.mixers, projs)] n_layers = stack.n_layers + linear_layers = [i for i, kind in enumerate(stack.kinds) if kind == "linear"] per_case: list[dict[str, Any]] = [] maps: dict[str, Any] = {} @@ -144,6 +152,8 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: } if abs(m_value) > 1e-6: row["dla_heads_share"] = round(float(flat.sum()) / m_value, 4) + if linear_layers: + row["dla_linear_share"] = round(float(dla[linear_layers].abs().sum()) / total_abs, 4) per_case.append(row) maps[case.id] = dla.numpy() if not first_top: @@ -166,6 +176,8 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: "skip_reasons": dict(skipped), "n_layers": n_layers, "n_heads": max(n_heads), + "heads_per_layer": n_heads, + "layer_kinds": stack.kinds, "target": self.target, "top_heads_first_case": first_top, "mean_heads_total": round(sum(r["dla_heads_total"] for r in per_case) / n, 4), diff --git a/evalrx/analyzers/patching/attention_knockout.py b/evalrx/analyzers/patching/attention_knockout.py index 389b5e29..fedf1030 100644 --- a/evalrx/analyzers/patching/attention_knockout.py +++ b/evalrx/analyzers/patching/attention_knockout.py @@ -13,11 +13,16 @@ a drop near zero says the answer did not come from the media at all — the language-prior signature. The band columns say WHERE the reading happened. -Mechanism: :mod:`evalrx.analyzers._internals.attn_tap` routes the decoder's -attention through an eager-equivalent function that masks the cut edges to -−inf before the softmax (an exact knockout, rows stay normalised). Before +Mechanism (:mod:`evalrx.analyzers._internals.edits`): in softmax layers the +attention tap masks the cut edges to −inf before the softmax (an exact +knockout, rows stay normalised); in the gated-delta-rule linear-attention +layers of a hybrid stack (Qwen3.5, Qwen3-Next), which have no attention matrix, +the text positions instead get the layer's output re-computed on the sequence +without the media tokens (:mod:`evalrx.analyzers._internals.linear_mix`). Before reporting anything the analyzer verifies, once per model, that the edit-free -tap engages on every attention layer and reproduces the native logits. +tap engages on every softmax layer and reproduces the native logits, and that +the media-absent re-run engages every linear layer and reproduces the outputs +before the first media token. References: - Dissecting Recall of Factual Associations in Auto-Regressive Language Models @@ -58,8 +63,9 @@ class AttentionKnockoutAnalyzer(Analyzer): ``"gold"`` (gold first token vs strongest other token). queries: ``"text"`` — every non-media position loses the media (Geva's setting); ``"last"`` — only the final position does. - fidelity_rtol: tolerance of the edit-free tap vs native logits (relative - to the largest logit) before the analyzer refuses. + fidelity_rtol: tolerance of the edit-free tap vs native logits, and of the + linear-layer re-run vs the real prefix outputs (relative to + the largest value) before the analyzer refuses. max_cases: label-stratified cap; 0 (the default) = every case. Cost: 5 logits-only forwards per case. """ @@ -96,15 +102,12 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: media_positions, torch_module, ) - from evalrx.analyzers._internals.attn_tap import attention_tap, knockout, require_fidelity + from evalrx.analyzers._internals.edits import media_flow_edit, require_edit_fidelity from evalrx.analyzers._internals.targets import build_margin from evalrx.core.model import CaptureSpec - check_editable_attention(model, self.name) stack = decoder_stack(torch_module(model, self.name)) - if not stack.fully_attentional: - raise ValueError(f"{self.name}: some decoder layers have no softmax attention " - "(hybrid stack); a knockout there would be partial.") + check_editable_attention(model, self.name, stack) bands = layer_bands(stack.n_layers) spec = CaptureSpec(to_cpu=False) @@ -128,15 +131,15 @@ def last_logits(inputs: Any) -> Any: skipped[why] += 1 continue if fidelity is None: - fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, - analyzer=self.name) + fidelity = require_edit_fidelity(model, case.inputs, media, + rtol=self.fidelity_rtol, analyzer=self.name) media_set = set(media) queries = ([clean.seq_len - 1] if self.queries == "last" else [i for i in range(clean.seq_len) if i not in media_set]) def cut(layers: "list[int] | None") -> Any: - edit = knockout(media, queries=queries, layers=layers) - with attention_tap(model, [edit], analyzer=self.name): + with media_flow_edit(model, media, queries=queries, layers=layers, weight=0.0, + analyzer=self.name): return last_logits(case.inputs) m_clean = margin.value(logits) @@ -167,6 +170,7 @@ def cut(layers: "list[int] | None") -> Any: "media_slot": slots.most_common(1)[0][0], "target": self.target, "queries": self.queries, + "layer_kinds": {"softmax": stack.n_softmax, "linear": stack.n_linear}, "tap_fidelity": fidelity, "mean_delta_full": round(sum(r["ko_delta_full"] for r in per_case) / n, 4), "frac_flips_top1": round(sum(bool(r["ko_flips_top1"]) for r in per_case) / n, 4), diff --git a/evalrx/analyzers/patching/attention_sensitivity.py b/evalrx/analyzers/patching/attention_sensitivity.py index ce677295..128956c5 100644 --- a/evalrx/analyzers/patching/attention_sensitivity.py +++ b/evalrx/analyzers/patching/attention_sensitivity.py @@ -17,6 +17,12 @@ for the gold token instead, i.e. whether amplifying the media would move the model TOWARD the right answer. +On a hybrid stack (Qwen3.5, Qwen3-Next) the gated-delta-rule linear-attention +layers have no attention weights to rescale; there the media's contribution to +the text positions' layer output is scaled instead — +``absent + w · (real − absent)``, with ``absent`` the layer re-run without the +media tokens (:mod:`evalrx.analyzers._internals.linear_mix`). + References: - Paying More Attention to Image: A Training-Free Method for Alleviating Hallucination in LVLMs (PAI) — Liu et al., ECCV 2024 — arXiv:2407.21771 @@ -89,16 +95,13 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: media_positions, torch_module, ) - from evalrx.analyzers._internals.attn_tap import attention_tap, require_fidelity, reweight + from evalrx.analyzers._internals.edits import media_flow_edit, require_edit_fidelity from evalrx.analyzers._internals.targets import build_margin from evalrx.analyzers.patching.attention_knockout import layer_bands from evalrx.core.model import CaptureSpec - check_editable_attention(model, self.name) stack = decoder_stack(torch_module(model, self.name)) - if not stack.fully_attentional: - raise ValueError(f"{self.name}: some decoder layers have no softmax attention " - "(hybrid stack); a reweight there would be partial.") + check_editable_attention(model, self.name, stack) layers = None if self.band == "all" else layer_bands(stack.n_layers)[self.band] spec = CaptureSpec(to_cpu=False) log_span = math.log(self.up) - math.log(self.down) @@ -118,16 +121,16 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: skipped[why] += 1 continue if fidelity is None: - fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, - analyzer=self.name) + fidelity = require_edit_fidelity(model, case.inputs, media, + rtol=self.fidelity_rtol, analyzer=self.name) media_set = set(media) queries = [i for i in range(clean.seq_len) if i not in media_set] top1 = int(logits[-1].argmax()) del clean def scaled(weight: float) -> Any: - edit = reweight(media, weight, queries=queries, layers=layers) - with attention_tap(model, [edit], analyzer=self.name): + with media_flow_edit(model, media, queries=queries, layers=layers, weight=weight, + analyzer=self.name): out = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) return out.require(Capability.LOGITS)[-1:].float() @@ -154,6 +157,7 @@ def scaled(weight: float) -> Any: "band": self.band, "weights": [self.down, self.up], "target": self.target, + "layer_kinds": {"softmax": stack.n_softmax, "linear": stack.n_linear}, "tap_fidelity": fidelity, "mean_slope": round(sum(r["rw_slope"] for r in per_case) / n, 4), "frac_negative_slope": round(sum(r["rw_slope"] < 0 for r in per_case) / n, 4), diff --git a/evalrx/eval_agent/prompts/diagnosis.py b/evalrx/eval_agent/prompts/diagnosis.py index 5f9a561f..b2e479a6 100644 --- a/evalrx/eval_agent/prompts/diagnosis.py +++ b/evalrx/eval_agent/prompts/diagnosis.py @@ -35,6 +35,8 @@ line is sent back for a rewrite.> FAILURE_MODE: CaseBatch: + return CaseBatch([ + FailureCase(inputs=Inputs(prompt=f"what is in picture {i}?", image=image(i)), + expected="cat", label=Label.FAIL if i % 2 else Label.PASS) + for i in range(n) + ]) + + +def _logits(model, inputs) -> "torch.Tensor": + return model.forward(inputs, capture={Capability.LOGITS}).logits + + +# ── stack discovery ─────────────────────────────────────────────────────────── + +def test_stack_sees_both_layer_kinds(hybrid): + stack = decoder_stack(hybrid._loaded[0]) + assert stack.kinds.count("softmax") == 2 and stack.kinds.count("linear") == 4 + assert stack.editable and not stack.fully_attentional + assert all(m is not None for m in stack.mixers) + + +def test_a_layer_with_an_unknown_mixer_is_refused(hybrid): + stack = decoder_stack(hybrid._loaded[0]) + broken = DecoderStack(layers=stack.layers, path=stack.path, attn=stack.attn, + parent=stack.parent, linear=[None] * stack.n_layers) + with pytest.raises(ValueError, match="neither softmax attention nor"): + check_editable_attention(hybrid, "test", broken) + + +# ── delta-rule mixing matrix ────────────────────────────────────────────────── + +def test_mixing_matrix_reproduces_the_kernel_and_is_causal(hybrid): + mod = next(m for m in decoder_stack(hybrid._loaded[0]).linear if m is not None) + kernel = _torch_kernel(mod) + torch.manual_seed(1) + b, t, h, dk, dv = 1, 70, 4, 8, 8 # 70 > one 64-token chunk: crosses a chunk boundary + q, k, v = torch.randn(b, t, h, dk), torch.randn(b, t, h, dk), torch.randn(b, t, h, dv) + g, beta = -torch.rand(b, t, h), torch.rand(b, t, h) + out, _ = kernel(q, k, v, g=g, beta=beta, use_qk_l2norm_in_kernel=True) + a = mixing_matrix(kernel, q, k, g, beta, use_qk_l2norm_in_kernel=True) + assert a.shape == (b, h, t, t) + assert torch.allclose(torch.einsum("bhtj,bjhd->bthd", a, v), out, atol=1e-5) + assert float(a.triu(1).abs().max()) == 0.0 # strictly causal + # Chefer's sum rule for a value-linear mixer: Σ_tj (∇A ⊙ A) = Σ_t ∂y/∂o_t · o_t + grad = torch.randn_like(out) + grad_a = torch.einsum("bthd,bjhd->bhtj", grad, v) + assert torch.allclose((grad_a * a).sum(), (grad * out).sum(), rtol=1e-4) + + +# ── media-absent re-run ─────────────────────────────────────────────────────── + +def test_linear_edit_at_weight_one_changes_nothing(hybrid): + inputs = Inputs(prompt="what?", image=image(0)) + with torch.no_grad(): + base = _logits(hybrid, inputs) + with linear_media_edit(hybrid, MEDIA, weight=1.0) as state: + same = _logits(hybrid, inputs) + assert torch.equal(base, same) and state.calls == [] + + +def test_linear_knockout_keeps_the_prefix_and_moves_the_answer(hybrid): + inputs = Inputs(prompt="what?", image=image(0)) + with torch.no_grad(): + base = _logits(hybrid, inputs) + with linear_media_edit(hybrid, MEDIA, weight=0.0) as state: + cut = _logits(hybrid, inputs) + assert sorted(set(state.calls)) == [i for i, k in enumerate(decoder_stack(hybrid._loaded[0]).kinds) + if k == "linear"] + assert max(state.prefix_rel_diff.values()) < 1e-4 # re-run reproduces pre-media outputs + assert torch.allclose(cut[:3], base[:3], atol=1e-6) # nothing before the image changes + assert not torch.allclose(cut[-1], base[-1]) + + +def test_linear_fidelity_verdict(hybrid): + hybrid._loaded[0].__dict__.pop("_evalrx_linear_fidelity", None) + verdict = check_linear_fidelity(hybrid, Inputs(prompt="x", image=image(1)), MEDIA) + assert verdict["linear_faithful"] is True + assert verdict["linear_layers_engaged"] == verdict["linear_layers_expected"] == 4 + + +# ── analyzers on the hybrid stack ───────────────────────────────────────────── + +def _contract(result, name: str, n: int) -> dict: + assert isinstance(result, Result) and result.analyzer == name + json.dumps(result.findings) + rows = result.findings["per_case"] + assert len(rows) == n + for row in rows: + for key, value in row.items(): + if isinstance(value, float): + assert math.isfinite(value), (key, value) + return result.findings + + +def test_attention_knockout_on_hybrid(hybrid): + from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer + + f = _contract(AttentionKnockoutAnalyzer().run(hybrid, _batch()), "attention_knockout", 4) + assert f["layer_kinds"] == {"softmax": 2, "linear": 4} + fid = f["tap_fidelity"] + assert fid["faithful"] and fid["tap_layers_engaged"] == 2 and fid["linear_layers_engaged"] == 4 + assert any(abs(r["ko_delta_full"]) > 0 for r in f["per_case"]) + + +def test_knockout_in_a_band_edits_only_that_band(hybrid): + # The early third of a 6-layer [lin, lin, soft] x2 stack is two linear + # layers: only the linear mechanism may engage, and the result must differ + # from the full-stack knockout. + from evalrx.analyzers._internals.edits import media_flow_edit + + inputs = Inputs(prompt="what?", image=image(0)) + text_q = [i for i in range(40) if i not in MEDIA] + with torch.no_grad(): + with media_flow_edit(hybrid, MEDIA, queries=text_q, layers=[0, 1]): + early = _logits(hybrid, inputs) + with media_flow_edit(hybrid, MEDIA, queries=text_q): + full = _logits(hybrid, inputs) + assert not torch.allclose(early[-1], full[-1]) + + +def test_attention_sensitivity_on_hybrid(hybrid): + from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer + + f = _contract(AttentionSensitivityAnalyzer().run(hybrid, _batch()), "attention_sensitivity", 4) + span = math.log(2.0) - math.log(0.5) + for row in f["per_case"]: + assert row["rw_slope"] == pytest.approx((row["rw_delta_up"] - row["rw_delta_down"]) / span, abs=2e-4) + + +def test_generic_attention_on_hybrid_reaches_every_layer(hybrid): + from evalrx.analyzers.attribution.generic_attn import GenericAttentionExplainability + + result = GenericAttentionExplainability().run(hybrid, _batch()) + f = _contract(result, "generic_attention", 4) + assert f["tap_fidelity"]["softmax_layers"] == 2 and f["tap_fidelity"]["linear_layers"] == 4 + assert all(0.0 <= r["rel_media_share"] <= 1.0 for r in f["per_case"]) + assert all(p.grad is None for p in hybrid._loaded[0].parameters()) + + +def test_head_attribution_on_hybrid_covers_linear_heads(hybrid): + from evalrx.analyzers.lens.head_attribution import HeadAttributionAnalyzer + + result = HeadAttributionAnalyzer().run(hybrid, _batch()) + f = _contract(result, "head_attribution", 4) + assert f["layer_kinds"].count("linear") == 4 + stack = decoder_stack(hybrid._loaded[0]) + expected = [m.num_v_heads if k == "linear" else m.config.num_attention_heads + for m, k in zip(stack.mixers, stack.kinds)] + assert f["heads_per_layer"] == expected + assert f["linearization_max_error"] < 1e-3 + assert all(0.0 <= r["dla_linear_share"] <= 1.0 for r in f["per_case"]) + + +@pytest.mark.parametrize("name", ["causal_trace", "gradcam", "tuned_lens"]) +def test_architecture_agnostic_analyzers_run_on_hybrid(hybrid, name): + from evalrx.analyzers.attribution.gradcam import GradCAMAnalyzer + from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer + from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer + + analyzer = {"causal_trace": CausalTraceAnalyzer(min_effect=0.0, n_points=3), + "gradcam": GradCAMAnalyzer(), "tuned_lens": TunedLensAnalyzer(steps=10)}[name] + _contract(analyzer.run(hybrid, _batch()), name, 4) diff --git a/tests/test_analyzers/test_intervention_analyzers.py b/tests/test_analyzers/test_intervention_analyzers.py index 099b8ff1..7a639476 100644 --- a/tests/test_analyzers/test_intervention_analyzers.py +++ b/tests/test_analyzers/test_intervention_analyzers.py @@ -57,13 +57,16 @@ def _batch(n: int = 4, with_gold: bool = True) -> CaseBatch: ]) -def _raw_inputs(seq_tail: list[int]) -> dict: +def _raw_inputs(seq_tail: list[int], module=None) -> dict: """Direct module inputs (bypassing the Model) for mask-level checks.""" - from tests.tiny_whitebox import GRID, _pixel_values + from tests.tiny_whitebox import GRID, _accepts, _pixel_values ids = [1, 2, VSTART] + [IMAGE_TOKEN] * 16 + [VEND] + seq_tail - return dict(input_ids=torch.tensor([ids]), pixel_values=_pixel_values(image(0)), - image_grid_thw=torch.tensor([[1, GRID, GRID]])) + enc = dict(input_ids=torch.tensor([ids]), pixel_values=_pixel_values(image(0)), + image_grid_thw=torch.tensor([[1, GRID, GRID]])) + if module is not None and _accepts(module, "mm_token_type_ids"): # transformers 5 M-RoPE + enc["mm_token_type_ids"] = (enc["input_ids"] == IMAGE_TOKEN).long() + return enc # ── attention tap ────────────────────────────────────────────────────────────── @@ -79,7 +82,7 @@ def test_tap_keeps_the_causal_mask(vlm): # transformers builds NO mask for an unknown attention name; the tap must # register one, or earlier positions would see a changed future token. module = vlm._loaded[0] - a, b = _raw_inputs([7, 8, 9, 10]), _raw_inputs([7, 8, 9, 11]) + a, b = _raw_inputs([7, 8, 9, 10], module), _raw_inputs([7, 8, 9, 11], module) with torch.no_grad(), attention_tap(vlm): la, lb = module(**a).logits, module(**b).logits assert torch.equal(la[0, :-1], lb[0, :-1]) @@ -100,7 +103,7 @@ def test_tap_retargets_only_the_decoder_and_restores_on_error(vlm): def test_knockout_touches_only_downstream_rows(vlm): module = vlm._loaded[0] - enc = _raw_inputs([7, 8, 9, 10]) + enc = _raw_inputs([7, 8, 9, 10], module) image_pos = list(range(3, 19)) text_q = [i for i in range(enc["input_ids"].shape[1]) if i not in image_pos] with torch.no_grad(): @@ -113,7 +116,7 @@ def test_knockout_touches_only_downstream_rows(vlm): def test_reweight_equivalences(vlm): module = vlm._loaded[0] - enc = _raw_inputs([7, 8, 9, 10]) + enc = _raw_inputs([7, 8, 9, 10], module) keys = list(range(3, 19)) with torch.no_grad(): base = module(**enc).logits @@ -256,12 +259,12 @@ def test_attention_knockout_refuses_text_only_batch(vlm): AttentionKnockoutAnalyzer().run(vlm, CaseBatch([FailureCase(inputs=Inputs(prompt="hi"))])) -def test_attention_knockout_refuses_hybrid_attention(): +def test_attention_knockout_refuses_latent_attention(): from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer - hybrid = TinyVLM("qwen2_5_vl", attn_semantics=AttnSemantics.HYBRID_SPARSE) - with pytest.raises(ValueError, match="hybrid_sparse"): - AttentionKnockoutAnalyzer().run(hybrid, _batch(1)) + latent = TinyVLM("qwen2_5_vl", attn_semantics=AttnSemantics.MLA_LATENT) + with pytest.raises(ValueError, match="mla_latent"): + AttentionKnockoutAnalyzer().run(latent, _batch(1)) def test_attention_summary_reports_media_mass(vlm): diff --git a/tests/test_eval_agent/test_probe_routing_whitebox.py b/tests/test_eval_agent/test_probe_routing_whitebox.py new file mode 100644 index 00000000..a91c2bc8 --- /dev/null +++ b/tests/test_eval_agent/test_probe_routing_whitebox.py @@ -0,0 +1,109 @@ +"""Routing of the white-box intervention analyzers (knockout, causal trace, gradients, lenses). + +Two guarantees: + +* every analyzer a routing table names is actually registered — a misspelt + name is silently dropped by ``select`` (the video list carried + ``"relative_attn"`` for months), so the table must be checked, not trusted; +* a white-box VLM/ALM gets the new analyzers in its priority list and when an + M3 hypothesis names the mechanism they test, while a black-box model never + sees them. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") # the attention analyzers register only with torch + +import evalrx.analyzers # noqa: F401,E402 - populate the registry +from evalrx.core.capability import Capability # noqa: E402 +from evalrx.core.case import CaseBatch, FailureCase, Inputs # noqa: E402 +from evalrx.core.registry import registry # noqa: E402 +from evalrx.eval_agent.stages.probe import ( # noqa: E402 + _FAILURE_MODE_TO_ANALYZERS, + _PRIORITY, + _SLOT_PRIORITY, + StrategyProbe, +) +from tests.conftest import FakeModel # noqa: E402 + +WRITE_SIDE = {"attention_knockout", "attention_sensitivity", "causal_trace", "gradcam", + "generic_attention", "head_attribution", "tuned_lens"} +_WHITE_BOX = {Capability.GENERATE, Capability.LOGITS, Capability.LOGPROBS, + Capability.ATTENTION, Capability.HIDDEN_STATES, Capability.GRADIENTS} + + +def _image_batch() -> CaseBatch: + return CaseBatch([FailureCase(inputs=Inputs(prompt="Is there a dog?", image="dog.png"))]) + + +def _audio_batch() -> CaseBatch: + return CaseBatch([FailureCase(inputs=Inputs(prompt="What do you hear?", audio="a.wav"))]) + + +def test_every_routed_analyzer_is_registered(): + registered = set(registry.analyzers.list()) + tables = { + **{f"failure_mode:{k}": v for k, v in _FAILURE_MODE_TO_ANALYZERS.items()}, + **{f"slot:{k}": v for k, v in _SLOT_PRIORITY.items()}, + **{f"kind:{k}": v for k, v in _PRIORITY.items()}, + } + unknown = {where: sorted(set(names) - registered) for where, names in tables.items() + if set(names) - registered} + assert not unknown, f"routing tables name unregistered analyzers: {unknown}" + + +def test_every_write_side_analyzer_is_routed_somewhere(): + routed = {name for names in _FAILURE_MODE_TO_ANALYZERS.values() for name in names} + routed |= {name for names in _SLOT_PRIORITY.values() for name in names} + assert WRITE_SIDE <= routed + + +def test_white_box_vlm_ranks_knockout_right_after_modality_ablation(): + vlm = FakeModel(capabilities=_WHITE_BOX, modalities={"text", "image"}) + ranked = StrategyProbe().select(vlm, data=_image_batch()) + assert ranked.index("attention_knockout") == ranked.index("modality_ablation") + 1 + assert WRITE_SIDE - {"tuned_lens", "head_attribution"} <= set(ranked[: len(_SLOT_PRIORITY["image"]) + 2]) + + +def test_white_box_alm_gets_the_audio_interventions(): + alm = FakeModel(capabilities=_WHITE_BOX, modalities={"text", "audio"}) + ranked = StrategyProbe().select(alm, data=_audio_batch()) + assert ranked.index("attention_knockout") == ranked.index("modality_ablation") + 1 + assert "gradcam" not in ranked # image-only: starved on an audio batch + + +@pytest.mark.parametrize("mode, first", [ + ("visual_neglect", "attention_knockout"), + ("language_prior_bias", "prompt_contrast"), + ("insufficient_visual_attention", "attention_sensitivity"), + ("visual_information_loss", "causal_trace"), + ("mislocalized_attention", "gradcam"), + ("late_suppression", "tuned_lens"), +]) +def test_m3_failure_modes_promote_the_analyzer_that_tests_them(mode, first): + vlm = FakeModel(capabilities=_WHITE_BOX, modalities={"text", "image"}) + ranked = StrategyProbe().select(vlm, data=_image_batch(), hint_failure_modes=[mode]) + assert ranked[0] == first + assert set(_FAILURE_MODE_TO_ANALYZERS[mode]) & set(ranked) <= set(ranked[: len(_FAILURE_MODE_TO_ANALYZERS[mode])]) + + +def test_language_prior_bias_now_reaches_the_media_interventions(): + assert {"attention_knockout", "attention_sensitivity"} <= set(_FAILURE_MODE_TO_ANALYZERS["language_prior_bias"]) + + +def test_black_box_model_is_never_offered_write_side_analyzers(): + api = FakeModel(capabilities={Capability.GENERATE, Capability.LOGPROBS}, + modalities={"text", "image"}) + ranked = StrategyProbe().select(api, data=_image_batch(), hint_failure_modes=["visual_neglect"]) + assert not WRITE_SIDE & set(ranked) + + +def test_m3_prompt_lists_the_white_box_tags(): + from evalrx.eval_agent.prompts import diagnosis + + text = " ".join(str(v) for v in vars(diagnosis).values() if isinstance(v, str)) + for tag in ("visual_neglect", "insufficient_visual_attention", "visual_information_loss", + "mislocalized_attention", "late_suppression"): + assert tag in text and tag in _FAILURE_MODE_TO_ANALYZERS diff --git a/tests/tiny_whitebox.py b/tests/tiny_whitebox.py index bfebb5fa..7ccd7844 100644 --- a/tests/tiny_whitebox.py +++ b/tests/tiny_whitebox.py @@ -9,6 +9,14 @@ positions in a TokenTypeMap, logits, optional attentions/hidden states, and gradients when ``CaptureSpec.with_grad`` asks for them. +Hybrid stacks (softmax + gated-delta-rule linear attention): + +* ``"qwen3_5"`` — a real Qwen3.5 VLM; needs transformers 5 (skipped on 4.x, + see :data:`HAS_QWEN3_5`); +* ``"qwen3_next"`` — the same layer kinds as a text-only model (transformers + 4.57+): an image is represented by its 16 placeholder tokens with no pixels, + which is all the analyzers' position logic needs. + Images are 16×16 PIL images: 8×8 patches of 2 px, merged 2×2 → 16 image tokens. """ @@ -33,6 +41,12 @@ GRID = IMG_PX // PATCH # 8 patches per side TOKENS_SIDE = GRID // MERGE # 4 merged tokens per side → 16 image tokens +HAS_QWEN3_5 = hasattr(transformers, "Qwen3_5ForConditionalGeneration") +#: Families whose stack mixes softmax and linear attention. +HYBRID_FAMILIES = ("qwen3_next",) + (("qwen3_5",) if HAS_QWEN3_5 else ()) +#: Families without a vision tower: images enter as placeholder tokens only. +TEXT_ONLY_FAMILIES = ("qwen3_next",) + # These models are tiny; on a many-core box torch's default one-thread-per-core # turns every matmul into scheduling contention (minutes instead of seconds). torch.set_num_threads(min(torch.get_num_threads(), 4)) @@ -65,11 +79,34 @@ def _build(family: str) -> Any: "mrope_interleaved": True}) cfg = transformers.Qwen3VLConfig(vision_config=vision, text_config=text, **ids) cls = transformers.Qwen3VLForConditionalGeneration + elif family == "qwen3_5": + vision.update(num_position_embeddings=64) + text.update(head_dim=16, layer_types=["linear_attention", "linear_attention", "full_attention"] * 2, + linear_num_value_heads=4, linear_num_key_heads=2, linear_key_head_dim=8, + linear_value_head_dim=8, linear_conv_kernel_dim=4, + rope_parameters={"rope_type": "default", "rope_theta": 10000.0, + "mrope_section": [2, 1, 1], "mrope_interleaved": True, + "partial_rotary_factor": 0.5}) + cfg = transformers.Qwen3_5Config(vision_config=vision, text_config=text, **ids) + cls = transformers.Qwen3_5ForConditionalGeneration + elif family == "qwen3_next": + cfg = transformers.Qwen3NextConfig( + **text, head_dim=8, full_attention_interval=3, linear_num_value_heads=4, + linear_num_key_heads=2, linear_key_head_dim=8, linear_value_head_dim=8, + linear_conv_kernel_dim=4, mlp_only_layers=list(range(6)), num_experts=2, + num_experts_per_tok=1) + cls = transformers.Qwen3NextForCausalLM else: # pragma: no cover - test helper misuse raise ValueError(family) return cls._from_config(cfg, attn_implementation="eager").eval() +def _accepts(module: Any, name: str) -> bool: + import inspect + + return name in inspect.signature(module.forward).parameters + + def _pixel_values(image: Any) -> torch.Tensor: """16×16 RGB → (64 patches, 3·2·2·2) in the processor's merge-window order.""" arr = np.asarray(image.convert("RGB").resize((IMG_PX, IMG_PX)), dtype=np.float32) / 127.5 - 1.0 @@ -93,6 +130,9 @@ def __init__(self, family: str = "qwen2_5_vl", caps.add(Capability.GRADIENTS) self.capabilities = frozenset(caps) self.modalities = frozenset({"text", "image"}) + if family in HYBRID_FAMILIES and attn_semantics is AttnSemantics.STANDARD: + attn_semantics = AttnSemantics.HYBRID_SPARSE # what the real specs declare + self.family = family self.spec = SimpleNamespace(key=f"tiny-{family}", attn_semantics=attn_semantics) self._hf = (_build(family), CharTokenizer()) self.n_forwards = 0 @@ -122,8 +162,9 @@ def forward(self, inputs: Any, capture: "set[Capability]", spec: Any = None) -> if image is not None: n_img = TOKENS_SIDE * TOKENS_SIDE ids = [1, 2, VSTART] + [IMAGE_TOKEN] * n_img + [VEND] + text_ids - enc["pixel_values"] = _pixel_values(image) - enc["image_grid_thw"] = torch.tensor([[1, GRID, GRID]]) + if self.family not in TEXT_ONLY_FAMILIES: + enc["pixel_values"] = _pixel_values(image) + enc["image_grid_thw"] = torch.tensor([[1, GRID, GRID]]) image_pos = list(range(3, 3 + n_img)) grids = [(1, TOKENS_SIDE, TOKENS_SIDE)] patches = [(0, 0, r, c) for r in range(TOKENS_SIDE) for c in range(TOKENS_SIDE)] @@ -131,6 +172,10 @@ def forward(self, inputs: Any, capture: "set[Capability]", spec: Any = None) -> ids = [1, 2] + text_ids image_pos, grids, patches = [], [], [] enc["input_ids"] = torch.tensor([ids]) + if "pixel_values" in enc and _accepts(module, "mm_token_type_ids"): + # transformers 5 computes M-RoPE from this processor output and + # refuses image inputs without it; 4.x models do not take it. + enc["mm_token_type_ids"] = (enc["input_ids"] == IMAGE_TOKEN).long() with_grad = bool(spec is not None and getattr(spec, "with_grad", False)) with torch.set_grad_enabled(with_grad): out = module(**enc, output_attentions=Capability.ATTENTION in capture, From a59bf6299c149793c13b415aa9354dac06a55702 Mon Sep 17 00:00:00 2001 From: Jiaqi Liu <32996169+jiaq-liu@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:10:52 -0700 Subject: [PATCH 7/9] add attribution_patching --- CHANGELOG.md | 33 +- docs/analyzers.md | 13 +- evalrx/analyzers/__init__.py | 4 +- evalrx/analyzers/_internals/residual.py | 8 + evalrx/analyzers/patching/__init__.py | 6 +- .../patching/attribution_patching.py | 480 ++++++++++++++++++ evalrx/analyzers/patching/causal_trace.py | 66 ++- evalrx/eval_agent/stages/probe.py | 14 +- evalrx/reporting/case_study.py | 6 +- .../test_analyzers/test_analyzer_contract.py | 1 + tests/test_analyzers/test_analyzer_docs.py | 1 + tests/test_analyzers/test_hybrid_analyzers.py | 15 +- .../test_intervention_analyzers.py | 189 +++++++ .../test_probe_routing_whitebox.py | 4 +- 14 files changed, 806 insertions(+), 34 deletions(-) create mode 100644 evalrx/analyzers/patching/attribution_patching.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c98ad35..696fbb4c 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,6 +64,35 @@ Verified on Qwen3.5-2B (transformers 5.14): tap 6/6 softmax layers exact, re-run 18/18 linear layers within 0.24% (bf16). Tests: a tiny Qwen3-Next (transformers ≥ 4.57) and a tiny Qwen3.5 VLM (transformers 5). +### Added — `attribution_patching`: causal tracing of single tokens, from gradients + +New `patching/attribution_patching` estimates the effect of restoring each +single token at each layer — `causal_trace`'s residual-stream site, corruptions +and noise draws — as `(h_clean − h_corrupt) · ∂margin/∂h`, with the gradient +averaged along the corrupt → clean media path (EAP-IG, `ig_steps=5`; +`ig_steps=1` = plain attribution patching, `direction="denoise" | "noise"`). +It reports the answer position's share of the media's effect per layer +(`ap_last_rise_frac`, `ap_last_share_mid`) and the strongest single relay token +(`ap_relay_token`, `ap_relay_share`, `ap_relay_depth_frac`), with a +cross-case tally of relay tokens. + +- Each case is checked against exact patches of the sites it reports (the last + position at `n_verify` layers, the relay token at its peak). Cases whose mean + gap exceeds `verify_tol` (0.25) are flagged `ap_verified=False` and kept out + of the aggregates. +- Measured on Qwen3-VL-2B and Qwen3.5-2B (8 deco_pope cases): the media tokens' + joint patch is not additive, and the sum of their gradient estimates does not + track it (r < 0.4 under every setting), so it is not reported. Single-token + estimates do track the exact patch where corrupting the media lowers the + margin (last position r 0.92–0.96), and fail where it makes the model more + confident (a saturated margin). With the shipped defaults, 2/6 (Qwen3-VL) and + 4/6 (Qwen3.5) traced cases pass their own check; 19 s / 55 s per 8 cases. +- `causal_trace`: the corrupted-variant construction is now + `corrupted_variants`, shared with the new analyzer (behaviour unchanged; the + exact checks reproduce `causal_trace`'s numbers digit for digit). +- `_internals/residual.py`: `keep_grad` also works with frozen weights (the + captured state becomes a leaf). + ### Changed — routing reaches the white-box analyzers `StrategyProbe` ranks the new analyzers for white-box image / audio / video @@ -74,7 +103,9 @@ media interventions, and new tags `visual_neglect`, `audio_neglect`, `mislocalized_attention`, `late_suppression` (also offered to the M3 prompt) route to the analyzer that tests them. The video list named the unregistered `relative_attn` (silently dropped); it is now `relative_attention`, and a test -fails on any unregistered name in the routing tables. +fails on any unregistered name in the routing tables. `attribution_patching` +follows `causal_trace` in the image / audio / video lists (the video list now +has both) and under `visual_information_loss` / `late_visual_integration`. ### Changed — `hf_local` provides `GRADIENTS` diff --git a/docs/analyzers.md b/docs/analyzers.md index 5e1ebbcf..9306a46f 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -24,11 +24,11 @@ model, not by a black-box / white-box label. | [Attention](#attention) | `attention/` | 4 | 0 | | [Attribution](#attribution) | `attribution/` | 2 | 0 | | [Lens](#lens) | `lens/` | 4 | 0 | -| [Patching](#patching) | `patching/` | 3 | 0 | +| [Patching](#patching) | `patching/` | 4 | 0 | | [Geometry](#geometry) | `geometry/` | 2 | 0 | | [Reasoning](#reasoning) | `reasoning/` | 7 | 0 | | [Agent](#agent) | `agent/` | 7 | 0 | -| **Total** | | **49** | **2** | +| **Total** | | **50** | **2** | How to read the columns: @@ -158,12 +158,21 @@ layer's text positions get its output re-run on the sequence without the media tokens (`_internals/linear_mix.py`). Latent / unverified attention is refused. Qwen3.5 needs transformers 5 (the `vllm35` environment), not the `[local]` extra. +`attribution_patching` estimates `causal_trace`'s patch for every single token +at every layer from gradients (`GRADIENTS`), averaged along the corrupt → clean +media path (EAP-IG). It reports single tokens only — the answer position, and +the question / template token that relays the media's information — because the +media tokens' joint patch is not additive and a sum of gradient estimates does +not track it. Each case is checked against exact patches of the sites it +reports, and only cases that pass enter the aggregates. + | Analyzer | Key | Requires | Modality | Paper | Status | |---|---|---|---|---|---| | Causal trace | `causal_trace` | `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Meng et al. 2022](https://arxiv.org/abs/2202.05262); [Basu et al. 2024](https://arxiv.org/abs/2406.04236) | ✓ | | Attention knockout | `attention_knockout` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Geva et al. 2023](https://arxiv.org/abs/2304.14767); [Zhang et al. 2025](https://arxiv.org/abs/2411.18620) | ✓ | | Attention sensitivity | `attention_sensitivity` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Liu et al. 2024 (PAI)](https://arxiv.org/abs/2407.21771) | ✓ | +| Attribution patching | `attribution_patching` | `GRADIENTS` + `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Syed et al. 2023](https://arxiv.org/abs/2310.10348); [Kramár et al. 2024 (AtP*)](https://arxiv.org/abs/2403.00745); [Hanna et al. 2024 (EAP-IG)](https://arxiv.org/abs/2403.17806) | ✓ | ### Geometry diff --git a/evalrx/analyzers/__init__.py b/evalrx/analyzers/__init__.py index cde5b2d7..87ed0620 100755 --- a/evalrx/analyzers/__init__.py +++ b/evalrx/analyzers/__init__.py @@ -13,8 +13,8 @@ attribution/ Grad-CAM / Chefer generic-attention — GRADIENTS lens/ logit-lens / tuned-lens / layer contrast / head logit attribution — HIDDEN_STATES - patching/ causal tracing / attention knockout / - attention sensitivity — HIDDEN_STATES / ATTENTION (write) + patching/ causal tracing / attribution patching / + attention knockout / attention sensitivity — HIDDEN_STATES / ATTENTION (write) geometry/ CKA / linear-probe — HIDDEN_STATES reasoning/ extraction/termination hygiene; arithmetic, self-repair, step-value, knowledge split, diff --git a/evalrx/analyzers/_internals/residual.py b/evalrx/analyzers/_internals/residual.py index f57bb845..e0fba543 100644 --- a/evalrx/analyzers/_internals/residual.py +++ b/evalrx/analyzers/_internals/residual.py @@ -49,7 +49,12 @@ def residual_hooks( Replacement happens before capture at the same layer. keep_grad: record the whole hidden tensor un-detached instead, so a caller can ``torch.autograd.grad`` a target with respect to it. + Under grad mode a state that does not require grad (frozen + weights upstream) is made a leaf that does — nothing upstream + needed a gradient, so no path is cut. """ + import torch + module = torch_module(model, analyzer) stack = decoder_stack(module) capture = dict(capture or {}) @@ -78,6 +83,9 @@ def hook(mod, args, kwargs): changed = True if idx in capture: if keep_grad: + if torch.is_grad_enabled() and not hidden.requires_grad: + hidden = hidden.detach().requires_grad_(True) + changed = True session.captured[idx] = hidden else: positions = capture[idx] diff --git a/evalrx/analyzers/patching/__init__.py b/evalrx/analyzers/patching/__init__.py index cfc6e078..a09427dc 100755 --- a/evalrx/analyzers/patching/__init__.py +++ b/evalrx/analyzers/patching/__init__.py @@ -1,7 +1,9 @@ -"""Causal intervention analyzers (causal tracing / attention knockout / attention reweighting).""" +"""Causal intervention analyzers (causal tracing and its gradient estimate / attention knockout / attention reweighting).""" from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer +from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer -__all__ = ["AttentionKnockoutAnalyzer", "AttentionSensitivityAnalyzer", "CausalTraceAnalyzer"] +__all__ = ["AttentionKnockoutAnalyzer", "AttentionSensitivityAnalyzer", "AttributionPatchingAnalyzer", + "CausalTraceAnalyzer"] diff --git a/evalrx/analyzers/patching/attribution_patching.py b/evalrx/analyzers/patching/attribution_patching.py new file mode 100644 index 00000000..9ca04ef9 --- /dev/null +++ b/evalrx/analyzers/patching/attribution_patching.py @@ -0,0 +1,480 @@ +"""Attribution patching — causal tracing of every single token at every layer, from gradients. + +``causal_trace`` patches GROUPS of positions exactly — all media tokens, or the +final prompt position — one forward pass per (layer, group). Attribution +patching estimates the effect of patching each SINGLE token ``t`` at each layer +``l`` at once, with a first-order expansion of the answer margin: + + effect(l, t) ≈ (h_clean[l, t] − h_corrupt[l, t]) · ḡ[l, t] + +where ``h[l, t]`` is the residual stream ENTERING decoder layer ``l`` (the site +``causal_trace`` patches, under the same corruptions and noise draws) and ``ḡ`` +is ``∂margin / ∂h[l, t]``: + +* ``ig_steps = 1`` — taken at one end: the CORRUPTED run for + ``direction="denoise"`` (Nanda: the estimate of restoring the clean token + into the corrupted run) or the CLEAN run for ``"noise"`` (Syed et al., + Kramár et al.: corrupting that one token of the clean run); +* ``ig_steps > 1`` (default) — averaged along the straight path that moves the + decoder's media input from its corrupted to its clean value (EAP-IG, Hanna et + al.): the media embeddings and, on DeepStack models, every re-injected copy of + the visual features. The pixels are not interpolated, so a gray image's path + does not run through the vision tower. + +Divided by the media effect ``m_clean − m_corrupt``, ``effect(l, t)`` is the +share of the media's effect that token carries at that depth. Reported, for +single tokens only: + +* ``last`` — the final prompt position, where the answer is read out: its + share per layer and the depth where it first carries half (read-out depth); +* ``relay`` — at each layer, the strongest single text token other than the + last one: the question word or template token through which the media's + information reaches the answer; across cases, which tokens those are. + +Not reported: the media tokens as a group. Restoring hundreds of image tokens +together is far from additive, and the sum of their single-token estimates +does not track the joint patch (on Qwen3-VL-2B its correlation with the exact +group patch was below 0.4 under every setting, against 0.9 and above for the +last position); ``causal_trace`` measures the group exactly. + +Each case checks itself: the reported sites — the last position at +``n_verify`` evenly spaced interior layers and the strongest relay token at its +layer — are also patched for real. A case whose mean gap exceeds +``verify_tol`` is kept, flagged ``ap_verified = False``, and left out of the +aggregate curves and the relay-token tally; ``verification`` pools estimate +against exact over all cases. The estimate fails mostly where corrupting the +media made the model MORE confident (a saturated margin, where gradients +under-read the effect — Kramár et al.). + +References: +- Attribution Patching: Activation Patching At Industrial Scale — Nanda, 2023 + (neelnanda.io/mechanistic-interpretability/attribution-patching) +- Attribution Patching Outperforms Automated Circuit Discovery + Syed, Rager & Conmy, BlackboxNLP 2024 — arXiv:2310.10348 +- AtP*: An efficient and scalable method for localizing LLM behaviour to + components — Kramár et al., 2024 — arXiv:2403.00745 +- Have Faith in Faithfulness: Going Beyond Circuit Overlap When Finding Model + Mechanisms (EAP-IG) — Hanna, Pezzelle & Belinkov, COLM 2024 — arXiv:2403.17806 +""" + +from __future__ import annotations + +from collections import Counter +from contextlib import contextmanager, nullcontext +from typing import TYPE_CHECKING, Any, Iterator + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + +DIRECTIONS = ("denoise", "noise") +SITES = ("last", "relay") + + +def _token_text(tokenizer: Any, token_id: int) -> "str | None": + """One token decoded, or ``None`` when the tokenizer cannot decode.""" + try: + text = tokenizer.decode([int(token_id)]) + except Exception: # noqa: BLE001 - no decoder → no token text, never a guess + return None + return text if isinstance(text, str) else None + + +@contextmanager +def _recorded(mergers: list, out: dict) -> Iterator[None]: + """Record each DeepStack merger's output into ``out[i]``.""" + def make(i: int): + def hook(mod, args, output): + out[i] = output.detach() + return hook + + handles = [m.register_forward_hook(make(i)) for i, m in enumerate(mergers)] + try: + yield + finally: + for handle in handles: + handle.remove() + + +def path_point(alpha: float, *, inputs: Any, media: "list[int]", clean_media: Any, + corrupt_media: Any, deltas: "dict | None") -> Any: + """The run *alpha* of the way from the corrupted (0) to the clean (1) media input. + + Runs on the clean *inputs*, writes ``clean + (1 − alpha)(corrupt − clean)`` + over the media embeddings at layer 0 and adds ``(1 − alpha)`` × each + DeepStack re-injection's corrupt-minus-clean difference (*deltas*). + """ + base = clean_media.float() + mixed = (base + (1.0 - alpha) * (corrupt_media.float() - base)).to(corrupt_media.dtype) + store = None if deltas is None else {i: d * (1.0 - alpha) for i, d in deltas.items()} + return inputs, {0: (media, mixed)}, (store, None) + + +def _pooled(pairs: "list[dict[str, Any]]") -> "dict[str, Any] | None": + """Estimate-vs-exact agreement over verified sites.""" + if not pairs: + return None + est = [p["estimate"] for p in pairs] + exact = [p["exact"] for p in pairs] + errors = [abs(a - b) for a, b in zip(est, exact)] + out: dict[str, Any] = { + "n_sites": len(pairs), + "mae": round(sum(errors) / len(errors), 4), + "max_abs_err": round(max(errors), 4), + } + if len(pairs) >= 3: + n = len(pairs) + me, mx = sum(est) / n, sum(exact) / n + cov = sum((a - me) * (b - mx) for a, b in zip(est, exact)) + var_e = sum((a - me) ** 2 for a in est) + var_x = sum((b - mx) ** 2 for b in exact) + if var_e > 0 and var_x > 0: + out["pearson"] = round(cov / (var_e * var_x) ** 0.5, 4) + return out + + +@register_analyzer("attribution_patching") +class AttributionPatchingAnalyzer(Analyzer): + """Gradient estimate of every single-token patch effect of the media, per case. + + Hyper-parameters: + corruption: ``"noise"`` (default, any media) or ``"gray"`` (images/video), + exactly as in ``causal_trace``. + noise_scale: noise std as a multiple of the media embeddings' std (Meng: 3). + n_noise: noise samples averaged per case (``"noise"`` only). + ig_steps: gradient points along the corrupt → clean path (midpoint rule); + 1 = plain attribution patching at one end (see ``direction``). + direction: ``"denoise"`` (restore clean into corrupt; default) or + ``"noise"`` (corrupt one token of the clean run): the + patch the estimate is checked against, and with + ``ig_steps=1`` the run whose gradient is used. + n_verify: interior layers where the last position's estimate is checked + by an exact patch (plus one check of the strongest relay + token); 0 = no exact patches (and every case is aggregated). + verify_tol: largest mean |estimate − exact| (in shares of the media + effect) at which a case counts as verified. + min_effect: smallest ``|m_clean - m_corrupt|`` (logits) worth attributing; + the estimates are divided by it (see ``causal_trace``). + target: ``"prediction"`` or ``"gold"``. + seed: noise seed — the same per-case noise as ``causal_trace``. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost per case and noise sample: 2 forwards, ``ig_steps`` + forward + backward passes, ``n_verify + 1`` forwards. + """ + + name = "attribution_patching" + description = ( + "Gradient estimate of causal tracing for every single token at every layer: how " + "much of the image/audio's effect the answer position and each question token " + "carry at each depth, checked against exact patches." + ) + requires = frozenset({Capability.GRADIENTS, Capability.HIDDEN_STATES, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "ap_margin_clean": ("Clean answer margin", "Logit margin of the target answer token with the media intact."), + "ap_media_effect": ("Margin lost to corrupt media", "How much the answer margin falls when the media tokens are corrupted; near zero = the answer did not depend on the media."), + "ap_media_effect_rel": ("Relative loss, corrupt media", "ap_media_effect as a fraction of the clean margin."), + "ap_effect_traceable": ("Media effect traceable", "Corrupting the media moved the margin by at least min_effect, so the effect could be attributed."), + "ap_last_rise_frac": ("Answer read-out depth (est.)", "Relative depth where the final prompt position first carries half of the media's effect, estimated at every layer; late = the media's content reaches the answer late."), + "ap_last_share_mid": ("Answer position share, mid", "Mean share of the media's effect carried by the final prompt position over the middle third of the layers."), + "ap_relay_share": ("Strongest relay token share", "Largest share of the media's effect carried by one text token other than the last, at any layer; high = the media reaches the answer through that token."), + "ap_relay_depth_frac": ("Strongest relay depth", "Relative depth at which that token carries it."), + "ap_relay_pos_from_end": ("Strongest relay position", "That token's distance from the end of the prompt, in tokens."), + "ap_relay_token": ("Strongest relay token", "That token's text (when the tokenizer can decode it)."), + "ap_verify_mae": ("Estimate vs exact patch", "Mean absolute gap between the estimated and the exactly patched share at the verified sites; large = do not read this case's curves."), + "ap_verified": ("Estimate verified", "ap_verify_mae is within verify_tol: this case's curves and relay token can be read (only these enter the aggregates)."), + } + + def __init__(self, corruption: str = "noise", noise_scale: float = 3.0, n_noise: int = 1, + ig_steps: int = 5, direction: str = "denoise", n_verify: int = 2, + verify_tol: float = 0.25, min_effect: float = 1.0, target: str = "prediction", + seed: int = 0, max_cases: int = 0) -> None: + from evalrx.analyzers.patching.causal_trace import CORRUPTIONS + + if corruption not in CORRUPTIONS: + raise ValueError(f"corruption must be one of {CORRUPTIONS}, got {corruption!r}") + if direction not in DIRECTIONS: + raise ValueError(f"direction must be one of {DIRECTIONS}, got {direction!r}") + if ig_steps < 1: + raise ValueError(f"ig_steps must be >= 1, got {ig_steps}") + super().__init__(corruption=corruption, noise_scale=noise_scale, n_noise=n_noise, + ig_steps=ig_steps, direction=direction, n_verify=n_verify, + verify_tol=verify_tol, min_effect=min_effect, target=target, seed=seed, + max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + from evalrx.analyzers._internals.access import ( + decoder_stack, + layer_points, + media_positions, + tokenizer_of, + torch_module, + ) + from evalrx.analyzers._internals.residual import residual_hooks + from evalrx.analyzers._internals.targets import build_margin + from evalrx.analyzers.patching.causal_trace import ( + _merge, + _noised, + corrupted_variants, + deepstack_mergers, + ) + from evalrx.core.model import CaptureSpec + + module = torch_module(model, self.name) + stack = decoder_stack(module) + mergers = deepstack_mergers(module) + n_layers = stack.n_layers + layers = list(range(n_layers)) + depth = [layer / max(1, n_layers - 1) for layer in layers] + mid = [layer for layer in layers if n_layers / 3 <= layer < 2 * n_layers / 3] or layers + verify_layers = layer_points(n_layers, self.n_verify + 2)[1:-1] if self.n_verify > 0 else [] + denoise = self.direction == "denoise" + integrated = self.ig_steps > 1 + alphas = [(j + 0.5) / self.ig_steps for j in range(self.ig_steps)] + tokenizer = tokenizer_of(model) + + def run(variant: Any, *, grad: bool = False, capture: bool = False, + restore: "dict | None" = None, record: "dict | None" = None) -> "tuple[Any, dict]": + inputs, plan, (store, gen) = variant + replace = _merge(plan, restore) if restore else plan + noised = _noised(mergers, store, self.noise_scale, gen) if store is not None else nullcontext() + recording = _recorded(mergers, record) if record is not None else nullcontext() + with recording, noised, residual_hooks( + model, capture=dict.fromkeys(layers) if capture else None, replace=replace, + keep_grad=grad, analyzer=self.name) as hooks: + trace = model.forward(inputs, capture={Capability.LOGITS}, + spec=CaptureSpec(with_grad=grad, to_cpu=False)) + return trace, hooks.captured + + def margin_of(margin: Any, variant: Any, restore: dict) -> float: + trace, _ = run(variant, restore=restore) + return margin.value(trace.require(Capability.LOGITS)[-1:]) + + def gradients(margin: Any, trace: Any, states: dict) -> "tuple[list, list]": + """``(∂margin/∂h[l], h[l])`` per layer, both ``(seq, dim)``.""" + hidden = [states[layer] for layer in layers] + if not all(h.requires_grad for h in hidden): + raise RuntimeError( + f"{self.name}: the forward pass kept no autograd graph — the backend " + "ignores CaptureSpec(with_grad=True), so there is nothing to differentiate." + ) + grads = torch.autograd.grad(margin(trace.require(Capability.LOGITS)[-1:]), hidden) + return [g[0] for g in grads], [h[0].detach() for h in hidden] + + def attributed(clean_h: list, corrupt_h: list, grads: list) -> Any: + """Σ_d (h_clean − h_corrupt) · ḡ, per (layer, token).""" + return torch.stack([ + ((clean_h[layer].float() - corrupt_h[layer].float()) * grads[layer].float()).sum(-1) + for layer in layers + ]) + + per_case: list[dict[str, Any]] = [] + curves: dict[str, dict[str, Any]] = {} + checks: dict[str, list[dict[str, Any]]] = {} + relay_tokens: Counter = Counter() + trusted: set = set() # traced cases whose estimates passed their exact check + skipped: Counter = Counter() + slots: Counter = Counter() + for index, case in enumerate(cases.stratified_head(self.max_cases)): + clean_variant = (case.inputs, {}, (None, None)) + # A gray image's DeepStack copies are interpolated too: record both runs' copies. + clean_rec: "dict | None" = {} if (integrated and mergers and self.corruption == "gray") else None + clean_grad_needed = not (denoise or integrated) + trace, states = run(clean_variant, grad=clean_grad_needed, capture=True, record=clean_rec) + slot, media = media_positions(trace) + if not media: + skipped["no_media_tokens"] += 1 + continue + logits = trace.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + m_clean = margin.value(logits) + if clean_grad_needed: + clean_grad, clean_h = gradients(margin, trace, states) + else: + clean_grad, clean_h = None, [states[layer] for layer in layers] + token_ids, last = trace.token_ids, trace.seq_len - 1 + del trace, states, logits + + variants, why = corrupted_variants( + case.inputs, self.corruption, slot=slot, media=media, embeddings=clean_h[0][media], + deepstack=bool(mergers), noise_scale=self.noise_scale, n_noise=self.n_noise, + seed=self.seed * 100003 + index) + if not variants: + skipped[why] += 1 + continue + + contrib: Any = None + m_corrupt, aligned = 0.0, True + corrupt_runs: list[list] = [] # "noise" checks write these corrupt states into the clean run + for variant in variants: + corrupt_rec: "dict | None" = {} if clean_rec is not None else None + trace, states = run(variant, grad=denoise and not integrated, capture=True, + record=corrupt_rec) + if trace.token_ids != token_ids: # a gray image that re-tokenised differently + aligned = False + break + m_corrupt += margin.value(trace.require(Capability.LOGITS)[-1:]) + if integrated: + corrupt_h = [states[layer] for layer in layers] + del trace, states + store = variant[2][0] + deltas = ({i: corrupt_rec[i] - clean_rec[i] for i in corrupt_rec} if corrupt_rec + else (dict(store) if store else None)) + term = None + for alpha in alphas: # ḡ averaged along the corrupt → clean media path + point = path_point(alpha, inputs=case.inputs, media=media, + clean_media=clean_h[0][media], + corrupt_media=corrupt_h[0][media], deltas=deltas) + trace, states = run(point, grad=True, capture=True) + grads, _ = gradients(margin, trace, states) + step = attributed(clean_h, corrupt_h, grads) / len(alphas) + term = step if term is None else term + step + del trace, states, grads, step + elif denoise: + grads, corrupt_h = gradients(margin, trace, states) + term = attributed(clean_h, corrupt_h, grads) + del trace, states, grads + else: + corrupt_h = [states[layer] for layer in layers] + term = attributed(clean_h, corrupt_h, clean_grad) + del trace, states + if not denoise: + corrupt_runs.append(corrupt_h) + contrib = term if contrib is None else contrib + term + del corrupt_h, term + if not aligned: + skipped["corruption_misaligned"] += 1 + continue + contrib = contrib / len(variants) + m_corrupt /= len(variants) + + effect = m_clean - m_corrupt + row: dict[str, Any] = { + "sample_id": case.id, + "ap_margin_clean": round(m_clean, 4), + "ap_media_effect": round(effect, 4), + "ap_effect_traceable": abs(effect) >= self.min_effect, + } + if m_clean > 1e-6: + row["ap_media_effect_rel"] = round(effect / m_clean, 4) + slots[slot] += 1 + if abs(effect) < self.min_effect: + per_case.append(row) + continue + + share = contrib / effect # (layers, seq) + text = torch.ones(share.shape[1], dtype=torch.bool, device=share.device) + text[media] = False + text[last] = False + last_curve = share[:, last].tolist() + row["ap_last_rise_frac"] = round(next((d for d, r in zip(depth, last_curve) if r >= 0.5), 1.0), 4) + row["ap_last_share_mid"] = round(sum(last_curve[layer] for layer in mid) / len(mid), 4) + sites = [(layer, last, "last", last_curve[layer]) for layer in verify_layers] + relay_curve: "list[float] | None" = None + relay_at: "list[int] | None" = None + if bool(text.any()): + best, where = share.masked_fill(~text, float("-inf")).max(1) + relay_curve, relay_at = best.tolist(), where.tolist() + peak = max(layers, key=lambda layer: relay_curve[layer]) + if relay_curve[peak] > 0: + pos = relay_at[peak] + row["ap_relay_share"] = round(relay_curve[peak], 4) + row["ap_relay_depth_frac"] = round(depth[peak], 4) + row["ap_relay_pos_from_end"] = last - pos + token = _token_text(tokenizer, token_ids[pos]) if tokenizer is not None else None + if token is not None: + row["ap_relay_token"] = token + if verify_layers: + sites.append((peak, pos, "relay", relay_curve[peak])) + + pairs: list[dict[str, Any]] = [] + for layer, pos, site, estimate in sites: + if denoise: + restore = {layer: ([pos], clean_h[layer][[pos]])} + patched = sum(margin_of(margin, v, restore) for v in variants) / len(variants) + exact = (patched - m_corrupt) / effect + else: + patched = sum(margin_of(margin, clean_variant, {layer: ([pos], run_h[layer][[pos]])}) + for run_h in corrupt_runs) / len(corrupt_runs) + exact = (m_clean - patched) / effect + pairs.append({"layer": layer, "site": site, + "estimate": round(estimate, 4), "exact": round(exact, 4)}) + if pairs: + row["ap_verify_mae"] = round( + sum(abs(p["estimate"] - p["exact"]) for p in pairs) / len(pairs), 4) + row["ap_verified"] = row["ap_verify_mae"] <= self.verify_tol + checks[case.id] = pairs + if row.get("ap_verified", True): + trusted.add(case.id) + if "ap_relay_token" in row: + relay_tokens[row["ap_relay_token"]] += 1 + per_case.append(row) + curves[case.id] = {"last": [round(r, 4) for r in last_curve]} + if relay_curve is not None: + curves[case.id]["relay"] = [round(r, 4) for r in relay_curve] + curves[case.id]["relay_pos_from_end"] = [last - p for p in relay_at] + del contrib, share, clean_h, clean_grad, corrupt_runs + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + label_of = {c.id: getattr(getattr(c, "label", None), "value", None) for c in cases} + + def mean_curve(label: str, site: str) -> "list[float] | None": + rows = [c[site] for cid, c in curves.items() + if cid in trusted and label_of.get(cid) == label and site in c] + if not rows: + return None + return [round(sum(col) / len(rows), 4) for col in zip(*rows)] + + every_pair = [p for pairs in checks.values() for p in pairs] + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_traced": len(curves), + "n_verified": len(trusted) if verify_layers else None, + "verify_tol": self.verify_tol, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": n_layers, + "media_slot": slots.most_common(1)[0][0] if slots else None, + "corruption": self.corruption, + "ig_steps": self.ig_steps, + "direction": self.direction, + "deepstack_layers": len(mergers), + "target": self.target, + "mean_media_effect": round(sum(r["ap_media_effect"] for r in per_case) / n, 4), + "frac_traceable": round(len(curves) / n, 4), + "verify_layers": verify_layers, + "verification": ({site: _pooled([p for p in every_pair if p["site"] == site]) for site in SITES} + | {"all": _pooled(every_pair)}) if every_pair else None, + "relay_tokens": relay_tokens.most_common(10), + "mean_share_curves": { + f"{label}_{site}": mean_curve(label, site) + for label in ("fail", "pass") for site in SITES + if mean_curve(label, site) is not None + }, + "per_case": per_case, + "_caveat": ( + "First-order (gradient) estimates of single-token residual-stream patch effects " + "on the next-token margin at the end of the prompt, as a share of the media's " + "effect" + (f", with the gradient averaged over {self.ig_steps} points of the " + "corrupt → clean media path" if integrated else "") + + ". The aggregate curves and relay tokens cover only cases whose estimates " + "passed their own exact check (ap_verified); the media tokens' joint effect " + "is not estimated here (it is not additive) — causal_trace measures it exactly." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, + artifacts={"share_curves": curves, "verification": checks}, + findings=findings) diff --git a/evalrx/analyzers/patching/causal_trace.py b/evalrx/analyzers/patching/causal_trace.py index f0acae82..82013a5c 100755 --- a/evalrx/analyzers/patching/causal_trace.py +++ b/evalrx/analyzers/patching/causal_trace.py @@ -118,6 +118,45 @@ def deepstack_mergers(module: Any) -> list: return mergers +def corrupted_variants( + inputs: "Inputs", + corruption: str, + *, + slot: str, + media: "list[int]", + embeddings: Any, + deepstack: bool, + noise_scale: float, + n_noise: int, + seed: int, +) -> "tuple[list[tuple[Any, dict, Any]], str]": + """The corrupted runs of one case: ``(variants, "")`` or ``([], skip_reason)``. + + A variant is ``(inputs, layer-0 replace plan, (DeepStack noise store, + generator))`` — one per noise sample (Gaussian noise of ``noise_scale`` × + the media embeddings' std written over *embeddings* at the *media* + positions; on DeepStack models the re-injected copies are noised from the + same generator on the variant's first forward, see :func:`_noised`), or the + single gray-image input. A gray variant still has to be checked for token + alignment by the caller. + """ + import torch + + if corruption == "gray": + grayed = gray_inputs(inputs) if slot == "image" else None + if grayed is None: + return [], "gray_needs_an_image" + return [(grayed, {}, (None, None))], "" + std = float(embeddings.float().std()) * noise_scale + gen = torch.Generator(device="cpu").manual_seed(seed) + variants = [] + for _ in range(max(1, n_noise)): + noise = torch.randn(embeddings.shape, generator=gen).to(embeddings.device, embeddings.dtype) * std + store = {} if deepstack and slot == "image" else None + variants.append((inputs, {0: (media, embeddings + noise)}, (store, gen))) + return variants, "" + + @contextmanager def _noised(mergers: list, store: dict, scale: float, gen: Any) -> Iterator[None]: """Add fixed noise to each merger output; drawn once per variant, then reused.""" @@ -187,8 +226,6 @@ def __init__(self, corruption: str = "noise", noise_scale: float = 3.0, n_noise: max_cases=max_cases) def _run(self, model: "Model", cases: "CaseBatch") -> Result: - import torch - from evalrx.analyzers._internals.access import ( decoder_stack, layer_points, @@ -241,27 +278,18 @@ def margin_with(variant: "tuple[Any, dict, Any]", margin: Any, restore: "dict | probe.logits = None # keep token ids for the alignment check, not the (seq, vocab) logits del clean, logits - # One corrupted variant per noise sample (or the single gray image): - # (inputs, layer-0 replace plan, (DeepStack noise store, generator)). - variants: list[tuple[Any, dict, Any]] = [] + variants, why = corrupted_variants( + case.inputs, self.corruption, slot=slot, media=media, embeddings=clean_media[0][1], + deepstack=bool(mergers), noise_scale=self.noise_scale, n_noise=self.n_noise, + seed=self.seed * 100003 + index) + if not variants: + skipped[why] += 1 + continue if self.corruption == "gray": - grayed = gray_inputs(case.inputs) if slot == "image" else None - if grayed is None: - skipped["gray_needs_an_image"] += 1 - continue - check = model.forward(grayed, capture={Capability.LOGITS}, spec=spec) + check = model.forward(variants[0][0], capture={Capability.LOGITS}, spec=spec) if check.token_ids != probe.token_ids: skipped["corruption_misaligned"] += 1 continue - variants.append((grayed, {}, (None, None))) - else: - emb = clean_media[0][1] - std = float(emb.float().std()) * self.noise_scale - gen = torch.Generator(device="cpu").manual_seed(self.seed * 100003 + index) - for _ in range(max(1, self.n_noise)): - noise = torch.randn(emb.shape, generator=gen).to(emb.device, emb.dtype) * std - store = {} if mergers and slot == "image" else None - variants.append((case.inputs, {0: (media, emb + noise)}, (store, gen))) m_corrupt = sum(margin_with(v, margin) for v in variants) / len(variants) effect = m_clean - m_corrupt diff --git a/evalrx/eval_agent/stages/probe.py b/evalrx/eval_agent/stages/probe.py index 73966445..ee3c21ee 100755 --- a/evalrx/eval_agent/stages/probe.py +++ b/evalrx/eval_agent/stages/probe.py @@ -91,16 +91,18 @@ class ModelKind(str, Enum): "logit_lens": ["logit_lens", "tuned_lens"], # White-box information-flow mechanisms (VLM / ALM). Each tag names what the # analyzers can falsify: is the media read at all (knockout / causal trace), - # would reading it more help (sensitivity), is the right region read - # (gradient maps), where does the answer get written or suppressed (heads / lens). + # where does its information travel, layer by layer and token by token + # (attribution patching), would reading it more help (sensitivity), is the + # right region read (gradient maps), where does the answer get written or + # suppressed (heads / lens). "visual_neglect": ["attention_knockout", "causal_trace", "modality_ablation", "attention_sensitivity"], "audio_neglect": ["attention_knockout", "causal_trace", "modality_ablation", "attention_sensitivity"], "modality_neglect": ["attention_knockout", "causal_trace", "modality_ablation"], "insufficient_visual_attention": ["attention_sensitivity", "attention_knockout", "attention"], - "visual_information_loss": ["causal_trace", "attention_knockout"], - "late_visual_integration": ["causal_trace", "attention_knockout"], + "visual_information_loss": ["causal_trace", "attribution_patching", "attention_knockout"], + "late_visual_integration": ["causal_trace", "attribution_patching", "attention_knockout"], "mislocalized_attention": ["gradcam", "generic_attention", "relative_attention"], "grounding_error": ["gradcam", "generic_attention", "relative_attention"], "late_suppression": ["tuned_lens", "logit_lens", "layer_contrast", "head_attribution"], @@ -150,6 +152,7 @@ class ModelKind(str, Enum): "attention_knockout", # white-box: is it READ, and at what depth? "relative_attention", # attention across frames "mm_shap", # per-slot reliance + "causal_trace", "attribution_patching", # where the video's effect hands off to the text ], "audio": [ # Grounding first, for the same reason the text list leads with hygiene: @@ -166,6 +169,7 @@ class ModelKind(str, Enum): "attention", "attention_sink", # where does it attend among audio tokens? "attention_sensitivity", # would attending more to the audio help? "causal_trace", # where the audio's effect hands off to the text + "attribution_patching", # ... estimated at every layer and token "verbalized_confidence", "calibration", ], "image": [ @@ -179,6 +183,7 @@ class ModelKind(str, Enum): "mm_shap", # text vs image reliance "gradcam", "generic_attention", # which image regions the answer depends on "causal_trace", # where the image's effect hands off to the text + "attribution_patching", # ... estimated at every layer and token "logprob_entropy", "self_consistency", ], "text": [], # filled below from the legacy LLM list @@ -198,6 +203,7 @@ class ModelKind(str, Enum): "mm_shap", # text vs image reliance "gradcam", "generic_attention", # which image regions the answer depends on "causal_trace", # where the image's effect hands off + "attribution_patching", # ... estimated at every layer and token "logprob_entropy", "self_consistency", ], ModelKind.AGENT: [ diff --git a/evalrx/reporting/case_study.py b/evalrx/reporting/case_study.py index c802005d..83b0915f 100644 --- a/evalrx/reporting/case_study.py +++ b/evalrx/reporting/case_study.py @@ -59,6 +59,7 @@ "attention_knockout": "internal_white_box", "attention_sensitivity": "internal_white_box", "causal_trace": "internal_white_box", + "attribution_patching": "internal_white_box", "mm_shap": "multimodal", "vl_shap": "multimodal", "modality_ablation": "multimodal", @@ -135,6 +136,7 @@ "multimodal_attribution": "Did the answer use the image or audio at all?", "attention_knockout": "Did the answer use the image or audio at all?", "causal_trace": "Did the answer use the image or audio at all?", + "attribution_patching": "Did the answer use the image or audio at all?", "attention_sensitivity": "Did the answer use the image or audio at all?", "gradcam": "Where was it looking?", "generic_attention": "Where was it looking?", @@ -219,6 +221,7 @@ "attention_knockout": "media attention knockout", "attention_sensitivity": "media attention rescaling", "causal_trace": "causal tracing", + "attribution_patching": "attribution patching", "mm_shap": "modality attribution", "vl_shap": "vision-language attribution", "modality_ablation": "modality ablation", @@ -243,7 +246,8 @@ "linear_probe", "gradcam", "generic_attention"}), ("logits & representations", {"logprob_entropy", "logit_lens", "layer_contrast", "tuned_lens", "head_attribution"}), - ("causal interventions", {"attention_knockout", "attention_sensitivity", "causal_trace"}), + ("causal interventions", {"attention_knockout", "attention_sensitivity", "causal_trace", + "attribution_patching"}), ], "multimodal": [ ("grounding & hallucination", {"pope", "chair", "opera", "vcd"}), diff --git a/tests/test_analyzers/test_analyzer_contract.py b/tests/test_analyzers/test_analyzer_contract.py index 9f129549..447042b5 100755 --- a/tests/test_analyzers/test_analyzer_contract.py +++ b/tests/test_analyzers/test_analyzer_contract.py @@ -227,6 +227,7 @@ def _long_prompt_batch() -> CaseBatch: # runs against a tiny real Qwen-VL in test_intervention_analyzers.py. "attention_knockout": "needs a live torch module; see test_intervention_analyzers.py", "causal_trace": "needs a live torch module; see test_intervention_analyzers.py", + "attribution_patching": "needs a live torch module; see test_intervention_analyzers.py", "attention_sensitivity": "needs a live torch module; see test_intervention_analyzers.py", "head_attribution": "needs a live torch module; see test_intervention_analyzers.py", } diff --git a/tests/test_analyzers/test_analyzer_docs.py b/tests/test_analyzers/test_analyzer_docs.py index 59280957..8cfc5aaf 100644 --- a/tests/test_analyzers/test_analyzer_docs.py +++ b/tests/test_analyzers/test_analyzer_docs.py @@ -37,6 +37,7 @@ def gen(): _SIBLING_ROWS = ( "| Attention knockout | `attention_knockout` | | | — | |\n" "| Attention sensitivity | `attention_sensitivity` | | | — | |\n" + "| Attribution patching | `attribution_patching` | | | — | |\n" ) diff --git a/tests/test_analyzers/test_hybrid_analyzers.py b/tests/test_analyzers/test_hybrid_analyzers.py index cfd669c0..64f92f40 100644 --- a/tests/test_analyzers/test_hybrid_analyzers.py +++ b/tests/test_analyzers/test_hybrid_analyzers.py @@ -204,12 +204,25 @@ def test_head_attribution_on_hybrid_covers_linear_heads(hybrid): assert all(0.0 <= r["dla_linear_share"] <= 1.0 for r in f["per_case"]) -@pytest.mark.parametrize("name", ["causal_trace", "gradcam", "tuned_lens"]) +@pytest.mark.parametrize("name", ["causal_trace", "attribution_patching", "gradcam", "tuned_lens"]) def test_architecture_agnostic_analyzers_run_on_hybrid(hybrid, name): from evalrx.analyzers.attribution.gradcam import GradCAMAnalyzer from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer analyzer = {"causal_trace": CausalTraceAnalyzer(min_effect=0.0, n_points=3), + "attribution_patching": AttributionPatchingAnalyzer(min_effect=0.0), "gradcam": GradCAMAnalyzer(), "tuned_lens": TunedLensAnalyzer(steps=10)}[name] _contract(analyzer.run(hybrid, _batch()), name, 4) + + +def test_attribution_patching_is_exact_in_the_small_corruption_limit_on_hybrid(hybrid): + # The gradient runs back through the gated-delta-rule kernel and its conv; + # with a tiny corruption the estimates must still meet the real patches. + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + for steps in (1, 5): + f = AttributionPatchingAnalyzer(noise_scale=1e-3, min_effect=0.0, ig_steps=steps, + n_verify=2).run(hybrid, _batch()).findings + assert f["verification"]["all"]["mae"] < 0.02 and f["verification"]["all"]["pearson"] > 0.99 diff --git a/tests/test_analyzers/test_intervention_analyzers.py b/tests/test_analyzers/test_intervention_analyzers.py index 7a639476..dc0e0e69 100644 --- a/tests/test_analyzers/test_intervention_analyzers.py +++ b/tests/test_analyzers/test_intervention_analyzers.py @@ -13,6 +13,7 @@ from __future__ import annotations +import contextlib import json import math @@ -423,6 +424,194 @@ def test_gray_inputs_keep_size_and_flatten_content(): assert gray_inputs(Inputs(prompt="x")) is None +# ── attribution patching ────────────────────────────────────────────────────── + +def test_residual_keep_grad_survives_frozen_weights(): + # A model loaded with requires_grad_(False) has no graph upstream of the + # residual stream; keep_grad makes the captured state a leaf so the gradient + # analyzers still work — and without grad mode it changes nothing. + from evalrx.core.model import CaptureSpec + + frozen = TinyVLM("qwen2_5_vl") + frozen._loaded[0].requires_grad_(False) + inputs = Inputs(prompt="q", image=image(0)) + with residual_hooks(frozen, capture={0: None, 3: None}, keep_grad=True) as s: + trace = frozen.forward(inputs, capture={Capability.LOGITS}, + spec=CaptureSpec(with_grad=True, to_cpu=False)) + grads = torch.autograd.grad(trace.logits[-1].sum(), [s.captured[0], s.captured[3]]) + assert all(bool(g.abs().sum() > 0) for g in grads) + with residual_hooks(frozen, capture={0: None}, keep_grad=True) as s: + frozen.forward(inputs, capture={Capability.LOGITS}) + assert not s.captured[0].requires_grad + + +@pytest.mark.parametrize("direction, ig_steps", [("denoise", 1), ("noise", 1), ("denoise", 5)]) +def test_attribution_patching_matches_exact_patches_for_a_small_corruption(vlm, direction, ig_steps): + # First order is exact in the limit: with a tiny corruption every reported + # single-token estimate must agree with the real patch of that token. + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + f = AttributionPatchingAnalyzer(noise_scale=1e-3, min_effect=0.0, direction=direction, + ig_steps=ig_steps, n_verify=2).run(vlm, _batch()).findings + checked = f["verification"] + assert checked["last"]["n_sites"] == 4 * 2 and checked["relay"]["n_sites"] >= 1 + assert checked["all"]["mae"] < 0.02 and checked["all"]["pearson"] > 0.99 + + +@pytest.mark.parametrize("corruption", ["noise", "gray"]) +def test_integration_path_runs_from_the_corrupted_to_the_clean_run(corruption): + # The path moves only the decoder's media input — the layer-0 media + # embeddings and every DeepStack re-injection — so its ends must be the + # corrupted run and the clean run exactly, pixels untouched. + from evalrx.analyzers.patching.attribution_patching import _recorded, path_point + from evalrx.analyzers.patching.causal_trace import ( + _noised, + corrupted_variants, + deepstack_mergers, + ) + + vlm = TinyVLM("qwen3_vl") + mergers = deepstack_mergers(vlm._loaded[0]) + inputs = Inputs(prompt="what?", image=image(1)) + + def forward(variant, record=None, capture=False): + run_inputs, plan, (store, gen) = variant + noised = _noised(mergers, store, 3.0, gen) if store is not None else contextlib.nullcontext() + recording = _recorded(mergers, record) if record is not None else contextlib.nullcontext() + with torch.no_grad(), recording, noised, residual_hooks( + vlm, capture={0: None} if capture else None, replace=plan) as s: + logits = vlm.forward(run_inputs, capture={Capability.LOGITS}).logits[-1] + return logits, s.captured + + clean_rec, corrupt_rec = {}, {} + clean, states = forward((inputs, {}, (None, None)), record=clean_rec, capture=True) + media = list(range(3, 19)) + variants, _ = corrupted_variants(inputs, corruption, slot="image", media=media, + embeddings=states[0][media], deepstack=True, + noise_scale=3.0, n_noise=1, seed=0) + corrupt, corrupt_states = forward(variants[0], record=corrupt_rec, capture=True) + store = variants[0][2][0] + deltas = ({i: corrupt_rec[i] - clean_rec[i] for i in corrupt_rec} if corruption == "gray" + else dict(store)) + assert len(deltas) == 2 and not torch.allclose(corrupt, clean) + ends = {alpha: forward(path_point(alpha, inputs=inputs, media=media, clean_media=states[0][media], + corrupt_media=corrupt_states[0][media], deltas=deltas))[0] + for alpha in (0.0, 1.0)} + assert torch.allclose(ends[0.0], corrupt, atol=1e-5) + assert torch.allclose(ends[1.0], clean, atol=1e-5) + + +def test_attribution_patching_exact_check_is_causal_traces_measurement(vlm): + # Same corruption, same noise draws: the exact patch of the last position it + # runs to verify itself is causal_trace's recovery at that layer, digit for digit. + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer + + batch = _batch() + ap = AttributionPatchingAnalyzer(min_effect=0.0, n_verify=1).run(vlm, batch) + ct = CausalTraceAnalyzer(min_effect=0.0, n_points=3).run(vlm, batch) + assert ap.findings["verify_layers"] == [2] and ct.findings["layers_probed"][1] == 2 + for cid, pairs in ap.artifacts["verification"].items(): + (last,) = [p for p in pairs if p["site"] == "last"] + assert last["exact"] == ct.artifacts["recovery_curves"][cid]["last"][1] + + +@pytest.mark.parametrize("corruption", ["noise", "gray"]) +def test_attribution_patching_contract(vlm, corruption): + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + result = AttributionPatchingAnalyzer(corruption=corruption, min_effect=0.0).run(vlm, _batch()) + f = _assert_result_contract(result, "attribution_patching", 4) + assert f["n_traced"] == 4 and f["direction"] == "denoise" and f["ig_steps"] == 5 + assert f["deepstack_layers"] == (2 if vlm.spec.key.endswith("qwen3_vl") else 0) + assert set(f["verification"]) == {"last", "relay", "all"} + assert set(f["mean_share_curves"]) == {f"{label}_{site}" for label in ("fail", "pass") + for site in ("last", "relay")} + assert all(len(curve) == 6 for curve in f["mean_share_curves"].values()) + assert f["relay_tokens"] == [] # the char tokenizer cannot decode: no token text made up + for row in f["per_case"]: + assert 0.0 <= row["ap_last_rise_frac"] <= 1.0 and row["ap_verify_mae"] >= 0.0 + assert "ap_relay_token" not in row + for curve in result.artifacts["share_curves"].values(): + # Nothing has left the media tokens yet at layer 0: no text token carries any of it. + assert curve["last"][0] == 0.0 and curve["relay"][0] == 0.0 + + +def test_attribution_patching_names_the_relay_token(vlm, monkeypatch): + from evalrx.analyzers._internals import access as access_mod + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + class Decoder: + def decode(self, ids): + return f"<{ids[0]}>" + + monkeypatch.setattr(access_mod, "tokenizer_of", lambda model: Decoder()) + batch = _batch() + f = AttributionPatchingAnalyzer(min_effect=0.0, n_verify=0).run(vlm, batch).findings + for case, row in zip(batch, f["per_case"]): + ids = vlm.forward(case.inputs, capture={Capability.LOGITS}).token_ids + assert row["ap_relay_token"] == f"<{ids[len(ids) - 1 - row['ap_relay_pos_from_end']]}>" + assert sum(count for _, count in f["relay_tokens"]) == 4 + + +def test_cases_failing_their_exact_check_stay_out_of_the_aggregates(vlm, monkeypatch): + from evalrx.analyzers._internals import access as access_mod + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + class Decoder: + def decode(self, ids): + return f"<{ids[0]}>" + + monkeypatch.setattr(access_mod, "tokenizer_of", lambda model: Decoder()) + strict = AttributionPatchingAnalyzer(min_effect=0.0, verify_tol=-1.0).run(vlm, _batch()).findings + assert strict["n_traced"] == 4 and strict["n_verified"] == 0 + assert all(row["ap_verified"] is False and "ap_relay_token" in row for row in strict["per_case"]) + assert strict["mean_share_curves"] == {} and strict["relay_tokens"] == [] + loose = AttributionPatchingAnalyzer(min_effect=0.0, verify_tol=1e9).run(vlm, _batch()).findings + assert loose["n_verified"] == 4 and sum(count for _, count in loose["relay_tokens"]) == 4 + + +def test_attribution_patching_small_effects_keep_the_row_but_skip_the_estimate(vlm): + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + result = AttributionPatchingAnalyzer().run(vlm, _batch(2)) # tiny model: effects << 1 logit + f = result.findings + assert f["n_cases_analyzed"] == 2 and f["n_traced"] == 0 and f["verification"] is None + for row in f["per_case"]: + assert row["ap_effect_traceable"] is False and "ap_last_rise_frac" not in row + assert result.artifacts["share_curves"] == {} + + +def test_attribution_patching_writes_no_parameter_grads(vlm): + from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer + + for direction, steps in (("denoise", 1), ("noise", 1), ("denoise", 3)): + AttributionPatchingAnalyzer(min_effect=0.0, direction=direction, ig_steps=steps, + n_verify=0).run(vlm, _batch(1)) + assert all(p.grad is None for p in vlm._loaded[0].parameters()) + + +def test_attribution_patching_argument_checks_and_token_text(): + from evalrx.analyzers.patching.attribution_patching import ( + AttributionPatchingAnalyzer, + _token_text, + ) + + with pytest.raises(ValueError, match="direction"): + AttributionPatchingAnalyzer(direction="sideways") + with pytest.raises(ValueError, match="corruption"): + AttributionPatchingAnalyzer(corruption="blur") + with pytest.raises(ValueError, match="ig_steps"): + AttributionPatchingAnalyzer(ig_steps=0) + + class Decoder: + def decode(self, ids): + return f"<{ids[0]}>" + + assert _token_text(Decoder(), 7) == "<7>" + assert _token_text(object(), 7) is None + + # ── head attribution ────────────────────────────────────────────────────────── def test_per_head_writes_sum_to_the_output_projection(vlm): diff --git a/tests/test_eval_agent/test_probe_routing_whitebox.py b/tests/test_eval_agent/test_probe_routing_whitebox.py index a91c2bc8..54a6908f 100644 --- a/tests/test_eval_agent/test_probe_routing_whitebox.py +++ b/tests/test_eval_agent/test_probe_routing_whitebox.py @@ -28,8 +28,8 @@ ) from tests.conftest import FakeModel # noqa: E402 -WRITE_SIDE = {"attention_knockout", "attention_sensitivity", "causal_trace", "gradcam", - "generic_attention", "head_attribution", "tuned_lens"} +WRITE_SIDE = {"attention_knockout", "attention_sensitivity", "causal_trace", "attribution_patching", + "gradcam", "generic_attention", "head_attribution", "tuned_lens"} _WHITE_BOX = {Capability.GENERATE, Capability.LOGITS, Capability.LOGPROBS, Capability.ATTENTION, Capability.HIDDEN_STATES, Capability.GRADIENTS} From d90500eaaecea8cb74f6beac081db62b1943a79b Mon Sep 17 00:00:00 2001 From: Jiaqi Liu <32996169+jiaq-liu@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:58:18 -0700 Subject: [PATCH 8/9] add path_patpatching and patchscope, fix bug in generic_attn --- CHANGELOG.md | 36 ++ docs/analyzers.md | 22 +- evalrx/analyzers/__init__.py | 5 +- evalrx/analyzers/_internals/heads.py | 85 ++++- evalrx/analyzers/_internals/residual.py | 39 ++ evalrx/analyzers/attribution/generic_attn.py | 3 +- evalrx/analyzers/lens/__init__.py | 7 +- evalrx/analyzers/lens/head_attribution.py | 9 +- evalrx/analyzers/lens/patchscope.py | 264 ++++++++++++++ evalrx/analyzers/patching/__init__.py | 5 +- evalrx/analyzers/patching/path_patching.py | 341 ++++++++++++++++++ evalrx/eval_agent/stages/probe.py | 20 +- evalrx/reporting/case_study.py | 9 +- evalrx/specs.py | 3 +- .../test_analyzers/test_analyzer_contract.py | 2 + tests/test_analyzers/test_analyzer_docs.py | 1 + tests/test_analyzers/test_hybrid_analyzers.py | 20 +- .../test_intervention_analyzers.py | 141 ++++++++ .../test_probe_routing_whitebox.py | 6 +- 19 files changed, 986 insertions(+), 32 deletions(-) create mode 100644 evalrx/analyzers/lens/patchscope.py create mode 100644 evalrx/analyzers/patching/path_patching.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 696fbb4c..a476068c 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -93,6 +93,42 @@ cross-case tally of relay tokens. - `_internals/residual.py`: `keep_grad` also works with frozen weights (the captured state becomes a leaf). +### Added — `path_patching` and `patchscope` + +- `patching/path_patching` (Wang et al., IOI): which heads carry the media's + effect directly to the answer. One head's output from the media-corrupted run + (`causal_trace`'s corruptions), every other token-mixing head frozen at its + clean output (the value heads of linear-attention layers included), MLPs and + norms recomputed. Heads are screened by direct logit attribution of their + media-dependent output and the `top_k` (8) patched exactly; per case the top + direct-path head, the joint direct / total share of the top-k, the screen's + rank agreement and (hybrids) the linear-attention share; across cases the + recurring heads. Patching every head at once reproduces the full effect at the + answer position — tested as an identity, on hybrids too. +- `lens/patchscope` (Ghandeharioun et al.): each layer's residual at the answer + position is patched into a token-identity prompt at the same layer and decoded + by the model: the depth where the answer first ranks first, its peak, and + `ps_late_drop` (peak minus the output probability; with `target="gold"`, a + right answer formed then suppressed). `media_layers > 0` also decodes every + image / audio token into a word ("what it sees"). +- Shared plumbing: `_internals/heads.mixer_hooks` (read / overwrite every + head's output at every position), `heads.readout_direction` (the frozen-norm + margin direction, now also used by `head_attribution`), + `residual.batched_patch` (a different residual patch per batch row). +- Measured on Qwen3.5-0.8B (16 POPE-adversarial cases): the screen ranks heads + like the exact patch (Spearman 0.86); the direct carriers recur in the last + two layers (`L23.H6` +0.33 of the effect in 8 cases; `L23.H0` −0.24 in 10 — + it consistently works against the image's effect). Patchscope's top-1 matches + the model's answer from layer 18 in 14/16 cases; the logit lens only at the + last layer, 6/16. Image tokens decode to scene words ("snow", "board", + "skiing", "slope" for a snowboarding photo). +- New spec `qwen3.5-0.8b-vl` (same 24-layer hybrid layout as the 2B; the fast + real-model check). +- Routing: `path_patching` after `attribution_patching` in the image / audio / + video lists and under `answer_head_conflict` / `visual_information_loss`; + `patchscope` after `tuned_lens` in the text list and under `late_suppression`, + `premature_layer_divergence`, `logit_lens`. + ### Changed — routing reaches the white-box analyzers `StrategyProbe` ranks the new analyzers for white-box image / audio / video diff --git a/docs/analyzers.md b/docs/analyzers.md index 9306a46f..e75d1bbd 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -23,12 +23,12 @@ model, not by a black-box / white-box label. | [Hallucination](#hallucination) | `hallucination/` | 5 | 2 | | [Attention](#attention) | `attention/` | 4 | 0 | | [Attribution](#attribution) | `attribution/` | 2 | 0 | -| [Lens](#lens) | `lens/` | 4 | 0 | -| [Patching](#patching) | `patching/` | 4 | 0 | +| [Lens](#lens) | `lens/` | 5 | 0 | +| [Patching](#patching) | `patching/` | 5 | 0 | | [Geometry](#geometry) | `geometry/` | 2 | 0 | | [Reasoning](#reasoning) | `reasoning/` | 7 | 0 | | [Agent](#agent) | `agent/` | 7 | 0 | -| **Total** | | **50** | **2** | +| **Total** | | **52** | **2** | How to read the columns: @@ -133,7 +133,11 @@ reads the decoder's image-token states rather than the vision tower. `evalrx/analyzers/lens/`: project intermediate hidden states through the unembedding (`HIDDEN_STATES`) to see at which layer a prediction forms, and where premature layers disagree with the final one; or project each attention head's own -write, to see which heads write the answer. +write, to see which heads write the answer. `patchscope` lets the model decode its +own states instead: each layer's residual at the answer position (and optionally +each image / audio token) is patched into a token-identity prompt at the same +layer. On Qwen3.5-0.8B it reads the final answer about five layers earlier than +the logit lens. | Analyzer | Key | Requires | Modality | Paper | Status | @@ -142,6 +146,7 @@ write, to see which heads write the answer. | Tuned lens | `tuned_lens` | `HIDDEN_STATES` | text + image | [Belrose et al. 2023](https://arxiv.org/abs/2303.08112) | ✓ | | Layer contrast | `layer_contrast` | `HIDDEN_STATES` | text + image | [Chuang et al. 2024](https://arxiv.org/abs/2309.03883); [Wang et al. 2025](https://arxiv.org/abs/2410.11779) | ✓ | | Head logit attribution | `head_attribution` | `HIDDEN_STATES` + `LOGITS` | text + image | [Elhage et al. 2021](https://transformer-circuits.pub/2021/framework/index.html); [Wang et al. 2023](https://arxiv.org/abs/2211.00593) | ✓ | +| Patchscope | `patchscope` | `HIDDEN_STATES` + `LOGITS` | text + image + audio + video | [Ghandeharioun et al. 2024](https://arxiv.org/abs/2401.06102) | ✓ | ### Patching @@ -166,6 +171,13 @@ media tokens' joint patch is not additive and a sum of gradient estimates does not track it. Each case is checked against exact patches of the sites it reports, and only cases that pass enter the aggregates. +`path_patching` asks which heads carry the media's effect DIRECTLY to the answer: +one head's output under corrupted media, every other head (linear-attention value +heads included) frozen at its clean output, MLPs recomputed. Heads are screened +by direct logit attribution and the top `top_k` patched exactly; patching every +head at once reproduces the whole effect at the answer position (a tested +identity), so the shares are exact fractions of it. + | Analyzer | Key | Requires | Modality | Paper | Status | |---|---|---|---|---|---| @@ -173,6 +185,7 @@ reports, and only cases that pass enter the aggregates. | Attention knockout | `attention_knockout` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Geva et al. 2023](https://arxiv.org/abs/2304.14767); [Zhang et al. 2025](https://arxiv.org/abs/2411.18620) | ✓ | | Attention sensitivity | `attention_sensitivity` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Liu et al. 2024 (PAI)](https://arxiv.org/abs/2407.21771) | ✓ | | Attribution patching | `attribution_patching` | `GRADIENTS` + `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Syed et al. 2023](https://arxiv.org/abs/2310.10348); [Kramár et al. 2024 (AtP*)](https://arxiv.org/abs/2403.00745); [Hanna et al. 2024 (EAP-IG)](https://arxiv.org/abs/2403.17806) | ✓ | +| Path patching | `path_patching` | `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Wang et al. 2023 (IOI)](https://arxiv.org/abs/2211.00593) | ✓ | ### Geometry @@ -287,6 +300,7 @@ chat template renders tools, which grants `TOOL_CALLS` on the local backend) and | `llama-4-scout` | `llama4` | text + image | MoE | — | | `llava-1.5-7b-hf` | `llava` | text + image | — | reference architecture for ViCrop | | `instructblip-vicuna-7b` | `instructblip` | text + image | — | reference Q-Former architecture | +| `qwen3.5-0.8b-vl` | `qwen3_5` | text + image | reasoning, tool calling | hybrid linear / full attention stack; smallest Qwen3.5 VLM (fast white-box checks) | ### Audio and omni models diff --git a/evalrx/analyzers/__init__.py b/evalrx/analyzers/__init__.py index 87ed0620..4c2edbd0 100755 --- a/evalrx/analyzers/__init__.py +++ b/evalrx/analyzers/__init__.py @@ -12,9 +12,10 @@ attention/ summary / rollout / sink / relative-attn — ATTENTION attribution/ Grad-CAM / Chefer generic-attention — GRADIENTS lens/ logit-lens / tuned-lens / layer contrast / - head logit attribution — HIDDEN_STATES + head logit attribution / patchscope — HIDDEN_STATES patching/ causal tracing / attribution patching / - attention knockout / attention sensitivity — HIDDEN_STATES / ATTENTION (write) + path patching / attention knockout / + attention sensitivity — HIDDEN_STATES / ATTENTION (write) geometry/ CKA / linear-probe — HIDDEN_STATES reasoning/ extraction/termination hygiene; arithmetic, self-repair, step-value, knowledge split, diff --git a/evalrx/analyzers/_internals/heads.py b/evalrx/analyzers/_internals/heads.py index a1e327dc..5acbf096 100644 --- a/evalrx/analyzers/_internals/heads.py +++ b/evalrx/analyzers/_internals/heads.py @@ -1,4 +1,5 @@ -"""Per-head attention outputs at one position, and the final norm frozen to a linear map. +"""Per-head token-mixer outputs (read at one position, or read / overwritten everywhere), +and the final norm frozen to a linear map. Every supported decoder token mixer — softmax attention, or the gated-delta-rule linear attention of a hybrid stack — ends in an output projection whose INPUT @@ -96,6 +97,74 @@ def grab_final(mod, args): handle.remove() +@dataclass +class MixerSession: + """``inputs[layer]`` — the output projection's whole input ``(seq, n_heads · head_dim)``; + ``final_residual`` — the final norm's input ``(seq, dim)``.""" + + inputs: "dict[int, Any]" = field(default_factory=dict) + final_residual: Any = None + + +@contextmanager +def mixer_hooks( + model: "Model", + *, + capture: bool = False, + replace: "dict[int, Any] | None" = None, + analyzer: str = "mixer hooks", +) -> Iterator[MixerSession]: + """Read and/or overwrite every token mixer's per-head outputs, all positions. + + ``replace[layer]`` — ``(seq, n_heads · head_dim)`` written over that layer's + output-projection input (heads concatenated, as the projection reads them), + so every head's write into the residual stream is set; the projection and + everything after it recompute. With *capture*, the inputs are recorded + (after any replacement) together with the final norm's input. + """ + from evalrx.models.backends.hf.discover import get_final_norm + + module = torch_module(model, analyzer) + stack = decoder_stack(module) + replace = dict(replace or {}) + session = MixerSession() + handles = [] + + def hook_for(layer: int): + def hook(mod, args): + x = args[0] + if x.dim() != 3 or x.shape[0] != 1: + raise RuntimeError(f"{analyzer}: expected a (1, seq, width) projection input, " + f"got {tuple(x.shape)}") + if layer in replace: + x = replace[layer].to(device=x.device, dtype=x.dtype).reshape(x.shape) + if capture: + session.inputs[layer] = x[0].detach() + return (x, *args[1:]) if layer in replace else None + return hook + + for layer, mixer in enumerate(stack.mixers): + proj = out_proj(mixer) if mixer is not None else None + if proj is None: + raise ValueError(f"{analyzer}: layer {layer} has no token-mixer output projection") + if capture or layer in replace: + handles.append(proj.register_forward_pre_hook(hook_for(layer))) + if capture: + norm = get_final_norm(module) + if norm is None: + raise ValueError(f"{analyzer}: no final norm found after the decoder layers") + + def grab_final(mod, args): + session.final_residual = args[0][0].detach() + + handles.append(norm.register_forward_pre_hook(grab_final)) + try: + yield session + finally: + for handle in handles: + handle.remove() + + def frozen_norm_scale(norm: Any, x: Any) -> "tuple[Any, bool]": """``(scale, centre)``: ``norm(x + c) ≈ norm(x) + scale ⊙ (c - centre·mean(c))``. @@ -113,3 +182,17 @@ def frozen_norm_scale(norm: Any, x: Any) -> "tuple[Any, bool]": return weight / torch.sqrt(xf.var(unbiased=False) + eps), True w_eff = norm(torch.ones_like(x)).float() * (1.0 + eps) ** 0.5 # undo rms(1) = sqrt(1 + eps) return w_eff / torch.sqrt(xf.pow(2).mean() + eps), False + + +def readout_direction(norm: Any, unembed: Any, x_final: Any, positive: int, negative: int) -> Any: + """``v`` with ``margin ≈ x · v`` around *x_final* (the final norm frozen there). + + A write ``c`` into the final residual moves the ``positive − negative`` logit + margin by ``c · v``: the frozen norm scale times the unembedding difference + (mean-removed for LayerNorm, whose centring is linear). + """ + scale, centre = frozen_norm_scale(norm, x_final) + w_u = unembed.weight + direction = (w_u[positive].float() - w_u[negative].float()).to(scale.device) + v = scale * direction + return v - v.mean() if centre else v diff --git a/evalrx/analyzers/_internals/residual.py b/evalrx/analyzers/_internals/residual.py index e0fba543..21195a2b 100644 --- a/evalrx/analyzers/_internals/residual.py +++ b/evalrx/analyzers/_internals/residual.py @@ -108,3 +108,42 @@ def hook(mod, args, kwargs): finally: for handle in handles: handle.remove() + + +@contextmanager +def batched_patch( + model: "Model", + patches: "Mapping[int, Sequence[tuple[int, int, Any]]]", + *, + analyzer: str = "batched patch", +) -> Iterator[None]: + """Overwrite single residual vectors in a BATCHED forward, row by row. + + ``patches[layer] = [(row, position, vector), ...]`` writes ``vector`` + (``(dim,)``) over ``hidden[row, position]`` entering decoder layer ``layer`` + — one forward can then carry a different patch (layer, source) per row. + """ + module = torch_module(model, analyzer) + stack = decoder_stack(module) + for idx in patches: + if not 0 <= idx < stack.n_layers: + raise ValueError(f"{analyzer}: layer {idx} outside [0, {stack.n_layers})") + + def make_hook(idx: int): + def hook(mod, args, kwargs): + hidden = args[0] if args else kwargs["hidden_states"] + hidden = hidden.clone() + for row, position, vector in patches[idx]: + hidden[row, position] = vector.to(device=hidden.device, dtype=hidden.dtype) + if args: + return (hidden, *args[1:]), kwargs + return args, {**kwargs, "hidden_states": hidden} + return hook + + handles = [stack.layers[idx].register_forward_pre_hook(make_hook(idx), with_kwargs=True) + for idx in sorted(patches) if patches[idx]] + try: + yield + finally: + for handle in handles: + handle.remove() diff --git a/evalrx/analyzers/attribution/generic_attn.py b/evalrx/analyzers/attribution/generic_attn.py index 3ce13605..18dc69cd 100755 --- a/evalrx/analyzers/attribution/generic_attn.py +++ b/evalrx/analyzers/attribution/generic_attn.py @@ -159,7 +159,8 @@ def hook(grad: Any, layer: int = layer, held: Any = held) -> None: # Asking for the first layer's mixing gradient sends the backward # pass through every layer's mixer, firing each layer's hook once. torch.autograd.grad(margin(logits), recorded[min(recorded)]) - del recorded, linear_out, logits + recorded.clear() # drop the graph-holding tensors; the closure keeps the dict + del linear_out, logits if len(abar) != n_layers: raise RuntimeError(f"{self.name}: relevance reached {len(abar)} of {n_layers} layers") diff --git a/evalrx/analyzers/lens/__init__.py b/evalrx/analyzers/lens/__init__.py index d0f54acb..0aac6a30 100755 --- a/evalrx/analyzers/lens/__init__.py +++ b/evalrx/analyzers/lens/__init__.py @@ -1,8 +1,11 @@ -"""Lens analyzers — project hidden states (or per-component writes) toward the vocabulary.""" +"""Lens analyzers — project hidden states (or per-component writes) toward the vocabulary, +or let the model decode them itself (Patchscope).""" from evalrx.analyzers.lens.head_attribution import HeadAttributionAnalyzer from evalrx.analyzers.lens.layer_contrast import LayerContrastAnalyzer from evalrx.analyzers.lens.logit_lens import LogitLensAnalyzer +from evalrx.analyzers.lens.patchscope import PatchscopeAnalyzer from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer -__all__ = ["LogitLensAnalyzer", "TunedLensAnalyzer", "LayerContrastAnalyzer", "HeadAttributionAnalyzer"] +__all__ = ["LogitLensAnalyzer", "TunedLensAnalyzer", "LayerContrastAnalyzer", "HeadAttributionAnalyzer", + "PatchscopeAnalyzer"] diff --git a/evalrx/analyzers/lens/head_attribution.py b/evalrx/analyzers/lens/head_attribution.py index fab2bdd2..624bf5e9 100644 --- a/evalrx/analyzers/lens/head_attribution.py +++ b/evalrx/analyzers/lens/head_attribution.py @@ -82,10 +82,10 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: from evalrx.analyzers._internals.access import decoder_stack, torch_module from evalrx.analyzers._internals.heads import ( - frozen_norm_scale, head_capture, n_heads_of, out_proj, + readout_direction, ) from evalrx.analyzers._internals.targets import build_margin from evalrx.models.backends.hf.discover import get_final_norm, get_unembed @@ -118,12 +118,7 @@ def _run(self, model: "Model", cases: "CaseBatch") -> Result: m_value = margin.value(logits) with torch.no_grad(): x_final = cap.final_residual - scale, centre = frozen_norm_scale(norm, x_final) - w_u = unembed.weight - direction = (w_u[margin.positive].float() - w_u[margin.negative].float()).to(scale.device) - v = scale * direction - if centre: - v = v - v.mean() + v = readout_direction(norm, unembed, x_final, margin.positive, margin.negative) readout = float(x_final.float().to(v.device) @ v) bias = getattr(unembed, "bias", None) if bias is not None: diff --git a/evalrx/analyzers/lens/patchscope.py b/evalrx/analyzers/lens/patchscope.py new file mode 100644 index 00000000..27edef1b --- /dev/null +++ b/evalrx/analyzers/lens/patchscope.py @@ -0,0 +1,264 @@ +"""Patchscope — read a hidden state by letting the model itself decode it. + +Logit lens and tuned lens push an intermediate residual straight to the +vocabulary. A Patchscope (Ghandeharioun et al.) instead writes the residual +into a second forward pass — a few-shot token-identity prompt, +``"cat -> cat; 1135 -> 1135; hello -> hello; ?"`` — at its placeholder +position and the SAME layer, and reads what the model predicts there: the +remaining layers do the decoding. Ghandeharioun et al. find this reads the +model's eventual next token better than either lens from mid depth on, with no +training. + +Two sources per case: + +* ``last`` (always) — the final prompt position, where the answer is formed, + at every layer (or ``n_points`` of them): the probability the decoder gives + the target token (the model's own answer, or the gold one with + ``target="gold"``) and whether it ranks first. The depth where the answer + first ranks first; the depth where it peaks; and the gap between that peak + and what the model finally outputs — positive = the answer was readable more + confidently mid-stack than at the output (with ``target="gold"`` on failing + cases: formed, then suppressed). +* ``media`` (``media_layers > 0``) — every image / audio token at a few + evenly spaced layers above the embedding, decoded the same way: the word the + model's own machinery reads off each media token ("what it sees"), kept as a + per-token word map and tallied (content words only) across cases. + +The patch is checked once per model: patching a target position with the +target's own residual must leave its logits unchanged. + +References: +- Patchscopes: A Unifying Framework for Inspecting Hidden Representations of + Language Models — Ghandeharioun et al., ICML 2024 — arXiv:2401.06102 +- interpreting GPT: the logit lens — nostalgebraist, 2020 (the baseline it improves on) +""" + +from __future__ import annotations + +from collections import Counter +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + +IDENTITY_PROMPT = "cat -> cat; 1135 -> 1135; hello -> hello; ?" +# Function words the decoder emits for almost any state; kept out of the word tally. +_STOP = frozenset("""a an and are as at be but by for from has have in into is it its of on or +that the their there this to was were with""".split()) + + +def _word(tokenizer: Any, token_id: int) -> "str | None": + try: + text = tokenizer.decode([int(token_id)]) + except Exception: # noqa: BLE001 - no decoder → no word, never a guess + return None + return text if isinstance(text, str) else None + + +@register_analyzer("patchscope") +class PatchscopeAnalyzer(Analyzer): + """Token-identity Patchscope readout of the answer position (and, optionally, the media tokens). + + Hyper-parameters: + target: ``"prediction"`` (the model's own top-1) or ``"gold"``. + target_prompt: the decoding prompt; its LAST token is the placeholder + that receives the patch (default: token identity, 3 shots). + n_points: layers read at the answer position; 0 (default) = all. + media_layers: layers at which every media token is decoded too, spread over + the stack above the embedding (which decodes to noise); 0 = off. + batch_rows: rows per batched decoding forward. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost per case: 1 forward + one batched forward of the + short target prompt (+ media tokens × media_layers rows). + """ + + name = "patchscope" + description = ( + "Decodes the answer position's hidden state at every layer with the model itself " + "(Patchscope, token-identity prompt): when the answer becomes readable, and whether " + "it was readable more strongly mid-stack than at the output; optionally decodes the " + "image/audio tokens into words." + ) + requires = frozenset({Capability.HIDDEN_STATES, Capability.LOGITS}) + applies_to_modalities = frozenset({"text", "image", "video", "audio"}) + signal_docs = { + "ps_final_prob": ("Answer prob. at the output", "Probability the model itself gives the target token."), + "ps_first_depth_frac": ("Answer readable from (depth)", "First relative depth at which the Patchscope readout of the answer position ranks the target token first (1.0 = never before the last layer)."), + "ps_peak_prob": ("Peak readable answer prob.", "Highest probability the Patchscope readout gives the target token at any depth."), + "ps_peak_depth_frac": ("Peak readout depth", "Relative depth of that peak."), + "ps_late_drop": ("Readable answer lost late", "Peak readout probability minus the model's final probability of the target token; positive = more confident mid-stack than at the output (suppressed late)."), + "ps_agree_mid": ("Mid-stack readout agrees", "Share of the middle third of the layers whose readout ranks the model's final answer first."), + } + + def __init__(self, target: str = "prediction", target_prompt: str = IDENTITY_PROMPT, + n_points: int = 0, media_layers: int = 0, batch_rows: int = 64, + max_cases: int = 0) -> None: + super().__init__(target=target, target_prompt=target_prompt, n_points=n_points, + media_layers=media_layers, batch_rows=batch_rows, max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + from evalrx.analyzers._internals.access import ( + decoder_stack, + layer_points, + media_positions, + tokenizer_of, + torch_module, + ) + from evalrx.analyzers._internals.residual import batched_patch, residual_hooks + from evalrx.analyzers._internals.targets import build_margin + from evalrx.core.model import CaptureSpec + + module = torch_module(model, self.name) + stack = decoder_stack(module) + n_layers = stack.n_layers + layers = layer_points(n_layers, self.n_points) + media_at = layer_points(n_layers, self.media_layers + 1)[1:] if self.media_layers > 0 else [] + mid = [layer for layer in layers if n_layers / 3 <= layer < 2 * n_layers / 3] or layers + depth = {layer: layer / max(1, n_layers - 1) for layer in range(n_layers)} + tokenizer = tokenizer_of(model) + special = set(getattr(tokenizer, "all_special_tokens", None) or []) + if tokenizer is None: + raise ValueError(f"{self.name}: needs the model's tokenizer to build the target prompt") + prompt_ids = tokenizer(self.target_prompt, add_special_tokens=False)["input_ids"] + if prompt_ids and isinstance(prompt_ids[0], (list, tuple)): + prompt_ids = prompt_ids[0] + if len(prompt_ids) < 2: + raise ValueError(f"{self.name}: target_prompt must tokenize to at least two tokens") + placeholder = len(prompt_ids) - 1 + device = next(module.parameters()).device + spec = CaptureSpec(to_cpu=False) + + def decode(sources: "list[tuple[int, Any]]") -> Any: + """Placeholder logits ``(rows, vocab)``: row ``i`` carries ``sources[i] = (layer, vector)``.""" + out = [] + for start in range(0, len(sources), max(1, self.batch_rows)): + chunk = sources[start:start + self.batch_rows] + patches: dict[int, list] = {} + for row, (layer, vector) in enumerate(chunk): + patches.setdefault(layer, []).append((row, placeholder, vector)) + ids = torch.tensor([prompt_ids] * len(chunk), device=device) + with torch.no_grad(), batched_patch(model, patches, analyzer=self.name): + logits = module(input_ids=ids).logits + out.append(logits[:, placeholder].float()) + return torch.cat(out) + + # The patch itself must be inert: writing the target's own residual back. + check_layer = layers[len(layers) // 2] + with torch.no_grad(): + own: dict[str, Any] = {} + + def grab(mod, args, kwargs): + own["h"] = (args[0] if args else kwargs["hidden_states"])[0, placeholder].detach().clone() + + handle = stack.layers[check_layer].register_forward_pre_hook(grab, with_kwargs=True) + try: + plain = module(input_ids=torch.tensor([prompt_ids], device=device)).logits[0, placeholder].float() + finally: + handle.remove() + patched = decode([(check_layer, own["h"])])[0] + patch_diff = float((patched - plain).abs().max()) + if patch_diff > 1e-3 * max(1.0, float(plain.abs().max())): + raise RuntimeError(f"{self.name}: re-writing a residual with itself changed the logits " + f"by {patch_diff:.3g}; the patch does not reach the layer input") + + per_case: list[dict[str, Any]] = [] + curves: dict[str, dict[str, Any]] = {} + media_words: dict[str, dict[int, list]] = {} + word_counts: Counter = Counter() + skipped: Counter = Counter() + for case in cases.stratified_head(self.max_cases): + with torch.no_grad(), residual_hooks(model, capture=dict.fromkeys(set(layers) | set(media_at)), + analyzer=self.name) as res: + trace = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + last = trace.seq_len - 1 + _, media = media_positions(trace) if media_at else (None, []) + logits = trace.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + final_probs = torch.softmax(logits[-1].float(), -1) + final_top = int(final_probs.argmax()) + target_id = margin.positive + del trace + + states = res.captured # layer -> (seq, dim) + readout = torch.softmax(decode([(layer, states[layer][last]) for layer in layers]), -1) + p_target = readout[:, target_id].tolist() + top1 = readout.argmax(-1).tolist() + peak = max(range(len(layers)), key=lambda i: p_target[i]) + p_final = float(final_probs[target_id]) + row: dict[str, Any] = { + "sample_id": case.id, + "ps_final_prob": round(p_final, 4), + "ps_first_depth_frac": round(next( + (depth[layer] for layer, t in zip(layers, top1) if t == target_id), 1.0), 4), + "ps_peak_prob": round(p_target[peak], 4), + "ps_peak_depth_frac": round(depth[layers[peak]], 4), + "ps_late_drop": round(p_target[peak] - p_final, 4), + "ps_agree_mid": round( + sum(top1[layers.index(layer)] == final_top for layer in mid) / len(mid), 4), + } + per_case.append(row) + curves[case.id] = {"p_target": [round(p, 4) for p in p_target], + "top1": [_word(tokenizer, t) or str(t) for t in top1]} + + if media_at and media: + rows = [(layer, states[layer][pos]) for layer in media_at for pos in media] + ids = decode(rows).argmax(-1).tolist() + words: dict[int, list] = {} + for (layer, _), token in zip(rows, ids): + word = _word(tokenizer, token) + words.setdefault(layer, []).append(word if word is not None else str(token)) + key = (word or "").strip().lower() + if word not in special and key.isalpha() and key not in _STOP: + word_counts[key] += 1 + media_words[case.id] = words + del res, states + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + label_of = {c.id: getattr(getattr(c, "label", None), "value", None) for c in cases} + + def mean_curve(label: str) -> "list[float] | None": + rows = [c["p_target"] for cid, c in curves.items() if label_of.get(cid) == label] + if not rows: + return None + return [round(sum(col) / len(rows), 4) for col in zip(*rows)] + + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": n_layers, + "layers_read": layers, + "target": self.target, + "target_prompt": self.target_prompt, + "patch_check_max_diff": round(patch_diff, 6), + "mean_late_drop": round(sum(r["ps_late_drop"] for r in per_case) / n, 4), + "mean_first_depth_frac": round(sum(r["ps_first_depth_frac"] for r in per_case) / n, 4), + "mean_p_target_curves": {label: mean_curve(label) for label in ("fail", "pass") + if mean_curve(label) is not None}, + "per_case": per_case, + "_caveat": ( + "Each layer's residual at the final prompt position is decoded by the model " + "itself, patched into the placeholder of the target prompt at the same layer. " + "A readout says what the model's remaining layers make of that state in a " + "copy-the-token context, not what the original run did with it." + ), + } + if media_at: + findings["media_layers_read"] = media_at + findings["media_word_counts"] = word_counts.most_common(20) + return Result(analyzer=self.name, model=repr(model), cases=cases, + artifacts={"readout_curves": curves, "media_words": media_words}, + findings=findings) diff --git a/evalrx/analyzers/patching/__init__.py b/evalrx/analyzers/patching/__init__.py index a09427dc..8ba81c75 100755 --- a/evalrx/analyzers/patching/__init__.py +++ b/evalrx/analyzers/patching/__init__.py @@ -1,9 +1,10 @@ -"""Causal intervention analyzers (causal tracing and its gradient estimate / attention knockout / attention reweighting).""" +"""Causal intervention analyzers (causal tracing and its gradient estimate / path patching / attention knockout / attention reweighting).""" from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer +from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer __all__ = ["AttentionKnockoutAnalyzer", "AttentionSensitivityAnalyzer", "AttributionPatchingAnalyzer", - "CausalTraceAnalyzer"] + "CausalTraceAnalyzer", "PathPatchingAnalyzer"] diff --git a/evalrx/analyzers/patching/path_patching.py b/evalrx/analyzers/patching/path_patching.py new file mode 100644 index 00000000..3d051d94 --- /dev/null +++ b/evalrx/analyzers/patching/path_patching.py @@ -0,0 +1,341 @@ +"""Path patching — which heads carry the media's effect DIRECTLY to the answer? + +Path patching (Wang et al., IOI, Appendix B) measures one head's effect along a +chosen path only. Here the path is a head → the answer logits: in the clean +run, the sender head's output is replaced by its value in the run whose media +was corrupted (``causal_trace``'s corruptions and noise draws), while EVERY +other token-mixing head — attention heads, and the value heads of a hybrid +stack's linear-attention layers — is frozen at its clean output. MLPs and norms +recompute. What the answer margin then loses travels from the sender to the +logits directly or through MLPs, never through another head: + + direct share(h) = (m_clean − m_patched) / (m_clean − m_corrupt) + +— the share of the media's effect on the answer that flows through ``h``'s +direct path. Without the freeze the same patch measures ``h``'s TOTAL effect +(through later heads too). + +Patching every head exactly costs a forward pass per head, so the heads are +screened first: each head's media-dependent output at the answer position is +read through the frozen final norm into the margin (direct logit attribution of +``o_corrupt − o_clean``, as in ``head_attribution``), and the ``top_k`` heads by +that estimate are patched for real. The screen ignores MLP mediation and the +norm's change; ``pp_screen_spearman`` says how well it ranked the heads that +were patched. + +Reported per case: the strongest direct-path head (exact), the joint direct +and total shares of the top-k set patched together, the screen's agreement, +and on hybrids how much of the direct effect runs through linear-attention +heads. Across cases: which heads recur as the media's direct carriers. + +References: +- Interpretability in the Wild: a Circuit for Indirect Object Identification in + GPT-2 small — Wang et al., ICLR 2023 — arXiv:2211.00593 +- Locating and Editing Factual Associations in GPT (corruption) — Meng et al., + NeurIPS 2022 — arXiv:2202.05262 +""" + +from __future__ import annotations + +from collections import Counter, defaultdict +from contextlib import nullcontext +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + + +def head_name(layer: int, head: int) -> str: + return f"L{layer}.H{head}" + + +def sender_inputs(base: "dict[int, Any]", donor: "dict[int, Any]", senders: "list[tuple[int, int]]", + n_heads: "list[int]", *, freeze: bool) -> "dict[int, Any]": + """Output-projection inputs for a patched run: the *senders*' head slices from + *donor*, everything else from *base* — at every layer when *freeze* (path + patching: no other head may respond), else only at the senders' layers + (the senders' total effect).""" + by_layer: dict[int, list[int]] = defaultdict(list) + for layer, head in senders: + by_layer[layer].append(head) + out = {} + for layer in (range(len(n_heads)) if freeze else sorted(by_layer)): + x = base[layer] + if by_layer.get(layer): + x = x.clone() + width = x.shape[-1] // n_heads[layer] + for head in by_layer[layer]: + x[:, head * width:(head + 1) * width] = donor[layer][:, head * width:(head + 1) * width] + out[layer] = x + return out + + +def _spearman(a: "list[float]", b: "list[float]") -> "float | None": + """Rank correlation (no tie correction — the inputs are continuous).""" + n = len(a) + if n < 3: + return None + + def ranks(xs: "list[float]") -> "list[float]": + order = sorted(range(n), key=lambda i: xs[i]) + out = [0.0] * n + for rank, i in enumerate(order): + out[i] = float(rank) + return out + + ra, rb = ranks(a), ranks(b) + mean = (n - 1) / 2 + cov = sum((x - mean) * (y - mean) for x, y in zip(ra, rb)) + var = sum((x - mean) ** 2 for x in ra) + return cov / var if var > 0 else None + + +@register_analyzer("path_patching") +class PathPatchingAnalyzer(Analyzer): + """Exact head → logits path patching of the media's effect, screened by direct logit attribution. + + Hyper-parameters: + corruption: ``"noise"`` (default, any media) or ``"gray"`` (images/video), + exactly as in ``causal_trace``. + noise_scale: noise std as a multiple of the media embeddings' std (Meng: 3). + top_k: heads patched exactly per case, chosen by the screen; + 0 = every head (one forward each). + min_effect: smallest ``|m_clean - m_corrupt|`` (logits) worth decomposing; + the shares are divided by it. + target: ``"prediction"`` or ``"gold"``. + seed: noise seed — the same per-case noise as ``causal_trace``. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost per case: ``top_k + 4`` logits-only forwards. + """ + + name = "path_patching" + description = ( + "Patches one head at a time with its output under corrupted image/audio, every " + "other head frozen, to find which heads carry the media's effect directly to the " + "answer (and how much of it the top heads carry together)." + ) + requires = frozenset({Capability.HIDDEN_STATES, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "pp_margin_clean": ("Clean answer margin", "Logit margin of the target answer token with the media intact."), + "pp_media_effect": ("Margin lost to corrupt media", "How much the answer margin falls when the media tokens are corrupted; near zero = the answer did not depend on the media."), + "pp_media_effect_rel": ("Relative loss, corrupt media", "pp_media_effect as a fraction of the clean margin."), + "pp_effect_traceable": ("Media effect traceable", "Corrupting the media moved the margin by at least min_effect, so it could be decomposed over heads."), + "pp_top_head": ("Top direct-path head", "The head (L.H) whose direct path to the answer carries the largest share of the media's effect (exact patch)."), + "pp_top_head_share": ("Top head's direct share", "Share of the media's effect lost when only that head's direct path sees the corrupted media."), + "pp_top_head_depth_frac": ("Top head depth", "Relative depth of that head's layer."), + "pp_top_head_linear": ("Top head is linear attention", "That head belongs to a linear-attention (gated delta rule) layer of a hybrid stack."), + "pp_topk_direct_share": ("Top-k heads, direct share", "Share of the media's effect lost when the top-k screened heads' direct paths are corrupted together, every other head frozen."), + "pp_topk_total_share": ("Top-k heads, total share", "The same patch without freezing the other heads: the top-k heads' total effect, through later heads too."), + "pp_screen_spearman": ("Screen vs exact (rank)", "Rank correlation between the direct-logit-attribution screen and the exact direct shares of the patched heads."), + "pp_linear_share": ("Direct effect via linear heads", "Share of the patched heads' absolute exact direct effect carried by linear-attention heads (hybrid stacks)."), + } + + def __init__(self, corruption: str = "noise", noise_scale: float = 3.0, top_k: int = 8, + min_effect: float = 1.0, target: str = "prediction", seed: int = 0, + max_cases: int = 0) -> None: + from evalrx.analyzers.patching.causal_trace import CORRUPTIONS + + if corruption not in CORRUPTIONS: + raise ValueError(f"corruption must be one of {CORRUPTIONS}, got {corruption!r}") + if top_k < 0: + raise ValueError(f"top_k must be >= 0, got {top_k}") + super().__init__(corruption=corruption, noise_scale=noise_scale, top_k=top_k, + min_effect=min_effect, target=target, seed=seed, max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + from evalrx.analyzers._internals.access import decoder_stack, media_positions, torch_module + from evalrx.analyzers._internals.heads import ( + mixer_hooks, + n_heads_of, + out_proj, + readout_direction, + ) + from evalrx.analyzers._internals.residual import residual_hooks + from evalrx.analyzers._internals.targets import build_margin + from evalrx.analyzers.patching.causal_trace import ( + _noised, + corrupted_variants, + deepstack_mergers, + ) + from evalrx.core.model import CaptureSpec + from evalrx.models.backends.hf.discover import get_final_norm, get_unembed + + module = torch_module(model, self.name) + stack = decoder_stack(module) + mergers = deepstack_mergers(module) + norm, unembed = get_final_norm(module), get_unembed(module) + projs = [out_proj(m) if m is not None else None for m in stack.mixers] + if any(p is None for p in projs): + raise ValueError(f"{self.name}: some decoder layers have no recognised token mixer " + "with an output projection; their heads could not be frozen.") + n_heads = [n_heads_of(m, p) for m, p in zip(stack.mixers, projs)] + n_layers = stack.n_layers + kinds = stack.kinds + all_heads = [(layer, head) for layer in range(n_layers) for head in range(n_heads[layer])] + spec = CaptureSpec(to_cpu=False) + + def forward(variant: Any, *, capture: bool = False, replace: "dict | None" = None, + residual0: bool = False) -> "tuple[Any, Any, dict]": + inputs, plan, (store, gen) = variant + noised = _noised(mergers, store, self.noise_scale, gen) if store is not None else nullcontext() + with torch.no_grad(), noised, \ + residual_hooks(model, capture={0: None} if residual0 else None, replace=plan, + analyzer=self.name) as res, \ + mixer_hooks(model, capture=capture, replace=replace, analyzer=self.name) as mix: + trace = model.forward(inputs, capture={Capability.LOGITS}, spec=spec) + return trace, mix, res.captured + + def with_senders(base: dict, donor: dict, senders: "list[tuple[int, int]]", + freeze: bool) -> dict: + return sender_inputs(base, donor, senders, n_heads, freeze=freeze) + + per_case: list[dict[str, Any]] = [] + tables: dict[str, list[dict[str, Any]]] = {} + recurring: dict[str, list[float]] = defaultdict(list) + screen_pairs: list[tuple[float, float]] = [] + skipped: Counter = Counter() + slots: Counter = Counter() + for index, case in enumerate(cases.stratified_head(self.max_cases)): + clean_variant = (case.inputs, {}, (None, None)) + trace, clean_mix, clean_res = forward(clean_variant, capture=True, residual0=True) + slot, media = media_positions(trace) + if not media: + skipped["no_media_tokens"] += 1 + continue + logits = trace.require(Capability.LOGITS)[-1:] + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + m_clean = margin.value(logits) + token_ids, last = trace.token_ids, trace.seq_len - 1 + variants, why = corrupted_variants( + case.inputs, self.corruption, slot=slot, media=media, + embeddings=clean_res[0][media], deepstack=bool(mergers), + noise_scale=self.noise_scale, n_noise=1, seed=self.seed * 100003 + index) + del trace, logits, clean_res + if not variants: + skipped[why] += 1 + continue + donor_variant = variants[0] + trace, donor_mix, _ = forward(donor_variant, capture=True) + if trace.token_ids != token_ids: # a gray image that re-tokenised differently + skipped["corruption_misaligned"] += 1 + continue + m_corrupt = margin.value(trace.require(Capability.LOGITS)[-1:]) + del trace + effect = m_clean - m_corrupt + row: dict[str, Any] = { + "sample_id": case.id, + "pp_margin_clean": round(m_clean, 4), + "pp_media_effect": round(effect, 4), + "pp_effect_traceable": abs(effect) >= self.min_effect, + } + if m_clean > 1e-6: + row["pp_media_effect_rel"] = round(effect / m_clean, 4) + slots[slot] += 1 + if abs(effect) < self.min_effect: + per_case.append(row) + continue + + clean_in, donor_in = clean_mix.inputs, donor_mix.inputs + # Screen: the media-dependent part of each head's output at the answer + # position, read through the frozen final norm into the margin. + with torch.no_grad(): + v = readout_direction(norm, unembed, clean_mix.final_residual[last], + margin.positive, margin.negative) + screen: dict[tuple[int, int], float] = {} + for layer, proj in enumerate(projs): + w_o = proj.weight.float() + delta = (donor_in[layer][last] - clean_in[layer][last]).float().to(w_o.device) + per_head = ((w_o.T @ v.to(w_o.device)) * delta).view(n_heads[layer], -1).sum(-1) + for head, value in enumerate(per_head.tolist()): + screen[(layer, head)] = -value / effect # margin LOST, as a share + k = len(all_heads) if self.top_k == 0 else min(self.top_k, len(all_heads)) + chosen = sorted(all_heads, key=lambda lh: -abs(screen[lh]))[:k] + + def share(replace: dict) -> float: + trace, _, _ = forward(clean_variant, replace=replace) + return (m_clean - margin.value(trace.require(Capability.LOGITS)[-1:])) / effect + + table = [] + for layer, head in chosen: + direct = share(with_senders(clean_in, donor_in, [(layer, head)], freeze=True)) + table.append({"head": head_name(layer, head), "layer": layer, "head_index": head, + "kind": kinds[layer], "screen": round(screen[(layer, head)], 4), + "direct": round(direct, 4)}) + screen_pairs.append((screen[(layer, head)], direct)) + top = max(table, key=lambda t: abs(t["direct"])) + row.update({ + "pp_top_head": top["head"], + "pp_top_head_share": top["direct"], + "pp_top_head_depth_frac": round(top["layer"] / max(1, n_layers - 1), 4), + "pp_top_head_linear": top["kind"] == "linear", + "pp_topk_direct_share": round(share(with_senders(clean_in, donor_in, chosen, freeze=True)), 4), + "pp_topk_total_share": round(share(with_senders(clean_in, donor_in, chosen, freeze=False)), 4), + }) + rho = _spearman([t["screen"] for t in table], [t["direct"] for t in table]) + if rho is not None: + row["pp_screen_spearman"] = round(rho, 4) + if "linear" in kinds: + mass = sum(abs(t["direct"]) for t in table) or 1.0 + row["pp_linear_share"] = round( + sum(abs(t["direct"]) for t in table if t["kind"] == "linear") / mass, 4) + per_case.append(row) + tables[case.id] = table + for t in sorted(table, key=lambda t: -abs(t["direct"]))[:3]: + recurring[t["head"]].append(t["direct"]) + del clean_mix, donor_mix, clean_in, donor_in + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + traced = [r for r in per_case if "pp_top_head" in r] + n = len(per_case) + + def mean(key: str) -> "float | None": + values = [r[key] for r in traced if key in r] + return round(sum(values) / len(values), 4) if values else None + + pooled = _spearman([a for a, _ in screen_pairs], [b for _, b in screen_pairs]) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_traced": len(traced), + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": n_layers, + "heads_per_layer": n_heads, + "layer_kinds": kinds, + "media_slot": slots.most_common(1)[0][0] if slots else None, + "corruption": self.corruption, + "top_k": self.top_k, + "target": self.target, + "mean_media_effect": round(sum(r["pp_media_effect"] for r in per_case) / n, 4), + "mean_topk_direct_share": mean("pp_topk_direct_share"), + "mean_topk_total_share": mean("pp_topk_total_share"), + "mean_screen_spearman": mean("pp_screen_spearman"), + "screen_spearman_pooled": round(pooled, 4) if pooled is not None else None, + "recurring_heads": [ + {"head": name, "n_cases": len(shares), "mean_direct_share": round(sum(shares) / len(shares), 4)} + for name, shares in sorted(recurring.items(), key=lambda kv: (-len(kv[1]), kv[0]))[:10] + ], + "per_case": per_case, + "_caveat": ( + "Exact path patches on the next-token margin at the end of the prompt: a head's " + "output under corrupted media, every other token-mixing head frozen at its clean " + "output, MLPs and norms recomputed. Only the top_k heads by a direct-logit-" + "attribution screen are patched, so a head the screen ranks low is never " + "measured — check pp_screen_spearman." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, + artifacts={"head_tables": tables}, findings=findings) diff --git a/evalrx/eval_agent/stages/probe.py b/evalrx/eval_agent/stages/probe.py index ee3c21ee..564e9cef 100755 --- a/evalrx/eval_agent/stages/probe.py +++ b/evalrx/eval_agent/stages/probe.py @@ -75,7 +75,7 @@ class ModelKind(str, Enum): "unfaithful_reasoning": ["cot_faithfulness"], "post_hoc_reasoning": ["cot_faithfulness"], "context_ignored": ["context_shap"], - "premature_layer_divergence": ["layer_contrast", "tuned_lens"], + "premature_layer_divergence": ["layer_contrast", "tuned_lens", "patchscope"], "confident_inconsistency": ["self_consistency", "verbalized_confidence"], "loop": ["loop_detect"], "ignored_obs": ["ignored_obs"], @@ -88,11 +88,12 @@ class ModelKind(str, Enum): "instruction_following": ["prompt_contrast"], "entropy": ["token_entropy", "logprob_entropy"], "perplexity": ["logprob_entropy"], - "logit_lens": ["logit_lens", "tuned_lens"], + "logit_lens": ["logit_lens", "tuned_lens", "patchscope"], # White-box information-flow mechanisms (VLM / ALM). Each tag names what the # analyzers can falsify: is the media read at all (knockout / causal trace), # where does its information travel, layer by layer and token by token - # (attribution patching), would reading it more help (sensitivity), is the + # (attribution patching), which heads carry it to the answer (path + # patching), would reading it more help (sensitivity), is the # right region read (gradient maps), where does the answer get written or # suppressed (heads / lens). "visual_neglect": ["attention_knockout", "causal_trace", "modality_ablation", @@ -101,12 +102,14 @@ class ModelKind(str, Enum): "attention_sensitivity"], "modality_neglect": ["attention_knockout", "causal_trace", "modality_ablation"], "insufficient_visual_attention": ["attention_sensitivity", "attention_knockout", "attention"], - "visual_information_loss": ["causal_trace", "attribution_patching", "attention_knockout"], + "visual_information_loss": ["causal_trace", "attribution_patching", "attention_knockout", + "path_patching"], "late_visual_integration": ["causal_trace", "attribution_patching", "attention_knockout"], "mislocalized_attention": ["gradcam", "generic_attention", "relative_attention"], "grounding_error": ["gradcam", "generic_attention", "relative_attention"], - "late_suppression": ["tuned_lens", "logit_lens", "layer_contrast", "head_attribution"], - "answer_head_conflict": ["head_attribution"], + "late_suppression": ["tuned_lens", "patchscope", "logit_lens", "layer_contrast", + "head_attribution"], + "answer_head_conflict": ["head_attribution", "path_patching"], "representational_collapse": ["cka"], "numerical_hallucination": ["self_consistency", "verbalized_confidence"], # Text-reasoning mechanisms (2026-08). The hygiene pair is listed under the @@ -153,6 +156,7 @@ class ModelKind(str, Enum): "relative_attention", # attention across frames "mm_shap", # per-slot reliance "causal_trace", "attribution_patching", # where the video's effect hands off to the text + "path_patching", # which heads carry it to the answer ], "audio": [ # Grounding first, for the same reason the text list leads with hygiene: @@ -170,6 +174,7 @@ class ModelKind(str, Enum): "attention_sensitivity", # would attending more to the audio help? "causal_trace", # where the audio's effect hands off to the text "attribution_patching", # ... estimated at every layer and token + "path_patching", # which heads carry it to the answer "verbalized_confidence", "calibration", ], "image": [ @@ -184,6 +189,7 @@ class ModelKind(str, Enum): "gradcam", "generic_attention", # which image regions the answer depends on "causal_trace", # where the image's effect hands off to the text "attribution_patching", # ... estimated at every layer and token + "path_patching", # which heads carry it to the answer "logprob_entropy", "self_consistency", ], "text": [], # filled below from the legacy LLM list @@ -204,6 +210,7 @@ class ModelKind(str, Enum): "gradcam", "generic_attention", # which image regions the answer depends on "causal_trace", # where the image's effect hands off "attribution_patching", # ... estimated at every layer and token + "path_patching", # which heads carry it to the answer "logprob_entropy", "self_consistency", ], ModelKind.AGENT: [ @@ -229,6 +236,7 @@ class ModelKind(str, Enum): "attention", "logit_lens", # interpretability "layer_contrast", # DoLa/DeCo divergence signal "tuned_lens", # unbiased per-layer readout + "patchscope", # the model decodes its own states "head_attribution", # which heads write / oppose the answer "token_entropy", "logprob_entropy", "attention_sink", "attention_rollout", diff --git a/evalrx/reporting/case_study.py b/evalrx/reporting/case_study.py index 83b0915f..c5bbed12 100644 --- a/evalrx/reporting/case_study.py +++ b/evalrx/reporting/case_study.py @@ -60,6 +60,8 @@ "attention_sensitivity": "internal_white_box", "causal_trace": "internal_white_box", "attribution_patching": "internal_white_box", + "path_patching": "internal_white_box", + "patchscope": "internal_white_box", "mm_shap": "multimodal", "vl_shap": "multimodal", "modality_ablation": "multimodal", @@ -137,6 +139,7 @@ "attention_knockout": "Did the answer use the image or audio at all?", "causal_trace": "Did the answer use the image or audio at all?", "attribution_patching": "Did the answer use the image or audio at all?", + "path_patching": "Did the answer use the image or audio at all?", "attention_sensitivity": "Did the answer use the image or audio at all?", "gradcam": "Where was it looking?", "generic_attention": "Where was it looking?", @@ -222,6 +225,8 @@ "attention_sensitivity": "media attention rescaling", "causal_trace": "causal tracing", "attribution_patching": "attribution patching", + "path_patching": "path patching", + "patchscope": "Patchscope readout", "mm_shap": "modality attribution", "vl_shap": "vision-language attribution", "modality_ablation": "modality ablation", @@ -245,9 +250,9 @@ ("attention & attribution", {"relative_attn", "relative_attention", "context_shap", "linear_probe", "gradcam", "generic_attention"}), ("logits & representations", {"logprob_entropy", "logit_lens", "layer_contrast", - "tuned_lens", "head_attribution"}), + "tuned_lens", "head_attribution", "patchscope"}), ("causal interventions", {"attention_knockout", "attention_sensitivity", "causal_trace", - "attribution_patching"}), + "attribution_patching", "path_patching"}), ], "multimodal": [ ("grounding & hallucination", {"pope", "chair", "opera", "vcd"}), diff --git a/evalrx/specs.py b/evalrx/specs.py index bd5c07cd..a7654de9 100755 --- a/evalrx/specs.py +++ b/evalrx/specs.py @@ -141,11 +141,12 @@ def list_specs() -> list[str]: # Qwen3.5 WITH its vision tower -- the same checkpoints as the text specs above, # loaded as ``Qwen3_5ForConditionalGeneration`` (``AutoModelForImageTextToText``) # for the image benchmarks (examples/m1_m5/*_qwen3_5_2b). Layout verified on -# transformers 5.15.0: ``model.language_model.layers`` (24 layers on the 2B, +# transformers 5.15.0: ``model.language_model.layers`` (24 layers on the 0.8B and 2B, # ``layer_types`` = [linear x3, full] x 6), ``model.visual.blocks``, # ``config.image_token_id`` top-level, ``vision_config.spatial_merge_size`` = 2, # processor emits ``image_grid_thw`` -- i.e. the Qwen3-VL shape. for _key, _repo, _n_layers in ( + ("qwen3.5-0.8b-vl", "Qwen/Qwen3.5-0.8B", 24), ("qwen3.5-2b-vl", "Qwen/Qwen3.5-2B", 24), ("qwen3.5-4b-vl", "Qwen/Qwen3.5-4B", 32), ("qwen3.5-9b-vl", "Qwen/Qwen3.5-9B", 32), diff --git a/tests/test_analyzers/test_analyzer_contract.py b/tests/test_analyzers/test_analyzer_contract.py index 447042b5..b8f15ebe 100755 --- a/tests/test_analyzers/test_analyzer_contract.py +++ b/tests/test_analyzers/test_analyzer_contract.py @@ -228,6 +228,8 @@ def _long_prompt_batch() -> CaseBatch: "attention_knockout": "needs a live torch module; see test_intervention_analyzers.py", "causal_trace": "needs a live torch module; see test_intervention_analyzers.py", "attribution_patching": "needs a live torch module; see test_intervention_analyzers.py", + "path_patching": "needs a live torch module; see test_intervention_analyzers.py", + "patchscope": "needs a live torch module; see test_intervention_analyzers.py", "attention_sensitivity": "needs a live torch module; see test_intervention_analyzers.py", "head_attribution": "needs a live torch module; see test_intervention_analyzers.py", } diff --git a/tests/test_analyzers/test_analyzer_docs.py b/tests/test_analyzers/test_analyzer_docs.py index 8cfc5aaf..edd47716 100644 --- a/tests/test_analyzers/test_analyzer_docs.py +++ b/tests/test_analyzers/test_analyzer_docs.py @@ -38,6 +38,7 @@ def gen(): "| Attention knockout | `attention_knockout` | | | — | |\n" "| Attention sensitivity | `attention_sensitivity` | | | — | |\n" "| Attribution patching | `attribution_patching` | | | — | |\n" + "| Path patching | `path_patching` | | | — | |\n" ) diff --git a/tests/test_analyzers/test_hybrid_analyzers.py b/tests/test_analyzers/test_hybrid_analyzers.py index 64f92f40..1b77a1c1 100644 --- a/tests/test_analyzers/test_hybrid_analyzers.py +++ b/tests/test_analyzers/test_hybrid_analyzers.py @@ -204,16 +204,19 @@ def test_head_attribution_on_hybrid_covers_linear_heads(hybrid): assert all(0.0 <= r["dla_linear_share"] <= 1.0 for r in f["per_case"]) -@pytest.mark.parametrize("name", ["causal_trace", "attribution_patching", "gradcam", "tuned_lens"]) +@pytest.mark.parametrize("name", ["causal_trace", "attribution_patching", "gradcam", "tuned_lens", + "patchscope"]) def test_architecture_agnostic_analyzers_run_on_hybrid(hybrid, name): from evalrx.analyzers.attribution.gradcam import GradCAMAnalyzer + from evalrx.analyzers.lens.patchscope import PatchscopeAnalyzer from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer analyzer = {"causal_trace": CausalTraceAnalyzer(min_effect=0.0, n_points=3), "attribution_patching": AttributionPatchingAnalyzer(min_effect=0.0), - "gradcam": GradCAMAnalyzer(), "tuned_lens": TunedLensAnalyzer(steps=10)}[name] + "gradcam": GradCAMAnalyzer(), "tuned_lens": TunedLensAnalyzer(steps=10), + "patchscope": PatchscopeAnalyzer(media_layers=2)}[name] _contract(analyzer.run(hybrid, _batch()), name, 4) @@ -226,3 +229,16 @@ def test_attribution_patching_is_exact_in_the_small_corruption_limit_on_hybrid(h f = AttributionPatchingAnalyzer(noise_scale=1e-3, min_effect=0.0, ig_steps=steps, n_verify=2).run(hybrid, _batch()).findings assert f["verification"]["all"]["mae"] < 0.02 and f["verification"]["all"]["pearson"] > 0.99 + + +def test_path_patching_freezes_and_patches_the_linear_heads_too(hybrid): + # Every head's direct path together must carry the whole media effect; a + # linear-attention layer left unfrozen or unpatched would break the identity. + from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer + + f = PathPatchingAnalyzer(min_effect=0.0, top_k=0).run(hybrid, _batch()).findings + stack = decoder_stack(hybrid._loaded[0]) + assert f["layer_kinds"] == stack.kinds and "linear" in f["layer_kinds"] + for row in f["per_case"]: + assert row["pp_topk_direct_share"] == pytest.approx(1.0, abs=1e-3) + assert 0.0 <= row["pp_linear_share"] <= 1.0 diff --git a/tests/test_analyzers/test_intervention_analyzers.py b/tests/test_analyzers/test_intervention_analyzers.py index dc0e0e69..44edd757 100644 --- a/tests/test_analyzers/test_intervention_analyzers.py +++ b/tests/test_analyzers/test_intervention_analyzers.py @@ -612,6 +612,147 @@ def decode(self, ids): assert _token_text(object(), 7) is None +# ── path patching ───────────────────────────────────────────────────────────── + +def _clean_and_donor(vlm, inputs): + """Projection inputs of the clean run and of the noise-corrupted run (seed 0).""" + from evalrx.analyzers._internals.heads import mixer_hooks + from evalrx.analyzers.patching.causal_trace import corrupted_variants + + with torch.no_grad(): + with residual_hooks(vlm, capture={0: None}) as res, mixer_hooks(vlm, capture=True) as clean: + base = vlm.forward(inputs, capture={Capability.LOGITS}).logits[-1] + media = list(range(3, 19)) + (variant,), _ = corrupted_variants(inputs, "noise", slot="image", media=media, + embeddings=res.captured[0][media], deepstack=False, + noise_scale=3.0, n_noise=1, seed=0) + with residual_hooks(vlm, replace=variant[1]), mixer_hooks(vlm, capture=True) as donor: + vlm.forward(inputs, capture={Capability.LOGITS}) + return base, clean.inputs, donor.inputs + + +def test_freezing_every_head_at_its_own_output_is_a_no_op(vlm): + from evalrx.analyzers._internals.heads import mixer_hooks + + inputs = Inputs(prompt="what?", image=image(1)) + base, clean, _ = _clean_and_donor(vlm, inputs) + with torch.no_grad(), mixer_hooks(vlm, replace=clean): + frozen = vlm.forward(inputs, capture={Capability.LOGITS}).logits[-1] + assert torch.equal(frozen, base) + + +def test_a_last_layer_heads_direct_path_is_its_whole_effect(vlm): + # No head follows the last layer, so freezing the others cannot change what + # its patch does: direct and total effects coincide exactly. + from evalrx.analyzers._internals.heads import mixer_hooks, n_heads_of, out_proj + from evalrx.analyzers.patching.path_patching import sender_inputs + + stack = access.decoder_stack(vlm._loaded[0]) + n_heads = [n_heads_of(m, out_proj(m)) for m in stack.mixers] + inputs = Inputs(prompt="what?", image=image(1)) + base, clean, donor = _clean_and_donor(vlm, inputs) + sender = [(stack.n_layers - 1, 0)] + out = {} + for freeze in (True, False): + with torch.no_grad(), mixer_hooks(vlm, replace=sender_inputs(clean, donor, sender, n_heads, freeze=freeze)): + out[freeze] = vlm.forward(inputs, capture={Capability.LOGITS}).logits[-1] + assert torch.equal(out[True], out[False]) and not torch.equal(out[True], base) + + +def test_all_heads_together_carry_the_whole_media_effect(vlm): + # MLPs are per-position, so the media reaches the answer position only + # through token mixers: patching every head's direct path at once must + # reproduce the corrupted run there, i.e. a share of exactly 1. + from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer + + f = PathPatchingAnalyzer(min_effect=0.0, top_k=0).run(vlm, _batch()).findings + for row in f["per_case"]: + assert row["pp_topk_direct_share"] == pytest.approx(1.0, abs=1e-3) + assert row["pp_topk_total_share"] == pytest.approx(1.0, abs=1e-3) + + +def test_path_patching_contract(vlm): + from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer + + result = PathPatchingAnalyzer(min_effect=0.0, top_k=5).run(vlm, _batch()) + f = _assert_result_contract(result, "path_patching", 4) + assert f["n_traced"] == 4 and f["heads_per_layer"] == [4] * 6 and f["top_k"] == 5 + assert 1 <= len(f["recurring_heads"]) <= 10 + for row in f["per_case"]: + assert row["pp_top_head"].startswith("L") and isinstance(row["pp_top_head_linear"], bool) + assert -1.0 <= row["pp_screen_spearman"] <= 1.0 and 0.0 <= row["pp_top_head_depth_frac"] <= 1.0 + for table in result.artifacts["head_tables"].values(): + assert len(table) == 5 and {"head", "screen", "direct", "kind"} <= set(table[0]) + screened = [abs(t["screen"]) for t in table] + assert screened == sorted(screened, reverse=True) # the top-k BY the screen were patched + + +def test_path_patching_small_effects_keep_the_row_but_skip_the_heads(vlm): + from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer + + f = PathPatchingAnalyzer().run(vlm, _batch(2)).findings # tiny model: effects << 1 logit + assert f["n_traced"] == 0 and f["recurring_heads"] == [] + assert all(row["pp_effect_traceable"] is False and "pp_top_head" not in row for row in f["per_case"]) + + +def test_path_patching_argument_checks_and_rank_correlation(): + from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer, _spearman + + with pytest.raises(ValueError, match="top_k"): + PathPatchingAnalyzer(top_k=-1) + with pytest.raises(ValueError, match="corruption"): + PathPatchingAnalyzer(corruption="blur") + assert _spearman([1, 2, 3, 4], [10, 20, 30, 40]) == pytest.approx(1.0) + assert _spearman([1, 2, 3, 4], [4, 3, 2, 1]) == pytest.approx(-1.0) + assert _spearman([1, 2], [1, 2]) is None + + +# ── patchscope ──────────────────────────────────────────────────────────────── + +def test_patchscope_layer0_readout_is_the_source_token_in_the_placeholder(vlm): + # Layer 0 reads the embedding, so decoding the answer position there is the + # target prompt with the source's last token written into the placeholder. + from evalrx.analyzers.lens.patchscope import IDENTITY_PROMPT, PatchscopeAnalyzer + from tests.tiny_whitebox import CharTokenizer + + case = _batch(1) + result = PatchscopeAnalyzer().run(vlm, case) + trace = vlm.forward(case[0].inputs, capture={Capability.LOGITS}) + target = int(trace.logits[-1].argmax()) + ids = CharTokenizer()(IDENTITY_PROMPT)["input_ids"] + ids[-1] = trace.token_ids[-1] + with torch.no_grad(): + plain = torch.softmax(vlm._loaded[0](input_ids=torch.tensor([ids])).logits[0, -1].float(), -1) + curve = result.artifacts["readout_curves"][case[0].id] + assert curve["p_target"][0] == pytest.approx(float(plain[target]), abs=1e-4) + assert curve["top1"][0] == str(int(plain.argmax())) # no decoder: the id, never a made-up word + assert result.findings["patch_check_max_diff"] == 0.0 + + +def test_patchscope_contract_with_media_words(vlm): + from evalrx.analyzers.lens.patchscope import PatchscopeAnalyzer + + result = PatchscopeAnalyzer(media_layers=2).run(vlm, _batch()) + f = _assert_result_contract(result, "patchscope", 4) + assert f["layers_read"] == list(range(6)) and f["media_layers_read"] == [2, 5] # not the embedding + for row in f["per_case"]: + assert 0.0 <= row["ps_first_depth_frac"] <= 1.0 and 0.0 <= row["ps_agree_mid"] <= 1.0 + assert row["ps_late_drop"] == pytest.approx(row["ps_peak_prob"] - row["ps_final_prob"], abs=2e-4) + for words in result.artifacts["media_words"].values(): + assert sorted(words) == [2, 5] and all(len(w) == 16 for w in words.values()) + assert set(f["mean_p_target_curves"]) == {"fail", "pass"} + + +def test_patchscope_reads_text_only_cases_and_skips_missing_gold(vlm): + from evalrx.analyzers.lens.patchscope import PatchscopeAnalyzer + + text_only = CaseBatch([FailureCase(inputs=Inputs(prompt="two plus two is"), expected="4"), + FailureCase(inputs=Inputs(prompt="no gold here"))]) + f = PatchscopeAnalyzer(target="gold", n_points=3).run(vlm, text_only).findings + assert f["n_cases_analyzed"] == 1 and f["skip_reasons"] == {"no_gold": 1} + assert f["layers_read"] == [0, 2, 5] + + # ── head attribution ────────────────────────────────────────────────────────── def test_per_head_writes_sum_to_the_output_projection(vlm): diff --git a/tests/test_eval_agent/test_probe_routing_whitebox.py b/tests/test_eval_agent/test_probe_routing_whitebox.py index 54a6908f..805ed2e9 100644 --- a/tests/test_eval_agent/test_probe_routing_whitebox.py +++ b/tests/test_eval_agent/test_probe_routing_whitebox.py @@ -29,7 +29,8 @@ from tests.conftest import FakeModel # noqa: E402 WRITE_SIDE = {"attention_knockout", "attention_sensitivity", "causal_trace", "attribution_patching", - "gradcam", "generic_attention", "head_attribution", "tuned_lens"} + "path_patching", "gradcam", "generic_attention", "head_attribution", "tuned_lens", + "patchscope"} _WHITE_BOX = {Capability.GENERATE, Capability.LOGITS, Capability.LOGPROBS, Capability.ATTENTION, Capability.HIDDEN_STATES, Capability.GRADIENTS} @@ -64,7 +65,8 @@ def test_white_box_vlm_ranks_knockout_right_after_modality_ablation(): vlm = FakeModel(capabilities=_WHITE_BOX, modalities={"text", "image"}) ranked = StrategyProbe().select(vlm, data=_image_batch()) assert ranked.index("attention_knockout") == ranked.index("modality_ablation") + 1 - assert WRITE_SIDE - {"tuned_lens", "head_attribution"} <= set(ranked[: len(_SLOT_PRIORITY["image"]) + 2]) + assert WRITE_SIDE - {"tuned_lens", "head_attribution", "patchscope"} <= set( + ranked[: len(_SLOT_PRIORITY["image"]) + 2]) def test_white_box_alm_gets_the_audio_interventions(): From 1d9d774b83801df8c2223417874827fa9a754331 Mon Sep 17 00:00:00 2001 From: Jiaqi Liu <32996169+jiaq-liu@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:11:15 -0700 Subject: [PATCH 9/9] add attention_temperature, embedLens --- CHANGELOG.md | 32 ++ docs/analyzers.md | 22 +- evalrx/analyzers/__init__.py | 4 +- evalrx/analyzers/_internals/attn_tap.py | 45 ++- evalrx/analyzers/lens/__init__.py | 6 +- evalrx/analyzers/lens/embed_lens.py | 284 ++++++++++++++++++ evalrx/analyzers/patching/__init__.py | 7 +- .../patching/attention_temperature.py | 214 +++++++++++++ evalrx/eval_agent/stages/probe.py | 19 +- evalrx/reporting/case_study.py | 11 +- .../test_analyzers/test_analyzer_contract.py | 2 + tests/test_analyzers/test_analyzer_docs.py | 1 + tests/test_analyzers/test_hybrid_analyzers.py | 35 +++ .../test_intervention_analyzers.py | 150 +++++++++ .../test_probe_routing_whitebox.py | 2 +- 15 files changed, 813 insertions(+), 21 deletions(-) create mode 100644 evalrx/analyzers/lens/embed_lens.py create mode 100644 evalrx/analyzers/patching/attention_temperature.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a476068c..52bff701 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -93,6 +93,38 @@ cross-case tally of relay tokens. - `_internals/residual.py`: `keep_grad` also works with frozen weights (the captured state becomes a leaf). +### Added — `attention_temperature` and `embed_lens` + +- `patching/attention_temperature`: divides the pre-softmax scores of the text + after the media by a temperature (`sharp` 0.5 / `flat` 2.0, by band) and + reports the margin's slope in log T. Default `scope="within"` tempers only the + media keys and keeps each row's total media weight — the model looks at the + image as much as before, only more focused or more spread out (FOCUS; the + AMOUNT is `attention_sensitivity`'s). `scope="all"` is plain temperature. + New tap edit kinds `temp` / `temp_within` (`attn_tap.temperature`), exact to + float precision (tested: `T=1` no-op, media mass kept, `p^(1/T)` inside, + `softmax(s/T)`). Softmax layers only: on Qwen3.5 6 of 24 layers are edited. +- `lens/embed_lens` (Fan et al., CVPR 2026): each media token at the decoder + input, read through its nearest input-embedding words (cosine). Tokens are + split into recurring (nearest word shared across ≥ `recurring_frac` of the + batch's images), high-norm (≥ `norm_ratio` × the image's median) and specific + — the paper's sink / dead / alive, inferred from the batch's DISTINCT images + instead of checkpoint-specific ids — and each prompt content word gets the + best cosine any media token reaches to it (grounding). `layer` reads any + depth (default 0, the decoder input, as in the paper). +- Measured on Qwen3.5-0.8B (POPE adversarial): + - `attention_temperature` (16 cases; tap exact; within-media mean slope + +0.33, no flips; plain temperature flips 19% when flattened). + - `embed_lens` (100 cases): at the input Qwen3.5's image tokens are not in + vocabulary space — best cosine 0.15, 90% recurring (space, digits, + `<|endoftext|>`), the queried object's word no closer when present + (AUROC 0.47). At `layer=0.9` 87% are image-specific (rocky, snow, wooden, + buses, ...) and object grounding reaches AUROC 0.76 (0.82 at the last layer). +- Routing: `attention_temperature` after `attention_sensitivity` and under + `attention_sink` / `insufficient_visual_attention` / `mislocalized_attention`; + `embed_lens` after the gradient maps and under `attention_sink` / + `visual_information_loss` / `grounding_error`. + ### Added — `path_patching` and `patchscope` - `patching/path_patching` (Wang et al., IOI): which heads carry the media's diff --git a/docs/analyzers.md b/docs/analyzers.md index e75d1bbd..4d2c4703 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -23,12 +23,12 @@ model, not by a black-box / white-box label. | [Hallucination](#hallucination) | `hallucination/` | 5 | 2 | | [Attention](#attention) | `attention/` | 4 | 0 | | [Attribution](#attribution) | `attribution/` | 2 | 0 | -| [Lens](#lens) | `lens/` | 5 | 0 | -| [Patching](#patching) | `patching/` | 5 | 0 | +| [Lens](#lens) | `lens/` | 6 | 0 | +| [Patching](#patching) | `patching/` | 6 | 0 | | [Geometry](#geometry) | `geometry/` | 2 | 0 | | [Reasoning](#reasoning) | `reasoning/` | 7 | 0 | | [Agent](#agent) | `agent/` | 7 | 0 | -| **Total** | | **52** | **2** | +| **Total** | | **54** | **2** | How to read the columns: @@ -137,7 +137,14 @@ write, to see which heads write the answer. `patchscope` lets the model decode i own states instead: each layer's residual at the answer position (and optionally each image / audio token) is patched into a token-identity prompt at the same layer. On Qwen3.5-0.8B it reads the final answer about five layers earlier than -the logit lens. +the logit lens. `embed_lens` reads the media tokens where they ENTER the language +model, through their nearest vocabulary embeddings: which carry image-specific +words, which repeat across every image (recurring / high-norm, the sink- and +dead-like tokens of Fan et al., inferred from the batch rather than from +checkpoint-specific ids), and how close the image comes to each prompt word. +On Qwen3.5-0.8B the input is not in vocabulary space (90% recurring, object +grounding at chance); `layer=0.9` reads the same lens deep, where 87% of image +tokens are specific and the queried object is grounded (POPE AUROC 0.76). | Analyzer | Key | Requires | Modality | Paper | Status | @@ -147,6 +154,7 @@ the logit lens. | Layer contrast | `layer_contrast` | `HIDDEN_STATES` | text + image | [Chuang et al. 2024](https://arxiv.org/abs/2309.03883); [Wang et al. 2025](https://arxiv.org/abs/2410.11779) | ✓ | | Head logit attribution | `head_attribution` | `HIDDEN_STATES` + `LOGITS` | text + image | [Elhage et al. 2021](https://transformer-circuits.pub/2021/framework/index.html); [Wang et al. 2023](https://arxiv.org/abs/2211.00593) | ✓ | | Patchscope | `patchscope` | `HIDDEN_STATES` + `LOGITS` | text + image + audio + video | [Ghandeharioun et al. 2024](https://arxiv.org/abs/2401.06102) | ✓ | +| EmbedLens | `embed_lens` | `HIDDEN_STATES` | **image** + **audio** + **video** | [Fan et al. 2026](https://arxiv.org/abs/2603.00510) | ✓ | ### Patching @@ -178,6 +186,11 @@ by direct logit attribution and the top `top_k` patched exactly; patching every head at once reproduces the whole effect at the answer position (a tested identity), so the shares are exact fractions of it. +`attention_temperature` divides the text's pre-softmax scores on the media by a +temperature while keeping each row's total media weight (`scope="within"`): the +model looks at the image as much as before, only more sharply or more broadly — +FOCUS, where `attention_sensitivity` changes the AMOUNT. Softmax layers only. + | Analyzer | Key | Requires | Modality | Paper | Status | |---|---|---|---|---|---| @@ -186,6 +199,7 @@ identity), so the shares are exact fractions of it. | Attention sensitivity | `attention_sensitivity` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | [Liu et al. 2024 (PAI)](https://arxiv.org/abs/2407.21771) | ✓ | | Attribution patching | `attribution_patching` | `GRADIENTS` + `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Syed et al. 2023](https://arxiv.org/abs/2310.10348); [Kramár et al. 2024 (AtP*)](https://arxiv.org/abs/2403.00745); [Hanna et al. 2024 (EAP-IG)](https://arxiv.org/abs/2403.17806) | ✓ | | Path patching | `path_patching` | `HIDDEN_STATES` + `LOGITS` | **image** + **audio** + **video** | [Wang et al. 2023 (IOI)](https://arxiv.org/abs/2211.00593) | ✓ | +| Attention temperature | `attention_temperature` | `ATTENTION` + `LOGITS` | **image** + **audio** + **video** | — (an intervention; cf. [Liu et al. 2024 (PAI)](https://arxiv.org/abs/2407.21771)) | ✓ | ### Geometry diff --git a/evalrx/analyzers/__init__.py b/evalrx/analyzers/__init__.py index 4c2edbd0..2aeef9ac 100755 --- a/evalrx/analyzers/__init__.py +++ b/evalrx/analyzers/__init__.py @@ -12,10 +12,10 @@ attention/ summary / rollout / sink / relative-attn — ATTENTION attribution/ Grad-CAM / Chefer generic-attention — GRADIENTS lens/ logit-lens / tuned-lens / layer contrast / - head logit attribution / patchscope — HIDDEN_STATES + head logit attribution / patchscope / embed-lens — HIDDEN_STATES patching/ causal tracing / attribution patching / path patching / attention knockout / - attention sensitivity — HIDDEN_STATES / ATTENTION (write) + attention sensitivity / temperature — HIDDEN_STATES / ATTENTION (write) geometry/ CKA / linear-probe — HIDDEN_STATES reasoning/ extraction/termination hygiene; arithmetic, self-repair, step-value, knowledge split, diff --git a/evalrx/analyzers/_internals/attn_tap.py b/evalrx/analyzers/_internals/attn_tap.py index 0c9f8299..1320fdc7 100644 --- a/evalrx/analyzers/_internals/attn_tap.py +++ b/evalrx/analyzers/_internals/attn_tap.py @@ -27,6 +27,12 @@ implemented exactly as ``+log w`` on the scores (``w = 0`` is ``block``). * ``scale_raw`` — multiply post-softmax weights by ``w`` without renormalising (the PAI-style amplification; rows no longer sum to one). +* ``temp`` — attention temperature ``T`` over the keys: their scores are + divided by ``T`` (``T < 1`` sharpens, ``T > 1`` flattens). +* ``temp_within`` — the same, but each row keeps the total weight it gave those + keys: only how that weight is distributed AMONG them changes, + ``p'_k ∝ p_k^(1/T)`` inside the set (a per-row shift of the divided scores). +Masked cells (future keys, padding) are never touched by a temperature. """ from __future__ import annotations @@ -42,7 +48,8 @@ from evalrx.core.model import Model TAP_NAME = "evalrx_tap" -EDIT_KINDS = ("block", "scale", "scale_raw") +EDIT_KINDS = ("block", "scale", "scale_raw", "temp", "temp_within") +TEMPERATURE_KINDS = ("temp", "temp_within") _ACTIVE: "list[TapState]" = [] _REGISTERED = False @@ -65,6 +72,8 @@ class AttentionEdit: def __post_init__(self) -> None: if self.kind not in EDIT_KINDS: raise ValueError(f"edit kind must be one of {EDIT_KINDS}, got {self.kind!r}") + if self.kind in TEMPERATURE_KINDS and not (self.weight > 0 and math.isfinite(self.weight)): + raise ValueError(f"attention temperature must be a finite number > 0, got {self.weight}") if self.weight < 0: raise ValueError(f"attention weight must be >= 0, got {self.weight}") @@ -93,6 +102,17 @@ def reweight(keys: Iterable[int], weight: float, *, queries: "Iterable[int] | No ) +def temperature(keys: Iterable[int], value: float, *, queries: "Iterable[int] | None" = None, + layers: "Iterable[int] | None" = None, within: bool = False) -> AttentionEdit: + """Attention temperature *value* over *keys* (``within``: keep each row's mass on them).""" + return AttentionEdit( + "temp_within" if within else "temp", tuple(int(k) for k in keys), + None if queries is None else tuple(int(q) for q in queries), + None if layers is None else frozenset(int(i) for i in layers), + float(value), + ) + + class TapState: """What the tap does while one ``attention_tap`` block is active.""" @@ -139,7 +159,10 @@ def edit_scores(self, layer: int, scores: Any) -> Any: if edit.kind == "scale_raw" or not edit.applies(layer): continue mask = self._mask(index, edit, q_len, k_len, scores.device) - if edit.kind == "block" or edit.weight == 0: + if edit.kind in TEMPERATURE_KINDS: + if edit.weight != 1.0: + scores = _tempered(scores, mask, edit.weight, within=edit.kind == "temp_within") + elif edit.kind == "block" or edit.weight == 0: scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min) elif edit.weight != 1.0: scores = scores + mask.to(scores.dtype) * math.log(edit.weight) @@ -158,6 +181,24 @@ def edit_probs(self, layer: int, probs: Any) -> Any: return probs +def _tempered(scores: Any, mask: Any, t: float, *, within: bool) -> Any: + """Divide the *mask* cells' scores by *t*; with *within*, shift each row so the + softmax mass on those cells is what it was (``logsumexp`` over them unchanged).""" + import torch + + floor = torch.finfo(scores.dtype).min + sel = mask & (scores > floor / 2) # never the masked (future / padding) cells + s = scores.float() + out = s / t + if within: + neg = torch.tensor(float("-inf"), device=s.device) + before = torch.logsumexp(torch.where(sel, s, neg), dim=-1, keepdim=True) + after = torch.logsumexp(torch.where(sel, out, neg), dim=-1, keepdim=True) + shift = torch.where(torch.isfinite(before), before - after, torch.zeros_like(before)) + out = out + shift + return torch.where(sel, out, s).to(scores.dtype) + + def _tap_attention(module: Any, query: Any, key: Any, value: Any, attention_mask: Any, scaling: "float | None" = None, dropout: float = 0.0, softcap: "float | None" = None, **kwargs: Any) -> "tuple[Any, Any]": diff --git a/evalrx/analyzers/lens/__init__.py b/evalrx/analyzers/lens/__init__.py index 0aac6a30..2d66cc39 100755 --- a/evalrx/analyzers/lens/__init__.py +++ b/evalrx/analyzers/lens/__init__.py @@ -1,6 +1,8 @@ """Lens analyzers — project hidden states (or per-component writes) toward the vocabulary, -or let the model decode them itself (Patchscope).""" +or let the model decode them itself (Patchscope), or read the media tokens at the +decoder input through their nearest vocabulary embeddings (EmbedLens).""" +from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer from evalrx.analyzers.lens.head_attribution import HeadAttributionAnalyzer from evalrx.analyzers.lens.layer_contrast import LayerContrastAnalyzer from evalrx.analyzers.lens.logit_lens import LogitLensAnalyzer @@ -8,4 +10,4 @@ from evalrx.analyzers.lens.tuned_lens import TunedLensAnalyzer __all__ = ["LogitLensAnalyzer", "TunedLensAnalyzer", "LayerContrastAnalyzer", "HeadAttributionAnalyzer", - "PatchscopeAnalyzer"] + "PatchscopeAnalyzer", "EmbedLensAnalyzer"] diff --git a/evalrx/analyzers/lens/embed_lens.py b/evalrx/analyzers/lens/embed_lens.py new file mode 100644 index 00000000..3b994c32 --- /dev/null +++ b/evalrx/analyzers/lens/embed_lens.py @@ -0,0 +1,284 @@ +"""EmbedLens — what do the media tokens mean when they enter the language model? + +Fan et al. read each projected visual token, at the language model's input, +through its nearest neighbours in the input-embedding table (cosine +similarity): many tokens sit next to concrete words — objects, colours, text in +the image — before any decoder layer has run, while others carry no +image-specific meaning. Here, for every media token (image / video / audio) at +the input of decoder layer 0: + +* its ``top_k`` nearest vocabulary tokens and the best cosine similarity; +* its activation norm. + +The paper sorts tokens into *sink* / *dead* / *alive* by checkpoint-specific +neighbour ids. Without such ids, the batch supplies the evidence: + +* **recurring** — the token's nearest word is the nearest word of media tokens + in at least ``recurring_frac`` of the images: a direction the projector emits + whatever the image shows, so it says nothing about this one (dead-/sink-like); +* **high-norm** — norm at least ``norm_ratio`` × the image's median token norm + (the massive-activation signature of sink tokens); +* **specific** — the rest: tokens whose nearest word is image-specific (alive-like). + +And one grounding readout per case: each content word of the prompt against +the media tokens — the best cosine similarity any media token reaches to that +word's embedding. A question about an object the image does contain should +find a media token close to the object's word. + +The paper reads the INPUT (``layer=0``, the default). Whether a model's +projector writes into vocabulary space there is itself the finding: on a +LLaVA-style model most media tokens sit near concrete words. On Qwen3.5-0.8B +they do not — the mean best cosine is 0.15, 90% of tokens share a recurring +nearest word, and the queried object's word is no closer to the image tokens +when the object is present than when it is absent (POPE adversarial, AUROC +≈ 0.5). The same lens at a late layer (``layer=0.85``–``1.0``) separates them +(AUROC 0.76–0.82) and 87% of the tokens turn image-specific (rocky, snow, +wooden, buses, ...): there the image tokens have been turned toward words. +``layer`` reads any depth. + +References: +- What Do Visual Tokens Really Encode? Uncovering Sparsity and Redundancy in + Multimodal Large Language Models — Fan et al., CVPR 2026 — arXiv:2603.00510 +""" + +from __future__ import annotations + +from collections import Counter +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + + +def content_word(tokenizer: Any, token_id: int, special: "set[str]") -> "str | None": + """The token as a lower-case content word, or ``None`` (punctuation, function word, special).""" + from evalrx.analyzers.lens.patchscope import _STOP, _word + + word = _word(tokenizer, token_id) + if word is None or word in special: + return None + key = word.strip().lower() + return key if key.isalpha() and len(key) >= 2 and key not in _STOP else None + + +@register_analyzer("embed_lens") +class EmbedLensAnalyzer(Analyzer): + """Nearest vocabulary words of every media token at the decoder input, per case. + + Hyper-parameters: + layer: residual stream read — ``0`` (default, the decoder input as in + the paper), an index, or a fractional depth in (0, 1). + top_k: nearest vocabulary tokens kept per media token. + recurring_frac: a nearest word shared by media tokens in at least this + fraction of the images is "recurring" (not image-specific); + needs at least ``min_images`` analysed images. + min_images: fewest images from which recurring words are inferred. + norm_ratio: a token is "high-norm" at ``norm ≥ norm_ratio × median``. + vocab_chunk: embedding rows compared per block (bounds memory). + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: 1 forward per case + one pass over the embedding table. + """ + + name = "embed_lens" + description = ( + "Reads each image/audio token, as it enters the language model, through its nearest " + "vocabulary words: how many carry image-specific meaning, which repeat across every " + "image, and whether a token sits near the words the question asks about." + ) + requires = frozenset({Capability.HIDDEN_STATES}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "el_n_media_tokens": ("Media tokens", "Number of media tokens entering the language model."), + "el_mean_max_sim": ("Nearest-word similarity", "Mean, over media tokens, of the cosine similarity to their nearest vocabulary embedding."), + "el_frac_specific": ("Image-specific tokens", "Share of media tokens whose nearest word is neither recurring across images nor carried by a high-norm token (alive-like)."), + "el_frac_recurring": ("Recurring-word tokens", "Share of media tokens whose nearest word recurs across most images: not specific to this one (dead-/sink-like)."), + "el_frac_high_norm": ("High-norm tokens", "Share of media tokens with norm ≥ norm_ratio × the image's median token norm (sink-like)."), + "el_prompt_best_word": ("Best-grounded prompt word", "The prompt content word some media token is closest to."), + "el_prompt_best_sim": ("Its grounding similarity", "Cosine similarity of that word's embedding to its closest media token."), + } + + def __init__(self, layer: "int | float" = 0, top_k: int = 5, recurring_frac: float = 0.5, + min_images: int = 3, norm_ratio: float = 4.0, vocab_chunk: int = 32768, + max_cases: int = 0) -> None: + if top_k < 1: + raise ValueError(f"top_k must be >= 1, got {top_k}") + if not 0 < recurring_frac <= 1: + raise ValueError(f"recurring_frac must be in (0, 1], got {recurring_frac}") + super().__init__(layer=layer, top_k=top_k, recurring_frac=recurring_frac, + min_images=min_images, norm_ratio=norm_ratio, vocab_chunk=vocab_chunk, + max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + from evalrx.analyzers._internals.access import ( + decoder_stack, + media_positions, + tokenizer_of, + torch_module, + ) + from evalrx.analyzers._internals.residual import residual_hooks + from evalrx.analyzers.attention.relative_attn import resolve_attention_layer + from evalrx.analyzers.lens.patchscope import _word + from evalrx.core.model import CaptureSpec + + module = torch_module(model, self.name) + n_layers = decoder_stack(module).n_layers + layer = resolve_attention_layer(self.layer, n_layers) % n_layers + table = module.get_input_embeddings().weight + tokenizer = tokenizer_of(model) + special = set(getattr(tokenizer, "all_special_tokens", None) or []) + with torch.no_grad(): + row_norms = table.float().norm(dim=-1).clamp(min=1e-12) # (vocab,) + + def nearest(x: Any) -> "tuple[Any, Any]": + """Top-k cosine neighbours of *x* ``(n, dim)``: ``(sims, ids)``, each ``(n, k)``.""" + x = x.float().to(table.device) + x = x / x.norm(dim=-1, keepdim=True).clamp(min=1e-12) + best_s = best_i = None + for start in range(0, table.shape[0], max(1, self.vocab_chunk)): + block = table[start:start + self.vocab_chunk].float() + sims = (x @ block.T) / row_norms[start:start + self.vocab_chunk] + k = min(self.top_k, sims.shape[1]) + s, i = sims.topk(k, dim=-1) + i = i + start + if best_s is None: + best_s, best_i = s, i + else: + merged_s = torch.cat([best_s, s], -1) + merged_i = torch.cat([best_i, i], -1) + s, order = merged_s.topk(min(self.top_k, merged_s.shape[1]), dim=-1) + best_s, best_i = s, merged_i.gather(-1, order) + return best_s, best_i + + rows: list[dict[str, Any]] = [] + per_token: dict[str, dict[str, Any]] = {} + prompt_sims: dict[str, dict[str, float]] = {} + images_with: Counter = Counter() # top-1 id -> number of DISTINCT images it is top-1 in + seen_media: set = set() # fingerprints of the media already counted + skipped: Counter = Counter() + spec = CaptureSpec(to_cpu=False) + for case in cases.stratified_head(self.max_cases): + with torch.no_grad(), residual_hooks(model, capture={layer: None}, analyzer=self.name) as res: + trace = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + _, media = media_positions(trace) + if not media: + skipped["no_media_tokens"] += 1 + continue + with torch.no_grad(): + x = res.captured[layer][media] # (n, dim) + sims, ids = nearest(x) + norms = x.float().norm(dim=-1) + media_set = set(media) + words = {} + for pos, tid in enumerate(trace.token_ids): + if pos not in media_set: + word = content_word(tokenizer, tid, special) if tokenizer is not None else None + if word is not None and word not in words: + words[word] = int(tid) + grounding = {} + if words: + w = table[list(words.values())].float() + w = w / w.norm(dim=-1, keepdim=True).clamp(min=1e-12) + xn = x.float().to(w.device) + xn = xn / xn.norm(dim=-1, keepdim=True).clamp(min=1e-12) + best = (xn @ w.T).max(0).values.tolist() + grounding = {word: round(s, 4) for word, s in zip(words, best)} + top1 = ids[:, 0].tolist() + # The media tokens precede the question, so several questions about one + # image give identical media states: count each image once. + fingerprint = (tuple(x.shape), round(float(x.float().sum()), 2), + tuple(round(v, 3) for v in x[0, :8].float().tolist())) + if fingerprint not in seen_media: + seen_media.add(fingerprint) + images_with.update(set(top1)) + median = float(norms.median()) or 1e-12 + per_token[case.id] = { + "top1": top1, + "top_ids": ids.tolist(), + "max_sim": [round(v, 4) for v in sims[:, 0].tolist()], + "high_norm": [bool(v >= self.norm_ratio * median) for v in norms.tolist()], + } + prompt_sims[case.id] = grounding + row: dict[str, Any] = { + "sample_id": case.id, + "el_n_media_tokens": len(media), + "el_mean_max_sim": round(float(sims[:, 0].mean()), 4), + "el_frac_high_norm": round(sum(per_token[case.id]["high_norm"]) / len(media), 4), + } + if grounding: + best_word = max(grounding, key=grounding.get) + row["el_prompt_best_word"] = best_word + row["el_prompt_best_sim"] = grounding[best_word] + rows.append(row) + del trace, res, x, sims, ids + + if not rows: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + n_cases, n_images = len(rows), len(seen_media) + recurring: set = set() + if n_images >= self.min_images: + recurring = {tid for tid, count in images_with.items() if count >= self.recurring_frac * n_images} + specific_words: Counter = Counter() + recurring_mass: Counter = Counter() + total_tokens = 0 + for row in rows: + tokens = per_token[row["sample_id"]] + n = len(tokens["top1"]) + total_tokens += n + high = tokens["high_norm"] + rec = [not h and tid in recurring for tid, h in zip(tokens["top1"], high)] + row["el_frac_recurring"] = round(sum(rec) / n, 4) + row["el_frac_specific"] = round(1.0 - row["el_frac_recurring"] - row["el_frac_high_norm"], 4) + for tid, h in zip(tokens["top1"], high): + if tid in recurring: + recurring_mass[tid] += 1 + elif not h and tokenizer is not None: + word = content_word(tokenizer, tid, special) + if word is not None: + specific_words[word] += 1 + + def label(tid: int) -> str: + word = _word(tokenizer, tid) if tokenizer is not None else None + return word if word is not None else str(tid) + + mean = lambda key: round(sum(r[key] for r in rows) / n_cases, 4) # noqa: E731 + findings: dict[str, Any] = { + "n_cases_analyzed": n_cases, + "n_distinct_media": n_images, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "top_k": self.top_k, + "layer_used": layer, + "n_layers": n_layers, + "recurring_inferred": n_images >= self.min_images, + "recurring_words": [ + {"word": label(tid), "token_id": int(tid), "n_images": images_with[tid], + "token_share": round(recurring_mass[tid] / max(1, total_tokens), 4)} + for tid in sorted(recurring, key=lambda t: -recurring_mass[t]) + ], + "mean_max_sim": mean("el_mean_max_sim"), + "mean_frac_specific": mean("el_frac_specific"), + "mean_frac_recurring": mean("el_frac_recurring"), + "mean_frac_high_norm": mean("el_frac_high_norm"), + "specific_word_counts": specific_words.most_common(20), + "per_case": rows, + "_caveat": ( + f"Nearest neighbours by cosine similarity in the input-embedding table, read at the " + f"input of decoder layer {layer}. 'Recurring' is inferred from this batch (a nearest word " + "shared across at least recurring_frac of the distinct images), not from " + "checkpoint-specific token ids; with fewer than min_images distinct images nothing " + "is called recurring. The rule is the paper's, for the input layer: read deep, a " + "scene word common to most of the batch's images can recur too." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, + artifacts={"tokens": per_token, "prompt_sims": prompt_sims}, + findings=findings) diff --git a/evalrx/analyzers/patching/__init__.py b/evalrx/analyzers/patching/__init__.py index 8ba81c75..b0e29167 100755 --- a/evalrx/analyzers/patching/__init__.py +++ b/evalrx/analyzers/patching/__init__.py @@ -1,10 +1,11 @@ -"""Causal intervention analyzers (causal tracing and its gradient estimate / path patching / attention knockout / attention reweighting).""" +"""Causal intervention analyzers (causal tracing and its gradient estimate / path patching / attention knockout / attention reweighting and temperature).""" from evalrx.analyzers.patching.attention_knockout import AttentionKnockoutAnalyzer from evalrx.analyzers.patching.attention_sensitivity import AttentionSensitivityAnalyzer +from evalrx.analyzers.patching.attention_temperature import AttentionTemperatureAnalyzer from evalrx.analyzers.patching.attribution_patching import AttributionPatchingAnalyzer from evalrx.analyzers.patching.causal_trace import CausalTraceAnalyzer from evalrx.analyzers.patching.path_patching import PathPatchingAnalyzer -__all__ = ["AttentionKnockoutAnalyzer", "AttentionSensitivityAnalyzer", "AttributionPatchingAnalyzer", - "CausalTraceAnalyzer", "PathPatchingAnalyzer"] +__all__ = ["AttentionKnockoutAnalyzer", "AttentionSensitivityAnalyzer", "AttentionTemperatureAnalyzer", + "AttributionPatchingAnalyzer", "CausalTraceAnalyzer", "PathPatchingAnalyzer"] diff --git a/evalrx/analyzers/patching/attention_temperature.py b/evalrx/analyzers/patching/attention_temperature.py new file mode 100644 index 00000000..347a9042 --- /dev/null +++ b/evalrx/analyzers/patching/attention_temperature.py @@ -0,0 +1,214 @@ +"""Attention temperature — does sharpening or spreading the model's look change the answer? + +Divides the pre-softmax attention scores of the text positions after the media +by a temperature ``T`` in a band of layers (``T < 1`` sharpens, ``T > 1`` +flattens) and measures the next-token margin at ``T = sharp`` and ``T = flat``. + +``scope="within"`` (default) tempers only the scores on the media keys AND +keeps each row's total weight on the media fixed: the model looks at the image +exactly as much as before, only more focused (``p_k ∝ p_k^(1/T)`` inside the +image) or more spread out. That separates FOCUS from AMOUNT — the amount is +what ``attention_sensitivity`` scales. The slope of the margin in ``log T``: + +* negative — sharper focus inside the image strengthens the answer: the model + reads the image too diffusely for this answer; +* positive — spreading the attention over the image strengthens it: the model + is over-focused (e.g. on a few sink tokens) and misses the rest; +* near zero — the answer does not depend on how the look is distributed. + +``scope="all"`` is plain attention temperature over every visible key (text +and media together), which also moves the media's share. + +A diagnostic only. With ``target="gold"`` the slope is for the gold token: would +focusing (or spreading) move the model TOWARD the right answer. Only softmax +attention has a temperature: on a hybrid stack (Qwen3.5) the linear-attention +layers are left as they are (``tt_edited_layers`` counts the edited ones). + +References: +- Attention temperature as a decoding-time intervention — no single + originating paper; related visual-attention interventions: PAI (Liu et al., + ECCV 2024, arXiv:2407.21771), VAR (Kang et al., ICLR 2025, arXiv:2503.03321). +""" + +from __future__ import annotations + +import math +from collections import Counter +from typing import TYPE_CHECKING, Any + +from evalrx.core.analyzer import Analyzer +from evalrx.core.capability import Capability +from evalrx.core.registry import register_analyzer +from evalrx.core.result import Result + +if TYPE_CHECKING: + from evalrx.core.case import CaseBatch + from evalrx.core.model import Model + +_BANDS = ("all", "early", "mid", "late") +SCOPES = ("within", "all") + + +@register_analyzer("attention_temperature") +class AttentionTemperatureAnalyzer(Analyzer): + """Margin response to sharpening / flattening the text→media attention, per case. + + Hyper-parameters: + sharp / flat: temperatures below / above 1; defaults 0.5 and 2.0. + scope: ``"within"`` (media keys only, media mass per row kept) or + ``"all"`` (every visible key). + band: layers edited — ``"all"`` or the ``"early"`` / ``"mid"`` / + ``"late"`` third of the stack (softmax layers only). + target: ``"prediction"`` or ``"gold"``. + fidelity_rtol: tolerance of the edit-free attention tap vs native logits. + max_cases: label-stratified cap; 0 (the default) = every case. + Cost: 3 logits-only forwards per case. + """ + + name = "attention_temperature" + description = ( + "Sharpens and flattens the text's attention over the image/audio tokens (same amount " + "of looking, different focus) and measures which way the answer's margin moves." + ) + requires = frozenset({Capability.ATTENTION, Capability.LOGITS}) + applies_to_modalities = frozenset({"image", "video", "audio"}) + requires_modalities = frozenset({"image", "video", "audio"}) + signal_docs = { + "tt_margin_clean": ("Clean answer margin", "Logit margin of the target answer token with attention untouched."), + "tt_delta_sharp": ("Margin change, sharper", "Change in the answer margin when the attention over the media is sharpened (T = sharp)."), + "tt_delta_flat": ("Margin change, flatter", "Change in the answer margin when the attention over the media is flattened (T = flat)."), + "tt_slope": ("Margin per log-temperature", "Slope of the answer margin in log T; negative = sharper focus strengthens the answer, positive = spreading the look strengthens it."), + "tt_flips_sharp": ("Answer flips, sharper", "Sharpening changes the model's top next token."), + "tt_flips_flat": ("Answer flips, flatter", "Flattening changes the model's top next token."), + "tt_media_entropy": ("Focus inside the media", "Normalized entropy of the answer position's attention distribution over the media tokens in the edited layers (1 = spread evenly, 0 = one token)."), + } + + def __init__(self, sharp: float = 0.5, flat: float = 2.0, scope: str = "within", + band: str = "all", target: str = "prediction", fidelity_rtol: float = 1e-2, + max_cases: int = 0) -> None: + if band not in _BANDS: + raise ValueError(f"band must be one of {_BANDS}, got {band!r}") + if scope not in SCOPES: + raise ValueError(f"scope must be one of {SCOPES}, got {scope!r}") + if not (0 < sharp < 1 < flat) or not math.isfinite(flat): + raise ValueError(f"need 0 < sharp < 1 < flat, got sharp={sharp}, flat={flat}") + super().__init__(sharp=sharp, flat=flat, scope=scope, band=band, target=target, + fidelity_rtol=fidelity_rtol, max_cases=max_cases) + + def _run(self, model: "Model", cases: "CaseBatch") -> Result: + import torch + + from evalrx.analyzers._internals.access import ( + check_editable_attention, + decoder_stack, + media_positions, + torch_module, + ) + from evalrx.analyzers._internals.attn_tap import ( + attention_tap, + require_fidelity, + temperature, + ) + from evalrx.analyzers._internals.targets import build_margin + from evalrx.analyzers.patching.attention_knockout import layer_bands + from evalrx.core.model import CaptureSpec + + stack = decoder_stack(torch_module(model, self.name)) + check_editable_attention(model, self.name, stack) + chosen = range(stack.n_layers) if self.band == "all" else layer_bands(stack.n_layers)[self.band] + layers = [i for i in chosen if stack.attn[i] is not None] + if not layers: + raise ValueError(f"{self.name}: no softmax attention layer in band {self.band!r} " + f"(layer kinds: {stack.kinds}); there is no temperature to set") + spec = CaptureSpec(to_cpu=False) + log_span = math.log(self.flat) - math.log(self.sharp) + + per_case: list[dict[str, Any]] = [] + skipped: Counter = Counter() + fidelity: "dict | None" = None + for case in cases.stratified_head(self.max_cases): + probe = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + _, media = media_positions(probe) + if not media: + skipped["no_media_tokens"] += 1 + continue + seq_len = probe.seq_len + del probe + queries = list(range(max(media) + 1, seq_len)) + if not queries: + skipped["no_text_after_media"] += 1 + continue + if fidelity is None: + fidelity = require_fidelity(model, case.inputs, rtol=self.fidelity_rtol, analyzer=self.name) + keys = media if self.scope == "within" else list(range(seq_len)) + entropies: list[float] = [] + + def focus(layer: int, probs: Any) -> None: + if layer in layers: + row = probs[0, :, -1, media].float() # (heads, n_media) + row = row / row.sum(-1, keepdim=True).clamp(min=1e-12) + ent = -(row * row.clamp(min=1e-12).log()).sum(-1) / math.log(max(2, len(media))) + entropies.append(float(ent.mean())) + + with torch.no_grad(), attention_tap(model, on_probs=focus, analyzer=self.name): + clean = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + logits = clean.require(Capability.LOGITS)[-1:].float() + margin, why = build_margin(self.target, logits, model=model, case=case) + if margin is None: + skipped[why] += 1 + continue + top1 = int(logits[-1].argmax()) + del clean + + def tempered(value: float) -> Any: + edit = temperature(keys, value, queries=queries, layers=layers, + within=self.scope == "within") + with torch.no_grad(), attention_tap(model, [edit], analyzer=self.name): + out = model.forward(case.inputs, capture={Capability.LOGITS}, spec=spec) + return out.require(Capability.LOGITS)[-1:].float() + + m_clean = margin.value(logits) + sharp, flat = tempered(self.sharp), tempered(self.flat) + m_sharp, m_flat = margin.value(sharp), margin.value(flat) + per_case.append({ + "sample_id": case.id, + "tt_margin_clean": round(m_clean, 4), + "tt_delta_sharp": round(m_sharp - m_clean, 4), + "tt_delta_flat": round(m_flat - m_clean, 4), + "tt_slope": round((m_flat - m_sharp) / log_span, 4), + "tt_flips_sharp": margin.top1_changed(sharp, top1), + "tt_flips_flat": margin.top1_changed(flat, top1), + "tt_media_entropy": round(sum(entropies) / len(entropies), 4) if entropies else None, + }) + if per_case[-1]["tt_media_entropy"] is None: + del per_case[-1]["tt_media_entropy"] + + if not per_case: + raise ValueError(f"{self.name}: no case could be analysed (skipped: {dict(skipped)})") + n = len(per_case) + findings: dict[str, Any] = { + "n_cases_analyzed": n, + "n_skipped": sum(skipped.values()), + "skip_reasons": dict(skipped), + "n_layers": stack.n_layers, + "band": self.band, + "scope": self.scope, + "temperatures": [self.sharp, self.flat], + "target": self.target, + "tt_edited_layers": len(layers), + "layer_kinds": {"softmax": stack.n_softmax, "linear": stack.n_linear}, + "tap_fidelity": fidelity, + "mean_slope": round(sum(r["tt_slope"] for r in per_case) / n, 4), + "frac_negative_slope": round(sum(r["tt_slope"] < 0 for r in per_case) / n, 4), + "frac_flips_sharp": round(sum(r["tt_flips_sharp"] for r in per_case) / n, 4), + "frac_flips_flat": round(sum(r["tt_flips_flat"] for r in per_case) / n, 4), + "per_case": per_case, + "_caveat": ( + "Diagnostic only — the next-token margin at the end of the prompt with the " + "attention of the text after the media tempered in the chosen band" + + (" over the media keys, each row keeping its total media weight" + if self.scope == "within" else " over every visible key") + + "; softmax layers only, no decoding was changed." + ), + } + return Result(analyzer=self.name, model=repr(model), cases=cases, findings=findings) diff --git a/evalrx/eval_agent/stages/probe.py b/evalrx/eval_agent/stages/probe.py index 564e9cef..b5408061 100755 --- a/evalrx/eval_agent/stages/probe.py +++ b/evalrx/eval_agent/stages/probe.py @@ -62,7 +62,7 @@ class ModelKind(str, Enum): # surface evidence for that mode. Used in cycle 2+ to focus the probe on what # the diagnosis agent flagged rather than running the same generic priority list. _FAILURE_MODE_TO_ANALYZERS: dict[str, list[str]] = { - "attention_sink": ["attention_sink"], + "attention_sink": ["attention_sink", "attention_temperature", "embed_lens"], "attention": ["attention", "attention_rollout", "generic_attention"], "hallucination": ["pope", "chair", "selfcheck_consistency", "attention_knockout", "attention_sensitivity"], @@ -101,12 +101,14 @@ class ModelKind(str, Enum): "audio_neglect": ["attention_knockout", "causal_trace", "modality_ablation", "attention_sensitivity"], "modality_neglect": ["attention_knockout", "causal_trace", "modality_ablation"], - "insufficient_visual_attention": ["attention_sensitivity", "attention_knockout", "attention"], + "insufficient_visual_attention": ["attention_sensitivity", "attention_knockout", "attention", + "attention_temperature"], "visual_information_loss": ["causal_trace", "attribution_patching", "attention_knockout", - "path_patching"], + "path_patching", "embed_lens"], "late_visual_integration": ["causal_trace", "attribution_patching", "attention_knockout"], - "mislocalized_attention": ["gradcam", "generic_attention", "relative_attention"], - "grounding_error": ["gradcam", "generic_attention", "relative_attention"], + "mislocalized_attention": ["gradcam", "generic_attention", "relative_attention", + "attention_temperature"], + "grounding_error": ["gradcam", "generic_attention", "relative_attention", "embed_lens"], "late_suppression": ["tuned_lens", "patchscope", "logit_lens", "layer_contrast", "head_attribution"], "answer_head_conflict": ["head_attribution", "path_patching"], @@ -157,6 +159,7 @@ class ModelKind(str, Enum): "mm_shap", # per-slot reliance "causal_trace", "attribution_patching", # where the video's effect hands off to the text "path_patching", # which heads carry it to the answer + "attention_temperature", "embed_lens", # focus of the look; frame tokens at the input ], "audio": [ # Grounding first, for the same reason the text list leads with hygiene: @@ -172,9 +175,11 @@ class ModelKind(str, Enum): "self_consistency", "logprob_entropy", "attention", "attention_sink", # where does it attend among audio tokens? "attention_sensitivity", # would attending more to the audio help? + "attention_temperature", # ... or attending more sharply / broadly? "causal_trace", # where the audio's effect hands off to the text "attribution_patching", # ... estimated at every layer and token "path_patching", # which heads carry it to the answer + "embed_lens", # what the audio tokens mean at the LLM input "verbalized_confidence", "calibration", ], "image": [ @@ -184,9 +189,11 @@ class ModelKind(str, Enum): "attention", "attention_rollout", # where is the model looking? "attention_sink", # sink collapse? "attention_sensitivity", # would looking more help (PAI-style repair)? + "attention_temperature", # ... or looking more sharply / more broadly? "prompt_contrast", # are failures prompt-repairable? "mm_shap", # text vs image reliance "gradcam", "generic_attention", # which image regions the answer depends on + "embed_lens", # what the image tokens mean at the LLM input "causal_trace", # where the image's effect hands off to the text "attribution_patching", # ... estimated at every layer and token "path_patching", # which heads carry it to the answer @@ -205,9 +212,11 @@ class ModelKind(str, Enum): "attention", "attention_rollout", # where is the model looking? "attention_sink", # sink collapse? "attention_sensitivity", # would looking more help? + "attention_temperature", # ... or looking more sharply / more broadly? "prompt_contrast", # are failures prompt-repairable? "mm_shap", # text vs image reliance "gradcam", "generic_attention", # which image regions the answer depends on + "embed_lens", # what the image tokens mean at the LLM input "causal_trace", # where the image's effect hands off "attribution_patching", # ... estimated at every layer and token "path_patching", # which heads carry it to the answer diff --git a/evalrx/reporting/case_study.py b/evalrx/reporting/case_study.py index c5bbed12..102147d4 100644 --- a/evalrx/reporting/case_study.py +++ b/evalrx/reporting/case_study.py @@ -61,6 +61,8 @@ "causal_trace": "internal_white_box", "attribution_patching": "internal_white_box", "path_patching": "internal_white_box", + "attention_temperature": "internal_white_box", + "embed_lens": "internal_white_box", "patchscope": "internal_white_box", "mm_shap": "multimodal", "vl_shap": "multimodal", @@ -140,6 +142,7 @@ "causal_trace": "Did the answer use the image or audio at all?", "attribution_patching": "Did the answer use the image or audio at all?", "path_patching": "Did the answer use the image or audio at all?", + "attention_temperature": "Where was it looking?", "attention_sensitivity": "Did the answer use the image or audio at all?", "gradcam": "Where was it looking?", "generic_attention": "Where was it looking?", @@ -227,6 +230,8 @@ "attribution_patching": "attribution patching", "path_patching": "path patching", "patchscope": "Patchscope readout", + "attention_temperature": "attention temperature", + "embed_lens": "EmbedLens readout", "mm_shap": "modality attribution", "vl_shap": "vision-language attribution", "modality_ablation": "modality ablation", @@ -250,9 +255,11 @@ ("attention & attribution", {"relative_attn", "relative_attention", "context_shap", "linear_probe", "gradcam", "generic_attention"}), ("logits & representations", {"logprob_entropy", "logit_lens", "layer_contrast", - "tuned_lens", "head_attribution", "patchscope"}), + "tuned_lens", "head_attribution", "patchscope", + "embed_lens"}), ("causal interventions", {"attention_knockout", "attention_sensitivity", "causal_trace", - "attribution_patching", "path_patching"}), + "attribution_patching", "path_patching", + "attention_temperature"}), ], "multimodal": [ ("grounding & hallucination", {"pope", "chair", "opera", "vcd"}), diff --git a/tests/test_analyzers/test_analyzer_contract.py b/tests/test_analyzers/test_analyzer_contract.py index b8f15ebe..e1a363e7 100755 --- a/tests/test_analyzers/test_analyzer_contract.py +++ b/tests/test_analyzers/test_analyzer_contract.py @@ -230,6 +230,8 @@ def _long_prompt_batch() -> CaseBatch: "attribution_patching": "needs a live torch module; see test_intervention_analyzers.py", "path_patching": "needs a live torch module; see test_intervention_analyzers.py", "patchscope": "needs a live torch module; see test_intervention_analyzers.py", + "attention_temperature": "needs a live torch module; see test_intervention_analyzers.py", + "embed_lens": "needs a live torch module; see test_intervention_analyzers.py", "attention_sensitivity": "needs a live torch module; see test_intervention_analyzers.py", "head_attribution": "needs a live torch module; see test_intervention_analyzers.py", } diff --git a/tests/test_analyzers/test_analyzer_docs.py b/tests/test_analyzers/test_analyzer_docs.py index edd47716..e66bf940 100644 --- a/tests/test_analyzers/test_analyzer_docs.py +++ b/tests/test_analyzers/test_analyzer_docs.py @@ -39,6 +39,7 @@ def gen(): "| Attention sensitivity | `attention_sensitivity` | | | — | |\n" "| Attribution patching | `attribution_patching` | | | — | |\n" "| Path patching | `path_patching` | | | — | |\n" + "| Attention temperature | `attention_temperature` | | | — | |\n" ) diff --git a/tests/test_analyzers/test_hybrid_analyzers.py b/tests/test_analyzers/test_hybrid_analyzers.py index 1b77a1c1..d245b203 100644 --- a/tests/test_analyzers/test_hybrid_analyzers.py +++ b/tests/test_analyzers/test_hybrid_analyzers.py @@ -242,3 +242,38 @@ def test_path_patching_freezes_and_patches_the_linear_heads_too(hybrid): for row in f["per_case"]: assert row["pp_topk_direct_share"] == pytest.approx(1.0, abs=1e-3) assert 0.0 <= row["pp_linear_share"] <= 1.0 + + +def test_attention_temperature_edits_only_the_softmax_layers(hybrid): + from evalrx.analyzers.patching.attention_temperature import AttentionTemperatureAnalyzer + + stack = decoder_stack(hybrid._loaded[0]) + f = AttentionTemperatureAnalyzer().run(hybrid, _batch(2)).findings + assert f["tt_edited_layers"] == stack.n_softmax < stack.n_layers + assert f["layer_kinds"] == {"softmax": stack.n_softmax, "linear": stack.n_linear} + early_softmax = [i for i in range(round(stack.n_layers / 3)) if stack.kinds[i] == "softmax"] + if not early_softmax: # a band of linear layers only has no temperature to set + with pytest.raises(ValueError, match="no softmax attention layer"): + AttentionTemperatureAnalyzer(band="early").run(hybrid, _batch(1)) + + +def test_embed_lens_runs_on_hybrid(hybrid): + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + + _contract(EmbedLensAnalyzer().run(hybrid, _batch()), "embed_lens", 4) + + +def test_embed_lens_reads_a_placeholder_token_as_itself(): + # The text-only Qwen3-Next fixture feeds an "image" as 16 copies of the + # placeholder token: each media token IS that token's embedding, so its + # nearest word is itself at cosine 1 — and every case shows the same + # "image", which counts once. + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + from tests.tiny_whitebox import IMAGE_TOKEN + + result = EmbedLensAnalyzer().run(TinyVLM("qwen3_next"), _batch()) + f = result.findings + assert f["n_cases_analyzed"] == 4 and f["n_distinct_media"] == 1 and f["recurring_inferred"] is False + for row in f["per_case"]: + assert row["el_mean_max_sim"] == pytest.approx(1.0, abs=1e-4) + assert all(set(t["top1"]) == {IMAGE_TOKEN} for t in result.artifacts["tokens"].values()) diff --git a/tests/test_analyzers/test_intervention_analyzers.py b/tests/test_analyzers/test_intervention_analyzers.py index 44edd757..7ee28686 100644 --- a/tests/test_analyzers/test_intervention_analyzers.py +++ b/tests/test_analyzers/test_intervention_analyzers.py @@ -824,6 +824,156 @@ def test_attention_sensitivity_validates_params(): AttentionSensitivityAnalyzer(down=2.0, up=3.0) +# ── attention temperature ───────────────────────────────────────────────────── + +def test_temperature_edits_are_exact(vlm): + # Layer-0 attention depends on the input only, so edited and clean rows compare directly. + from evalrx.analyzers._internals.attn_tap import temperature + + inputs = Inputs(prompt="what is here?", image=image(0)) + media = list(range(3, 19)) + probs: dict = {} + + def grab(tag): + def record(layer, p): + if layer == 0: + probs[tag] = p.detach()[0].float() + return record + + with torch.no_grad(): + with attention_tap(vlm, on_probs=grab("clean")): + base = vlm.forward(inputs, capture={Capability.LOGITS}).logits + with attention_tap(vlm, [temperature(media, 0.5, within=True)], on_probs=grab("within")): + vlm.forward(inputs, capture={Capability.LOGITS}) + with attention_tap(vlm, [temperature(range(64), 2.0)], on_probs=grab("all")): + vlm.forward(inputs, capture={Capability.LOGITS}) + with attention_tap(vlm, [temperature(media, 1.0, within=True)]): + same = vlm.forward(inputs, capture={Capability.LOGITS}).logits + assert torch.equal(same, base) # T = 1 is no edit + p, w, a = probs["clean"], probs["within"], probs["all"] + last = p.shape[-1] - 1 + mass = p[:, last, media].sum(-1) + assert torch.allclose(w[:, last, media].sum(-1), mass, atol=1e-6) # how much it looks: kept + sharpened = p[:, last, media] ** 2 + assert torch.allclose(w[:, last, media], sharpened / sharpened.sum(-1, keepdim=True) * mass[:, None], atol=1e-6) + assert torch.allclose(w[:, last, 19:], p[:, last, 19:], atol=1e-6) # text keys untouched + flat = torch.softmax(torch.log(p[:, last, : last + 1]) / 2.0, -1) + assert torch.allclose(a[:, last, : last + 1], flat, atol=1e-6) + + +def test_temperature_must_be_positive_and_finite(): + from evalrx.analyzers._internals.attn_tap import temperature + + for bad in (0.0, -1.0, float("inf")): + with pytest.raises(ValueError, match="temperature"): + temperature([1, 2], bad) + + +@pytest.mark.parametrize("scope", ["within", "all"]) +def test_attention_temperature_contract(vlm, scope): + from evalrx.analyzers.patching.attention_temperature import AttentionTemperatureAnalyzer + + f = _assert_result_contract(AttentionTemperatureAnalyzer(scope=scope).run(vlm, _batch()), + "attention_temperature", 4) + assert f["tap_fidelity"]["tap_faithful"] is True and f["scope"] == scope + assert f["tt_edited_layers"] == 6 and f["temperatures"] == [0.5, 2.0] + for row in f["per_case"]: + assert isinstance(row["tt_flips_sharp"], bool) and 0.0 <= row["tt_media_entropy"] <= 1.0 + assert row["tt_slope"] == pytest.approx( + (row["tt_delta_flat"] - row["tt_delta_sharp"]) / math.log(4.0), abs=2e-4) + + +def test_attention_temperature_validates_params(vlm): + from evalrx.analyzers.patching.attention_temperature import AttentionTemperatureAnalyzer + + with pytest.raises(ValueError, match="sharp"): + AttentionTemperatureAnalyzer(sharp=1.5) + with pytest.raises(ValueError, match="scope"): + AttentionTemperatureAnalyzer(scope="text") + with pytest.raises(ValueError, match="band"): + AttentionTemperatureAnalyzer(band="top") + assert AttentionTemperatureAnalyzer(band="late").run(vlm, _batch(1)).findings["tt_edited_layers"] == 2 + + +# ── embed lens ──────────────────────────────────────────────────────────────── + +class _LetterDecoder: + """Decodes every id to a distinct three-letter word, so content words exist.""" + + def decode(self, ids): + i = int(ids[0]) + return chr(ord("a") + i % 26) + chr(ord("a") + (i // 26) % 26) + "x" + + +def test_embed_lens_neighbours_match_brute_force(vlm): + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + + batch = _batch(2) + result = EmbedLensAnalyzer(top_k=3, vocab_chunk=7).run(vlm, batch) # many blocks to merge + table = vlm._loaded[0].get_input_embeddings().weight.float() + for case in batch: + with torch.no_grad(), residual_hooks(vlm, capture={0: None}) as res: + vlm.forward(case.inputs, capture={Capability.LOGITS}) + x = res.captured[0][3:19].float() + cos = torch.nn.functional.normalize(x, dim=-1) @ torch.nn.functional.normalize(table, dim=-1).T + want = cos.topk(3, dim=-1) + got = result.artifacts["tokens"][case.id] + assert got["top_ids"] == want.indices.tolist() + assert got["max_sim"] == pytest.approx([round(v, 4) for v in want.values[:, 0].tolist()], abs=1e-4) + + +def test_embed_lens_contract_and_grounding(vlm, monkeypatch): + from evalrx.analyzers._internals import access as access_mod + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + + monkeypatch.setattr(access_mod, "tokenizer_of", lambda model: _LetterDecoder()) + result = EmbedLensAnalyzer().run(vlm, _batch()) + f = _assert_result_contract(result, "embed_lens", 4) + assert f["recurring_inferred"] is True + for row in f["per_case"]: + assert row["el_n_media_tokens"] == 16 + parts = row["el_frac_specific"] + row["el_frac_recurring"] + row["el_frac_high_norm"] + assert parts == pytest.approx(1.0, abs=1e-3) + sims = result.artifacts["prompt_sims"][row["sample_id"]] + assert row["el_prompt_best_sim"] == max(sims.values()) and row["el_prompt_best_word"] in sims + + +def test_embed_lens_reads_any_depth(vlm): + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + + batch = _batch(2) + deep = EmbedLensAnalyzer(layer=0.6, top_k=2).run(vlm, batch) + assert deep.findings["layer_used"] == 3 # round(0.6 * 5) + table = vlm._loaded[0].get_input_embeddings().weight.float() + with torch.no_grad(), residual_hooks(vlm, capture={3: None}) as res: + vlm.forward(batch[0].inputs, capture={Capability.LOGITS}) + cos = (torch.nn.functional.normalize(res.captured[3][3:19].float(), dim=-1) + @ torch.nn.functional.normalize(table, dim=-1).T) + assert deep.artifacts["tokens"][batch[0].id]["top_ids"] == cos.topk(2, dim=-1).indices.tolist() + + +def test_embed_lens_counts_each_image_once(vlm): + # Four questions about ONE image are one image's worth of evidence: every + # nearest word of it "recurs" in 100% of the distinct images, but only one + # image backs that, below min_images, so nothing is called recurring. + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + + same = CaseBatch([FailureCase(inputs=Inputs(prompt=f"question {i}?", image=image(0))) for i in range(4)]) + f = EmbedLensAnalyzer().run(vlm, same).findings + assert f["n_cases_analyzed"] == 4 and f["n_distinct_media"] == 1 + assert f["recurring_inferred"] is False and f["recurring_words"] == [] + + +def test_embed_lens_infers_nothing_recurring_from_too_few_images(vlm): + from evalrx.analyzers.lens.embed_lens import EmbedLensAnalyzer + + f = EmbedLensAnalyzer(min_images=10).run(vlm, _batch()).findings + assert f["recurring_inferred"] is False and f["recurring_words"] == [] + assert all(row["el_frac_recurring"] == 0.0 for row in f["per_case"]) + with pytest.raises(ValueError, match="recurring_frac"): + EmbedLensAnalyzer(recurring_frac=0.0) + + # ── tuned lens ──────────────────────────────────────────────────────────────── def test_tuned_lens_fit_beats_the_identity_translator_and_leaves_weights_alone(vlm, tmp_path): diff --git a/tests/test_eval_agent/test_probe_routing_whitebox.py b/tests/test_eval_agent/test_probe_routing_whitebox.py index 805ed2e9..4f66d6c5 100644 --- a/tests/test_eval_agent/test_probe_routing_whitebox.py +++ b/tests/test_eval_agent/test_probe_routing_whitebox.py @@ -30,7 +30,7 @@ WRITE_SIDE = {"attention_knockout", "attention_sensitivity", "causal_trace", "attribution_patching", "path_patching", "gradcam", "generic_attention", "head_attribution", "tuned_lens", - "patchscope"} + "patchscope", "attention_temperature", "embed_lens"} _WHITE_BOX = {Capability.GENERATE, Capability.LOGITS, Capability.LOGPROBS, Capability.ATTENTION, Capability.HIDDEN_STATES, Capability.GRADIENTS}