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: $$ \hat{O} = V_L = P_L \cdot P_{L-1} \cdots P_1 \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} = \mathbf{\mathcal{O}(N \cdot (\log N) \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

📚 Installation and Usage

An open idea to researchers