microsoft/Mage-ViT

1---2license: mit3library_name: transformers4pipeline_tag: image-feature-extraction5tags:6- mage-vl7- vision-encoder8- codec-vit9- video-understanding10---11 12<h1 align="center">Mage-ViT<br><span style="font-size: 0.55em; font-weight: normal;">A Codec-Native Visual Encoder Trained from Scratch</span></h1>13 14<p align="center">15  <a href="https://microsoft.github.io/Mage"><img alt="Project Page" src="https://img.shields.io/badge/%F0%9F%8C%90-Project%20Page-blue" height="22" /></a>16  <a href="https://github.com/microsoft/Mage"><img src="https://img.shields.io/badge/Code-GitHub-181717?logo=github" alt="GitHub" height="22"></a>17  <a href="https://huggingface.co/microsoft/Mage-VL"><img alt="Mage-VL" src="https://img.shields.io/badge/%F0%9F%A4%97-Mage--VL-yellow" height="22" /></a>18  <a href="https://huggingface.co/microsoft/Mage-ViT"><img alt="Mage-ViT" src="https://img.shields.io/badge/%F0%9F%A4%97-Mage--ViT-yellow" height="22" /></a>19  <a href="https://opensource.org/licenses/MIT"><img src="https://img.shields.io/badge/License-MIT-green" alt="License: MIT" height="22"></a>20</p>21 22---23 24**Mage-ViT** is the visual encoder at the core of **[Mage-VL](https://huggingface.co/microsoft/Mage-VL)**. It is a *Codec-ViT* built primarily for video, where a single image is simply the degenerate one-frame case. Mage-ViT follows a **codec-aligned sparsity** principle: visual tokens should be spent where a video codec spends bits, since codec bit-allocation is a natural proxy for spatio-temporal importance. On a shared `16×16` patch grid it keeps every anchor (I-frame) patch and only the motion-salient predicted (P-frame) patches, while a shared **3D rotary position encoding** preserves spatio-temporal structure even after large fractions of the grid are dropped.25 26This repository is the **ViT-pre-trained checkpoint only** — it has not gone through the joint VLM training with the language model. Use it as a data-efficient, codec-native visual encoder, or as a drop-in ViT for your own multimodal training.27 28## ✨ Highlights29 30- **Codec-driven patchifier.** Patches are selected by a per-patch importance map derived from the codec — motion vectors + P-frame residual energy for HEVC/H.265, or the learned rate map of the neural codec DCVC-RT. For a 64-frame clip it keeps all I-frame patches plus the top-*k* P-frame patches within a **4096-token budget (~75% token reduction)**. Chunk-wise and collage patchification are also supported.31- **Codec-agnostic.** The same interface accepts a traditional codec (HEVC/H.265) or a neural codec (DCVC-RT) with no architecture or retraining change.32- **Trained from scratch.** No billion-scale image-text ViT initialization — Mage-ViT is optimized with a large-scale **cluster-discrimination** objective on **~100M unlabeled images/videos**.33 34## 🏗️ Architecture35 36A **24-layer pre-norm Vision Transformer** trunk (hidden size `1024`, `16` attention heads, GELU MLP at 4× expansion) processes the variable-length token sequence produced by the codec-driven patchifier, followed by a multi-head attention pooling head. The standalone encoder in this repo consumes `pixel_values` (and optional `patch_positions`); codec-driven patch **selection** is applied upstream in the Mage-VL data pipeline.37 38Pre-training is a two-stage, from-scratch recipe in bf16: **(1)** variable-resolution image pre-training (224–448), then **(2)** joint image + video pre-training (video at resolution 256, 64 frames per clip, 4096-token budget).39 40## 📦 Installation41 42Mage-ViT follows the Mage-VL runtime:43 44```bash45pip install "transformers>=5.7" torch torchvision pillow46# optional, for the fastest GPU attention path:47# pip install flash-attn --no-build-isolation48```49 50Attention is dispatched internally to `sdpa` (default) → `flash_attention_2` → `eager`, so **flash-attn is optional** and the model runs on CPU or GPU out of the box.51 52## 🚀 Encoding an image53 54```python55import torch56from PIL import Image57from transformers import AutoModel, AutoImageProcessor58 59model = AutoModel.from_pretrained(60    "microsoft/Mage-ViT", trust_remote_code=True61).to("cuda").eval()62processor = AutoImageProcessor.from_pretrained(63    "microsoft/Mage-ViT", trust_remote_code=True64)65 66image = Image.open("your_image.jpg")67pixel_values = processor(images=image, return_tensors="pt")["pixel_values"].to("cuda")68 69with torch.no_grad():70    out = model(pixel_values)              # pixel_values: [B, 3, H, W]71 72patch_features = out.last_hidden_state     # [B, num_patches, 1024]73pooled_feature = out.pooler_output         # [B, 1024]74```75 76Pass `attn_implementation="flash_attention_2"` (or `"sdpa"` / `"eager"`) to `from_pretrained` to pick the attention backend.77 78## 🎞️ Encoding a video79 80Stack the preprocessed frames into `[B, C, T, H, W]` and pass a `patch_positions` tensor of shape `[B, T * tokens_per_frame, 3]` giving the `(t, h, w)` grid coordinate of every patch:81 82```python83import torch84from PIL import Image85 86PATCH = 1687 88def build_patch_positions(num_frames, target_frames, grid_h, grid_w, device="cuda"):89    # temporal index for each frame, spread across the target timeline90    t = torch.linspace(0, target_frames - 1, num_frames, device=device).long()91    t = t.repeat_interleave(grid_h * grid_w)                    # [T * H * W]92    h = torch.arange(grid_h, device=device).repeat_interleave(grid_w).repeat(num_frames)93    w = torch.arange(grid_w, device=device).repeat(grid_h).repeat(num_frames)94    return torch.stack([t, h, w], dim=-1).unsqueeze(0)          # [1, T*H*W, 3]95 96frames = [Image.open(f"frame_{i}.jpg") for i in range(16)]      # your sampled frames97pv = processor(images=frames, return_tensors="pt")["pixel_values"]   # [T, C, H, W]98video = pv.unsqueeze(0).permute(0, 2, 1, 3, 4).to("cuda")            # [1, C, T, H, W]99 100gh, gw = video.shape[-2] // PATCH, video.shape[-1] // PATCH101patch_positions = build_patch_positions(num_frames=16, target_frames=64,102                                        grid_h=gh, grid_w=gw)103 104with torch.no_grad():105    out = model(video, patch_positions=patch_positions)106```107 108## 📤 Outputs109 110| Field | Shape | Description |111| ----- | ----- | ----------- |112| `last_hidden_state` | `[B, num_patches, 1024]` | per-patch features after the final layer norm |113| `pooler_output`     | `[B, 1024]`              | global feature from the multi-head attention pooling head |114 115## 📋 Specifications116 117| | |118| --- | --- |119| Architecture | Codec-ViT (pre-norm ViT, SigLIP-style MLP) |120| Parameters | ~316M |121| Hidden size / MLP | 1024 / 4096 |122| Layers / heads | 24 / 16 |123| Patch size | 16 |124| Position encoding | shared 3D rotary (4:6:6 split over T:H:W) |125| Pooling | learned-probe multi-head attention head |126| Pre-training resolution | images 224–448 (variable), video 256 |127| Weights dtype | bfloat16 |128| License | MIT |129 130## 📄 License131 132Released under the [MIT License](https://opensource.org/licenses/MIT).133