transformer_lens.benchmarks.backward_gradients module

Backward gradient benchmarks for TransformerBridge.

transformer_lens.benchmarks.backward_gradients.benchmark_backward_hooks(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None, abs_tolerance: float = 0.2, rel_tolerance: float = 0.0003) → BenchmarkResult

Benchmark all backward hooks for gradient matching.

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing

  • reference_model – Optional HookedTransformer reference model

  • abs_tolerance – Absolute tolerance for gradient comparison

  • rel_tolerance – Relative tolerance for gradient comparison

Returns:

BenchmarkResult with backward hook comparison details

transformer_lens.benchmarks.backward_gradients.benchmark_critical_backward_hooks(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None, abs_tolerance: float = 0.2, rel_tolerance: float = 0.0003) → BenchmarkResult

Benchmark critical backward hooks for gradient matching.

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing

  • reference_model – Optional HookedTransformer reference model

  • abs_tolerance – Absolute tolerance for gradient comparison

  • rel_tolerance – Relative tolerance for gradient comparison

Returns:

BenchmarkResult with critical backward hook comparison details

transformer_lens.benchmarks.backward_gradients.benchmark_gradient_computation(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None, atol: float = 0.001) → BenchmarkResult

Benchmark basic gradient computation.

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing

  • reference_model – Optional HookedTransformer reference model

  • atol – Absolute tolerance for gradient comparison

Returns:

BenchmarkResult with gradient computation comparison details

transformer_lens.benchmarks.backward_gradients.gradient_mismatch_is_numerical_noise(rel_l2: float, over_count: int) → bool

True when a gradient mismatch is diffuse and tiny rather than a divergence.

Elementwise worst-case cannot separate the two: one element of 55k crossing the tolerance scores the same as a head scaled by 1%. rel_l2 separates them by 58x or more, and the element COUNT guards the localized case rel_l2 would dilute. A count (not a fraction) keeps the band reachable on small tensors: detection guarantees count >= 1, so a fractional guard of 1e-4 was arithmetically unsatisfiable below 10,000 elements (gemma-3-270m’s MQA hook_rot_k is 6,912).

transformer_lens.benchmarks.backward_gradients.gradient_mismatch_stats(bridge_finite: Tensor, reference_finite: Tensor, abs_tolerance: float, rel_tolerance: float) → dict

Scale-aware statistics for grading one recorded gradient mismatch.

A zero reference with a nonzero bridge gradient is the maximally divergent case, not perfect agreement, so rel_l2 is inf there rather than 0.

transformer_lens.benchmarks.backward_gradients.needs_fp32_gradients(dtype: dtype | None) → bool

Reduced-precision gradients cannot be graded against the fp32-calibrated band — bf16’s rounding floor alone is ~2e-3 rel_l2, inside the bug band.