thinkingvit-swin-s-imagenet1k

ThinkingViT-Swin / Swin-S ImageNet-1K

This repository contains the ImageNet-1K EMA weights for ThinkingViT-Swin / Swin-S ImageNet-1K from ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference.

• Paper: https://arxiv.org/abs/2507.10800
• Code: https://github.com/ds-kiel/ThinkingViT
• Project page: https://ds-kiel.github.io/ThinkingViT-project-page/
• Exported checkpoint key: state_dict_ema
• Weight format: model.safetensors

Usage

import torch
from timm.models import create_model

# Run from the ThinkingViT repository root, or put this repository on PYTHONPATH.
model = create_model("hf-hub:alihjt/thinkingvit-swin-s-imagenet1k", pretrained=True)
model.eval()

x = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    logits, stage = model(x, threshold=1.0)
print(logits.shape, stage)

This is a custom timm-based architecture. Use the code from the ThinkingViT repository when loading this model.

Threshold Behavior

The entropy threshold controls early exit. Lower thresholds send more samples to the full Swin-S round; higher thresholds exit after the reduced-head round.

ImageNet-1K Results

Citation

Please cite the ThinkingViT paper if you use this model: https://arxiv.org/abs/2507.10800