[论文解读] White-Box Transformers via Sparse Rate Reduction: Compression Is All There Is?

Yaodong Yu, Sam Buchanan+8 more|arXiv (Cornell University)|Nov 22, 2023

Image and Signal Denoising Methods被引用 4

一句话总结

本文提出 CRATE，一种源自稀疏率压缩的白盒 Transformer 架构——在表征中同时优化压缩与稀疏性。通过将多头自注意力解释为对编码率的梯度下降，将 MLP 解释为对稀疏性的近端梯度下降，CRATE 在视觉与语言任务上的性能与当前最先进模型相当，表明压缩是有效表征学习的核心。

ABSTRACT

In this paper, we contend that a natural objective of representation learning is to compress and transform the distribution of the data, say sets of tokens, towards a low-dimensional Gaussian mixture supported on incoherent subspaces. The goodness of such a representation can be evaluated by a principled measure, called sparse rate reduction, that simultaneously maximizes the intrinsic information gain and extrinsic sparsity of the learned representation. From this perspective, popular deep network architectures, including transformers, can be viewed as realizing iterative schemes to optimize this measure. Particularly, we derive a transformer block from alternating optimization on parts of this objective: the multi-head self-attention operator compresses the representation by implementing an approximate gradient descent step on the coding rate of the features, and the subsequent multi-layer perceptron sparsifies the features. This leads to a family of white-box transformer-like deep network architectures, named CRATE, which are mathematically fully interpretable. We show, by way of a novel connection between denoising and compression, that the inverse to the aforementioned compressive encoding can be realized by the same class of CRATE architectures. Thus, the so-derived white-box architectures are universal to both encoders and decoders. Experiments show that these networks, despite their simplicity, indeed learn to compress and sparsify representations of large-scale real-world image and text datasets, and achieve performance very close to highly engineered transformer-based models: ViT, MAE, DINO, BERT, and GPT2. We believe the proposed computational framework demonstrates great potential in bridging the gap between theory and practice of deep learning, from a unified perspective of data compression. Code is available at: https://ma-lab-berkeley.github.io/CRATE .

研究动机与目标

• 建立一个基于数据压缩与稀疏性的原理性、理论驱动的表征学习框架。
• 通过稀疏率压缩的统一目标视角，解释 Transformer 成功的原因：即迭代优化过程。
• 从第一性原理出发，推导出一个完全可解释、数学基础坚实的深度网络架构（CRATE）。
• 通过可逆的、结构化的去噪与扩散，统一编码器与解码器功能于单一架构中。
• 通过实证验证 CRATE 在视觉与语言任务中学习到有意义的、语义结构化的表征。

提出的方法

• 将表征学习形式化为优化稀疏率压缩——一种结合内在信息增益与外在稀疏性的度量。
• 将多头自注意力推导为输入 token 编码率上的近似梯度下降步。
• 将多层感知机建模为特征稀疏编码的近端梯度下降步。
• 通过交替堆叠自注意力与稀疏编码模块构建 CRATE，确保数学可解释性。
• 通过使用相同 CRATE 模块进行结构化去噪-扩散，反转压缩过程，实现可逆解码。
• 在 PyTorch 中实现 CRATE，采用模块化组件：MSSA（多头自注意力）与 ISTA（用于稀疏性的迭代软阈值化）。

实验结果

研究问题

• RQ1 Transformer 的行为能否通过根植于信息论与稀疏性的统一优化目标来解释？
• RQ2 在多大程度上，可通过直接从原理性压缩目标推导，设计出完全白盒、可解释的深度网络？
• RQ3 基于稀疏率压缩推导出的网络是否能达到与高度工程化的 Transformer 模型相媲美性能？
• RQ4 同一架构能否通过可逆的、结构化的去噪实现编码与解码的双重功能？
• RQ5 CRATE 中的注意力图与学习到的特征是否表现出涌现的语义结构，如基于对象的分割？

主要发现

• CRATE 在图像与文本基准测试中，性能与 ViT、MAE、DINO、BERT 和 GPT2 等最先进模型相比仅低 1-2%。
• CRATE 层中的注意力图表现出涌现的语义分割，注意力高度集中于物体边界与显著特征处。
• 分层分析表明，CRATE 学习到分层的、结构化的表征，早期层捕捉低级模式，深层则捕捉高级语义。
• CRATE 编码器的逆过程可通过结构化去噪-扩散实现，证明同一架构可同时支持编码与解码。
• 消融研究显示模型性能稳定，证实了稀疏率压缩目标的鲁棒性。
• 学习到的表征简洁且结构化，特征集中于非相干子空间，与理论压缩目标一致。