transformer_lens.benchmarks.vision module

Vision benchmarks for TransformerBridge (Phase 9).

Tests that vision encoder models (ViT, DeiT) correctly handle pixel inputs through forward(), run_with_cache(), and produce stable representations — the hook/cache coverage that Phase 1 (HF parity on one forward) doesn’t give non-text models. The audio analog is Phase 8 (audio.py).

transformer_lens.benchmarks.vision.benchmark_vision_cache(bridge: TransformerBridge, test_pixels: Tensor) → BenchmarkResult

Benchmark run_with_cache() for vision models.

Verifies that critical vision hooks fire and produce valid tensors: the patch embeddings, final layernorm, classifier head (when present), and the first and last block.

Parameters:

• bridge – TransformerBridge model to test

• test_pixels – Pixel tensor [batch, channels, height, width]

transformer_lens.benchmarks.vision.benchmark_vision_classification_decode(bridge: TransformerBridge) → BenchmarkResult

Benchmark image-classification decoding on a real image.

Loads the cats-image fixture, preprocesses it with the bridge’s image processor, and reports the top predicted labels — the vision analog of the audio CTC decode, exercising the processor wiring that synthetic pixel tensors don’t. Skipped for bare encoders (no classifier head), tiny-random models, and when no processor/datasets are available.

Parameters:

bridge – TransformerBridge model to test

transformer_lens.benchmarks.vision.benchmark_vision_embeddings(bridge: TransformerBridge, test_pixels: Tensor) → BenchmarkResult

Verify patch-embedding hook outputs.

Checks that embed.hook_out produces [batch, seq, d_model] tensors with non-degenerate values — the vision analog of the audio feature-extractor check. Seq length is architecture-dependent (ViT: patches + CLS; DeiT: patches + CLS + distillation token), so only lower-bounded here.

Parameters:

• bridge – TransformerBridge model to test

• test_pixels – Pixel tensor [batch, channels, height, width]

transformer_lens.benchmarks.vision.benchmark_vision_forward(bridge: TransformerBridge, test_pixels: Tensor, reference_model: Module | None = None) → BenchmarkResult

Benchmark forward pass with pixel input.

Compares bridge output against the HF native model on the same pixels. Bare encoders (ViTModel) compare last_hidden_state; classification heads compare logits.

Parameters:

• bridge – TransformerBridge model to test

• test_pixels – Pixel tensor [batch, channels, height, width]

• reference_model – Optional HF reference model for comparison

transformer_lens.benchmarks.vision.benchmark_vision_representation_stability(bridge: TransformerBridge, test_pixels: Tensor) → BenchmarkResult

Benchmark representation stability under small pixel perturbations.

Parameters:

• bridge – TransformerBridge model to test

• test_pixels – Pixel tensor [batch, channels, height, width]

transformer_lens.benchmarks.vision.run_vision_benchmarks(bridge: TransformerBridge, test_pixels: Tensor | None = None, verbose: bool = True) → List[BenchmarkResult]

Run all vision benchmarks.

Parameters:

• bridge – TransformerBridge model to test

• test_pixels – Optional pixel tensor. If None, generates a synthetic input shaped for this architecture from the HF config.

• verbose – Whether to print progress

Returns:

List of BenchmarkResult objects