Visual Encoders (ViT, DeiT) Support Rollout

transformer_lens/config/transformer_bridge_config.py

@@ -90,6 +90,8 @@ def __init__(
         attn_implementation: Optional[str] = None,
         # Audio model configuration
         is_audio_model: bool = False,
+        # Vision model (ViT, DeiT) configuration
+        is_visual_model: bool = False,
         # Stateful model configuration (e.g., Mamba SSMs use cache_params,
         # not past_key_values, so generation delegates to hf_generate)
         is_stateful: bool = False,
@@ -185,6 +187,8 @@ def __init__(
         self.attn_implementation = attn_implementation
         # Audio model configuration
         self.is_audio_model = is_audio_model
+        # Vision model (ViT, DeiT) configuration
+        self.is_visual_model = is_visual_model
         # Stateful model configuration
         self.is_stateful = is_stateful
         # Multimodal configuration

transformer_lens/factories/architecture_adapter_factory.py

@@ -147,6 +147,7 @@
     T5Gemma2ArchitectureAdapter,
     T5GemmaArchitectureAdapter,
     VaultGemmaArchitectureAdapter,
+    ViTArchitectureAdapter,
     XGLMArchitectureAdapter,
     YoutuArchitectureAdapter,
     Zamba2ArchitectureAdapter,
@@ -310,6 +311,10 @@
     "MinGPTForCausalLM": MingptArchitectureAdapter,
     "GPTNeoForCausalLM": NeoArchitectureAdapter,
     "GPTNeoXForCausalLM": NeoxArchitectureAdapter,
+    "ViTModel": ViTArchitectureAdapter,
+    "ViTForImageClassification": ViTArchitectureAdapter,
+    "DeiTModel": ViTArchitectureAdapter,
+    "DeiTForImageClassification": ViTArchitectureAdapter,  # "DeiTForImageClassificationWithTeacher" unsupported for now
 }
 
 

transformer_lens/model_bridge/bridge.py

@@ -844,7 +844,7 @@ def enable_compatibility_mode(
         if not getattr(self.adapter, "supports_compatibility_mode", True):
             raise RuntimeError(
                 f"{type(self.adapter).__name__} does not support compatibility mode: "
-                "its stored-processed-weights forward is known to diverge from the "
+                "its stored-processed-weights  is known to diverge from the "
                 "reference model. Use the default bridge forward instead."
             )
 
@@ -1760,9 +1760,10 @@ def forward(
                 (fragile across HF versions) or exception-based layer skipping (corrupts
                 model state). Raises NotImplementedError if a non-None value is passed.
             stop_at_layer: Layer to stop forward pass at
-            pixel_values: Optional image tensor for multimodal models (e.g., LLaVA, Gemma3).
+            pixel_values: Optional image tensor for multimodal models (e.g., LLaVA, Gemma3)
+                and vision models (eg. ViT, DeiT).
                 The tensor is passed directly to the underlying HuggingFace model.
-                Only valid when cfg.is_multimodal is True.
+                Only valid when cfg.is_multimodal is True or cfg.is_visual_model is True.
             input_values: Optional audio waveform tensor for audio models (e.g., HuBERT).
                 The tensor is passed directly to the underlying HuggingFace model.
                 Only valid when cfg.is_audio_model is True.
@@ -1822,6 +1823,7 @@ def forward(
             isinstance(input, list)
             and len(input) > 1
             and not getattr(self.cfg, "is_audio_model", False)
+            and not getattr(self.cfg, "is_visual_model", False)
         )
 
         try:
@@ -1831,6 +1833,11 @@ def forward(
                         "Audio models require tensor input (raw waveform), not text. "
                         "Pass a torch.Tensor or use the input_values parameter."
                     )
+                if getattr(self.cfg, "is_visual_model", False):
+                    raise ValueError(
+                        "Visual models require tensor input (pixel values), not text. "
+                        "Pass a torch.Tensor or use the pixel_values parameter."
+                    )
                 if _is_batched_list and padding_side is None:
                     # Force left-padding so real tokens are flush-right.
                     _orig_padding_side = self.tokenizer.padding_side
@@ -1918,10 +1925,13 @@ def forward(
 
             # Handle pixel_values for multimodal models
             if pixel_values is not None:
-                if not getattr(self.cfg, "is_multimodal", False):
+                if not (
+                    getattr(self.cfg, "is_multimodal", False)
+                    or getattr(self.cfg, "is_visual_model", False)
+                ):
                     raise ValueError(
-                        "pixel_values can only be passed to multimodal models "
-                        "(cfg.is_multimodal must be True)"
+                        "pixel_values can only be passed to multimodal or vision models "
+                        "(cfg.is_multimodal or cfg.is_visual_model must be True)"
                     )
                 kwargs["pixel_values"] = pixel_values
 
@@ -1946,6 +1956,19 @@ def forward(
                         "Audio models require tensor input (raw waveform). "
                         "Pass a torch.Tensor or use input_values parameter."
                     )
+            elif getattr(self.cfg, "is_visual_model", False):
+                # "pixel_values" may already be in kwargs from the "if pixel_values is not None:"
+                # gate above (explicit pixel_values=... call); otherwise treat `input` itself as
+                # the image tensor, matching how the audio branch treats `input` as the waveform.
+                if "pixel_values" not in kwargs:
+                    if isinstance(input, torch.Tensor):
+                        kwargs["pixel_values"] = input
+                    else:
+                        raise ValueError(
+                            "Visual models require tensor input (pixel values). "
+                            "Pass a torch.Tensor as `input` or use the pixel_values parameter."
+                        )
+                output = self.original_model(**kwargs)
             elif _is_inputs_embeds:
                 output = self.original_model(inputs_embeds=input_ids, **kwargs)
             else:
@@ -1957,6 +1980,13 @@ def forward(
                 logits = output.logits
             elif isinstance(output, tuple) and len(output) > 0:
                 logits = output[0]
+            elif hasattr(output, "last_hidden_state"):
+                # Bare encoder models (ViTModel, DeiTModel, BertModel, etc. without
+                # a task head) return e.g. BaseModelOutput/BaseModelOutputWithPooling,
+                # which has neither `.logits` nor tuple semantics. Fall back to
+                # `last_hidden_state` so return_type="logits" still yields a plain
+                # tensor rather than silently handing back the raw HF output object.
+                logits = output.last_hidden_state
             else:
                 logits = output
             if return_type == "logits":
@@ -1970,6 +2000,14 @@ def forward(
                         "Audio models do not support return_type='loss'. "
                         "CTC loss requires aligned frame-level labels."
                     )
+                if getattr(self.cfg, "is_visual_model", False):
+                    raise ValueError(
+                        "Vision classification models do not support return_type='loss' "
+                        "via this path (no next-token LM target exists for image "
+                        "classification). Compute cross-entropy against `labels` "
+                        "yourself from the returned logits, or use hf_generate()-style "
+                        "direct access to self.original_model for HF's own loss."
+                    )
                 if _is_inputs_embeds:
                     raise ValueError(
                         "Cannot compute loss with inputs_embeds — token IDs required for labels."

transformer_lens/model_bridge/generalized_components/__init__.py

@@ -137,6 +137,12 @@
 from transformer_lens.model_bridge.generalized_components.vision_projection import (
     VisionProjectionBridge,
 )
+from transformer_lens.model_bridge.generalized_components.vision_embeddings import (
+    VisionEmbeddingsBridge,
+)
+from transformer_lens.model_bridge.generalized_components.vision_classifier_head import (
+    VisionClassifierHeadBridge,
+)
 
 __all__ = [
     "AttentionBridge",
@@ -193,4 +199,6 @@
     "SSMBlockBridge",
     "SSMMixerBridge",
     "VisionProjectionBridge",
+    "VisionEmbeddingsBridge",
+    "VisionClassifierHeadBridge",
 ]

transformer_lens/model_bridge/generalized_components/vision_classifier_head.py

@@ -0,0 +1,37 @@
+"""CLS-token classifier head bridge component.
+
+HF's ViTForImageClassification.forward() / DeiTForImageClassification.forward()
+already pool the CLS token before calling self.classifier:
+
+    sequence_output = outputs.last_hidden_state
+    pooled_output = sequence_output[:, 0, :]
+    logits = self.classifier(pooled_output)
+
+So this component only ever receives an already-pooled (batch, hidden) tensor.
+No slicing needed here — this is a thin, hook-named pass-through.
+
+Deliberately NOT covering DeiTForImageClassificationWithTeacher (dual
+cls_classifier + distillation_classifier head) — see vit.py's docstring.
+"""
+
+from typing import Any
+
+from torch import Tensor
+
+from transformer_lens.model_bridge.generalized_components.base import (
+    GeneralizedComponent,
+)
+
+
+class VisionClassifierHeadBridge(GeneralizedComponent):
+    """Wraps the classifier nn.Linear that HF calls with an already-pooled CLS token."""
+
+    def forward(self, pooled_output: Tensor, **kwargs: Any) -> Tensor:
+        original_component = self.original_component
+        if original_component is None:
+            raise RuntimeError(
+                f"Original component not set for {self.name}. Call set_original_component() first."
+            )
+        pooled_output = self.hook_in(pooled_output)
+        logits = original_component(pooled_output)
+        return self.hook_out(logits)

transformer_lens/model_bridge/generalized_components/vision_embeddings.py

@@ -0,0 +1,64 @@
+"""Vision (ViT-style) embeddings bridge component.
+
+Wraps a HF `ViTEmbeddings` / `DeiTEmbeddings` module directly and forwards through
+it unmodified. The patch-conv projection, CLS-token (and, for DeiT, distillation-
+token) concatenation, and position-embedding addition are HF's math — we don't
+reimplement any of it, we just point at the real module and hook its input/output.
+
+Why one class covers both ViT and DeiT: `ViTEmbeddings.forward` and
+`DeiTEmbeddings.forward` have an *identical* signature
+    (pixel_values, bool_masked_pos=None, interpolate_pos_encoding=False)
+The only structural difference between them (DeiT prepends a distillation token
+in addition to CLS, and sizes `position_embeddings` for +2 slots instead of +1)
+lives entirely inside whichever module `set_original_component` resolves to — it's
+invisible to this wrapper, so no DeiT-specific branching is needed here.
+
+Verified against the real generalized_components/base.py (not extrapolated):
+`GeneralizedComponent.forward(*args, **kwargs)` already does exactly what we need
+— it hooks the input via `self.hook_in`, casts it to match the wrapped module's
+own parameter dtype (equivalent to HF's own `pixel_values.to(expected_dtype)`
+"kept for BC" cast in ViTModel.forward — both land on the same dtype in the
+overwhelming common case of a non-mixed-precision checkpoint), calls
+`self.original_component(*args, **kwargs)`, and hooks the output via
+`self.hook_out`. So this subclass only needs to guarantee `pixel_values` reaches
+`super().forward()` *positionally* — the base class's own kwarg-name sniffing
+list (`input`, `hidden_states`, `input_ids`, `query_input`, `x`, `inputs_embeds`)
+doesn't include `pixel_values`, so if this component were ever called
+all-keyword (`self.embeddings(pixel_values=x, ...)`, unlike the two HF call
+sites we've confirmed) the base class's hook_in would silently never fire.
+Re-emitting positionally here closes that gap regardless of how *we* were
+called.
+"""
+
+from typing import Any, Optional
+
+import torch
+from torch import Tensor
+
+from transformer_lens.model_bridge.generalized_components.base import (
+    GeneralizedComponent,
+)
+
+
+class VisionEmbeddingsBridge(GeneralizedComponent):
+    """Bridge for ViTEmbeddings / DeiTEmbeddings.
+
+    Point `name=` at the embeddings module directly, e.g.:
+        VisionEmbeddingsBridge(name="embeddings")        # bare ViTModel/DeiTModel
+        VisionEmbeddingsBridge(name="vit.embeddings")     # ViTForImageClassification
+        VisionEmbeddingsBridge(name="deit.embeddings")    # DeiTForImageClassification
+    """
+
+    def forward(
+        self,
+        pixel_values: Tensor,
+        bool_masked_pos: Optional[torch.BoolTensor] = None,
+        interpolate_pos_encoding: Optional[bool] = None,
+        **kwargs: Any,
+    ) -> Tensor:
+        return super().forward(
+            pixel_values,
+            bool_masked_pos=bool_masked_pos,
+            interpolate_pos_encoding=interpolate_pos_encoding,
+            **kwargs,
+        )

transformer_lens/model_bridge/sources/transformers.py

@@ -157,6 +157,19 @@ def map_default_transformer_lens_config(hf_config):
         tl_config.n_ctx = source_config.max_length
     elif hasattr(source_config, "seq_length"):
         tl_config.n_ctx = source_config.seq_length
+    elif hasattr(source_config, "image_size") and hasattr(source_config, "patch_size"):
+        # Vision Transformers calculate sequence length dynamically:
+        # (image_size / patch_size)^2 + 1 (for the CLS token)
+        image_size = source_config.image_size
+        patch_size = source_config.patch_size
+
+        # HF configs allow these to be integers or tuples/lists
+        img_h = image_size[0] if isinstance(image_size, (list, tuple)) else image_size
+        img_w = image_size[1] if isinstance(image_size, (list, tuple)) else image_size
+        patch_h = patch_size[0] if isinstance(patch_size, (list, tuple)) else patch_size
+        patch_w = patch_size[1] if isinstance(patch_size, (list, tuple)) else patch_size
+
+        tl_config.n_ctx = (img_h // patch_h) * (img_w // patch_w) + 1
     elif hasattr(source_config, "max_sequence_length"):
         tl_config.n_ctx = source_config.max_sequence_length
     else:
@@ -367,6 +380,8 @@ def get_hf_model_class_for_architecture(architecture: str):
         MASKED_LM_ARCHITECTURES,
         MULTIMODAL_ARCHITECTURES,
         SEQ2SEQ_ARCHITECTURES,
+        VISION_ARCHITECTURES,
+        VISION_CLASSIFICATION_ARCHITECTURES,
     )
 
     if architecture in SEQ2SEQ_ARCHITECTURES or architecture in AUDIO_TEXT_ARCHITECTURES:
@@ -388,6 +403,14 @@ def get_hf_model_class_for_architecture(architecture: str):
             return AutoModelForCTC
         from transformers import AutoModel
 
+        return AutoModel
+    elif architecture in VISION_ARCHITECTURES:
+        if architecture in VISION_CLASSIFICATION_ARCHITECTURES:
+            from transformers import AutoModelForImageClassification
+
+            return AutoModelForImageClassification
+        from transformers import AutoModel
+
         return AutoModel
     else:
         return AutoModelForCausalLM
@@ -813,7 +836,8 @@ def boot(
     use_fast = getattr(adapter.cfg, "use_fast", True)
     # Audio models use feature extractors, not text tokenizers
     _is_audio = getattr(adapter.cfg, "is_audio_model", False)
-    if _is_audio and tokenizer is None:
+    _is_visual = getattr(adapter.cfg, "is_visual_model", False)
+    if (_is_audio or _is_visual) and tokenizer is None:
         tokenizer = None  # Skip tokenizer loading for audio models
     elif tokenizer is not None:
         tokenizer = setup_tokenizer(tokenizer, default_padding_side=default_padding_side)

transformer_lens/model_bridge/supported_architectures/__init__.py

@@ -252,6 +252,9 @@
 from transformer_lens.model_bridge.supported_architectures.switch_transformers import (
     SwitchTransformersArchitectureAdapter,
 )
+from transformer_lens.model_bridge.supported_architectures.vit import (
+    ViTArchitectureAdapter,
+)
 from transformer_lens.model_bridge.supported_architectures.t5 import T5ArchitectureAdapter
 from transformer_lens.model_bridge.supported_architectures.t5gemma import T5GemmaArchitectureAdapter
 from transformer_lens.model_bridge.supported_architectures.t5gemma2 import (
@@ -396,6 +399,8 @@
     "Starcoder2ArchitectureAdapter",
     "T5ArchitectureAdapter",
     "T5GemmaArchitectureAdapter",
+    "T5Gemma2ArchitectureAdapter",
+    "ViTArchitectureAdapter",
     "XGLMArchitectureAdapter",
     "Zamba2ArchitectureAdapter",
     "HrmTextArchitectureAdapter",