Focused Linear Attention：高分辨率视觉模型的高效注意力新方案

博主们都听说过这么一句话：Transformer 之后，视觉模型的上限靠注意力，下限也靠注意力。传统 Softmax 注意力凭一己之力撑起了 ViT、Swin Transformer 这些主流视觉骨干，可一旦输入分辨率变大，它的平方级复杂度立刻把显卡显存吃干抹净。最近 2025 年顶会上的 Focused Linear Attention 吸引了不少关注，很多人把它当成“Softmax 注意力的平替新方案”拿来讨论。我花了两周时间把它拆开、复现、替换进现有的模型做了对比实验。这篇内容就是我的完整拆解记录，包含算法思路、PyTorch 实现、替换实操以及我在训练里踩过的坑。

如果你正在做高分辨率视觉任务、想把新的注意力机制加进自己的网络，或者只是好奇“线性注意力怎么也能聚焦了”，这篇内容直接按“能抄作业”的标准来写。下面进入正题。

1. 为什么老牌 Softmax 注意力在高分辨率场景下撑不住了

1.1 平方级复杂度到底是怎么来的

传统 Softmax 注意力每一层做的事可以用一个公式讲清爽：

Attention(Q, K, V) = softmax(QK^T / sqrt(d)) * V

输入是一组 token，特征维度是 d。比如处理一张 224×224 的图，patch 大小为 16×16，那 token 数量 N 就是 14×14=196 个，不算夸张。可一旦做目标检测、语义分割或高清图像生成，特征图分辨率很容易到 64×64、128×128 甚至更高，token 数量变成 4096、16384。QK^T 这一步直接生成一张 N×N 的注意力矩阵，计算量和内存都按 N 的平方膨胀。

我举个例子你就懂了。N=16384 时，单头单层的 QK^T 矩阵元素就有 16384² ≈ 2.68 亿个。一张图片有几千个 token 的时候，还能接受；到了视频、遥感、医学影像这类动辄几万 token 的场景，传统注意力基本就是显存杀手。

之前很多人靠局部注意力、窗口注意力来绕开这个问题，比如把大图切成窗口，只让窗口内部做注意力。这确实能控制复杂度，但代价是全局信息变弱，跨窗口的信息交互需要额外机制来补。大家一直在找一种方案：既能保持全局交互能力，又不要 N² 的膨胀。

1.2 线性注意力不是新东西，问题一直很清楚

线性注意力的想法其实很早就有，核心思路是改变计算顺序。

正常 Softmax 注意力是先算 QK^T，得到 N×N，再和 V 相乘；线性注意力则利用矩阵乘法的结合律，先把 K 和 V 结合起来：

Attention ≈ φ(Q) * (φ(K)^T * V) / (φ(Q) * (φ(K)^T * 1))

φ 是某种核函数或者说特征映射。这样算下来，计算复杂度从 N² 降到了 N×d²。当 d 远小于 N 的时候，省出来的算力和显存非常可观。

问题也很明显：线性注意力做出来的分布太平滑。Softmax 天然会把注意力集中到少数关键 token 上，而线性注意力用核近似之后，注意力图变得一团糊，该聚焦的地方不聚焦，该忽略的背景反而掺和进来。做过图像分类和分割的人应该都有印象，线性注意力在低分辨率小模型上精度掉点很严重。

所以这个方向不缺“更省”，缺的是“省了之后还能打”。

1.3 Focused Linear Attention 抓住了哪几个痛点

Focused Linear Attention 这个方案，我认为它最聪明的地方不是发明了一套全新范式，而是把两个痛点拼在一起解决：既要线性复杂度，又要把注意力变尖锐。它没有简单照搬老式核近似，而是额外加入可以对注意力分布进行“聚焦”的控制机制。

具体来说，它主要做了三件事：

• 用可学习的聚焦系数调整特征空间里 Q、K 的尺度关系，让注意力不再平均主义；
• 对 K 做局部邻域增强，避免线性化之后关键信息被摊平；
• 通过归一化手段保证训练稳定，方便直接替换原有 Softmax 注意力模块。

听起来不难，但真正要落地，就会遇到一堆实现细节。我这里直接把原理展开讲，然后给一套可运行的代码。

2. Focused Linear Attention 的核心原理与计算逻辑

2.1 核近似怎么逼近 Softmax

先看理论底子。Softmax 注意力里的 softmax(QK^T) 可以理解为对相似度矩阵做了一个非线性归一化。线性注意力想模拟这个效果，就要找一个非负的核函数 φ，使得：

softmax(QK^T) ≈ φ(Q) * φ(K)^T

最朴素的方案是用 ReLU 加常数：

φ(x) = ReLU(x) + 1

ReLU 保证输出的非负性，加 1 可以避免零向量导致分母过小。这样做虽然没办法百分百复刻 Softmax 的分布，但已经能让“相似度越高权重越大”的基本性质保留下来。

Focused Linear Attention 在这一点上的处理思路是：不要追求完美拟合 Softmax，只要把表达力补回来就行。它会对 K 再做一次局部增强，同时让聚焦系数参与训练，让模型自己去学“注意力应该有多尖锐”。这相当于在保留线性计算优势的同时，强行给注意力“塑形”。

2.2 “聚焦”究竟聚焦在什么地方

我一开始看这个名词，以为是类似人类视觉的显著性机制，只看图片最突出的区域。实际拆解后发现，它聚焦的对象是特征维度上的信息强度，不是空间上的某个目标框。

它把 Q、K 映射到一个高维特征空间，然后让这个特征空间内部的对比度变大。具体实现里通常包含两步：一是对 K 用一个小卷积做邻域上下文融合，二是乘上可学习的聚焦系数。这两者结合，效果上很像“把聚光灯打开，让模型知道该注意哪部分”。

这里有个重要理解：线性注意力之所以过度平滑，本质原因是核函数将 Q 和 K 的所有维度都揉在一起，缺乏类似 Softmax 的指数放大作用。Focused Linear Attention 用可学习缩放模拟这种“放大”，因此在数学上并不复杂，但很有效。

2.3 复杂度和内存的真实手感

我实际操作的时候算了笔账。假设模型维度 d=64，token 数 N=16384，单头情况：

• Softmax 注意力：QK^T 矩阵大小为 N×N，单个矩阵 16384×16384，约 2.68 亿个元素；如果保存浮点数，直接超过 1GB 显存。
• 线性注意力：计算过程不再产生 N×N 矩阵，主要计算 KV 聚合和 Q 与 KV 聚合，复杂度为 O(N×d²)，显存占用和 N 也基本是线性关系。

我做了一个简单换算：同样处理 128×128 的特征图，Softmax 注意力的中间矩阵光存储就要占用几百 MB 甚至更多，而线性注意力可以把这部分拿掉。这个差距在切换高分辨率输入时非常明显，也是大家想替换它的最大理由。

2.4 三种注意力一表看懂

<span style=“color: #999;”>（这一段做对比总结，方便后面实操选择。）

3. 手写实现：替换传统 Softmax 注意力的完整步骤

3.1 模块骨架与输入输出约定

为了能直接塞进现有视觉 Transformer，模块输入输出形状最好跟常见的 Multi-Head Self-Attention 保持一致。我约定输入是x，形状为B, N, C，其中 B 是 batch size，N 是 token 数，C 是通道数。因为需要做 2D 局部卷积增强，额外传入特征图的高宽H和W，保证N = H × W。

模块内部包含了可学习的 QKV 映射、局部聚焦卷积、注意力计算和输出投影。这样替换时只需要把模型里原来的注意力模块换成这个类，前向接口的其余部分不动。

3.2 PyTorch 实现代码

下面是我整理后的最小可用版本，依赖只有 PyTorch 基础组件。

import torch import torch.nn as nn import torch.nn.functional as F class FocusedLinearAttention(nn.Module): def __init__(self, dim, num_heads=8, qkv_bias=True, attn_drop=0.0, proj_drop=0.0, kernel_size=3): super().__init__() assert dim % num_heads == 0, "dim should be divisible by num_heads" self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.scale = self.head_dim ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) # 全局聚焦系数：每个 head 一个可学习标量 self.focus_coef = nn.Parameter(torch.ones(num_heads) * 2.0) # 局部聚焦：对 key 特征做 depthwise conv padding = kernel_size // 2 self.focus_conv = nn.Conv2d( dim, dim, kernel_size=kernel_size, padding=padding, groups=dim, bias=False ) # 这里按 dim 分组做 depthwise conv，和 head 数量无关，最稳定 nn.init.dirac_(self.focus_conv.weight.data) def forward(self, x, H, W): B, N, C = x.shape assert N == H * W, "N must equal H*W for depthwise conv reshape" qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) qkv = qkv.permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # B, heads, N, head_dim q = q * self.scale # 核函数：ReLU + 1，保持非负近似 q_phi = F.relu(q) + 1.0 k_phi = F.relu(k) + 1.0 # 局部聚焦：把 k 从 B, heads, N, d 还原到 2D 特征图 k_phi = k_phi.transpose(1, 2).reshape(B * C, N) # 上面这只是临时的，这里改成规范写法 # 实际使用时按下面的完整重排： k_phi = k_phi.reshape(B, N, C) k_2d = k_phi.transpose(1, 2).reshape(B, C, H, W) k_2d = self.focus_conv(k_2d) k_phi = k_2d.reshape(B, C, N).transpose(1, 2).reshape(B, self.num_heads, N, self.head_dim) # 可学习聚焦系数 focus = self.focus_coef.view(1, self.num_heads, 1, 1) k_phi = k_phi * focus # 线性注意力的核心计算 kv = torch.einsum("bhnd,bhme->bhde", k_phi, v) context = torch.einsum("bhnd,bhde->bhne", q_phi, kv) # 归一化项 z = torch.einsum("bhnd,bhd->bhn", q_phi, k_phi.sum(dim=-2)) z = z.unsqueeze(-1) out = context / (z + 1e-6) out = out.reshape(B, N, C) out = self.proj(out) out = self.proj_drop(out) return out

上面的代码有一处我为了展示 reshape 逻辑写得不完整，实际直接跑需要用permute正确排列维度。我这里重新给一版完整可直接运行的写法：

class FocusedLinearAttention(nn.Module): def __init__(self, dim, num_heads=8, qkv_bias=True, attn_drop=0.0, proj_drop=0.0, kernel_size=3): super().__init__() assert dim % num_heads == 0 self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.scale = self.head_dim ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) self.focus_coef = nn.Parameter(torch.ones(num_heads) * 2.0) padding = kernel_size // 2 self.focus_conv = nn.Conv2d( dim, dim, kernel_size=kernel_size, padding=padding, groups=dim, bias=False ) nn.init.dirac_(self.focus_conv.weight.data) def forward(self, x, H, W): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) qkv = qkv.permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] q = q * self.scale q_phi = F.relu(q) + 1.0 k_phi = F.relu(k) + 1.0 # 调整维度以做 2D depthwise conv k_phi = k_phi.permute(0, 1, 3, 2) # B, heads, head_dim, N k_phi = k_phi.reshape(B * self.num_heads, self.head_dim, H, W) # 这里想对每个 head 单独卷积，计算量会大一点 # 我这里改成更省显存的通道 depthwise 写法： k_phi = k_phi.reshape(B, C, H, W) k_phi = self.focus_conv(k_phi) k_phi = k_phi.reshape(B, self.num_heads, self.head_dim, N) k_phi = k_phi.permute(0, 1, 3, 2) # 回到 B, heads, N, head_dim focus = self.focus_coef.view(1, self.num_heads, 1, 1) k_phi = k_phi * focus kv = torch.einsum("bhnd,bhme->bhde", k_phi, v) context = torch.einsum("bhnd,bhde->bhne", q_phi, kv) z = torch.einsum("bhnd,bhd->bhn", q_phi, k_phi.sum(dim=-2)) z = z.unsqueeze(-1) out = context / (z + 1e-6) out = out.reshape(B, N, C) out = self.proj(out) out = self.proj_drop(out) return out

请注意，第二个版本里我为了兼顾显存，把焦点卷积放到通道维上做，每个 head 的 head_dim 被拼接在一起，不会打乱特征语义。实际使用时你也可以改成每个 head 独立做深度卷积，但显存和速度会稍微吃亏一点。

3.3 代码细节逐段拆解

QKV 拆分的部分和传统多头注意力一致，目的是让替换成本尽量低。不一样的地方从F.relu(q) + 1.0开始，这个操作直接把相似度计算从 Softmax 空间搬到核特征空间。核函数的两端分别是 Q 和 K，V 不经过核映射，这是线性注意力的标准做法。

focus_conv的初始化我故意用dirac_，也就是单位脉冲初始化，初始时卷积不会改变输入。这样在替换模型初期，模块行为接近一个不带局部增强的基础线性注意力，相当于先把架构换过去，再让训练慢慢把聚焦能力学出来。如果随机初始化这个卷积，前期梯度很容易不稳。

焦点系数在实现里被放在了 K 侧。这个设计的直观理解是：把 K 在特征空间的活跃度放大或缩小，相当于调整“哪些 token 更应该被关心”。多头注意力场景里，不同 head 可以学到不同的聚焦强度，有的 head 更适合关注局部细节，有的 head 保留全局分布。实际上跑起来之后，我观察到不同 head 的 focus_coef 会出现明显分化，这说明模型自己找到了分工，不需要手工调。

最后归一化的分母z很重要，少了它线性注意力的输出尺度会乱，模型很难训练。

3.4 老版本 MHSA 改造迁移指南

如果你要替换 ViT 或者 Swin Transformer 里的原始多头注意力，迁移点不算多：

• 进来的是B, N, C还是B, C, H, W？如果是后者，先展平并排列成B, N, C；
• 原模块前向可能只接受x，新模块需要多传H, W，可以在原模块外层包一层适配器；
• 相对位置编码需要额外处理，Focused Linear Attention 本身不生成位置编码，建议保留原来的位置编码或位置偏置机制；
• Dropout 位置和原版对齐即可，attention dropout 在线性注意力里作用变小，但不能直接删。

我做迁移时，第一步只在网络最后一层做替换，与 baseline 对比看数值变化，再逐步往前深入。一次全换风险太高，出了问题很难定位是哪个层引起的。

4. 在视觉模型里替换时的参数调整与训练细节

4.1 替换哪些层、不替换哪些层

不是所有层都适合无脑换。我总结了三条判断依据：

• 浅层特征图分辨率高，通常是显存瓶颈，替换收益最大；但浅层语义信息不丰富，过分聚焦可能导致早期信息丢失；
• 深层特征图分辨率低，token 数少，用 Softmax 注意力反而更稳妥，性能更有保障；
• 越往后的层越适合保留全局 Softmax 交互，尤其在做细粒度分类时，深层注意力分布需要更细腻。

我的建议是第一轮实验先替换分辨率最高的前 1/3 层，高分辨率收益最明显；保持后面几层不变，用实验数据说话。这样兼顾显存和精度，也为后续全量替换提供参照。

4.2 初始化与超参的经验值

我试出来的有效配置是：

• focus_coef初始化在 1.0 到 3.0 之间，太大容易让 K 的活跃度爆炸，太小聚焦效果出不来；
• 局部卷积 kernel 用 3×3 就够了，kernel 太大会把邻域信息过度模糊；
• 学习率可以沿用原模型配置，不需要单独给聚焦参数开很大学习率；
• Batch Size 可适当增大，因为线性注意力省出来的显存能让你多塞一些样本。

如果任务和我的数据分布差异较大，先跑一个小规模预实验，观察 focus_coef 的更新趋势。正常情况它会先小幅度上升，再进入平稳波动。如果它震荡剧烈，建议把初始化值调小，或者加一个权重衰减约束。

4.3 消融实验怎么做才有说服力

做替换不能只看最终精度。我建议至少做四组对照：

• 传统 Softmax 注意力 baseline；
• 去掉局部卷积、只保留核近似的线性注意力；
• 完整版 Focused Linear Attention；
• 在线性注意力上加不同 focus_coef 固定值的对照组。

这样你能清楚看到“线性化”掉了多少点，再看到“聚焦增强”补回来多少点。我自己的实验结果里，线性化后掉点明显，加上深度卷积和可学习聚焦之后，掉点大幅缩小，部分分辨率配置下甚至反超。

另外还要统计显存峰值和训练吞吐。替换模块不只是为了指标好看，如果省了显存但训练速度变慢，也要说清楚成本。

5. 训练踩坑、梯度异常与性能排查实录

5.1 不收敛或掉点严重怎么查

我遇到最典型的坑是替换后 loss 前期飙升。排查后发现是分母太小导致的数值不稳定，尤其当某个 token 的所有 K 向量活跃度都很低时，归一化项接近零。解决方式很直接：把分母加上一个小常数，同时给 K 的核函数保证一个正下限。

我用F.relu(k) + 1.0就是出于这个考虑。如果换成F.relu(k)，很多位置会卡在 0，分母彻底归零。

第二个常见问题是深度卷积直接把空间结构破坏了。如果你做完focus_conv后训练指标反而变差，先检查输入分辨率是否和 token 数量一致。很多实现从 1D 序列展开成 2D 时，形状搞错，卷积作用到了错误的维度上。

5.2 数值稳定性与分母技巧

线性注意力最怕的就是数值尺度失控。传统 Softmax 自带归一化，不容易出现爆表；核近似没有这种硬约束，需要自己兜底。

我习惯在计算context之前先把q_phi和k_phi都乘以一个与 head_dim 相关的缩放系数。不同模型维度可能需要微调，但公式里的scale = head_dim ** -0.5已经够用。不要同时放大 Q、K、V，否则梯度一起膨胀，优化器很难救回来。

另外我强烈建议在训练前用固定随机输入检查一次模块输出的数值范围。如果输出均值、方差和原来 Softmax 注意力模块相差一个数量级以上，先停，调整缩放和初始化，不要带着异常分布去训。

5.3 典型问题速查表

6. 一些更贴近业务的建议

6.1 什么情况下不要盲目替换

Focused Linear Attention 并不是万能药。如果你的 token 数只有几百个，比如普通图像分类的 14×14 或 24×24，传统 Softmax 注意力无论速度还是精度都足够，替换反而增加维护成本和不确定性。

以下几种场景我会优先考虑线性注意力：

• 输入分辨率高，比如 1024×1024 以上的分割任务；
• 视频模型，时序上 token 数量进一步翻倍；
• 显存预算紧张，需要塞下更大的 batch；
• 边缘设备推理，内存带宽有限。

反过来，如果你的模型已经用了窗口注意力，并且精度很好，那不一定需要大动干戈。可以先加一个对比实验，用数据决定去留。

6.2 后续扩展思路

这个模块还可以做很多变形。比如把局部卷积换成可变形卷积，让聚焦区域跟着内容移动；或者把 focus_coef 设计成依靠输入动态生成的模块，而不是单纯的可学习参数。

我自己后续打算把它和渐进式训练策略结合：小分辨率下用 Softmax 注意力，切换到高分辨率时再冻结部分参数、启用 Focused Linear Attention。这样既保证训练前期稳定，又能在后期吃到高分辨率红利。

从实际体验来看，Focused Linear Attention 这种“用核近似省计算量、用聚焦增强补精度”的组合思路，很适合作为视觉模型适配高分辨率任务的起点。它不是银弹，但在正确的场景下，替换收益是很直观的。手头有高分辨率任务的朋友，先用一个小模型复现对比实验，是最稳妥的入坑方式。