Beyond Self Attention: A Subquadratic Fourier Wavelet Transformer with Multi Modal Fusion

TL;DR

The paper tackles long‑sequence abstractive summarization by replacing traditional attention with a subquadratic spectral approach. It develops a full encoder–decoder Transformer that fuses global Fourier token mixing with local wavelet filtering (MDFWA) and enforces autoregressive decoding via causal spectral masking. Key contributions include a rigorous mathematical formulation, explicit complexity guarantees of $O(N\log N)$ time and $O(N)$ memory, gradient expressions for learned frequencies, and extensions to adaptive scales and multi‑modal fusion. Empirical results on PubMed 200K RCT demonstrate improvements over Fourier‑based and sparse‑attention baselines, highlighting MDFWA’s efficacy for long documents and multi‑modal contexts.

Abstract

We revisit the use of spectral techniques to replaces the attention mechanism in Transformers through Fourier Transform based token mixing, and present a comprehensive and novel reformulation of this technique in next generation transformer models. We provide expanded literature context, detailed mathematical formulations of Fourier mixing and causal masking, and introduce a novel MultiDomain Fourier Wavelet Attention(MDFWA) that integrates frequency and time localized transforms to capture both global and local dependencies efficiently. We derive the complexity bounds, gradient formulas, and show that MDFWA achieves sub quadratic time and memory cost while improving expressive power. We validate our design on an abstractive summarization task using PubMed dataset, by enhancing the proposed approach with learned frequency bases, adaptive scale selection, and multi-modal extensions.

Figures (1)

• Figure 1: Overview of the MDFWA Transformer. The input embeddings $X\in\mathbb{R}^{N\times d}$ are first combined with positional encodings $P$ to form $X^{(0)}=X+P$. Each of the $L$ encoder and decoder layers applies an MDFWA block, in which the Fourier branch computes $X'^{(\ell)} \;=\;\Re\bigl(\mathrm{FFT}(X^{(\ell-1)})\bigr),$ and the wavelet branch computes $\widetilde{X}^{(\ell)} \;=\;\mathrm{DWT}(X^{(\ell-1)})\,P^\top$. These are fused by $Y^{(\ell)}=\sigma\bigl(X'^{(\ell)}F_{F}+\widetilde{X}^{(\ell)}F_{W}+b\bigr)$, then combined with a residual connection and layer normalization: $X^{(\ell)}=X^{(\ell-1)}+\mathrm{LayerNorm}(Y^{(\ell)})$. In the decoder, a causal spectral mask restricts each inverse FFT sum to $k\le n$, preserving autoregressivity. Cross‑mixing replaces conventional encoder–decoder attention via $C=\Re\bigl(\mathrm{FFT}(\mathrm{concat}(Q,K))\bigr)\,V^\top$, thereby conditioning global source and target representations without $O(N^2)$ dot‑products. Finally, the decoder outputs are passed through a linear layer and softmax to produce token probabilities.