Research map / Architectures and efficiency / RNNs and normalization
Masked Completion via Structured Diffusion with White-Box Transformers
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.
Derives CRATE, a white-box transformer-like masked autoencoder whose encoder and decoder layers are interpreted as structured denoising-diffusion steps (compression ≈ denoising), using ~30% of MAE parameters.
Problem
Existing white-box (interpretable, unrolled-optimization) transformer networks work only for supervised learning; building matching decoders for unsupervised autoencoding requires inverting their layers.
Approach
Shows that denoising and lossy compression coincide under a low-dimensional Gaussian mixture signal model, links discretized probability-flow ODEs to unrolled compression networks, and derives an encoder (MSSA + ISTA) and an analytically inverted decoder (anti-compression + dictionary) for masked autoencoding.
Key results
- CRATE-Base uses about 30% of the parameters of ViT-MAE-Base (44.6M vs 143.8M) with the same configuration.
- Fine-tuned CRATE-B reaches 96.8 on CIFAR10 and 80.3 on CIFAR100, versus 98.5 and 87.0 for ViT-MAE-B.
- Linear probing CRATE-B gives 80.9 on CIFAR10 (vs 87.9 for ViT-MAE-B) and 46.2 on Oxford-IIIT Pets (vs 80.1).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · CIFAR10 fine-tuning | top-1 accuracy: 98.5 -> 96.8 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · CIFAR100 fine-tuning | top-1 accuracy: 87.0 -> 80.3 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · Oxford Flowers-102 fine-tuning | top-1 accuracy: 92.5 -> 78.5 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · Oxford-IIIT-Pets fine-tuning | top-1 accuracy: 90.3 -> 76.7 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · CIFAR10 linear probing | top-1 accuracy: 87.9 -> 80.9 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · CIFAR100 linear probing | top-1 accuracy: 68.0 -> 60.1 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · Oxford Flowers-102 linear probing | top-1 accuracy: 66.4 -> 61.8 | Hurt |
| Use CRATE instead of ViT-MAE | ImageNet-1K pre-trained representations · Oxford-IIIT-Pets linear probing | top-1 accuracy: 80.1 -> 46.2 | Hurt |
| Use CRATE-Base instead of ViT-MAE-Base | ImageNet-1K masked autoencoding · ImageNet-1K | parameter count: 143.8M -> 44.6M | Helped |
Surprising findings
- Linear probing CRATE-B on CIFAR10 (80.9) slightly exceeds ViT-MAE-S (79.9) despite being a model of similar size, though it trails ViT-MAE-B.
Limitations
- Performance trails ViT-MAE on fine-tuning and linear probing, especially on Oxford Flowers and Pets.
- Results are not optimally engineered and use far less compute than state-of-the-art MAEs.
- Theoretical connection between compression and denoising holds only under simplified uncorrelated Gaussian mixture signal models.
Experimental setup
Datasets: ImageNet-1K CIFAR10 CIFAR100 Oxford Flowers-102 Oxford-IIIT Pets
Benchmarks: CIFAR10 CIFAR100 Oxford Flowers-102 Oxford-IIIT Pets
Metrics: top-1 accuracy reconstruction loss
Baselines: ViT-MAE-Small ViT-MAE-Base
Builds on: CRATE masked autoencoder probability flow ODE ISTA sparse rate reduction
Closest papers by meaning
- Scaling White-Box Transformers for Vision (NeurIPS 2024 · similarity 0.89)
- Masked Autoencoders Are Effective Tokenizers for Diffusion Models (ICML 2025 · similarity 0.86)
- Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking (NeurIPS 2025 · similarity 0.85)
- Factorized Diffusion Autoencoder for Unsupervised Disentangled Representation Learning (AAAI 2024 · similarity 0.85)
- Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity (ICLR 2025 · similarity 0.85)
- Improved Baselines with Representation Autoencoders (NeurIPS 2026 · similarity 0.85)
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.