白盒CRATE具有可扩展性，NeurIPS2024

论文名称：Scaling White-Box Transformers for Vision

GitHub主页： https://rayjryang.github.io/CRATE-alpha/

论文地址：https://arxiv.org/pdf/2405.20299

顶会收录：NeurIPS2024

论文简介：

在论文《Scaling White-Box Transformers for Vision》中，作者探讨了白盒Transformer架构CRATE在视觉任务中的可扩展性。CRATE以其固有的数学可解释性，通过学习压缩和稀疏表示，提供了与传统视觉Transformer（ViTs）不同的替代方案。尽管在语言和视觉Transformer的扩展性方面已有广泛研究，但CRATE的可扩展性仍是一个未解的问题。为此，作者提出了CRATE-α，通过对CRATE架构中的稀疏编码模块进行战略性且最小化的修改，以及设计轻量级的训练方案，以提高CRATE的可扩展性。通过大量实验，作者证明了CRATE-α可以有效地随着模型规模和数据集的增大而扩展。例如，CRATE-α-B在ImageNet分类任务中显著超越了之前最好的CRATE-B模型，准确率提高了3.7%，达到了83.2%。同时，进一步扩展的CRATE-α-L在ImageNet分类中取得了85.1%的准确率。值得注意的是，这些模型性能的提升是在保持甚至增强CRATE模型可解释性的同时实现的。作者通过展示训练后的CRATE-α模型在无监督对象分割任务中的高质量表现，证明了这一点。

项目主页为https://rayjryang.github.io/CRATE-alpha/

主要内容概述：

1. 引言：

   • 背景：Transformer架构在自然语言处理、图像处理和视觉计算等领域取得了主导地位。然而，其设计主要基于经验，缺乏严格的数学解释，这在一定程度上阻碍了新型Transformer变体在效率和可解释性方面的发展。白盒Transformer模型CRATE通过对稀疏率减少表示学习目标的优化，提供了一个简化且完全可解释的Transformer模块。具体而言，CRATE的自注意力模块通过对特征进行降噪，实现压缩；MLP模块则被替代为特征的增量稀疏化处理。此前的研究表明，CRATE不仅在数学上具有可解释性，其学习到的特征在语义上也比传统Transformer更具可解释性，例如在仅进行分类训练的情况下，CRATE能够自然地实现图像的零样本分割。

2. CRATE的可扩展性挑战：

   • 问题：尽管CRATE在小规模模型上表现出色，但其在大规模模型和数据集上的可扩展性尚未得到验证。相比之下，标准的ViT模型已成功扩展至数十亿参数的规模。

   • 目标：探索并验证CRATE在不同模型规模（如Tiny、Small、Base、Large、Huge）和大型数据集上的性能表现。

3. CRATE-α的提出：

   • 架构修改：作者识别出CRATE中的ISTA模块是限制其扩展性的关键因素。为此，提出了以下改进：

     • 通道扩展：显著增加通道数量，以提升模型容量。

     • 矩阵解耦：将关联矩阵解耦，以减少计算复杂度。

     • 残差连接：引入残差连接，增强模型的训练稳定性。

   • 训练策略：设计了轻量级的训练方案，包括优化器选择、学习率调度和数据增强策略，以适应大规模训练需求。

   • 

4. 实验验证：

   • 数据集：在ImageNet-1K和ImageNet-21K数据集上进行训练和评估。

   • 模型规模：评估了从Tiny到Huge不同规模的CRATE-α模型。

   • 性能提升：实验结果显示，CRATE-α在各个模型规模上均取得了显著的性能提升。例如，CRATE-α-B在ImageNet-1K上的准确率达到83.2%，相比之前的CRATE-B提高了3.7%；CRATE-α-L的准确率则达到85.1%。

   • 计算效率：在保持或减少计算复杂度的情况下，CRATE-α实现了性能的提升，展示了其高效性。

5. 可解释性分析：

   • 语义分割：作者通过无监督的对象分割任务，评估了CRATE-α模型的可解释性。结果表明，随着模型规模的增大，CRATE-α学习到的特征在分割任务中表现出更高的质量。

   • 特征可视化：对不同层次的特征进行可视化，展示了CRATE-α在捕捉图像语义信息方面的有效性。

6. 结论：

   • 总结：CRATE-α通过对原始CRATE架构的关键性改进，实现了在大规模模型和数据集上的有效扩展，同时保持了模型的数学和语义可解释性。

   • 未来工作：计划进一步优化CRATE-α的训练策略，探索其在更多视觉任务和更大规模数据集上的应用潜力。

参考文献：

• Yang, J., Li, X., Pai, D., Zhou, Y., Ma, Y., Yu, Y., & Xie, C. (2024). Scaling White-Box Transformers for Vision. arXiv preprint arXiv:2405.20299. https://arxiv.org/abs/2405.20299

附注：

• 该论文的项目主页提供了更多信息和资源：https://rayjryang.github.io/CRATE-alpha/