[论文解读] White-Box Transformers via Sparse Rate Reduction

Yaodong Yu, Sam Buchanan+6 more|arXiv (Cornell University)|Jun 1, 2023

Advanced Neural Network Applications被引用 23

一句话总结

该论文将 transformer-like 层统一为展开的步骤，优化稀疏率降低目标，产生一个完全可解释的白盒架构（CRATE），它在压缩和稀疏化标记表示方面具有竞争力，并在大规模视觉数据上与工程化 Transformer 相媲美。

ABSTRACT

In this paper, we contend that the objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a mixture of low-dimensional Gaussian distributions supported on incoherent subspaces. The quality of the final representation can be measured by a unified objective function called sparse rate reduction. From this perspective, popular deep networks such as transformers can be naturally viewed as realizing iterative schemes to optimize this objective incrementally. Particularly, we show that the standard transformer block can be derived from alternating optimization on complementary parts of this objective: the multi-head self-attention operator can be viewed as a gradient descent step to compress the token sets by minimizing their lossy coding rate, and the subsequent multi-layer perceptron can be viewed as attempting to sparsify the representation of the tokens. This leads to a family of white-box transformer-like deep network architectures which are mathematically fully interpretable. Despite their simplicity, experiments show that these networks indeed learn to optimize the designed objective: they compress and sparsify representations of large-scale real-world vision datasets such as ImageNet, and achieve performance very close to thoroughly engineered transformers such as ViT. Code is at \url{https://github.com/Ma-Lab-Berkeley/CRATE}.

研究动机与目标

• 将表示学习的动机描述为将数据分布压缩为低维子空间的混合并对表示进行稀疏化。
• 引入一个统一的稀疏率降低目标，将有损编码率与稀疏性相结合，以学习紧凑的标记表示。
• 将 transformer-like 层推导为展开的优化步骤，为注意力和 MLP 块提供数学上的可解释性。
• 提出 CRATE（Coding RAte TransformEr）作为一个白盒架构，具有逐层的分布和字典的概率模型，参数从数据中学习。
• 证明 CRATE 能在大规模视觉数据上学习压缩和稀疏表示，并接近 ViT 风格的性能。

提出的方法

• 定义一个统一目标：最大化稀疏率降低，将率降低与 ℓ0 稀疏惩罚相结合，设 Z = f(X)。
• 将标记分布建模为低维子空间的混合，在每一层具有学习到的基 U[K]。
• 将自注意力风格的更新推导为梯度步骤，以最小化对子空间混合的编码率（MSSA）。
• 通过对学习到的字典 D 进行 ISTA 风格的更新来表示稀疏化，从而在 Z 中促进稀疏性。
• 通过堆叠执行基于 MSSA 的压缩、随后执行基于 ISTA 的稀疏化、并在端到端学习的层特定 U[K] 与 D 的方式构建 CRATE。
• 提供可复现的代码链接：https://github.com/Ma-Lab-Berkeley/CRATE

实验结果

研究问题

• RQ1 带有稀疏性的率降低目标是否能够产生紧凑、可解释的标记表示？
• RQ2 从展开优化推导出的白盒 Transformer 层是否在大规模视觉任务中实现具有竞争力的性能？
• RQ3 自注意力和 MLP 块是否可以在一个统一框架内被重新解释为去噪/压缩和稀疏编码步骤？
• RQ4 逐层学习的子空间基与字典对表示质量和可迁移性有何影响？
• RQ5 在训练过程中，所提出的 MSSA 和 ISTA 块与预期的优化目标的一致性有多高？

主要发现

• CRATE 层实现递增的优化，将标记分布压缩至子空间混合并稀疏化表示。
• MSSA 组件对应于一种梯度步长样的操作，类似自注意力，但来自针对子空间的降噪的率降低推导。
• ISTA 基于的稀疏化层在学习到的字典上促进稀疏性，使得对基于率的多样性的近似可处理。
• 在 ImageNet-1K 上的实验表明 CRATE 学会压缩和稀疏化表示，性能接近像 ViT 这样的工程化 Transformer。
• 按层分析表明压缩和稀疏化在各层都在改进，支持了旨在目标驱动的设计。