Emergence of Segmentation with Minimalistic White-Box Transformers
Yaodong Yu Tianzhe Chu Shengbang Tong Ziyang Wu Druv Pai Sam Buchanan Yi Ma
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}.