transformer_lens.benchmarks.weight_processing module¶

Weight processing benchmarks for TransformerBridge.

transformer_lens.benchmarks.weight_processing.benchmark_attention_output_centering(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark attention output centering - W_O should have mean ≈ 0.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with attention output centering verification details

transformer_lens.benchmarks.weight_processing.benchmark_layer_norm_folding(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark layer norm folding - norm weights should be identity after folding.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with layer norm folding verification details

transformer_lens.benchmarks.weight_processing.benchmark_mlp_output_centering(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark MLP output centering - MLP output weights should have mean ≈ 0.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with MLP output centering verification details

transformer_lens.benchmarks.weight_processing.benchmark_no_nan_inf(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark that weights contain no NaN or Inf values.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with NaN/Inf verification details

transformer_lens.benchmarks.weight_processing.benchmark_unembed_centering(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark unembed centering - unembed matrix should have mean ≈ 0.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with unembed centering verification details

transformer_lens.benchmarks.weight_processing.benchmark_value_bias_folding(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark value bias folding - b_V should be zero after folding.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with value bias folding verification details

transformer_lens.benchmarks.weight_processing.benchmark_weight_magnitudes(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark that weight magnitudes are in reasonable ranges.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with weight magnitude verification details

transformer_lens.benchmarks.weight_processing.benchmark_weight_modification(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark that weight modifications propagate correctly.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model (not used)

Returns:

BenchmarkResult with weight modification verification details

transformer_lens.benchmarks.weight_processing.benchmark_weight_processing(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None) → BenchmarkResult¶

Benchmark weight processing (folding, centering) application.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model

Returns:

BenchmarkResult with weight processing verification details

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

Benchmark weight sharing and modification effects.

Parameters:

bridge – TransformerBridge model to test
test_text – Input text for testing
reference_model – Optional HookedTransformer reference model
atol – Absolute tolerance for effect comparison

Returns:

BenchmarkResult with weight sharing verification details