用多项式混合机制替代注意力,计算量线性增长却保持性能。
PoM: A Linear-Time Replacement for Attention with the Polynomial Mixer
- 用可学习的多项式函数聚合令牌,实现线性复杂度的上下文混合。
- 在五类任务中性能接近注意力模型,长序列时计算成本大幅降低。
- 适合需要高效处理长序列的应用,如大规模文本或高分辨率图像生成。
本文提出多项式混合器(PoM),一种线性复杂度的新型令牌混合机制,可直接替代自注意力。PoM通过可学习的多项式函数将输入令牌压缩为紧凑表征,每个令牌从中提取上下文信息。我们证明了PoM满足上下文映射性质,确保配备PoM的Transformer仍具备通用序列到序列近似能力。在文本生成、手写识别、图像生成、3D建模和地球观测五个不同领域中,用PoM替换标准自注意力,性能与注意力模型相当,但长序列下计算开销显著下降。代码已公开于https://github.com/davidpicard/pom。
原文摘要 · Abstract (English)
This paper introduces the Polynomial Mixer (PoM), a novel token mixing mechanism with linear complexity that serves as a drop-in replacement for self-attention. PoM aggregates input tokens into a compact representation through a learned polynomial function, from which each token retrieves contextual information. We prove that PoM satisfies the contextual mapping property, ensuring that transformers equipped with PoM remain universal sequence-to-sequence approximators. We replace standard self-attention with PoM across five diverse domains: text generation, handwritten text recognition, image generation, 3D modeling, and Earth observation. PoM matches the performance of attention-based models while drastically reducing computational cost when working with long sequences. The code is available at https://github.com/davidpicard/pom.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。