chaiya1962/fft-ia-transformer

---

tags: - transformer - attention-mechanism - long-sequence - n-log-n-complexity - research license: apache-2.0

🚀 FFT-Inspired Attention (FFT-IA): $\mathcal{O}(N \log N)$ Scalable Transformer

💡 Overview

The FFT-Inspired Attention ($\text{FFT}$-$\text{IA}$) framework is a novel architecture designed to fundamentally solve the $\mathcal{O}(N^2)$ quadratic complexity bottleneck of the Multi-Head Self-Attention (MHSA) mechanism. Inspired by the computational principles of the Fast Fourier Transform (FFT) and the Cooley-Tukey algorithm, $\text{FFT}$-$\text{IA}$ achieves an $\mathbf{\mathcal{O}(N \log N)}$ asymptotic complexity by structurally factoring the attention matrix.

Unlike approximation methods, $\text{FFT}$-$\text{IA}$ achieves efficiency through fixed structural pruning while maintaining Softmax Fidelity and content dynamism.

⚙️ Architectural Mechanics

The architecture replaces the single, dense attention matrix with a cascade of $L = \log_2 N$ sequential sparse projection factors, $P_i$.

1. The Hierarchical Factorization

The attention operation is decomposed into $L$ sequential stages: $$\widehat{O}=V_L=P_L\cdot P_{L-1}\cdots P_1\cdot V\backslash hat\{O\}\; =\; V\_L\; =\; P\_L\; \backslash cdot\; P\_\{L-1\}\; \backslash cdots\; P\_1\; \backslash cdot\; V$$

Each factor $P_i$ is a fixed, sparse matrix defined by a radix-2 butterfly connection pattern.

2. The Butterfly-Attention Block (BAB)

The core unit, the BAB, performs a local attention operation between two tokens $j$ and $k$ (where $k = j \pm 2^{i-1}$) in $\mathbf{\mathcal{O}(d^2)}$ time per token.

• Dynamic Re-Projection: Query ($Q_i$) and Key ($K_i$) are dynamically re-projected from the intermediate value state $V_{i-1}$ at every stage to ensure content-dependent scoring: $$ Q_i = W_{Q, i} V_{i-1}, \quad K_i = W_{K, i} V_{i-1} $$
• Local Softmax Fidelity: The exact attention score is computed and normalized locally over the fixed connection set $\mathcal{C}j = {j, k}$: $$ P_i[j, k'] = \text{softmax}{k' \in \mathcal{C}j} \left(\frac{Q{i, j} K_{i, k'}^\top}{\sqrt{d_k}}\right) $$ The local Softmax acts as a normalized adaptive pooling step.

📊 Performance and Complexity

The framework guarantees $\mathbf{\mathcal{O}(N \log N)}$ asymptotic complexity in sequence length $N$. The total computational cost is: $$\text{Total FLOPs}=\mathcal{O}(\mathbf{N}\cdot (\log \mathbf{N})\cdot ({\mathbf{d}}^{\mathbf{2}}+{\mathbf{d}}_{\mathbf{k}}))\backslash text\{Total\; FLOPs\}\; =\; \backslash mathbf\{\backslash mathcal\{O\}(N\; \backslash cdot\; (\backslash log\; N)\; \backslash cdot\; (d^2\; +\; d\_k))\}$$

Critical Implementation Note: Kernel Fusion

The primary practical barrier to achieving wall-clock speedup is the $\mathbf{\mathcal{O}(N d^2 \log N)}$ overhead from $\log_2 N$ sequential $Q/K$ re-projection and kernel launches. Dedicated Kernel Fusion (e.g., in CUDA or Triton) is the necessary condition to realize the full theoretical speed advantage.

🔑 Key Distinctions

Feature	Standard MHSA ($\mathcal{O}(N^2)$)	$\text{FFT}$-$\text{IA}$ ($\mathcal{O}(N \log N)$)
Complexity	Quadratic $\mathcal{O}(N^2)$	Near-linear $\mathcal{O}(N \log N)$
Efficiency Mechanism	None (Dense matrix)	Fixed Structural Pruning (FFT-inspired)
Softmax	Global $\mathcal{O}(N)$ scope	Local $\mathcal{O}(1)$ scope (Softmax Fidelity)
Dynamism	Direct content-based interaction	Dynamic Q/K re-projection on fixed graph

📚 Installation and Usage

An open idea to researchers