基于Transformer的视觉基础模型在各种下游任务，如分割和检测中都展现出了非常强大的性能，并且DINO等模型经过自监督训练后已经涌现出了语义的分割属性。
马毅教授领导的研究团队开发了CRATE模型，推动了神经网络可解释研究！完全数学可解释的白盒Transformer，性能不输ViT
过去十多年，AI的飞速发展主要是工程实践上的进步，AI理论并没有起到指导算法开发的作用，经验设计的神经网络依然是一个黑盒。
而随着ChatGPT的爆火，AI的能力也被不断夸大、炒作，甚至到了威胁、绑架社会的地步，让Transformer架构设计变透明已刻不容缓！
最近，马毅教授团队发布了最新研究成果，设计了一个完全可用数学解释的白盒Transformer模型CRATE，并在真实世界数据集ImageNet-1K上取得了接近ViT的性能。
代码链接：https://github.com/Ma-Lab-Berkeley/CRATE
论文链接：https://arxiv.org/abs/2306.01129
在这篇论文中，研究人员认为，表示学习的目标是压缩和转换数据（例如token集合）的分布，以支持在不相干子空间（incoherent subspace）上的低维高斯分布混合，最终表征的质量可以通过稀疏率降低（sparse rate reduction）的统一目标函数来度量。
从这个角度来看，流行的深度网络模型，如Transformer等可以很自然地被认为是实现迭代方案（realizing iterative schemes）以逐步优化该目标。
特别是，研究结果表明标准Transformer块可以从对该目标的互补部分的交替优化中派生出：多头自注意力运算符可以被视为通过最小化有损编码率来压缩token集合的梯度下降步骤，而随后的多层感知器可以被视为尝试稀疏化token的表示。
这一发现也促进设计了一系列在数学上完全可解释的白盒Transformer类深度网络架构，尽管设计上很简单，但实验结果表明，这些网络确实学会了优化设计目标：压缩和稀疏化了大规模真实世界视觉数据集（如ImageNet）的表示，并实现了接近高度工程化Transformer模型（ViT）的性能。
图灵奖得主Yann LeCun对马毅教授的工作也表示赞同，认为Transformer使用LISTA（Learned Iterative Shrinkage and Thresholding Algorithm）类似的方法增量地优化稀疏压缩。
白盒CRATE的涌现属性
DINO的分割涌现能力
智能系统中的表征学习旨在将世界的高维、多模态感官数据（图像、语言、语音）转换为更紧凑的形式，同时保留其基本的低维结构，实现高效的识别（比如分类）、分组（比如分割）和追踪。
深度学习模型的训练通常采用数据驱动的方式，输入大规模数据，以自监督的方式进行学习。
在视觉基础模型中，自监督Transformer架构的DINO模型展现出了令人惊讶的涌现能力，即使没有经过有监督分割训练，ViTs也能识别出显式的语义分割信息。
后续有工作研究了如何在DINO模型中利用这种分割信息，并在下游任务中，如分割、检测等实现了最先进的性能，也有工作证明了用DINO训练的ViTs中的倒数第二层特征与视觉输入中的显著性信息强烈相关，如区分前景、背景和物体边界，从而提升图像分割和其他任务的性能。
为了让分割属性涌现，DINO需要在训练期间将自监督学习、知识蒸馏和权重平均巧妙地结合起来。
目前还不清楚DINO中引入的每个组件是否对于分割遮罩的涌现来说必不可缺，尽管DINO也采用ViT架构作为其主干，但在分类任务上训练的普通有监督ViT模型中，并没有观察到分割涌现行为。
让Transformer变白盒
这篇论文的主要目的在于用一个更统一的框架以设计类似Transformer的网络结构，从而实现数学上的可解释性和良好的实际性能。
为此，研究人员提出学习一个增量映射（incremental mappings）序列，以获得输入数据（token集合）的最小压缩和最稀疏的表征，优化一个统一的目标函数，即稀疏率降低。
这个框架统一了「Transformer模型和自注意力」、「扩散模型和降噪」、「结构化查找和率降低」（Structure-seeking models and rate reduction）三种看似不同的方法，并表明类似Transformer的深层网络层可以自然地从展开迭代优化（unrolling iterative optimization）方案中导出， 以增量地优化稀疏率降低目标。
映射的目标
Self-Attention via Denoising Tokens Towards Multiple Subspaces
研究人员使用一个理想化的token分布模型表明，如果朝着低维子空间系列迭代去噪，相关的评分函数就会呈现出类似于Transformer中的自注意力操作符的显式形式。
Self-Attention via Compressing Token Sets through Optimizing Rate Reduction
研究人员将多头自注意力层推导为一个展开的梯度下降步，以最小化速率降低的有损编码率部分，从而展现了将自注意力层解释为压缩token表征的另一种解释方法。
MLP via Iterative Shrinkage-Thresholding Algorithms (ISTA) for Sparse Coding
研究人员展示了在Transformer块中紧随多头自注意力层后面的多层感知机可以被解释为（并且可以被替换为）一个层，该层通过构建token表征稀疏编码来逐步优化稀疏率降低目标剩余部分。
CRATE
结合上述理解，研究人员创建了一个全新的的白盒Transformer架构CRATE（Coding RAte reduction TransformEr），学习目标函数、深度学习架构和最终学习到的表征都完全可以用数学解释，其中每一层执行交替最小化算法（alternating minimization algorithm）的一个步骤，以优化稀疏率降低目标。
可以注意到，CRATE在构建的每个阶段都选择了尽可能最简单的构建方式，只要新构建的部分保持相同的概念角色，就可以直接替换，并获得一个新的白盒架构。
基于DINO的成功案例，研究人员想要探究，复杂的自监督学习管道对于获得类似Transformer的视觉模型中的涌现属性是否是必要的。
研究人员认为，在Transformer模型中促进分割属性的一种有前途的方法是，在考虑输入数据结构的情况下设计Transformer模型架构，也代表了表征学习经典方法与现代数据驱动的深度学习框架的结合。
与目前主流的Transformer模型对比，这种设计方法也可以叫做白盒Transformer模型。
基于马毅教授组之前的工作，研究人员对白盒架构的CRATE模型进行了广泛的实验，证明了CRATE的白盒设计是自注意力图中分割属性涌现的原因。定性评估
研究人员采用基于[CLS] token的注意力图方法对模型进行解释和可视化，结果发现CRATE中的query-key-value矩阵都是相同的。
定性评估
研究人员采用基于[CLS] token的注意力图方法对模型进行解释和可视化，结果发现CRATE中的query-key-value矩阵都是相同的。
可以观察到CRATE模型的自注意力图（self-attention map）可以对应到输入图像的语义上，模型的内部网络对每个图像都进行了清晰的语义分割，实现了类似DINO模型的效果。相比之下，在有监督分类任务上训练的普通ViT却并没有表现出类似的分割属性。
遵循之前关于可视化图像学习的逐块深度特征的工作，研究人员对CRATE和ViT模型的深度token表征进行主成分分析（PCA）研究。
可以发现，CRATE可以在没有分割监督训练的情况下，依然可以捕捉到图像中物体的边界。
并且，主成分（principal compoenents）也说明了token和物体中相似部分的特征对齐，例如红色通道对应马腿。
而有监督ViT模型的PCA可视化结构化程度相当低。
定量评估
研究人员使用现有的分割和对象检测技术对CRATE涌现的分割属性进行评估。
从自注意力图可以看到，CRATE用清晰的边界显式地捕获了对象级的语义，为了定量测量分割的质量，研究人员利用自注意力图生成分割遮罩（segmentation mask），对比其与真实遮罩之间的标准mIoU（平均交并比）。
从实验结果中可以看到，CRATE在视觉和mIOU评分上都显著优于ViT，表明CRATE的内部表征对于分割遮罩任务生成来说要更有效。
对象检测和细粒度分割
为了进一步验证和评估CRATE捕获的丰富语义信息，研究人员采用了一种高效的对象检测和分割方法MaskCut，无需人工标注即可获得自动化评估模型，可以基于CRATE学到的token表征从图像中提取更细粒度的分割。
在COCO val2017上的分割结果中可以看到，有CRATE的内部表征在检测和分割指标上都要好于有监督ViT，有监督ViT特征的MaskCut在某些情况下甚至完全不能产生分割遮罩。
CRATE分割能力的白盒分析
深度在CRATE中的作用
CRATE的每一层设计都遵循相同的概念目的：优化稀疏速率降低，并将token分布转换为紧凑和结构化的形式。
假设CRATE中语义分割能力的涌现类似于「表征Z中属于相似语义类别token的聚类」，预期CRATE的分割性能可以随着深度的增加而提高。
为了测试这一点，研究人员利用MaskCut管道来定量评估跨不同层的内部表征的分割性能；同时应用PCA可视化来理解分割是如何随深度加深而涌现的。
从实验结果中可以观察到，当使用来自更深层的表征时，分割分数提高了，与CRATE的增量优化设计非常一致。
相比之下，即使ViT-B/8的性能在后面的层中略有提高，但其分割分数明显低于CRATE，PCA结果显示，从CRATE深层提取的表征会逐渐更关注前景对象，并且能够捕捉纹理级别的细节。
CRATE的消融实验
CRATE中的注意力块（MSSA）和MLP块（ISTA）都不同于ViT中的注意力块。
为了了解每个组件对CRATE涌现分割属性的影响，研究人员选取了三个CRATE变体：CRATE, CRATE-MHSA, CRATE-MLP，分别表示ViT中的注意块（MHSA）和MLP块。
研究人员在ImageNet-21k数据集上应用相同的预训练设置，然后应用粗分割评估和遮罩分割评估来定量对比不同模型的性能。
实验结果显示，CRATE在所有任务中都明显优于其他模型架构，可以发现，尽管MHSA和MSSA之间的架构差异很小，但只需要简单地用CRATE中的MSSA替换ViT中的MHSA，可以显著改善ViT的粗分割性能（即VOC Seg），证明了白盒设计的有效性。
识别注意头的语义属性
[CLS] token和图像块token之间的自注意力图可以看到清晰的分段掩码，根据直觉，每个注意力头应该都可以捕捉到数据的部分特征。
研究人员首先将图像输入到CRATE模型，然后由人来检查、选择四个似乎具有语义含义的注意力头；然后在其他输入图像上在这些注意力头上进行自注意力图可视化。
可以发现，每个注意力头都捕捉到了物体的不同部分，甚至不同的语义：例如第一列中显示的注意力头可以捕捉到不同动物的腿，最后一列中显示的注意力头捕捉的是耳朵和头部。
自从可形变部件模型（deformable part model）和胶囊网络发布以来，这种将视觉输入解析为部分-整体层次结构的能力一直是识别架构的目标，白盒设计的CRATE模型也具有这种能力。
实验部分
研究人员的实验目标不仅仅是在使用基本设计的情况下与其他精心设计的Transformer竞争，还包括：
1、与通常仅在端到端性能上评估的经验设计的黑盒网络不同，白盒设计的网络可以查看深层架构的内部，并验证学习网络的层是否确实执行其设计目标， 即对目标进行增量优化。
2、尽管CRATE架构很简单，但实验结果应当验证该架构的巨大潜力，即可以在大规模真实世界的数据集和任务上取得与高度工程化Transformer模型相匹配的性能。
模型架构
通过变化token维度、头数和层数，研究人员创建了四个不同规模的CRATE模型，表示为CRATE-Tiny，CRATE-Small，CRATE-Base和CRATE-Large
数据集和优化
文中主要考虑ImageNet-1K作为测试平台，使用Lion优化器来训练具有不同模型规模的CRATE模型。
同时还评估了CRATE的迁移学习性能：在ImageNet-1K上训练的模型作为预训练模型，然后在几个常用的下游数据集（CIFAR10/100、Oxford Flowers、Oxford-IIT-Pets）上对CRATE进行微调。
CRATE的层实现设计目标了吗？
随着层索引的增加，可以看到CRATE-Small模型在大多数情况下的压缩和稀疏化项都得到了提升，最后一层稀疏性度量的增加是由于用于分类的额外线性层。
结果表明，CRATE与原始的设计目标非常契合：一旦学习完毕，基本上通过其层逐渐学习对表示进行压缩和稀疏化。
在其他规模的CRATE模型以及中间模型检查点上测量压缩和稀疏化项后可以发现，实验结果依然非常一致，具有更多层的模型往往能更有效地优化目标，验证了之前对每个层角色的理解。
性能对比
通过测量ImageNet-1K上的最高准确率以及在几个广泛使用的下游数据集上的迁移学习性能来研究所提出的网络的经验性能。
由于设计的架构在注意力块（MSSA）和MLP块（ISTA）中都利用了参数共享，所以CRATE-Base模型（2208万）与ViT-Small（2205万）的参数数量相似。
可以看到，在模型参数数量相似的情况下，文中提出的网络实现了与ViT相似的ImageNet-1K和迁移学习性能，但CRATE的设计更简单，可解释性强。
此外，在相同的训练超参数下，CRATE还可以继续扩展，即通过扩大模型的规模不断提高性能，而在ImageNet-1K上直接扩大ViT的规模并不总是能带来一致的性能改善。
也就是说，CRATE网络尽管简单，但已经可以在大规模的真实世界数据集上学习所需的压缩和稀疏表示，并在各种任务（如分类和迁移学习）上取得与更工程化Transformer网络（如ViT）相当的性能。
参考资料：
https://arxiv.org/abs/2306.01129
