（ICLR 2025）TSSA：基于 token 统计的自注意力机制，即插即用，涨点起飞

主要观点总结

本文介绍了一篇名为TOKEN STATISTICS TRANSFORMER的论文，提出了一种新的注意力机制Token Statistics Self-Attention（TSSA），能够降低注意力复杂度从常规的二次为线性。文章还介绍了一个新的MCR²变分形式，以及基于TSSA模块构建的Transformer架构Token Statistics Transformer（TOST）。实验证明，该架构在图像分类、语言建模、长序列建模等任务上性能优异，同时显著节省计算和内存。文章还给出了TSSA注意力机制的具体实现代码。

关键观点总结

关键观点1: Token Statistics Self-Attention（TSSA）机制

TSSA机制不依赖传统的token间两两相似度计算，而是基于输入token的统计特性（例如二阶矩）来构建注意力操作，将注意力复杂度从常规的二次降低为线性。

关键观点2: MCR²变分形式

基于最大编码率压缩（Maximal Coding Rate Reduction, MCR²）原则，引入新的变分上界方法，使得注意力机制可以通过梯度下降进行高效优化。

关键观点3: Token Statistics Transformer（TOST）架构

TOST架构基于TSSA模块构建，具备理论可解释性，时间和空间复杂度为线性，且无需近似任何pairwise attention矩阵。

关键观点4: 实验验证

在图像分类、语言建模、长序列建模等任务上，TOST架构替换注意力模块即可实现与标准Transformer相当或更优的性能，同时显著节省计算和内存。

关键观点5: TSSA模块的作用

TSSA模块是一个可替换传统attention的高效子模块，尤其适用于长序列、高清图像和资源受限场景下的Transformer架构，能显著降低计算与内存成本，同时保留甚至提升模型性能。

免责声明