transformer_lens.model_bridge.supported_architectures.vit module

ViT / DeiT architecture adapter.

Supports HF ViTModel, ViTForImageClassification, DeiTModel, DeiTForImageClassification (single CLS-token classifier head). Encoder blocks are structurally near-identical between ViT and DeiT — same field names (layernorm_before/after, attention.{q,k,v,o}_proj, mlp.{fc1,fc2}) — differing only in the embeddings module used (DeiT’s carries an extra distillation token, which is invisible to this adapter — see vision_embeddings.py).

NOT covered: DeiTForImageClassificationWithTeacher (dual cls+distillation head, averaged). See vision_classifier_head.py’s docstring for why, and prepare_model() below raises loudly if you load one anyway rather than silently producing wrong logits.

ViT/DeiT blocks are pre-LN (LayerNorm applied before attention/MLP, residual added after — same shape as Llama/GPT2), unlike BERT’s post-LN. That’s why supports_fold_ln = True here where BertArchitectureAdapter sets it False.

NOTE (transformers >= the ViT/DeiT flattening refactor): as of this version of modeling_vit.py, HF removed the separate ViTEncoder wrapper — the blocks now live directly at <prefix>.layers on the model, not <prefix>.encoder.layer. Attention was flattened too: ViTAttention now owns q_proj/k_proj/v_proj/ o_proj directly (no more nested attention.attention.{query,key,value} + output.dense). And ViTLayer already exposes a flat .mlp submodule (ViTMLP with .fc1/.fc2) instead of the old intermediate/output split. Because of this, the old ViTMLPWrapper shim, the block-forward tuple-unwrapping monkey-patch (ViTLayer.forward returns a plain tensor now, not a tuple), and the hf_model.encoder_layer = hf_model.encoder.layer aliasing hack are all gone — the component mapping below points straight at the real attributes.

class transformer_lens.model_bridge.supported_architectures.vit.ViTArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for ViT and (non-distilled-head) DeiT vision models.

applicable_phases: list[int] = [1]

component_mapping: ComponentMapping | None

prepare_model(hf_model: Any) → None

Detect ViTForImageClassification vs DeiTForImageClassification vs a bare *Model, and add/omit the classifier head + prefix accordingly.

No structural patching of the HF model is needed any more: current transformers ViT/DeiT blocks already expose a flat .mlp (fc1/fc2) and return plain tensors from forward, and the blocks live directly on <prefix>.layers rather than behind a now-removed .encoder wrapper. This method only has to figure out the right prefix/classifier and build the component mapping to point at those real attributes.

supports_generation: bool = False

uses_split_attention: bool

weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None