Native Multi-Dimensional Subquadratic Operators via Input Dependent Long Convolutions¶

会议: NeurIPS 2026
arXiv: 2607.19378
OpenReview: https://openreview.net/forum?id=YF3dDMGYTt
代码: https://github.com/HazyResearch/hyena
领域: 其他
关键词: 亚二次方复杂度算子、多维长卷积、HyenaND、输入依赖核、nSubQ

一句话总结¶

本文提出针对多维数据的亚二次方全局算子 HyenaND，通过显式保持多维几何结构的坐标隐式神经网络与基于寄存器 token 的 FiLM 样本级动态卷积核生成，结合软硬件协同优化的 nSubQ CUDA 算子库，在保持 \(\mathcal{O}(L\log L)\) 计算复杂度的同时突破了注意力机制的二次方瓶颈与递归模型的空间光栅化失真。

研究背景与动机¶

深度学习模型在应对高分辨率图像、三维医学体素扫描、时空偏微分方程（PDE）动力学场以及数百万 token 的长基因组序列等多维数据时，面临计算复杂度与几何保真度的根本性权衡。标准 Transformer 的多头自注意力机制虽然原生支持任意多维结构，但其 \(\mathcal{O}(L^2)\) 的时间与显存复杂度在序列长度膨胀时导致计算成本迅速失控。计算机视觉与科学计算中普遍采用的补丁化（Patchification）操作虽能将 token 数量压缩至显存允许范围，但经验缩放定律表明模型的表征精度会随着补丁尺寸的缩小（即向像素级 token 逼近）而单调递增，而注意力机制的二次方墙恰恰阻碍了这一极限精度的获取。

在注意力机制的替代方案中，状态空间模型（SSM，如 S4）及其演化出的选择性递归模型（如 Mamba）凭借 \(\mathcal{O}(L)\) 的线性训练复杂度和 \(\mathcal{O}(1)\) 的单步自回归推理崭露头角。然而，递归模型在数学形式上本质属于因果一维时序算子：在双向上下文处理中必须沿正反向运行两次递归，更关键的是递归机制缺乏向高维欧氏空间的自然拓扑扩展。现有视觉 Mamba 或多维 Mamba 变体均依赖人工指定的 1D 空间光栅化扫描路径（如行列扫描、锯齿扫描或四向交叉扫描），不可避免地割裂了多维数据内在的空间邻域关联与各向同性几何结构。此外，递归模型依赖固定容量的隐状态进行序列压缩，其状态瓶颈在多维空间中随感受野维度的指数级扩展而急剧恶化。

长卷积模型作为另一类非递归亚二次方架构，通过频域快速傅里叶变换（FFT）能以 \(\mathcal{O}(L\log L)\) 复杂度直接作用于输入的原生多维坐标系，且无任何递归隐状态瓶颈。然而，以往的长卷积模型普遍受限于静态卷积核的输入无关性，无法像 Mamba 或自注意力那样根据输入动态调整权重。核心 idea：利用可学习寄存器 token 聚合全局语义并经由 FiLM 动态调制多维隐式坐标网络，在完全保留原生多维几何结构、线性时不变（LTI）频域 FFT 路径与 \(\mathcal{O}(L\log L)\) 亚二次方扩展性的同时，赋予全局长卷积样本级动态输入依赖能力。

方法详解¶

整体框架¶

HyenaND 旨在为任意 \(N\) 维数据输入 \(\mathbf{x}\in\mathbb{R}^{L_1\times\cdots\times L_N\times H}\) 构建原生多维空间下的长程交互通道。在算子内部，输入特征经过线性投影与极短的局部深度卷积，分解为三路表征：查询投影 \(\mathbf{q}\)、键投影 \(\mathbf{k}\) 与值投影 \(\mathbf{v}\)。为了控制张量幅值并在长序列甚至数百万 token 场景下抑制数值溢出，算子对 \(\mathbf{q}\) 施加通道维度的 \(\ell_2\) 归一化，与经由 SiLU 激活的 \(\mathbf{k}\) 逐元素相乘构成内门控信号，再经过 RMSNorm 稳定输入进多维频域卷积网路。

该门控信号随后与动态生成的全多维全局卷积核 \(K(\mathbf{x})\) 执行原生 \(N\) 维快速傅里叶变换长卷积。卷积输出通过 Sigmoid 激活的值分支 \(\mathbf{v}\) 进行逐通道软衰减外门控，最后经由残差连接前的 RMSNorm 输出为 \(\mathbf{y}\)。算子的完整前向数学形式表达如下：

\[\mathbf{y} = \mathrm{RMSNorm}\bigl(\mathrm{Sigmoid}(\mathbf{v})\odot (K(\mathbf{x})\ast \mathrm{RMSNorm}(\ell_2(\mathbf{q})\odot \mathrm{SiLU}(\mathbf{k})))\bigr)\]

整体算子可无缝替换标准 Transformer 中的自注意力层或 Mamba 块，嵌入到预归一化残差模块堆叠中。整个前向过程及输入依赖卷积核的合成数据流如下图所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多维输入网格 + 寄存器"] --> B["多维隐式坐标内核合成<br/>SIREN + FiLM调制 + 高斯窗口"]
    A --> C["三路投影与内门控稳定化<br/>l2归一化q ⊙ SiLU(k) + RMSNorm"]
    B --> D["原生多维频域长卷积<br/>nSubQ 融合 N-D FFT 路径"]
    C --> D
    A --> E["外门控选择与残差输出<br/>Sigmoid(v) 调制 + RMSNorm"]
    D --> E
    E --> F["输出到下一残差层"]

关键设计¶

1. 多维隐式坐标内核合成：突破可分离轴分解的联合多维连续表征

传统多维卷积（如 S4ND）通常采用各轴独立的一维卷积核乘积 \(f_\theta(\mathbf{c}_i)=\prod_{n=1}^N f_{\theta,n}(c_{i,n})\) 来近似多维核，这使得生成的卷积核在空间秩上退化为秩 1 结构，无法刻画斜向边缘、各向异性扩散或复杂的高维流体漩涡。HyenaND 抛弃轴可分离假设，采用联合映射的连续隐式正弦周期神经网络（SIREN）\(f_\theta:\mathbb{R}^N\to\mathbb{R}^H\) 直接作用于归一化网格坐标 \(\mathbf{c}_i\in[-1,1]^N\)。为了解决多维输入下激活方差随维度 \(N\) 线性膨胀导致正弦函数初始饱和与梯度弥散的痛点，论文给出了严谨的方差校准证明：第一层权重采用 \(W_0\sim\mathcal{U}(-1/\sqrt{N},1/\sqrt{N})\) 初始化，将方差精确限制在 \(1/9\)。同时，隐式核的频率缩放参数 \(\bm{\omega_0}\) 结构化为 \(B\) 个分块，在线性频率空间中覆盖多频段，并在外部覆盖具有可学习通道与空间各轴独立方差的零均值高斯包络窗口 \(w(\mathbf{c}_i)=\exp(-\frac{1}{2}\sum_d (r_d/\sigma_{d,c})^2)\)，有效克服了传统指数衰减窗口在多维空间中随维度加倍而剧烈衰退的几何缺陷。

2. 基于寄存器与 FiLM 的样本级动态调制：兼顾输入依赖性与单一 FFT 高速通道

逐 token 变化的卷积核会彻底破坏线性时不变性（LTI），迫使算子回退为昂贵的通用矩阵乘法或一维自回归扫描。为了在保留 \(N\) 维 FFT 高速卷积优势的同时引入注意力机制的核心能力——内容自适应选择性，HyenaND 设计了基于寄存器 token（Register Tokens）的样本级控制机制。在输入多维张量的顶层空间维度前置 \(R\) 个可学习寄存器向量，使其与空间 token 在各层联合交互并汇聚全局上下文。各层通过对寄存器激活值进行可学习 Softmax 加权池化提取样本级控制向量 \(\mathbf{z}(\mathbf{x})\in\mathbb{R}^H\)，并通过两层 MLP \(f_{\mathrm{FiLM}}\) 预测 SIREN 隐层的逐层仿射尺度与平移量 \((\bm{\gamma}_j,\bm{\beta}_j)\)：

\[\mathbf{h}_j\leftarrow \bm{\gamma}_j\odot \mathbf{h}_j + \bm{\beta}_j\]

该设计使得卷积核 \(K(\mathbf{x})\) 针对每个样本动态特化，而在单个样本内部的空间位置上完全共享，从而在批次处理中仅需为每个样本实例化一个全局多维核，严格维持了单次多维 FFT 卷积的计算路径。初始化时将 FiLM 投影层权重置零且偏置置零，确保初始状态等价于静态平稳核，保障了深度堆叠训练初期的数值稳定性。

3. nSubQ 软硬件协同优化算子库：打通多维 FFT 的片上显存墙瓶颈

尽管多维 FFT 卷积在理论上具有 \(\mathcal{O}(L\log L)\) 的渐近复杂度优势，但在硬件实现层面，多维非循环边界所必需的零填充（在 \(D\) 维欧氏空间下带来约 \(2^D\) 倍的数据体量膨胀）、实部到复数的转换以及频繁的内核调用，往往导致显存（HBM）读写流量急剧攀升，使得渐近优势被巨大的访存常数淹没。为彻底释放多维长卷积的硬件潜力，作者开发了基于 CuTe 与 cuFFTDx 的定制 CUDA 算子库 nSubQ。nSubQ 将输入填充、正向 N-D FFT、频谱逐元素相乘复数调制、逆向 N-D IFFT 以及输出有效区域截断融合成单一 CUDA Kernel，将全部中间频域复数激活保留在片上高速共享内存（SRAM）与寄存器中。此外，针对隐式网络自身生成卷积核耗时占比高达 30% 的系统瓶颈，nSubQ 构建了专用的隐式核熔合生成流水线，在模块级别将卷积核生成延迟降低了 40 倍，显存流量缩减 32 倍，为 16M 超长序列的分布式上下文并行训练奠定了工程基础。

损失函数 / 训练策略¶

HyenaND 算子本身属于通用序列/空间混合器，支持直接接入主流无监督预训练或有监督任务的目标函数。在超长基因组预训练实验中，模型采用自回归交叉熵损失，并在 GB200 集群上采用 FP8 精度与上下文并行（Context Parallelism, CP）技术，实现了从 8K 平滑扩展至 16M token 窗口的渐进式长度泛化训练。在 ImageNet-1K 视觉分类实验中，模型采用 ViT-5-Small 骨干（22M 参数），使用 LAMB 优化器配合 Cosine 衰减调度训练 800 个 epoch，并辅以 Mixup、CutMix、3-Augment 和 LayerScale。在偏微分方程建模中，模型以均方误差（MSE）作为训练目标，并在评估阶段采用标准化均方根误差（VRMSE）作为核心基准指标。

实验关键数据¶

主实验¶

论文在跨维度受控空间召回基准、ImageNet-1K 图像分类以及 The Well 物理偏微分方程模拟上开展了系统性评测。下表展示了受控空间记忆探测任务（无补丁化、原生分辨率）的均方误差对比，验证不同算子在空间路由与内容联想检索上的能力：

评估任务维度	评测子项	随机均值基线	标准 Attention (RoPE)	双向 Mamba (bidir)	HyenaND (本文)
1D 序列 (\(L=4096\))	simple_copy	0.5490	0.1500	0.0850	0.000052
1D 序列 (\(L=4096\))	color_cond	0.2840	0.2740	0.0065	0.0062
2D 平面 (\(64\times 64\))	simple_copy	0.5490	0.0920	0.3490	0.00013
2D 平面 (\(64\times 64\))	color_cond	0.2840	0.2530	0.0083	0.0025
3D 时空体 (\(32^3\))	simple_copy	0.4820	OOM (显存溢出)	0.4510	0.00012
3D 时空体 (\(32^3\))	color_cond	0.2570	OOM (显存溢出)	0.00010	0.000057

在 ImageNet-1K 图像分类任务中，采用 ViT-5-Small 相同架构骨干（22M 参数，不使用额外 patch-merge），不同补丁尺寸 \(p\times p\) 下的性能与计算量对比如下：

模型架构类别	算子配置与结构	补丁大小 (\(p\times p\))	Token 数量	Top-1 准确率 (%)	计算量 (GFLOPs)
标准 Attention 骨干	ViT-5-Small 基线	\(16\times 16\)	196	81.8	9.41
标准 Attention 骨干	ViT-5-Small 基线	\(8\times 8\)	784	84.3	45.51
标准 Attention 骨干	ViT-5-Small 基线	\(4\times 4\)	3,136	85.1	317.34
标准 Attention 骨干	ViT-5-Small 基线	\(2\times 2\)	12,544	85.3	3,443.93
纯卷积亚二次方算子	HyenaND-S (纯算子)	\(16\times 16\)	196	81.5	9.97
纯卷积亚二次方算子	HyenaND-S (纯算子)	\(8\times 8\)	784	83.7	39.04
纯卷积亚二次方算子	HyenaND-S (纯算子)	\(4\times 4\)	3,136	84.0	155.32
纯卷积亚二次方算子	HyenaND-S (纯算子)	\(2\times 2\)	12,544	84.0	623.02
混合交织架构	HyenaND-S \((HA)\times 6\)	\(16\times 16\)	196	82.1	9.69
混合交织架构	HyenaND-S \((HA)\times 6\)	\(4\times 4\)	3,136	85.0	236.33
混合交织架构	HyenaND-S \((HA)\times 6\)	\(2\times 2\)	12,544	85.2	2,033.48
递归模型基线	Mamba-S \((MA)\times 6\)	\(16\times 16\)	196	81.8	12.46 (30M参)
递归模型基线	Mamba-S \((MMMA)\times 3\)	\(16\times 16\)	196	81.0	13.99 (34M参)

消融实验与多领域分析¶

针对三维医学图像分割（PanTS 数据集，输入 ROI 大小为 \(96^3\)，测试集 901 例独立腹部 CT 体素，28 类分割目标），论文将 SwinUNETR 的四级分层编码器中的窗口注意力按不同组合替换为三维 HyenaND 算子，实验对比各配置的分割精度（Dice）与显存开销：

编码器四级结构配置	算子说明	平均 Dice 准确率	显存峰值 (GB)	显存降低比例	前向计算量 (GFLOPs)
AAAA 基线	全自注意力 (SwinUNETR)	0.7496 ± 0.0053	65.79	基准 (0.0%)	638
HHHH 纯卷积	全 HyenaND 替换	0.7510 ± 0.0110	58.69	降低 10.8%	621
HAHA 交错混合	条带化 Hyena-Attention 混合	0.7535 ± 0.0054	58.70	降低 10.8%	627
HHAA 层次混合	前两级 Hyena，后两级 Attention	0.7559 ± 0.0052	58.67	降低 10.8%	623

关键发现¶

几何结构完好性决定多维空间建模能力：在受控空间记忆实验中，由于 Mamba 必须将 2D/3D 输入人为压平为一维扫描链，破坏了空间连续性，在 3D simple_copy 任务上的误差高达 0.4510，几乎退化至随机预测水平（0.4820）；而纯正多维长卷积 HyenaND 能直接在原生欧几里得网格上对齐多维空间偏移，将误差降至 \(1.2\times 10^{-4}\)，性能差异达四个数量级。
混合架构在极限微补丁下实现精度与算力的最优帕累托前沿：在 ImageNet 实验中，当图像补丁缩小至 \(p=2\)（对应 12,544 个 token）时，全注意力机制的计算开销暴增至 3,443.9 GFLOPs，而纯 HyenaND 仅需 623.0 GFLOPs（节约 82% 计算量）。交错混合配置 \((HA)\times 6\) 达到了 85.2% 的 Top-1 准确率，与纯注意力基线（85.3%）差距仅在 0.1% 以内，同时计算量锐减约 41%（从 3,444 GFLOPs 降至 2,033 GFLOPs）。
空间分辨率梯度与算子特性的天然契合：在三维医疗体素分割中，HHAA 阶梯分层配置取得了最高的 0.7559 Dice 表现。原因在于网络浅层特征分辨率极高（token 数超 11 万），全局 FFT 卷积提供了比有限局部窗口注意力更开阔的全局解剖结构视野，且规避了高分辨率注意力矩阵显存暴涨的弊端；而在深层网格缩减后，注意力机制的全局内容寻址能力重新发挥关键优势。

亮点与洞察¶

厘清亚二次方算子的理论谱系：论文清晰界定了状态空间模型（SSM）、广义递归模型（如 Mamba）以及广义长卷积模型（如 HyenaND）的本质边界，指出现有研究将亚二次方算子泛称为 SSM 是认知上的混淆。为处理多维物理信号与视觉信号，保留线性时不变性（LTI）以激活多维 FFT 快速通道，远比一维逐 token 因果选择性更加关键。
寄存器 token 巧妙解耦动态输入依赖与快速 FFT 卷积：过去认为长卷积缺乏动态自适应能力的缺陷，在此被转化为“样本级动态性”的优雅实现。利用附加的全局寄存器抽取样本语义控制向量并经由 FiLM 调节隐式坐标网络的参数，使网络能够针对不同输入图像生成各异的动态滤波器，同时又在样本内部保持频域卷积定理的严格适用性。
颠覆补丁化作为视觉预处理的刻板认知：论文指出现代视觉架构中的 Patchification 并非物理特征提取所必须，而是注意力二次方计算量逼迫下的妥协产物。HyenaND 证明了亚二次方多维算子能够跨越算力墙，直接在极高分辨率网格甚至原生像素/体素级别进行无损全局混合，为构建下一代像素级原生视觉与科学基础模型指明了方向。

局限与展望¶

全局卷积核生成存在常数级计算开销：尽管 nSubQ 将隐式核生成进行了深度融合与加速，但相比于静态参数卷积或无参数 FFT，在每个 batch 动态调用 SIREN 生成高维张量核依然存在不可忽略的时间常数，尤其在小批量和浅层特征推理时开销更为突出。
高频自适应初始化的跨模态泛化敏感性：偏微分方程实验表明，不同物理动力学过程（如剪切流与反应扩散系统）对隐式网络初始化频率 \(\bm{\omega_0}\) 的分布范围敏感度各异，固定的频率超参数在个别数据集上表现不及调整后的补丁尺寸，未来需要引入自适应频谱初始化或频域自动微调机制。
从预处理切片向原生高维分词器的演进契机：作者指出，当前 HyenaND 仍主要作为深层 Transformer 的替代混合器，未来极具潜力的演进路径是将浅层 HyenaND 直接作为可微分的原生像素分词器（Learned Tokenizer），在极低开销下自底向上聚合高密度信息，赋能下游大模型架构。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙融合隐式连续多维坐标表征、寄存器 FiLM 动态调制与统一 nSubQ CUDA 库，彻底攻克多维亚二次方算子的几何割裂与静态核缺陷。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1D 长上下文基因组、2D ImageNet 分类、3D 医学腹部 CT 分割以及 8 个复杂多维 PDE 动力学数据集，对比基线全面且包含大量微补丁与极限消融。
写作质量: ⭐⭐⭐⭐⭐ 概念界定严谨透彻，从理论分析到系统实现层层递进，软硬件瓶颈解剖精辟，图表设计极其翔实。
价值: ⭐⭐⭐⭐⭐ 为高分辨率视觉基础模型、百万级生物基因大模型以及三维大规模科学仿真提供了关键的高效底层算子，具有重大的学术与工业应用价值。