[Paper Review] Emergence of Segmentation with Minimalistic White-Box Transformers

Yaodong Yu, Tianzhe Chu+5 more|arXiv (Cornell University)|Aug 30, 2023

Advanced Memory and Neural ComputingEngineering3 citations

TL;DR

This paper demonstrates that segmentation properties emerge in a minimalistic, white-box transformer architecture called CRATE through standard supervised training—without complex self-supervised objectives. The model's mathematically interpretable design, based on sparse coding and unrolled optimization, enables emergent object-level and part-level segmentation in self-attention maps, establishing a new path for interpretable foundation models.

ABSTRACT

Transformer-like models for vision tasks have recently proven effective for a wide range of downstream applications such as segmentation and detection. Previous works have shown that segmentation properties emerge in vision transformers (ViTs) trained using self-supervised methods such as DINO, but not in those trained on supervised classification tasks. In this study, we probe whether segmentation emerges in transformer-based models solely as a result of intricate self-supervised learning mechanisms, or if the same emergence can be achieved under much broader conditions through proper design of the model architecture. Through extensive experimental results, we demonstrate that when employing a white-box transformer-like architecture known as CRATE, whose design explicitly models and pursues low-dimensional structures in the data distribution, segmentation properties, at both the whole and parts levels, already emerge with a minimalistic supervised training recipe. Layer-wise finer-grained analysis reveals that the emergent properties strongly corroborate the designed mathematical functions of the white-box network. Our results suggest a path to design white-box foundation models that are simultaneously highly performant and mathematically fully interpretable. Code is at \url{https://github.com/Ma-Lab-Berkeley/CRATE}.

Motivation & Objective

• To investigate whether segmentation properties in vision transformers arise solely from complex self-supervised training or can emerge from architectural design alone.
• To evaluate whether a white-box transformer architecture can achieve segmentation emergence without self-supervised pretraining or segmentation-specific supervision.
• To demonstrate that minimalistic supervised training on a mathematically interpretable architecture can yield segmentation-aware representations.
• To establish a framework for designing foundation models that are both highly performant and fully interpretable through explicit modeling of data structure.

Proposed method

• The CRATE architecture is derived via unrolling optimization on a sparse rate reduction objective, explicitly modeling low-dimensional structures in the data distribution.
• Each layer of CRATE performs structured compression using a learnable dictionary and sparsification, with components designed to be mathematically interpretable.
• The model uses a multi-head self-attention mechanism with a [CLS] token and patch embeddings, where attention maps are de-rasterized to visualize segmentation behavior.
• The architecture incorporates both a multi-scale self-attention (MSSA) block and an ISTA-based block for sparse coding, with a soft-thresholding activation in one variant.
• Training is performed using standard supervised classification with cross-entropy loss on ImageNet, without any segmentation annotations or self-supervised components.
• Layer-wise analysis and ablation studies are conducted to correlate emergent attention patterns with the model’s designed mathematical functions.

Experimental results

Research questions

• RQ1 Can segmentation properties emerge in a vision transformer without self-supervised pretraining, solely through architectural design and supervised fine-tuning?
• RQ2 To what extent does the mathematical interpretability of a white-box transformer architecture contribute to the emergence of semantic segmentation in attention maps?
• RQ3 How do the attention heads in a white-box transformer like CRATE correspond to distinct semantic parts of objects?
• RQ4 Does the emergent segmentation behavior in CRATE persist under noisy input conditions, indicating robustness to input perturbations?
• RQ5 What architectural components are essential for the emergence of segmentation, and how do they compare to standard ViT components?

Key findings

• Self-attention maps in CRATE, trained only with supervised classification, exhibit clear object-level and part-level segmentation, matching qualitative behavior seen in DINO.
• Layer-wise analysis reveals that individual attention heads in CRATE specialize in different semantic parts of objects, such as edges, textures, or object components.
• The model maintains segmentation-aware attention even when up to 50% of input pixels are corrupted with Gaussian noise, demonstrating robustness.
• Ablation studies show that the MSSA block is essential for segmentation emergence, while replacing ReLU with soft thresholding in the ISTA block preserves segmentation behavior.
• CRATE outperforms standard ViT in attention map quality on COCO val2017, with segmentation performance improving steadily over training epochs.
• The penultimate-layer features of CRATE correlate strongly with saliency and object boundaries, confirming that segmentation properties emerge from the model’s design rather than training recipe.