[论文解读] Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

John Guibas, Morteza Mardani+4 more|arXiv (Cornell University)|Nov 24, 2021

CCD and CMOS Imaging Sensors被引用 78

一句话总结

本文提出自适应傅里叶神经算子（AFNO），作为视觉变换器的高效令牌混合器，将基于傅里叶的全局卷积与自适应、稀疏、块状通道混合相结合，实现准线性复杂度，并在下游任务上取得改进，尤其是在高分辨率与小样本分割任务中。

ABSTRACT

Vision transformers have delivered tremendous success in representation learning. This is primarily due to effective token mixing through self attention. However, this scales quadratically with the number of pixels, which becomes infeasible for high-resolution inputs. To cope with this challenge, we propose Adaptive Fourier Neural Operator (AFNO) as an efficient token mixer that learns to mix in the Fourier domain. AFNO is based on a principled foundation of operator learning which allows us to frame token mixing as a continuous global convolution without any dependence on the input resolution. This principle was previously used to design FNO, which solves global convolution efficiently in the Fourier domain and has shown promise in learning challenging PDEs. To handle challenges in visual representation learning such as discontinuities in images and high resolution inputs, we propose principled architectural modifications to FNO which results in memory and computational efficiency. This includes imposing a block-diagonal structure on the channel mixing weights, adaptively sharing weights across tokens, and sparsifying the frequency modes via soft-thresholding and shrinkage. The resulting model is highly parallel with a quasi-linear complexity and has linear memory in the sequence size. AFNO outperforms self-attention mechanisms for few-shot segmentation in terms of both efficiency and accuracy. For Cityscapes segmentation with the Segformer-B3 backbone, AFNO can handle a sequence size of 65k and outperforms other efficient self-attention mechanisms.

研究动机与目标

• 为了在高分辨率视觉变换器中超越平方自注意力复杂度，激发高效的令牌混合机制。
• 通过将傅里叶神经算子应用于处理具有不连续性和高分辨率的图像，来开发 AFNO。
• 引入架构修改（块对角权重、Adaptive 权重共享、稀疏性）以提高效率和泛化性。
• 展示 AFNO 在上采样/修复预训练、小样本分割，以及 Cityscapes/高分辨率分割中的有效性。
• 提供代码和与现有混合器的经验比较。

提出的方法

• 将令牌混合框定为通过算子学习的连续全局卷积。
• 采用带离散化 FFT 实现的傅里叶神经算子（FNO）进行令牌混合。
• 在通道混合权重上施加块对角结构以降低参数量。
• 通过一个跨令牌的共享两层多层感知器引入自适应权重共享。
• 对频率模态应用软阈值（稀疏性），以聚焦于信息量大的分量。
• 允许自适应权重掩蔽和残差局部项以处理非周期边界和局部特征。

实验结果

研究问题

• RQ1 AFNO 是否在视觉变换器上实现与自注意力具有竞争力甚至优越的准确性，同时显著降低计算成本？
• RQ2 强制块对角通道结构、权重共享和稀疏性是否有助于高分辨率图像任务的泛化和效率？
• RQ3 在预训练任务（修复）、下游的少样本分割以及高分辨率 Cityscapes 分割中，AFNO 相较于现有混合器的表现如何？
• RQ4 AFNO 是否具有分辨率不变性，且在非常长的令牌序列上也能实现可扩展性且不会带来难以承受的内存消耗？

主要发现

• AFNO 在图像修复中实现了具有竞争力的 PSNR/SSIM，同时使用的 GFLOPs 比自注意力低约 30%。
• 在少样本分割中，AFNO 在 2/3 的数据集上超越自注意力，在其他数据集上与之并驾齐驱，且 FLOPs 更低。
• 在 Cityscapes 的 SegFormer-B3 设置下，AFNO-100% 达到 mIoU 80.9，优于在可比设置下的 Efficient Self-Attention、LS 与 GFN 变体。
• ImageNet-1K 分类（ViT-S/4）显示 AFNO 的 Top-1 准确率为 80.89%（自注意力为 80.87%），参数量与 FLOPs 相近。
• 消融研究表明稀疏性（lambda ~0.01）、8 块及自适应权重共同提升了 AFNO 相对于 FNO 和非自适应变体的性能。
• AFNO 展示了分辨率不变性与高效通道混合，为高分辨率视觉任务提供了传统自注意力的有力替代方案。