Upload model files

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,110 @@
+---
+license: mit
+library_name: transformers
+tags:
+  - vision-transformer
+  - image-classification
+  - efficient-transformer
+  - selective-attention
+  - knowledge-distillation
+  - computer-vision
+pipeline_tag: image-classification
+---
+
+# Soft-Masked Selective Vision Transformer
+
+## Model Description
+
+Soft-Masked Selective Vision Transformer is an efficient **Vision Transformer (ViT)** model designed to reduce the computational overhead of self-attention while maintaining competitive accuracy.  
+The model introduces a **patch-selective attention mechanism** that enables the transformer to focus on the most salient image regions and dynamically disregard less informative patches. This selective strategy significantly reduces the quadratic complexity typically associated with full self-attention, making the model particularly suitable for **high-resolution vision tasks** and **resource-constrained environments**.
+
+To further improve performance, the model leverages **knowledge distillation**, transferring representational knowledge from a stronger teacher network to enhance the accuracy of lightweight transformer variants.
+
+---
+
+## Intended Use
+
+This model is intended for:
+
+- Image classification tasks
+- Deployment in **compute- or memory-constrained environments**
+- High-resolution image processing where standard ViTs are prohibitively expensive
+- Research on efficient attention mechanisms and transformer compression
+
+### Example Use Cases
+
+- Edge or embedded vision systems
+- Large-scale image analysis with reduced inference cost
+- Efficient backbones for downstream vision tasks
+
+---
+
+## Training Details
+
+- **Training Objective:** Cross-entropy loss with optional distillation loss
+- **Distillation:** Teacher–student framework
+- **Optimization:** AdamW
+- **Training Dataset:** ILSVRC 2012
+- **Evaluation Metrics:** Top-1 accuracy, FLOPs, parameter count
+
+---
+
+## Usage
+
+### Image Classification Example
+
+```python
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
+import requests
+
+# Load image
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
+# Load processor and model
+processor = AutoImageProcessor.from_pretrained(
+    "XAFT/SM-Selective-ViT-Base-224-Distilled",
+    trust_remote_code=True,
+)
+
+model = AutoModelForImageClassification.from_pretrained(
+    "XAFT/SM-Selective-ViT-Base-224-Distilled",
+    trust_remote_code=True,
+)
+
+# Preprocess
+inputs = processor(
+    images=image,
+    return_tensors="pt",
+)
+
+# Forward pass
+outputs = model(**inputs)
+logits = outputs.logits
+predicted_class = logits.argmax(-1).item()
+
+print("Predicted class index:", predicted_class)
+```
+
+---
+
+## Evaluation Results
+
+| Model                  | Top-1 Acc. | Top-5 Acc. | # Params | Avg. GFLOPs |
+|------------------------|------------|------------|----------|-------------|
+| Base                   | 80.350%    | 94.980%    | 86.60M   | 9.61        |
+| Base (distilled)       | 80.990%    | 95.386%    | 87.37M   | 9.21        |
+| Small                  | 78.662%    | 94.454%    | 22.06M   | 3.12        |
+| Small (distilled)      | 79.000%    | 94.494%    | 22.45M   | 3.05        |
+| Tiny tall              | 74.802%    | 92.794%    | 11.07M   | 1.64        |
+| Tiny tall (distilled)  | 75.676%    | 92.988%    | 11.26M   | 1.64        |
+| Tiny                   | 71.056%    | 90.192%    | 5.72M    | 0.95        |
+| Tiny (distilled)       | 72.618%    | 91.338%    | 5.92M    | 0.93        |
+
+
+---
+
+## Acknowledgments
+
+We thank the TPU Research Cloud program for providing cloud TPUs that were used to build and train the models for our extensive experiments.

__init__.py

@@ -0,0 +1,7 @@
+from .configuration_selectivevit import SMSelectiveViTConfig
+from .modeling_selectivevit import SMSelectiveViTModelForClassification
+
+__all__ = [
+    "SMSelectiveViTConfig",
+    "SMSelectiveViTModelForClassification",
+]

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "SMSelectiveViTModelForClassification"
+  ],
+  "atten_dim": 768,
+  "attention_scale": 8.0,
+  "auto_map": {
+    "AutoConfig": "configuration_selectivevit.SMSelectiveViTConfig",
+    "AutoModelForImageClassification": "modeling_selectivevit.SMSelectiveViTModelForClassification"
+  },
+  "channels": 3,
+  "depth": 12,
+  "drop_path": 0.2,
+  "dropout": 0.0,
+  "dtype": "float32",
+  "embed_dim": 768,
+  "ffn_groups": null,
+  "image_size": 224,
+  "mask_threshold": 0.05,
+  "mlp_dim": 3072,
+  "model_type": "softmasked_selective_vit",
+  "num_classes": 1000,
+  "num_groups": 12,
+  "num_heads": 12,
+  "patch_drop": 0.05,
+  "patch_size": 16,
+  "transformers_version": "4.57.3",
+  "use_distil_token": true
+}

configuration_selectivevit.py

@@ -0,0 +1,43 @@
+# configuration_my_model.py
+from transformers import PretrainedConfig
+
+class SMSelectiveViTConfig(PretrainedConfig):
+    model_type = "softmasked_selective_vit"
+
+    def __init__(
+        self,
+        image_size=224,
+        patch_size=16,
+        num_classes=1000,
+        embed_dim=768,
+        atten_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_dim=3072,
+        channels=3,
+        dropout=0.0,
+        drop_path=0.0,
+        attention_scale=0.0,
+        mask_threshold=0.0,
+        patch_drop=0.0,
+        use_distil_token=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        # store everything as attributes (HF will save them in config.json)
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_classes = num_classes
+        self.embed_dim = embed_dim
+        self.atten_dim = atten_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.channels = channels
+        self.dropout = dropout
+        self.drop_path = drop_path
+        self.attention_scale = attention_scale
+        self.mask_threshold = mask_threshold
+        self.patch_drop = patch_drop
+        self.use_distil_token = use_distil_token
+

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cd1a78510e92f164589ff458c23d69309413d14687029d5bcfccbc102ac06c6
+size 349491616

modeling_selectivevit.py

@@ -0,0 +1,91 @@
+# modeling_my_model.py
+from dataclasses import dataclass
+from typing import Optional, Tuple, List
+import inspect
+import torch
+from torch import nn
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from .configuration_selectivevit import SMSelectiveViTConfig
+
+from .selective_vit import VisionTransformer
+
+
+@dataclass
+class ImageClassifierWithMasksOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    distil_logits: Optional[torch.FloatTensor] = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    masks: Optional[List[torch.FloatTensor]] = None
+
+
+
+class SMSelectiveViTModelForClassification(PreTrainedModel ):
+    config_class = SMSelectiveViTConfig
+    base_model_prefix = "backbone"
+
+    def __init__(self, config: SMSelectiveViTConfig):
+        super().__init__(config)
+
+        # Build kwargs for your original class from config fields
+        cfg_dict = config.to_dict()
+
+        # Filter to only args your MyModel.__init__ accepts
+        sig = inspect.signature(VisionTransformer.__init__)
+        allowed = set(sig.parameters.keys())
+        allowed.discard("self")
+        model_kwargs = {k: v for k, v in cfg_dict.items() if k in allowed}
+
+        self.backbone = VisionTransformer(**model_kwargs)
+
+        self.post_init()  # important: ties into HF weight init hooks
+
+    def forward(
+        self,
+        pixel_values=None,
+        labels=None,
+        full=False,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict
+            if return_dict is not None
+            else self.config.use_return_dict
+        )
+
+        last_hidden, all_hidden, masks = self.backbone.forward_features(
+            pixel_values,
+            full=full,
+            output_hidden_states=output_hidden_states,
+        )
+
+        logits, distil_logits = self.backbone.forward_classifier(last_hidden)
+
+        loss = None
+        if labels is not None:
+            loss_fn = nn.CrossEntropyLoss()
+            loss = loss_fn(logits, labels)
+
+        if not return_dict:
+            return (loss, logits, distil_logits, last_hidden, all_hidden, masks)
+
+        return ImageClassifierWithMasksOutput(
+            loss=loss,
+            logits=logits,
+            distil_logits=distil_logits,
+            last_hidden_state=last_hidden,
+            hidden_states=all_hidden,
+            masks=masks,
+        )
+
+

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "DeiTImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 236,
+    "width": 236
+  }
+}

selective_vit.py

@@ -0,0 +1,412 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.ops import StochasticDepth
+import math
+
+class SoftMaskedMultiheadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True,
+                 add_bias_kv=True, kdim=None, vdim=None,
+                 scale=8., device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.scale = scale
+
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+
+        self.head_dim = embed_dim // num_heads
+
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs)
+        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias and add_bias_kv, **factory_kwargs)
+        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias and add_bias_kv, **factory_kwargs)
+
+        self.dropout_layer = nn.Dropout(dropout)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.k_proj.weight)
+        nn.init.xavier_uniform_(self.v_proj.weight)
+
+        if self.q_proj.bias is not None:
+            nn.init.constant_(self.q_proj.bias, 0.)
+        if self.k_proj.bias is not None:
+            nn.init.constant_(self.k_proj.bias, 0.)
+        if self.v_proj.bias is not None:
+            nn.init.constant_(self.v_proj.bias, 0.)
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.v_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.)
+
+    def forward(self, query, key, value, key_padding_mask=None,
+                need_weights=True, attn_mask=None, average_attn_weights=True):
+        """
+        query, key, value: shape (L, N, E)
+        where L is the sequence length, N is the batch size, E is the embedding dimension.
+        """
+        batch_size, tgt_len, embed_dim = query.size()
+        batch_size, src_len, _ = key.size()
+
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(value)
+
+        # Reshape q, k, v for multihead attention
+        q = q.view(batch_size, tgt_len, self.num_heads, self.head_dim).transpose(1,2)
+        k = k.view(batch_size, src_len, self.num_heads, self.head_dim).transpose(1,2)
+        v = v.view(batch_size, src_len, self.num_heads, self.head_dim).transpose(1,2)
+
+        # Compute scaled dot-product attention scores
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        # scores shape: (batch_size, num_heads, tgt_len, src_len)
+
+        # Apply the soft [0, 1] mask
+        if attn_mask is not None:
+            # Ensure attn_mask values are in (0, 1] to avoid log(0)
+            # attn_mask shape [b, l]
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
+            if not self.training:
+                scores = scores.masked_fill((attn_mask == 0.), float('-inf'))
+            eps = 1e-6
+            attn_mask = attn_mask.clip(min=eps).log()
+            # attn_mask shape [b, 1, 1, l]
+            scores = scores + self.scale * attn_mask
+
+        # Apply key padding mask
+        if key_padding_mask is not None:
+            key_padding_mask = key_padding_mask.view(batch_size, 1, 1, src_len)
+            scores = scores.masked_fill(key_padding_mask, float('-inf'))
+
+        # Compute attention weights
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout_layer(attn_weights)
+
+        # Compute attention output
+        attn_output = torch.matmul(attn_weights, v)
+        # attn_output shape: (batch_size, num_heads, tgt_len, head_dim)
+
+        # Concatenate heads and project
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, tgt_len, embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+
+        if need_weights:
+            # Optionally average attention weights over heads
+            if average_attn_weights:
+                attn_weights = attn_weights.mean(dim=1)
+            else:
+                attn_weights = attn_weights
+        else:
+            attn_weights = None
+
+        return attn_output, attn_weights
+
+def get_ffn(input_dim, output_dim, middle_dim, dropout=0.1):
+    fc1 = nn.Linear(input_dim, middle_dim)
+    fc2 = nn.Linear(middle_dim, output_dim)
+    fc3 = nn.Identity()
+    return nn.Sequential(
+            fc1,
+            nn.GELU(),
+            nn.Dropout(dropout),
+            fc2,
+            nn.Dropout(dropout),
+            fc3
+            )
+# Assuming SoftMaskedMultiheadAttention is already defined as provided earlier
+class EncoderBlock(nn.Module):
+    def __init__(self, input_dim, embed_dim, num_heads, mlp_dim, dropout=0.1, drop_path=0.0, patch_drop=0.0, attention_scale=2., mask_threshold=0.05):
+        super().__init__()
+        self.mask_threshold = mask_threshold
+        self.self_attn = SoftMaskedMultiheadAttention(
+            embed_dim, num_heads, dropout=dropout, scale=attention_scale
+        )
+        if attention_scale > 0:
+            self.linear_mask = nn.Linear(input_dim, 1)  # Linear layer to compute mask scores
+            self.patch_drop = nn.Dropout(patch_drop)
+        else:
+            self.linear_mask = None
+        if input_dim != embed_dim:
+            raise ValueError("embed_dim must equal atten_dim but {input_dim}!={embed_dim}")
+        else:
+            self.embed = nn.Identity()
+            self.project = nn.Identity()
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        # Feed-forward network (MLP)
+        self.mlp = get_ffn(embed_dim, embed_dim, mlp_dim, dropout=dropout)
+        self.path_drop = StochasticDepth(drop_path, mode='row')
+        self.norm3 = nn.LayerNorm(input_dim)
+
+    def _reset_parameters(self):
+        for n, m in self.named_modules():
+            if n.startswith('self_attn'):
+                continue
+            if isinstance(m, (nn.Linear, GroupedLinear)):
+                nn.init.trunc_normal_(m.weight.data, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias.data)
+        nn.init.ones_(self.norm1.weight)
+        nn.init.zeros_(self.norm1.bias)
+        nn.init.ones_(self.norm2.weight)
+        nn.init.zeros_(self.norm2.bias)
+        nn.init.zeros_(self.norm3.weight)
+        nn.init.zeros_(self.norm3.bias)
+
+    def forward_common(self, x, mask):
+        """
+        x: shape (batch_size, seq_len, embed_dim)
+        """
+        # Compute mask scores: (batch_size, seq_len, 1)
+        x1 = x
+        x = self.embed(x)
+        x = self.norm1(x)
+        # Apply attention mechanism
+        attn_output, _ = self.self_attn(x, x, x, attn_mask=mask)
+        # Add & Norm
+        x = x + self.path_drop(attn_output)
+        x = self.norm2(x)
+        # Feed-forward network
+        mlp_output = self.mlp(x)
+        # Add & Norm
+        x = self.path_drop(self.project(x + mlp_output))
+        x = self.norm3(x)
+        if mask is not None:
+            x = x * mask.unsqueeze(-1)
+        x = x1 + x
+        return x
+
+    def get_groups(self, mask, full=False):
+        n_items, index = (mask != 0.0).sum(-1).cpu().sort(descending=True)
+        n_items, index = n_items.tolist(), index.tolist()
+        groups = []
+        t = 1.0 if full else 1.2
+        for ni, ii in zip(n_items, index):
+            if ni == 0:
+                break
+            if len(groups) == 0 or groups[-1][1] / ni > t:
+                groups.append(([], ni))
+            groups[-1][0].append(ii)
+        return groups
+
+    def infer_forward(self, x, mask, full=False):
+        """
+        The “sparse‐inference” path: for each group of batch‐samples that have the same
+        number n of tokens ≥ mask_threshold, gather only those top‐n tokens (in original order),
+        run forward_common on the smaller (b’, n, dim) tensor, then scatter the results back.
+        Fully masked tokens are left untouched.
+        """
+        # Step 1: Threshold the mask without in-place ops
+        mask_thresholded = mask * (mask >= self.mask_threshold)
+        # Step 2: Prepare output tensor (copy of x)
+        x_out = x.clone()
+        # Step 3: Group samples by number of kept tokens
+        groups = self.get_groups(mask_thresholded, full)
+        # Step 4: Process each group
+        for batch_indices, n_keep in groups:
+            x_sel = x[batch_indices]                     # (Bg, seq_len, input_dim)
+            mask_sel = mask_thresholded[batch_indices]   # (Bg, seq_len)
+            # Top-k selection and sorting
+            topk_vals, topk_idx_unsorted = torch.topk(mask_sel, k=n_keep, dim=1, sorted=False)
+            topk_idx_sorted, _ = topk_idx_unsorted.sort(dim=1)
+            # Gather tokens in sorted order
+            idx_expanded = topk_idx_sorted.unsqueeze(-1).expand(-1, -1, x_sel.size(-1))
+            X_topk = torch.gather(x_sel, dim=1, index=idx_expanded)
+            mask_topk = torch.gather(mask_sel, dim=1, index=topk_idx_sorted)
+            # Run forward pass
+            results = self.forward_common(X_topk, mask_topk)
+            # Scatter results into a new x_sel tensor
+            x_sel_updated = x_sel.clone()
+            x_sel_updated = x_sel_updated.scatter(1, idx_expanded, results)
+            # Write the updated batch slice into the new output tensor
+            x_out[batch_indices] = x_sel_updated
+        return x_out
+
+    def forward(self, x, full=False):
+        if self.linear_mask is not None:
+            attn_mask = self.patch_drop(self.linear_mask(x).sigmoid().squeeze(-1))
+        else:
+            attn_mask = None
+        if not self.training and not attn_mask is None and self.mask_threshold >= 0:
+            x = self.infer_forward(x, attn_mask, full)
+        else:
+            x = self.forward_common(x, attn_mask)
+        return x, attn_mask
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+        self,
+        image_size=256,
+        patch_size=16,
+        num_classes=1000,
+        embed_dim=768,
+        atten_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_dim=768,
+        channels=3,
+        dropout=0.1,
+        drop_path=0.1,
+        patch_drop=0.1,
+        attention_scale=2.,
+        mask_threshold=0.05,
+        use_distil_token=False
+    ):
+        super().__init__()
+        assert image_size % patch_size == 0, "Image dimensions must be divisible by the patch size."
+        num_patches = (image_size // patch_size) ** 2
+
+        # Patch embedding layer
+        self.patch_embed = nn.Conv2d(
+            in_channels=channels,
+            out_channels=embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size
+        )
+
+        # Class token
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+
+        # Positional embedding
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1 + (1 if use_distil_token else 0), embed_dim))
+
+        self.dropout = nn.Dropout(dropout)
+
+        # Encoder blocks
+        self.encoder_layers = nn.ModuleList([
+            EncoderBlock(
+                embed_dim, atten_dim,
+                num_heads, mlp_dim,
+                dropout, drop_path * i / (depth - 1),
+                patch_drop=patch_drop,
+                attention_scale=attention_scale,
+                mask_threshold=mask_threshold,
+                )
+            for i in range(depth)
+        ])
+
+        # Classification head
+        self.post_norm = nn.LayerNorm(embed_dim)
+        self.head = nn.Linear(embed_dim, num_classes)
+        if use_distil_token:
+            self.dis_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+            self.dis_head = nn.Linear(embed_dim, num_classes)
+        else:
+            self.dis_token = None
+
+        # Initialize weights
+        self._init_weights()
+
+    def _init_weights(self):
+        for n, m in self.named_modules():
+            if n.startswith('encoder_layers'):
+                continue
+            if isinstance(m, (nn.Linear, nn.Conv2d)):
+                nn.init.trunc_normal_(m.weight.data, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias.data)
+            if isinstance(m, nn.LayerNorm):
+                nn.init.ones_(m.weight.data)
+                nn.init.zeros_(m.bias.data)
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        if self.cls_token is not None:
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+        if self.dis_token is not None:
+            nn.init.trunc_normal_(self.dis_token, std=0.02)
+
+    def forward_features(
+        self,
+        pixel_values,
+        full=False,
+        output_hidden_states=False,
+    ):
+        """
+        Args:
+            pixel_values: (B, C, H, W)
+        Returns:
+            last_hidden_state: (B, N, D)
+            all_hidden_states: tuple or None
+            masks: Tensor or None
+        """
+        batch_size = pixel_values.size(0)
+        hidden_states = []
+
+        # Patch embedding
+        x = self.patch_embed(pixel_values)
+        x = x.flatten(2).transpose(1, 2)
+
+        # Distillation token
+        if self.dis_token is not None:
+            dis_tokens = self.dis_token.expand(batch_size, -1, -1)
+            x = torch.cat((dis_tokens, x), dim=1)
+
+        # CLS token
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # Position + dropout
+        x = x + self.pos_embed
+        x = self.dropout(x)
+
+        masks = []
+
+        for layer in self.encoder_layers:
+            x, mask = layer(x, full)
+
+            if output_hidden_states:
+                hidden_states.append(x)
+
+            if mask is not None:
+                masks.append(mask)
+
+        x = self.post_norm(x)
+
+        if output_hidden_states:
+            hidden_states = tuple(hidden_states)
+        else:
+            hidden_states = None
+
+        if len(masks) > 0:
+            masks = tuple(masks)
+        else:
+            masks = None
+
+        return x, hidden_states, masks
+
+
+    def forward_classifier(self, hidden_states):
+        """
+        Args:
+            hidden_states: (B, N, D)
+        Returns:
+            logits: (B, num_classes)
+            dis_logits: (B, num_classes) or None
+        """
+        cls_token = hidden_states[:, 0]
+        logits = self.head(cls_token)
+
+        dis_logits = None
+        if self.dis_token is not None:
+            dis_cls_token = hidden_states[:, 1]
+            dis_logits = self.dis_head(dis_cls_token)
+
+            # Inference-time averaging (same as original)
+            if not self.training:
+                logits = (logits + dis_logits) / 2
+
+        return logits, dis_logits
+
+    def forward(self, x, full=False):
+        last_hidden_states, hidden_states, masks = self.forward_features(x, full)
+        logits, dis_logits = self.forward_classifier(last_hidden_states)
+        return logits, dis_logits, masks