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