Émergence de la segmentation avec des Transformers minimalistes en boîte blanche

Emergence of Segmentation with Minimalistic White-Box Transformers

August 30, 2023

Auteurs: Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, Yi Ma

cs.AI

Résumé

Les modèles de type Transformer pour les tâches de vision ont récemment démontré leur efficacité pour un large éventail d'applications en aval, telles que la segmentation et la détection. Des travaux antérieurs ont montré que les propriétés de segmentation émergent dans les vision transformers (ViTs) entraînés à l'aide de méthodes auto-supervisées comme DINO, mais pas dans ceux entraînés sur des tâches de classification supervisée. Dans cette étude, nous explorons si la segmentation émerge dans les modèles basés sur Transformer uniquement en raison de mécanismes d'apprentissage auto-supervisé complexes, ou si la même émergence peut être obtenue dans des conditions beaucoup plus larges grâce à une conception appropriée de l'architecture du modèle. À travers des résultats expérimentaux approfondis, nous démontrons qu'en utilisant une architecture de type Transformer transparente appelée CRATE, dont la conception modélise et recherche explicitement des structures de faible dimension dans la distribution des données, les propriétés de segmentation, à la fois au niveau global et des parties, émergent déjà avec une recette d'entraînement supervisé minimaliste. Une analyse plus fine couche par couche révèle que les propriétés émergentes corroborent fortement les fonctions mathématiques conçues du réseau transparent. Nos résultats suggèrent une voie pour concevoir des modèles de base transparents qui sont à la fois très performants et entièrement interprétables mathématiquement. Le code est disponible à l'adresse https://github.com/Ma-Lab-Berkeley/CRATE.

English

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 https://github.com/Ma-Lab-Berkeley/CRATE.