Research map / Architectures and efficiency / Attention and transformers
Scaling White-Box Transformers for Vision
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CRATE-α, a scaled white-box transformer for vision with an overcomplete dictionary, decoupled dictionary, and residual connection in the sparse coding block, plus an improved training recipe, showing it scales from Base to Huge on ImageNet classification and zero-shot tasks.
Problem
Scalability of CRATE, a mathematically interpretable white-box transformer, has been limited to small sizes and lags far behind standard ViTs in accuracy.
Approach
Replace CRATE's complete-dictionary ISTA block with an overcomplete (4x) dictionary, a decoupled output dictionary, and a residual connection, while keeping the MSSA compression block and an improved training recipe, so the model scales while remaining interpretable.
Key results
- CRATE-α-B/8 reaches 83.2% ImageNet-1K top-1 vs. 79.5% for prior best CRATE-B, a 3.7% gain.
- CRATE-α-L/8 reaches 85.1% ImageNet-1K top-1 after IN-21K pre-training; scaling Base to Large gives +1.9% on this setting.
- CRATE-α-CLIPA-H/14 achieves 72.3% zero-shot ImageNet-1K accuracy, 2.5% above the Large model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the complete dictionary with an overcomplete dictionary (ablation) | CRATE · ImageNet-1K | top-1 accuracy: 5.3% | Helped |
| Use a decoupled dictionary in the last sparse-coding step (ablation) | CRATE with an overcomplete dictionary · ImageNet-1K | top-1 accuracy: 2.0% | Helped |
| Add a residual connection to the sparse-coding block (ablation) | CRATE with overcomplete and decoupled dictionaries · ImageNet-1K | top-1 accuracy: 0.7% | Helped |
| Apply all three sparse-coding modifications | CRATE · ImageNet-1K | top-1 accuracy: 8% | Helped |
| Scale from CRATE-B to CRATE-L without ImageNet-21K pre-training | CRATE · ImageNet-1K | top-1 accuracy: +0.5%, from 70.8% to 71.3% | Helped |
| Scale CRATE-α from Base to Large with ImageNet-21K pre-training and smaller patch size | CRATE-α · ImageNet-1K | top-1 accuracy: from 83.2% to 85.1% | Helped |
| Scale from CRATE-α-CLIPA-B/16 to CRATE-α-CLIPA-L/14 | CRATE-α-CLIPA · ImageNet-1K zero-shot | accuracy: 11.3% during pre-training; 9.0% during fine-tuning | Helped |
| Scale from CRATE-α-CLIPA-L/14 to CRATE-α-CLIP-H/14 | CRATE-α-CLIP · ImageNet-1K zero-shot | accuracy: 3.1% during pre-training; 2.5% during fine-tuning | Helped |
| Pre-train with CRATE-α-L/32 instead of CRATE-α-L/14, then fine-tune with L/14 | CRATE-α-L · ImageNet-1K | top-1 accuracy: 83.7% vs 83.9% | No clear effect |
| Pre-train with CRATE-α-L/32 instead of CRATE-α-L/8, then fine-tune with L/8 | CRATE-α-L · ImageNet-1K | top-1 accuracy: 84.2% compared to 85.1% | Hurt |
| Use CRATE-α-L/32 for pre-training and L/8 for fine-tuning instead of L/14 for both stages | CRATE-α-L · ImageNet-1K | top-1 accuracy: 84.2% vs 83.9% | Helped |
| Replace CRATE with CRATE-α | Image classification on downstream datasets · CIFAR-10 | accuracy: 97.22 -> 98.17 | Helped |
Surprising findings
- Vanilla CRATE gains only 0.5% (70.8% to 71.3%) going from Base to Large on ImageNet-1K, showing diminishing returns that CRATE-α avoids.
- A CRATE-α-L/32 model pre-trained on ImageNet-21K and fine-tuned at higher token length uses about 35% of the compute of L/14 training while reaching comparable accuracy (83.7% vs 83.9%).
Limitations
- Requires substantial computational resources, which may be inaccessible to many researchers.
- Still trails ViT in some settings: CRATE-α-L/16 at 336 resolution reaches 84.6% vs ViT-L/16 at 85.2%.
- Language modeling experiments were incomplete due to time and resource limits (80% and 55% of iterations).
Experimental setup
Datasets: ImageNet-21K ImageNet-1K DataComp1B CIFAR-10 CIFAR-100 Oxford Flowers-102 Oxford-IIIT Pets ADE20K OpenWebText COCO val2017
Benchmarks: ImageNet-1K CIFAR-10 CIFAR-100 Oxford Flowers-102 Oxford-IIIT Pets ADE20K OpenWebText COCO val2017
Metrics: top-1 accuracy zero-shot accuracy mIoU CE val loss average precision throughput FLOPs
Baselines: CRATE ViT GPT-2 CLIPA
Builds on: CRATE CLIPA ViT MaskCut
Closest papers by meaning
- Masked Completion via Structured Diffusion with White-Box Transformers (ICLR 2024 · similarity 0.89)
- Win-Win: Training High-Resolution Vision Transformers from Two Windows (ICLR 2024 · similarity 0.87)
- Vision Transformers Need Registers (ICLR 2024 · similarity 0.86)
- Progressive Token Length Scaling in Transformer Encoders for Efficient Universal Segmentation (ICLR 2025 · similarity 0.86)
- Slicing Vision Transformer for Flexible Inference (NeurIPS 2024 · similarity 0.86)
- The Need for Speed: Pruning Transformers with One Recipe (ICLR 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.