新智元报道
编辑：LRS
CRATE-α是一种新型 Transformer 架构变体，通过设计改进提升了模型的可扩展性、性能和可解释性，CRATE-α-Base 在 ImageNet 分类任务上的性能显著超过了之前最好的 CRATE-B 模型，其性能会随着模型和数据集规模扩大而继续提升。
在过去的几年里，Transformer 架构在自然语言处理（NLP）、图像处理和视觉计算领域的深度表征学习中取得了显著的成就，几乎成为了 AI 领域的主导技术。
然而，虽然 Transformer 架构及其众多变体在实践中取得了巨大成功，但其设计大多是基于经验的，并没有严格的数学解释，也在一定程度上限制了研究人员的思路，无法开发出更高效、更具可解释性的 Transformer 新变体。
为了填补这一空白，马毅教授团队曾发布过白盒 Transformer 模型 CRATE，其架构的每一层都是通过数学推导得到的，可以完全解释为展开的梯度下降迭代；此外，CRATE 学习到的模型和特征在语义上也比传统的 Transformer 模型具有更好的可解释性，例如，即使模型仅在分类任务上进行训练，可视化图像的特征也能自然地形成该图像的零样本分割。
然而，到目前为止，CRATE 的应用规模仍然相对有限，CRATE-Large 只包含 77.6M 参数，与标准 Vision Transformer（ViTs）的 22B 参数量形成了鲜明对比。
最近，加利福尼亚大学圣克鲁斯分校和伯克利分校的研究团队联合提出了 CRATE-α，首次探索了不同规模的 CRATE 用于视觉任务（从 Tiny 到 Huge）时的模型性能，研究人员在 CRATE 架构设计中对稀疏编码块进行了策略性但最小化的（strategic yet minimal）修改，并设计了一种轻量级的训练方法，以提高 CRATE 的可扩展性。
论文链接：https://arxiv.org/pdf/2405.20299
项目链接：https://rayjryang.github.io/CRATE-alpha/
具体来说，CRATE 中的 ISTA 模块是限制进一步扩展的因素，为了克服这一限制，CRATE-α主要做了三个修改：
1. 大幅扩展了通道，对稀疏编码块进行过参数化（overparameterized），使用过完备字典（overcomplete dictionary）对 token 表征进行稀疏化。
2. 解耦了关联矩阵，在稀疏编码块的最后一部中引入一个解耦字典（decoupled dictionary）
3. 添加了残差连接。
实验结果证明，CRATE-α能够随着模型尺寸和训练数据集的增大而扩展，性能可以持续提升。
例如，CRATE-α-B 在 ImageNet 分类任务上的性能显著超过了之前最好的 CRATE-B 模型，准确率提高了 3.7%，达到了 83.2%；进一步对模型进行扩展时，CRATE-α-L 在 ImageNet 分类任务上达到了 85.1% 的准确率。
值得注意的是，模型性能的提升是在保持甚至增强了 CRATE 模型可解释性的同时实现的，因为更大尺寸的 CRATE-α模型学到的 token 表征能够生成更高质量的无监督图像分割。
实验结果
从基础尺寸（base）到大尺寸（large）
ImageNet-21K 是一个广泛用于图像识别和分类任务的大型数据集，文中用于训练的数据集版本包含 19,000 个类别和大约 1300 万张图片，由于数据丢失，比标准数据集（包含 21,000 个类别和大约 1400 万张图片）的数据量要少一点。
在预训练时，从数据集中随机选取1% 作为验证集。
预训练完成后，在 ImageNet-1K 数据集上对模型进行微调，其中 ImageNet-1K 是一个更小的子集，包含 1000 个类别，通常用于模型的最终评估。在微调阶段，模型会针对这 1000 个类别进行更精细的训练，以提高其在特定任务上的性能。
最后，在 ImageNet-1K 的验证集上评估模型的性能。
研究人员对比了在 32、16 和 8 像素块大小下的 CRATE-α-B 和 CRATE-α-L，从实验结果中可以看到，CRATE-α-L 在所有像素块大小上都取得了显著的改进，但从 CRATE-B 增加到 CRATE-L 只能带来 0.5% 的性能提升，表明了收益递减的情况，证明了 CRATE-α模型的可扩展性显著优于普通 CRATE
同时，预训练阶段的训练损失显示，随着模型容量的增加，训练损失的趋势可预测地得到改善。
从大（large）到巨大（huge）
多模态数据集 DataComp1B 包含 14 亿图文对，可以提供足够的数据来训练和扩展模型。
研究人员采用对比学习的方法来训练 CRATE-α，不仅能够利用上庞大的图文对数据集，还能在模型尺寸从大到巨大的提升过程中，观察到显著的性能提升。
然而，直接训练一个类似 CLIP 的模型需要巨大的计算资源，研究人员采用了优化后的 CLIPA 协议，可以在减少计算资源消耗的同时，可以保持与 CLIP 相当的性能。
最后，为了评估 CRATE-α模型的性能，研究人员采用了零样本学习的方法，在 ImageNet-1K 数据集上测试模型的准确率，该方法可以有效地评估模型在面对未见过类别数据时的泛化能力，提供了一个衡量模型可扩展性和实用性的重要指标。
从实验结果中可以看到，
1. 模型尺寸的影响：CRATE-α-CLIPA-L/14 在预训练和微调阶段的 ImageNet-1K 零样本准确率上，分别比 CRATE-α-CLIPA-B/16 高出 11.3% 和 9.0%，表明学习到的表征质量可能受到模型尺寸的限制，即增加模型尺寸可以利用上更多数据。
2. 扩展模型尺寸的益处：当继续增加模型尺寸时，可以观察到 CRATE-α-CLIP-H/14 从更大的训练数据集中继续获益，在预训练和微调阶段的 ImageNet-1K 零样本准确率上，分别比 CRATE-α-CLIP-L/14 高出 3.1% 和 2.5%，证明了 CRATE-α模型的强大可扩展性。
3. 性能上限的探索：为了探索性能的上限，研究人员从头开始训练了一个标准的 ViT-CLIPA-H/14，并观察到了性能的提升。
节省计算资源的扩展策略
在追求模型扩展的效率和计算资源的优化方面，研究人员发现，通过调整预训练阶段的图像 token 序列长度，可以在极大减少计算资源消耗的同时，保持模型性能。
具体来说，研究人员尝试了一种新的方法：在预训练时使用较长序列长度的 CRATE-α-L/32，在微调时切换到较短序列长度的 CRATE-α-L/14 或 CRATE-α-L/8，不仅大幅度降低了预训练阶段的计算成本，而且在微调后，模型在 ImageNet-1K 数据集上的准确率仍然非常接近全尺寸模型的性能。
例如，使用 CRATE-α-L/32 进行预训练，然后微调到 CRATE-α-L/14，可以节省约 70% 的计算资源，而准确率只是略有下降；更进一步，当从 CRATE-α-L/32 预训练后微调到 CRATE-α-L/8 时，仅使用了原模型所需训练时间的 10%，准确率依然达到了 84.2%，与全尺寸模型的 85.1% 相差无几。
上述结果表明，通过精心设计预训练和微调阶段的策略，可以在资源有限的情况下，有效地扩展 CRATE-α模型。
CRATE-α的语义可解释性得到提升
除了可扩展性，文中还研究了不同模型大小的 CRATE-α的可解释性，使用 MaskCut 来验证和评估模型捕获的丰富语义信息，包括定性和定量结果。
为 CRATE-α、CRATE 和 ViT 在 COCO val2017 上提供了分割可视化后，可以发现，CRATE-α模型保持甚至提高了 CRATE 的（语义）可解释性优势。
在 COCO val2017 上的定量评估结果显示，当为 CRATE-α扩展模型大小时，大型模型在目标检测和分割方面比 base 模型有所提高。
参考资料：
