功能注意力:从成对亲和度到功能对应关系 Functional Attention: From Pairwise Affinities to Functional Correspondences
将注意力重新解释为函数空间之间的线性算子,实现高效且分辨率无关的算子学习
前置知识
算子学习
算子学习是学习无限维函数空间之间映射的机器学习范式,目标是学习一个算子 O: F 到 G,其中 F 和 G 是函数空间。与传统的向量到向量映射不同,算子学习的输入和输出都是连续函数,这使得学习目标独立于特定的离散化方式,能够在不同分辨率间泛化。典型应用包括偏微分方程求解、物理仿真等。
本文核心贡献就是提出了一种新的注意力机制用于算子学习,理解算子学习的基本概念对于理解本文的问题设定和方法动机至关重要。
函数映射
函数映射是几何处理中用于表示流形之间对应关系的框架,其核心思想是将困难的点对点匹配问题转化为函数空间之间的线性算子问题。给定两个流形 M 和 N,以及它们各自的一组基函数,函数映射用一个 k 乘 k 的矩阵 C 表示,其中 k 远小于点数 n。这个矩阵描述了如何将 M 上的函数映射到 N 上的函数。
本文直接借鉴了函数映射的思想,将注意力机制重新解释为函数空间之间的线性算子,这是理解本文方法的关键理论基础。
频谱基
频谱基是指通过求解拉普拉斯-贝尔特拉米算子的特征函数得到的一组正交基函数。在几何和物理问题中,这些基函数具有很好的性质:低频基函数捕捉全局平滑变化,高频基函数捕捉局部细节。傅里叶基是最常见的频谱基之一。
本文将函数投影到学习到的基上表示,理解频谱基的概念有助于理解本文如何在压缩空间中计算注意力,以及这种表示如何捕获全局依赖关系。
研究动机
现有的算子学习方法面临两个核心问题。首先,基于变换器的方法通常采用逐token的注意力机制,将连续场视为离散token集合,这种方法忽略了函数的全局结构。具体来说,标准注意力需要 O(n²) 的计算复杂度,其中 n 是表示函数所需的采样点数,这使得计算成本随着分辨率增加而急剧增长。其次,传统神经算子方法如 Fourier Neural Operator (FNO) 使用固定的频谱基(如傅里叶基),这些基函数虽然计算高效,但对规则网格结构和周期性边界条件有很强假设,限制了在复杂几何形状上的表现能力。例如,在具有尖锐局部特征的领域(如带缺口的三角形域),固定基函数无法有效对齐数据的内在结构。
本文的目标是本文的目标是提出一种新的注意力机制,能够直接在函数空间层面建模,而不是在离散token层面。具体来说,作者希望设计一种既能够捕捉全局函数结构,又能够适应不同几何形状和数据特征的算子学习方法。该方法应该具有以下特性:计算复杂度与采样点数 n 线性相关,而非二次相关;能够学习适应数据的基函数,而不依赖于固定的人为选择的基;保持分辨率不变性,即在粗网格上训练的模型能够泛化到细网格;在各种算子学习任务(包括PDE求解、3D分割、回归)上达到最先进的性能。
与已有工作不同的是,本文的独特切入角度是借鉴函数映射理论,将注意力机制重新解释为函数空间之间的功能对应关系,而不是token之间的成对亲和度。虽然已有工作如 Galerkin Transformer 也注意到了 Q、K、V 矩阵的列可以看作是希尔伯特空间中函数的离散化,但它们没有提供如何在实际中计算这些基和系数的具体方法。与之不同的是,本文提出了一种完整的框架,通过学习到的基将输入函数投影到压缩的频谱空间,在该空间中通过最小二乘问题计算最优线性算子,然后再反投影回原始空间。这种设计在理论上具有更好的稳定性(通过正则化参数 lambda 控制 Lipschitz 常数),在实验中也表现出更强的泛化能力和样本效率。
核心方法
方法的核心思想是将注意力机制从token空间提升到函数空间。直观上,标准注意力计算的是token之间的成对相似度(一个 n 乘 n 的矩阵),而本文的方法计算的是函数空间之间的线性算子(一个 k 乘 k 的矩阵,其中 k 远小于 n)。技术路线如下:首先,为查询空间和键值空间各学习一组基函数 Phi 和 Psi,这些基通过简单的线性层加softmax归一化得到;然后,将查询 Q、键 K、值 V 投影到各自的频谱系数空间,得到 Q_tilde = Phi_dagger 乘 Q、K_tilde = Psi_dagger 乘 K、V_tilde = Psi_dagger 乘 V;接着,在频谱空间中求解一个Tikhonov正则化的最小二乘问题,得到最优线性算子 C_star = Q_tilde 乘 K_tilde_dagger 乘 (K_tilde 乘 K_tilde_dagger + lambda 乘 I_k) 的负一次方;最后,用这个算子将值的频谱系数投影到查询空间,并反投影得到输出 FUNCATTN(Q, K, V) = Phi 乘 C_star 乘 V_tilde。整个流程避免了计算显式的 n 乘 n 亲和度矩阵,而是在紧凑的 k 乘 k 空间中进行操作。
核心创新点在于用结构化的线性算子替代softmax亲和度,这使得注意力机制能够显式地捕获全局函数结构。具体而言,标准注意力计算 Softmax(Q 乘 K_dagger 除以 sqrt(d_k)) 乘 V 可以被重新解释为找到一个最优的线性算子 C,使得 Q_tilde 约等于 C 乘 K_tilde(在频谱空间中)。这种解释的关键优势在于:它将离散的点对点匹配问题转化为连续的函数空间对应问题;通过学习到的基函数,能够适应数据的内在几何、语义和物理结构;最小二乘求解天然支持正则化,通过参数 lambda 可以控制解的稳定性;整个计算过程是连续的,对分辨率变化具有鲁棒性。与 Transolver 等方法的区别在于,Transolver 学习的是物理感知的基函数,然后在这些基上应用标准注意力,而 FUNCATTN 直接在函数空间层面定义注意力,不依赖于物理特定的token化或域依赖的切片。
方法步骤详情
方法包含以下完整步骤:基函数计算阶段,给定输入 X 属于 R 的 n 乘 d 维,通过线性变换 B = Softmax(Linear(X)) 得到基函数矩阵,其中 Phi 和 Psi 分别对应查询空间和键值空间,它们是 n 乘 k 的矩阵,每一列代表一个基函数在 n 个离散点上的评估值。基函数的学习通过简单的全连接层实现,softmax 操作确保了权重有界且和为一。频谱系数投影阶段,计算查询、键、值的频谱系数 Q_tilde = Phi_dagger 乘 K、K_tilde = Psi_dagger 乘 K、V_tilde = Psi_dagger 乘 V。在实现中使用转置 Phi_dagger 代替 Moore-Penrose 伪逆,以获得更好的数值稳定性和运行时复杂度。算子求解阶段,在频谱空间中求解Tikhonov正则化的最小二乘问题,目标是找到 C 使得 Q_tilde 减去 C 乘 K_tilde 的 F 范数平方加上 lambda 乘 C 的 F 范数平方最小,得到闭式解 C_star = Q_tilde 乘 K_tilde_dagger 乘 (K_tilde 乘 K_tilde_dagger + lambda 乘 I_k) 的负一次方。正则化参数 lambda 大于 0 用于数值稳定。传输与反投影阶段,用算子 C_star 将值的频谱系数投影到查询空间,然后反投影得到最终输出 FUNCATTN(Q, K, V) = Phi 乘 C_star 乘 V_tilde。这些步骤可以堆叠成多个 FUNCATTN 块,形成深层的网络架构。
技术新颖性
技术新颖性体现在三个方面:首先,本文建立了标准注意力公式与函数映射框架之间的原则性连接,这是首次将注意力机制明确表述为函数空间之间的线性算子估计问题。其次,提出了基于学习的自适应基函数,这些基函数通过简单的线性层加 softmax 归一化得到,既保持了计算效率,又能够适应数据的内在结构。理论上证明了这种学习方法可以视为广义的 P0(分段常数)元素,当温度参数 tau 趋于 0 时退化为经典的硬分区基函数。第三,提供了严格的理论分析,证明了 FUNCATTN 的局部 Lipschitz 连续性,其 Lipschitz 常数由正则化参数 lambda 控制,这为方法的稳定性提供了理论保证。此外,证明了 FUNCATTN 等价于可学习的积分算子,建立了与传统神经算子方法的联系。
实验结果
实验结果在多个任务上验证了 FUNCATTN 的有效性。在正弦回归任务中,使用仅 4 个观测点时,本文方法比标准注意力方法的误差低三个数量级,比 Transolver 低一个数量级。更重要的是,本文方法表现出优异的泛化能力:在不同观测数量下,误差始终最低,即使只有 5 个观测点时,本文方法的误差也比标准注意力方法用 40 个观测点时更低。在 PDE 求解的六个基准测试中(Elasticity、Airfoil、Darcy、Pipe、Navier-Stokes、Plasticity),本文方法在五个任务上达到最优性能。具体而言,与 Transolver 相比,相对误差提升在 6% 到 26.3% 之间。在 3D RNA 分割任务中,本文方法达到 89.0% 的准确率,比 Transolver 高 1.5 个百分点,比 PointNet++ 高 14.6 个百分点。在带缺口的三角形域上的 Darcy 流测试中,本文方法的相对 L2 误差为 0.64%,比专门为复杂几何 PDE 设计的 WNO 方法相对提升 30.9%,而基于网格的频谱方法(如 dgFNO+)表现显著较差(7.82%),说明固定笛卡尔基函数不适合具有尖锐局部特征的非矩形域。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| PDE求解 | 相对L2误差 (乘100, 向下) | 0.50 | Transolver: 0.64, FNO: 1.08 | 相比Transolver提升21.9% |
| PDE求解 | 相对L2误差 (乘100, 向下) | 0.43 | Transolver: 0.53, LNO: 0.54 | 相比Transolver提升18.9% |
| PDE求解 | 相对L2误差 (乘100, 向下) | 0.42 | Transolver: 0.57, LNO: 0.60 | 相比Transolver提升26.3% |
| PDE求解 | 相对L2误差 (乘100, 向下) | 0.29 | Transolver: 0.31, LNO: 0.25 | 相比Transolver提升6.5% |
| PDE求解 | 相对L2误差 (乘100, 向下) | 8.00 | Transolver: 9.44, LNO: 8.45 | 相比Transolver提升15.3% |
| PDE求解 | 相对L2误差 (乘100, 向下) | 0.11 | Transolver: 0.13, LSM: 0.25 | 相比Transolver提升15.4% |
| RNA分割 | 准确率 (向上) | 89.0% | Transolver: 87.5%, PointNet++: 74.4% | 相比Transolver提升1.5个百分点 |
| 零样本超分辨率 | 相对L2误差 (乘1e3, 向下) | 1.081 | FNO: 1.195, Transolver: 1.243 | 相比FNO提升9.5% |
| 复杂几何PDE | 相对L2误差 (向下) | 0.64% | WNO: 0.92%, dgFNO+: 7.82% | 相比WNO提升30.9% |
局限与改进
作者承认的局限性包括:学习到的基函数使用简单的 softmax 投影,探索更具表达力或结构化的设计仍是开放方向;虽然功能注意力显示出对算子学习有利的归纳偏置,但缺乏严格的理论分析,如近似保证或泛化界;正式建立压缩比 k 除以 n 与近似误差之间的关系将进一步加强理论基础。此外,本文方法的局限性还包括:当基函数数量 k 增加时,计算复杂度会以 O(k³) 增长,虽然 k 通常设置得较小(如 64),但在需要高表达力的场景下可能成为瓶颈;对于具有极高频特征的函数,可能需要较大的 k 才能充分表达,这与计算效率之间存在权衡;虽然实验表明学习到的基函数表现良好,但缺乏对不同数据分布下基函数学习行为的深入分析。
独立分析的弱点
独立分析的弱点包括:在处理具有极端高频特征的函数时(如湍流中的小尺度涡旋),较小的基函数数量可能无法充分捕捉细节,而增加 k 又会带来计算负担。改进方向可以是探索层次化的基函数设计,类似小波变换的多分辨率表示。当前基函数的学习仅基于简单的线性层和 softmax,缺乏对基函数正交性或其他几何结构的显式约束,这在某些需要特定归纳偏置的任务中可能限制表现。改进方向可以是引入软约束或正则化项来引导基函数的学习。理论上证明了 Lipschitz 连续性,但缺乏对泛化误差的定量分析,这使得在实践中选择超参数(如正则化参数 lambda、基函数数量 k)时缺乏理论指导。改进方向是建立更完善的理论框架,提供泛化界或误差估计。方法在算子学习任务上表现出色,但在其他领域(如自然语言处理)的适用性尚未验证,这可能限制了方法的广泛应用。改进方向是探索函数空间解释在非传统算子学习任务中的应用。
未来方向
作者提出的未来研究方向包括:探索更具表达力或结构化的基函数设计,如层次化基、稀疏基或具有特定几何结构的基;建立更完善的理论分析,包括近似保证、泛化界和压缩比 k 除以 n 与近似误差的关系;探索其他正则化方法(如 L1 惩罚)在特定应用中的潜在改进;研究功能注意力在函数空间解释不那么直接的领域(如自然语言处理)中的应用。基于本文成果可以延伸的研究方向包括:将功能注意力与其他注意力机制结合,形成混合架构,在不同层次使用不同的注意力方式;探索自适应的基函数数量选择机制,根据输入函数的复杂度动态调整 k;研究功能注意力在反向传播中的梯度性质,优化训练过程;扩展到更广泛的算子学习任务,如随机偏微分方程、时空演化等。
复现评估
复现评估:作者表示项目页面在 https://github.com/xjffff/FUNCATTN,表明代码开源。实验均在单个 Nvidia A40 GPU 上进行,重复三次以确保稳定性。论文提供了详细的实现细节(附录 C),包括网络架构、超参数设置和训练过程。基函数数量 k 的默认值为 64,在各种数据集上表现稳健(在最佳结果的 5% 以内)。对于较平滑的场(Darcy、Pipe),k = 32-64 足够;对于高频场(Elasticity、Navier-Stokes),k = 128-256 可进一步提升性能。正则化参数 lambda 的选择影响数值稳定性和条件数,作者在附录中提供了敏感性分析。总体而言,本文的复现难度中等,主要挑战在于理解函数空间投影的概念和实现基函数学习的正确性。
论文图表