transformer_lens.benchmarks.hook_registration module

Hook registration and behavior benchmarks for TransformerBridge.

transformer_lens.benchmarks.hook_registration.benchmark_critical_forward_hooks(bridge: TransformerBridge, test_text: str, reference_activations: Dict[str, Tensor] | None = None, tolerance: float = 0.02) → BenchmarkResult

Benchmark critical forward hooks commonly used in interpretability research.

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing (must match the snapshot’s prompt)

  • reference_activations – Optional reference activations keyed by hook name (e.g. a golden fixture snapshot). Capture-only self-check if None.

  • tolerance – Tolerance for activation comparison

Returns:

BenchmarkResult with critical hook comparison details

transformer_lens.benchmarks.hook_registration.benchmark_forward_hooks(bridge: TransformerBridge, test_text: str, reference_activations: Dict[str, Tensor] | None = None, tolerance: float = 0.5, prepend_bos: bool | None = None) → BenchmarkResult

Benchmark all forward hooks for activation matching.

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing (must match the snapshot’s prompt)

  • reference_activations – Optional reference activations keyed by hook name (e.g. a golden fixture snapshot). Fire-only self-check if None.

  • tolerance – Tolerance for activation matching (fraction of mismatches allowed)

  • prepend_bos – Whether to prepend BOS token. If None, uses model default.

Returns:

BenchmarkResult with hook activation comparison details

transformer_lens.benchmarks.hook_registration.benchmark_gated_hooks_fire(bridge: TransformerBridge, test_text: str = 'The quick brown fox', prepend_bos: bool | None = None) → BenchmarkResult

Verify each cfg-gated attention hook fires when its flag is enabled.

Hooks like hook_result, hook_q_input, hook_attn_in exist unconditionally on the attention bridge but are only populated when the corresponding config flag is set (keeping default-path cost at zero). This benchmark toggles each flag in turn, runs a short forward, and asserts at least one layer’s matching hook actually captured an activation.

use_attn_in and use_split_qkv_input are mutually exclusive, so each flag runs in its own forward pass. Plain AttentionBridge (non-PEA/JPEA) adapters raise NotImplementedError from the setter — recorded as skipped rather than failed, since the applicability gate is intentional.

transformer_lens.benchmarks.hook_registration.benchmark_hook_functionality(bridge: TransformerBridge, test_text: str, reference_effect: float | None = None, atol: float = 0.002) → BenchmarkResult

Benchmark hook system functionality through ablation effects.

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing (must match the reference’s prompt)

  • reference_effect – Optional reference ablation loss-delta (e.g. from a golden fixture). Effect-only self-check if None.

  • atol – Absolute tolerance for effect comparison

Returns:

BenchmarkResult with hook functionality comparison details

transformer_lens.benchmarks.hook_registration.benchmark_hook_registry(bridge: TransformerBridge, reference_hooks: Iterable[str] | None = None) → BenchmarkResult

Benchmark hook registry completeness.

Parameters:

  • bridge – TransformerBridge model to test

  • reference_hooks – Optional reference hook-name collection (e.g. the keys of a golden hook manifest). Structural self-check only if None.

Returns:

BenchmarkResult with registry comparison details