FFTNet实战：5分钟教你用快速傅里叶变换优化ViT模型（附PyTorch代码）

原创

于 2026-02-27 03:40:45 发布 · 599 阅读

标签

FFTNet实战：5分钟教你用快速傅里叶变换优化ViT模型（附PyTorch代码）

最近在复现一些长序列视觉任务时，我又一次被Transformer那惊人的显存占用给“教育”了。一张高分辨率图片，patch序列轻松破万，传统的自注意力机制瞬间就成了算力黑洞。就在我对着O(n²)的复杂度曲线发愁时，一篇新鲜出炉的论文进入了视野——它把信号处理领域的老将快速傅里叶变换（FFT） 请了回来，在频域里重新设计了注意力机制。今天，我就带你绕过那些复杂的公式，直接上手，看看如何把这种名为FFTNet的模块，像乐高积木一样，“咔哒”一声塞进你现有的Vision Transformer项目里，换来效率的显著提升。

这篇文章面向的是已经熟悉ViT基本结构，但在实际部署中受困于计算和内存开销的工程师和研究者。我们不会停留在理论探讨，而是聚焦于实操：从核心代码块的解读，到完整替换ViT中自注意力层的步骤，再到我本人在本地环境下的性能实测对比。你会发现，这种“即插即用”的改造，远比想象中简单。

1. 理解核心：为什么是频域？FFTNet解决了什么痛点？

在深入代码之前，我们得先搞明白一个根本问题：为什么要把视觉Transformer的注意力计算搬到频域去？这背后直指传统自注意力机制的一个阿喀琉斯之踵——二次方复杂度。

想象一下，你有一张1024x1024的图片，采用16x16的patch划分，你会得到4096个视觉标记（token）。标准的多头自注意力需要计算这4096个标记中每一对之间的关联度，这个计算量随着序列长度n的增长是O(n²)。当n达到数千甚至上万时（例如处理医学图像或长文档），无论是训练时的显存占用还是推理时的延迟，都变得难以承受。

而快速傅里叶变换（FFT） 提供了一条“捷径”。它能在O(n log n)的时间内，将一个序列从时域（或空间域）转换到频域。在频域里，信号的全局特性被分解到不同的频率分量上。FFTNet的核心思想就在于：与其在原始空间里费力地计算所有标记对之间的交互，不如在频域里对这些频率分量进行自适应滤波。这种滤波操作是逐元素（element-wise） 的，复杂度是线性的O(n)。一次FFT变换（O(n log n)）加上一次频域滤波（O(n)），整体复杂度依然控制在O(n log n)量级，对于长序列来说，这是数量级的优势。

注意：这里有一个关键但常被误解的点。FFTNet并非直接“计算”注意力权重，而是通过可学习的滤波器在频域对信号进行调制，这相当于隐式地实现了全局信息混合。它牺牲了标准注意力那种显式的、可解释的“谁关注谁”的权重矩阵，换来了极高的计算效率。对于许多视觉任务，这种全局混合能力已经足够。

那么，这种转换会丢失信息吗？根据帕塞瓦尔定理，信号在时域的总能量等于其在频域的总能量（仅差一个常数因子）。这意味着，从数学上讲，这种变换是保能量的，信息本身没有丢失，只是换了一种更容易进行某些操作（如滤波）的表示形式。

2. 核心模块拆解：手把手实现FFTNet Block

理论说得再多，不如一行代码来得实在。FFTNet最吸引人的就是其模块化设计，我们可以先从一个最基础的、可独立运行的块开始理解。下面这个 FFTNetBlock 类，就是整个体系的基石。

import torch
import torch.nn as nn
import torch.nn.functional as F

class ModReLU(nn.Module):
    """
    复数域激活函数。标准ReLU无法直接处理复数，ModReLU通过调整复数的幅度来实现非线性。
    它保持相位不变，仅对幅度进行ReLU式的阈值处理。
    """
    def __init__(self, features):
        super().__init__()
        # 一个可学习的偏置参数，用于调整相位阈值
        self.b = nn.Parameter(torch.Tensor(features))
        nn.init.uniform_(self.b, -0.1, 0.1)

    def forward(self, z):
        # z是复数张量
        magnitude = torch.abs(z)
        phase = torch.angle(z)
        # 核心操作：对幅度应用ReLU，但阈值受到相位偏置b的影响
        magnitude_activated = F.relu(magnitude + self.b)
        # 用激活后的幅度和原始相位重构复数
        return magnitude_activated * torch.exp(1j * phase)

class FFTNetBlock(nn.Module):
    """
    基础的FFTNet模块。
    输入: [batch_size, sequence_length, feature_dim]
    输出: [batch_size, sequence_length, feature_dim]
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        # 频域滤波器：两个独立的线性层分别处理实部和虚部（也可用复数线性层）
        self.filter_real = nn.Linear(dim, dim, bias=False)
        self.filter_imag = nn.Linear(dim, dim, bias=False)
        self.modrelu = ModReLU(dim)

    def forward(self, x):
        # 1. 时空域 -> 频域
        # torch.fft.fft 默认在最后一个维度做FFT，我们需在序列维度（dim=1）操作
        x_fft = torch.fft.fft(x, dim=1)

        # 2. 频域自适应滤波
        # 分别对实部和虚部进行线性变换
        real_part = self.filter_real(x_fft.real) - self.filter_imag(x_fft.imag)
        imag_part = self.filter_real(x_fft.imag) + self.filter_imag(x_fft.real)
        x_filtered = torch.complex(real_part, imag_part)

        # 3. 频域非线性激活
        x_activated = self.modrelu(x_filtered)

        # 4. 频域 -> 时空域
        x_out = torch.fft.ifft(x_activated, dim=1).real  # 取实部作为输出
        return x_out

我们来逐段解析这个 forward 过程：