🤖 AI Summary
Existing subquadratic alternatives to attention struggle to simultaneously achieve global receptive fields, input-dependent modeling, and preservation of geometric structure in multidimensional data. This work proposes HyenaND, the first natively multidimensional subquadratic global convolution operator, which leverages implicitly parameterized long convolutional kernels to enable input-dependent modeling without flattening or disrupting spatial structure. By integrating FFT-based acceleration with a fused CUDA kernel (nSubQ), HyenaND attains a computational complexity of O(L log L). Experiments demonstrate that pure HyenaND models match strong attention-based baselines across diverse domains—including genomics, computer vision, medical imaging, and partial differential equation modeling—while hybrid architectures significantly outperform existing attention-recurrent hybrids.
📝 Abstract
Subquadratic alternatives to attention require compromises when applied to multi-dimensional data: standard convolutions lack global receptive fields and input dependency, while recurrent models require rasterizing data such as images, volumes, and partial differential equation (PDE) into an ad-hoc $1\rm D$ scan order that violates their spatial structure. We introduce \textit{HyenaND}, a subquadratic, global, input-dependent operator that acts directly on the native geometry of multidimensional data through convolutions with implicitly parametrized global, input-dependent multi-dimensional convolutional kernels. Our CUDA implementation, \texttt{nSubQ}, fuses the FFT-convolution path to turn HyenaND's $\mathcal{O}(L \log L)$ scaling into wall-clock speedups. Across long-context genomics, computer vision, medical imaging, and PDE modeling, pure HyenaND stacks match the accuracy of strong attention baselines, while hybrid configurations that interleave HyenaND and attention layers outperform both pure attention and strong recurrence-based hybrids.