[论文解读] Emergence of Segmentation with Minimalistic White-Box Transformers

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

Advanced Memory and Neural ComputingEngineering被引用 3

一句话总结

该论文表明，通过标准的监督训练，一种称为 CRATE 的极简白盒 Transformer 架构中会涌现出分割特性——无需复杂的自监督目标。该模型基于稀疏编码和展开优化的数学可解释性设计，使自注意力图中自然出现对象级和部件级的分割，为可解释的基础模型开辟了新路径。

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}.

研究动机与目标

• 探究视觉 Transformer 中的分割特性是否仅源于复杂的自监督训练，还是可仅通过架构设计本身涌现。
• 评估白盒 Transformer 架构是否可在无需自监督预训练或分割特定监督的情况下实现分割的涌现。
• 证明在数学可解释性架构上进行极简监督训练，可生成具备分割感知能力的表征。
• 建立一种设计基础模型的框架，通过显式建模数据结构，实现高性能与完全可解释性的统一。

提出的方法

• CRATE 架构通过在稀疏率降低目标上展开优化推导而来，显式建模数据分布中的低维结构。
• CRATE 的每一层通过可学习字典和稀疏化实现结构化压缩，其组件设计具有数学可解释性。
• 模型采用多头自注意力机制，包含 [CLS] 标记和图像块嵌入，其中注意力图被去光栅化以可视化分割行为。
• 架构同时包含多尺度自注意力（MSSA）模块和基于 ISTA 的稀疏编码模块，其中一个变体中使用了软阈值激活。
• 训练采用标准监督分类，使用 ImageNet 上的交叉熵损失，未引入任何分割标注或自监督组件。
• 通过逐层分析和消融研究，将涌现的注意力模式与模型设计的数学函数相关联。

实验结果

研究问题

• RQ1 在不使用自监督预训练的情况下，仅通过架构设计和监督微调，视觉 Transformer 中的分割特性是否可以涌现？
• RQ2 白盒 Transformer 架构的数学可解释性在注意力图中促进语义分割涌现的程度如何？
• RQ3 在类似 CRATE 的白盒 Transformer 中，注意力头如何对应于对象的不同语义部件？
• RQ4 CRATE 中的涌现分割行为在输入存在噪声时是否依然保持，表明其对输入扰动的鲁棒性？
• RQ5 哪些架构组件对分割的涌现至关重要？它们与标准 ViT 组件相比有何差异？

主要发现

• 仅通过监督分类训练的 CRATE 自注意力图中，清晰展现出对象级和部件级的分割，其定性行为与 DINO 中观察到的一致。
• 逐层分析表明，CRATE 中的各个注意力头会专门关注对象的不同语义部件，如边缘、纹理或对象组件。
• 即使输入像素高达 50% 被高斯噪声破坏，模型仍能保持具备分割感知的注意力，表现出强鲁棒性。
• 消融研究显示，MSSA 模块对分割的涌现至关重要，而将 ISTA 模块中的 ReLU 替换为软阈值激活可保持分割行为。
• 在 COCO val2017 上，CRATE 的注意力图质量优于标准 ViT，且分割性能随训练轮次稳步提升。
• CRATE 的倒数第二层特征与显著性图和物体边界高度相关，证实分割特性源于模型架构本身，而非训练策略。