基于Token Statistics Transformer的线性复杂度自注意力机制创新模型ToST，在多模态学习中实现变分编码率缩减，荣获ICLR Spotlight高分认可

机器之心学术专栏：AIxiv 推动全球学术交流
机器之心的 AIxiv 栏目致力于发布和报道前沿的学术和技术。自成立以来，该栏目已报道了超过2000篇来自全球顶尖高校和企业实验室的研究成果，极大地促进了学术界的交流与传播。如果您有优秀的工作想要分享，欢迎通过以下邮箱投稿或联系：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com。

Token Statistics Transformer (ToST)：革新注意力机制
加州大学伯克利分校三年级博士生吴梓阳及其导师马毅教授提出了一种全新的注意力机制——Token Statistics Transformer (ToST)。该工作由多所知名高校和机构共同完成，包括加州大学伯克利分校、宾夕法尼亚大学、密歇根大学、清华大学等。马毅教授将在今年四月的 ICLR 大会上就相关研究进行主题报告。

Transformer 架构在计算机视觉、自然语言处理等领域取得了显著成就，但其核心组件「自注意力机制」的计算复杂度随输入 token 数量呈二次方增长，导致资源消耗巨大。ToST 提出了一种线性时间复杂度的注意力机制，通过统计建模序列特征，显著提高了处理效率。该研究已被选为 ICLR 2025 大会的 Spotlight 论文。

论文Token Statistics Transformer： Linear-Time Attention via Variational Rate Reduction
论文地址：[https：//arxiv.org/abs/2412.17810](https：//arxiv.org/abs/2412.17810)
项目主页：[https：//robinwu218.github.io/ToST/](https：//robinwu218.github.io/ToST/)
开源代码：[https：//github.com/RobinWu218/ToST](https：//github.com/RobinWu218/ToST)

研究背景与动机
传统的自注意力机制依赖于两两相似性的计算，尽管有效但资源开销巨大。当输入 token 数量极大时，传统方法的计算复杂度和内存使用成为瓶颈。为此，本文提出了基于统计学特征的 Token Statistics Self-Attention (TSSA)，通过避免两两相似性计算，仅依赖 token 特征的统计量，显著降低了计算复杂度。

核心方法
ToST 的核心方法是通过特定的概率分布函数对输入序列进行建模，减少冗余信息并提取关键特征。具体步骤包括：
1. 统计特征提取：对序列中的每个 token 提取其统计特征。
2. 变分编码率缩减：利用 VRR 框架对特征进行压缩，减少信息冗余。
3. 线性复杂度实现：通过一系列优化，将计算复杂度从 O(n²) 降低为 O(n)。

技术细节与创新点
1. 线性时间注意力机制：TSSA 通过白盒设计方法从最大编码率减少（Maximal Coding Rate Reduction, MCR²）的变分形式中推导而来，实现了线性时间和内存复杂度。
2. 创新性的网络结构：ToST 替代标准的自注意力模块，不仅提升了效率，还增强了模型的可解释性。
3. 理论支撑与数学推导：基于 MCR² 的变分形式，提出了一种新颖的压缩项公式，通过低秩投影保留关键信息并消除冗余方向。

实验验证与性能分析
实验覆盖了自然语言处理（NLP）、计算机视觉（CV）等多个领域的任务，结果表明 ToST 在保证模型性能的同时大幅降低了计算资源消耗。具体表现为：
1. 计算和内存的线性复杂度：与现有注意力机制相比，TSSA 的时间和内存复杂度更低，实现了线性扩展。
2. 视觉任务性能：在 ImageNet-1k 等数据集上，ToST 的性能可与传统 Transformer 相媲美，同时减少了模型参数量和计算开销。
3. 长序列任务和语言建模：ToST 展现出优异的长距离建模能力，并适用于多种任务场景，如因果语言建模。

应用前景
1. 大模型的高效化：ToST 为实现线性复杂度的大模型提供了可能性，推动了高效化应用。
2. 普适化应用：高效的注意力机制使得 ToST 能够广泛应用于资源受限场景，如边缘计算、实时系统等。
3. 多模态融合：低复杂度机制为处理多模态长序列任务提供了新的技术框架。
4. 跨学科应用：结合数学理论与工程实现，ToST 可能推动其在量子计算、生物信息学等新兴领域中的应用。

本文来源：