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) → BenchmarkResult

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

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with attention output centering verification details

transformer_lens.benchmarks.weight_processing.benchmark_layer_norm_folding(bridge: TransformerBridge, test_text: str) → BenchmarkResult

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

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with layer norm folding verification details

transformer_lens.benchmarks.weight_processing.benchmark_mlp_output_centering(bridge: TransformerBridge, test_text: str) → BenchmarkResult

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

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with MLP output centering verification details

transformer_lens.benchmarks.weight_processing.benchmark_no_nan_inf(bridge: TransformerBridge, test_text: str) → BenchmarkResult

Benchmark that weights contain no NaN or Inf values.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with NaN/Inf verification details

transformer_lens.benchmarks.weight_processing.benchmark_unembed_centering(bridge: TransformerBridge, test_text: str) → BenchmarkResult

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

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with unembed centering verification details

transformer_lens.benchmarks.weight_processing.benchmark_value_bias_folding(bridge: TransformerBridge, test_text: str) → BenchmarkResult

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

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with value bias folding verification details

transformer_lens.benchmarks.weight_processing.benchmark_weight_magnitudes(bridge: TransformerBridge, test_text: str) → BenchmarkResult

Benchmark that weight magnitudes are in reasonable ranges.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with weight magnitude verification details

transformer_lens.benchmarks.weight_processing.benchmark_weight_modification(bridge: TransformerBridge, test_text: str) → BenchmarkResult

Benchmark that weight modifications propagate correctly.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with weight modification verification details