PoM: Trading Quadratic Attention for Linear Polynomial Mixing

Paper: 2604.06129 Authors: David Picard, Nicolas Dufour, Lucas Degeorge, Arijit Ghosh, Davide Allegro, Tom Ravaud, Yohann Perron, Corentin Sautier, Zeynep Sonat Baltaci, Fei Meng Categories: cs.CV, cs.AI

The Gap

Self-attention has dominated sequence modeling since 2017, but it carries a dirty secret: O(n²) complexity. For every token you add, the computational cost grows quadratically. This isn’t just a theoretical concern—it’s a wall that practitioners hit when working with long documents, high-resolution images, or extended video sequences. Prior attempts to fix this fall into two camps: sparse attention patterns (Longformer, BigBird) that sacrifice global context, and linear approximations (Performer, RWKV) that often underperform standard attention on benchmarks. The field has been stuck in a trade-off: either pay the quadratic tax or accept degraded performance.

Problem: Self-attention O(n^2) blocks long sequences
                    |
                    v
Assumption: Can we aggregate tokens into compact form
            then retrieve context in O(n)?
                    |
                    v
Method: Polynomial Mixer (PoM)
        - Aggregate via learned polynomial
        - Retrieve per-token context
                    |
                    v
Evidence: Match attention performance across 5 domains
          (text, handwriting, images, 3D, Earth obs)
          with linear complexity
                    |
                    v
Conclusion: Linear mixing can replace quadratic attention
            without sacrificing universal approximation

The Increment

One sentence: Before PoM, you chose between quadratic attention or degraded linear alternatives; after PoM, you get attention-level performance at linear cost with theoretical guarantees.

Core Mechanism

PoM operates in two phases. First, it compresses all input tokens into a single compact representation using a learned polynomial function. Think of this as creating a “summary vector” that captures global context. Second, each token queries this summary to retrieve its contextualized representation. The polynomial acts as a parameterized aggregation function—instead of computing pairwise interactions between all tokens (the attention bottleneck), you compute each token’s contribution to a shared summary, then broadcast that summary back.

The polynomial itself is a weighted sum of powers: f(x) = w₀ + w₁x + w₂x² + … + wₖxᵏ. The weights are learned parameters. When you feed token embeddings through this polynomial, you’re essentially asking: “How should this token contribute to the global context?” The beauty is that polynomial evaluation is linear in sequence length—you process each token once to build the summary, then once more to retrieve context.

Input tokens:  [x1] [x2] [x3] ... [xn]
                 |    |    |        |
                 v    v    v        v
            +------------------------+
            | Polynomial Aggregation |  <-- f(x) = w0 + w1*x + w2*x^2 + ...
            +------------------------+
                       |
                       v
                  [Summary S]  <-- Compact global context
                       |
         +-------------+-------------+
         |             |             |
         v             v             v
    Retrieve(x1,S) Retrieve(x2,S) Retrieve(xn,S)
         |             |             |
         v             v             v
      [y1]          [y2]          [yn]  <-- Contextualized outputs

Here’s the structural metaphor: PoM works like a town hall meeting. In standard attention, every person (token) whispers to every other person—that’s n² conversations. In PoM, everyone speaks into a microphone (polynomial aggregation) that feeds into a PA system (the summary vector). Then each person listens to the PA system to understand the room’s consensus. You’ve gone from n² whispers to n speeches plus n listening sessions—linear complexity. The polynomial weights determine whose voice gets amplified in the PA mix.

Key Concepts

• Contextual Mapping Property: A token mixer satisfies this if it can map any input sequence to any output sequence given enough parameters. Self-attention has this property, which is why transformers are universal approximators. The paper proves PoM also satisfies it—meaning you’re not losing theoretical expressiveness by going linear. This matters because many prior linear alternatives couldn’t make this guarantee. Think of it like this: if attention is a Swiss Army knife that can cut, saw, and open bottles, PoM proves it’s also a Swiss Army knife, just with a different blade mechanism. You’re not downgrading to a butter knife.

• Polynomial Aggregation: Why polynomials specifically? Because they’re universal function approximators (Weierstrass theorem) and computationally cheap. When you aggregate tokens via f(x) = Σ wᵢxⁱ, you’re learning a flexible weighting scheme. Low-degree terms (w₀, w₁) capture linear relationships; higher-degree terms capture non-linear interactions. The key insight: you don’t need to compute all pairwise interactions explicitly—the polynomial implicitly encodes them through its learned coefficients. It’s like how a Taylor series can approximate any smooth function without explicitly storing every point on the curve.

• Summary Vector: This is the bottleneck through which all information flows. It’s a fixed-size representation (typically same dimension as token embeddings) that compresses the entire sequence. The risk: information loss. The mitigation: the polynomial’s non-linearity and learned weights ensure that important patterns get preserved. Think of it like JPEG compression—you’re throwing away some data, but if you tune the compression algorithm well (via learning), you keep what matters for downstream tasks.

Framework Shift

Before (Self-Attention):              After (PoM):

[x1]--\                               [x1]--\
[x2]---+--[Pairwise]--[Weighted]      [x2]---+--[Polynomial]--[Summary]
[x3]--/    [n^2 ops]   [Sum]          [x3]--/    [n ops]       [Vector]
 ...                     |              ...                        |
[xn]--------------------+             [xn]-----------------------+
         |                                        |
         v                                        v
    [Attention                              [Broadcast &
     Matrix]                                 Retrieve]
         |                                        |
         v                                        v
    [y1...yn]                                [y1...yn]

Cost: O(n^2)                          Cost: O(n)

From pairwise token interactions to collective aggregation, the core shift is replacing explicit all-to-all communication with implicit global summarization.

Expert Assessment

Problem choice: This is a real gap, not manufactured. The O(n²) wall is the single biggest obstacle to scaling transformers to longer contexts. The problem sits at the intersection of theory (can we maintain expressiveness?) and practice (can we actually run this on real hardware?). It’s a high-value target.

Method maturity: The polynomial choice is elegant but not obviously necessary—why not other aggregation functions? The paper doesn’t deeply explore this design space. That said, the theoretical proof of contextual mapping is solid work. The method feels like a clever insight (use polynomials as learnable aggregators) rather than brute force, though I’d want to see ablations on polynomial degree and alternative aggregation schemes.

Experimental integrity: Testing across five diverse domains (text, handwriting, images, 3D, Earth observation) is commendable—most papers cherry-pick one favorable domain. The baselines appear fair. However, the paper doesn’t deeply probe failure modes. Where does PoM underperform attention? On what sequence lengths does the crossover happen? The numbers hold up, but I’d want more analysis of the summary vector’s capacity limits.

Writing quality: The introduction is crisp, but Section 3 (method) buries the lead. The polynomial aggregation mechanism should be front and center with a visual diagram, not hidden in equations. The related work section is thorough but could be trimmed. If I were rewriting, I’d expand Section 4.5 (analysis of what the polynomial learns) into a full section—that’s where the real insight lives.

Verdict: Weak accept — solid contribution with strong empirical validation, but the design space exploration is shallow and the theoretical novelty is incremental (applying known polynomial approximation theory to token mixing).

Takeaways

Practitioners can steal the core pattern: aggregate globally, retrieve locally. This isn’t specific to polynomials—you could use other learnable aggregation functions (MLPs, kernels, even learned hash functions). The key is breaking the O(n²) barrier by forcing information through a fixed-size bottleneck. If you’re building custom architectures for long sequences, try this: (1) compress your sequence into a fixed-size summary via a learnable function, (2) let each token query that summary for context. Test whether your task’s information can survive the compression. For domains where global context matters but pairwise interactions don’t (e.g., document classification, time series forecasting), this pattern is immediately applicable.