PoM: 다항식 혼합기를 이용한 어텐션의 선형 시간 복잡도 대체 모델

PoM: A Linear-Time Replacement for Attention with the Polynomial Mixer

본 논문에서는 다항식 혼합기(Polynomial Mixer, PoM)를 소개합니다. PoM은 자기 어텐션의 대체 모델로 사용될 수 있는, 선형 복잡도를 가지는 새로운 토큰 혼합 메커니즘입니다. PoM은 학습된 다항 함수를 통해 입력 토큰을 압축된 표현으로 통합하고, 각 토큰은 이를 통해 문맥 정보를 얻습니다. 우리는 PoM이 문맥 매핑 속성을 만족함을 증명하여, PoM을 탑재한 트랜스포머가 여전히 범용적인 시퀀스-투-시퀀스 근사 모델임을 보장합니다. 우리는 텍스트 생성, 손글씨 텍스트 인식, 이미지 생성, 3D 모델링, 지구 관측을 포함한 다섯 가지 다양한 영역에서 기존의 자기 어텐션을 PoM으로 대체했습니다. PoM은 어텐션 기반 모델과 동등한 성능을 제공하면서, 특히 긴 시퀀스를 처리할 때 계산 비용을 크게 절감합니다. 코드 및 관련 정보는 다음 링크에서 확인할 수 있습니다: https://github.com/davidpicard/pom.

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.