# Masked Completion via Structured Diffusion with White-Box Transformers

arXiv: https://arxiv.org/abs/2404.02446 · HTML: https://mlsanity.com/paper/2404.02446/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 61 references that are in the map.
- First posted on arXiv 3 Apr 2024.

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Quantifying Attention Flow in Transformers](https://mlsanity.com/paper/2005.00928/) (2020)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Deep ViT Features as Dense Visual Descriptors](https://mlsanity.com/paper/2112.05814/) (2021)
- [Test-Time Training with Masked Autoencoders](https://mlsanity.com/paper/2209.07522/) (2022)
- [Score Approximation, Estimation and Distribution Recovery of Diffusion Models on Low-Dimensional Data](https://mlsanity.com/paper/2302.07194/) (2023)
- [Diffusion Models are Minimax Optimal Distribution Estimators](https://mlsanity.com/paper/2303.01861/) (2023)
- [Learning Diverse and Discriminative Representations via the Principle of Maximal Coding Rate Reduction](https://mlsanity.com/paper/2006.08558/) (2020)
- [Denoising Diffusion Autoencoders are Unified Self-supervised Learners](https://mlsanity.com/paper/2303.09769/) (2023)
- [Deconstructing Denoising Diffusion Models for Self-Supervised Learning](https://mlsanity.com/paper/2401.14404/) (2024)
- [ReduNet: A White-box Deep Network from the Principle of Maximizing Rate Reduction](https://mlsanity.com/paper/2105.10446/) (2021)
- [White-Box Transformers via Sparse Rate Reduction](https://mlsanity.com/paper/2306.01129/) (2023)
- [On the Principles of Parsimony and Self-Consistency for the Emergence of Intelligence](https://mlsanity.com/paper/2207.04630/) (2022)
- [Understanding DDPM Latent Codes Through Optimal Transport](https://mlsanity.com/paper/2202.07477/) (2022)

## Built on by (most cited first, 9 of 9)

- [U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation](https://mlsanity.com/paper/2406.02918/) (2024)
- [Simplifying DINO via Coding Rate Regularization](https://mlsanity.com/paper/2502.10385/) (2025)
- [Scaling White-Box Transformers for Vision](https://mlsanity.com/paper/2405.20299/) (2024)
- [Attention-Only Transformers via Unrolled Subspace Denoising](https://mlsanity.com/paper/2506.03790/) (2025)
- [Path-SAM2: Transfer SAM2 for digital pathology semantic segmentation](https://mlsanity.com/paper/2408.03651/) (2024)
- [Rethinking Decoders for Transformer-based Semantic Segmentation: A Compression Perspective](https://mlsanity.com/paper/2411.03033/) (2024)
- [Geometry of Neural Reinforcement Learning in Continuous State and Action Spaces](https://mlsanity.com/paper/2507.20853/) (2025)
- [Unifying Learning Dynamics and Generalization in Transformers Scaling Law](https://mlsanity.com/paper/2512.22088/) (2025)
- [Latent Zoning Network: A Unified Principle for Generative Modeling, Representation Learning, and Classification](https://mlsanity.com/paper/2509.15591/) (2025)

## Abstract (arXiv)

Modern learning frameworks often train deep neural networks with massive amounts of unlabeled data to learn representations by solving simple pretext tasks, then use the representations as foundations for downstream tasks. These networks are empirically designed; as such, they are usually not interpretable, their representations are not structured, and their designs are potentially redundant. White-box deep networks, in which each layer explicitly identifies and transforms structures in the data, present a promising alternative. However, existing white-box architectures have only been shown to work at scale in supervised settings with labeled data, such as classification. In this work, we provide the first instantiation of the white-box design paradigm that can be applied to large-scale unsupervised representation learning. We do this by exploiting a fundamental connection between diffusion, compression, and (masked) completion, deriving a deep transformer-like masked autoencoder architecture, called CRATE-MAE, in which the role of each layer is mathematically fully interpretable: they transform the data distribution to and from a structured representation. Extensive empirical evaluations confirm our analytical insights. CRATE-MAE demonstrates highly promising performance on large-scale imagery datasets while using only ~30% of the parameters compared to the standard masked autoencoder with the same model configuration. The representations learned by CRATE-MAE have explicit structure and also contain semantic meaning. Code is available at https://github.com/Ma-Lab-Berkeley/CRATE .