点击下方卡片，关注“AI前沿速递”公众号
点击下方卡片，关注“AI前沿速递”公众号
各种重磅干货，第一时间送达
各种重磅干货，第一时间送达
标题：Fast Vision Transformers with HiLo Attention
论文链接：https://arxiv.org/abs/2205.13213
代码链接：https://github.com/ziplab/LITv2
创新点
HiLo Attention机制
高频/低频分离：提出了一种新的自注意力机制HiLo，将图像的高频和低频信息分离处理。高频信息（Hi-Fi）通过局部窗口自注意力捕捉细节，低频信息（Lo-Fi）通过全局注意力处理整体结构。这种分离方式能够更高效地处理图像特征，同时减少计算复杂度。
高效设计：HiLo通过分离自注意力头（heads），将高频和低频信息分别处理，显著降低了计算复杂度（FLOPs），并且在GPU和CPU上表现出更快的速度。
LITv2模型架构
基于直接速度评估的设计原则：与以往基于FLOPs（浮点运算次数）作为效率指标不同，LITv2直接在目标平台上评估模型的速度（throughput），以实现更实际的效率优化。
改进的Transformer架构：在LITv1的基础上，LITv2引入HiLo注意力机制，并通过移除早期的多头自注意力层（MSA），改为深度可分离卷积（Depthwise Convolution）来隐式编码位置信息，进一步提升了模型的效率和性能。
高效的频率处理
高频和低频的独立处理：HiLo通过局部窗口注意力处理高频信息，通过全局注意力处理低频信息，这种设计充分利用了不同频率信息在图像中的作用（高频关注细节，低频关注整体结构），从而在保持性能的同时显著提高效率。
动态窗口大小调整：论文提出通过调整窗口大小来优化HiLo在高分辨率图像上的效率，窗口大小的增加可以进一步降低计算复杂度，同时保持性能。
零填充卷积的位置编码
隐式位置信息编码：为了避免传统位置编码（如相对位置编码）在不同分辨率下的插值问题，LITv2在每个前馈网络（FFN）中引入3×3深度可分离卷积层，通过零填充隐式学习位置信息，同时扩大早期多层感知机（MLP）的感知野。
整体结构
LITv2模型整体架构
Patch Embedding（分块嵌入）
输入图像被分割成固定大小的patch（如16×16像素）。 每个patch通过线性层嵌入到高维空间，形成初始特征表示。 作用：将图像从像素空间转换为特征空间。
LITv2主体由多个阶段（Stage）组成，每个阶段包含：
Deformable Token Merging（可变形Token合并）：
1、对特征图进行下采样，减少空间维度，增加通道数。
2、设计思想：类似金字塔结构，逐步降低分辨率以提取高级特征。
Transformer Block（Transformer块）：
包含两个核心部分：
1、HiLo Attention（HiLo注意力机制）：
分离处理高频（Hi-Fi）和低频（Lo-Fi）信息。 高频通过局部窗口自注意力捕捉细节；低频通过全局注意力处理整体结构。 分离自注意力头，降低计算复杂度。
2、ConvFFN（改进的前馈网络）：
在FFN中引入3×3深度可分离卷积。 优势：隐式编码位置信息，扩大感受野。
最终分类头（Classification Head）
最后一个Transformer阶段后，特征图通过全局平均池化生成全局特征。 全局特征通过线性分类器输出类别预测。
HiLo模块的作用
高效处理高分辨率图像HiLo通过分离高频（Hi-Fi）和低频（Lo-Fi）信息，能够显著减少计算复杂度和内存占用。在高分辨率图像处理任务中（如图像分类、目标检测和语义分割），HiLo可以高效处理局部细节和全局结构，提升推理速度。 提升模型效率和性能HiLo在GPU和CPU上表现出色，能够加快推理速度，特别适用于低延迟应用场景（如无人机图像处理和自动驾驶中的实时感知系统）。此外，HiLo在FLOPs、吞吐量和内存消耗方面优于现有的注意力机制。 增强模型鲁棒性在复杂的场景中，HiLo可以帮助模型更好地适应光照、噪声和遮挡等变化。通过关注不同频率的信息，HiLo能够减少噪声和干扰的影响，提高模型的鲁棒性。 易于集成和扩展HiLo模块具有良好的即插即用特性，可以轻松替换或补充现有的注意力机制。它能够与各种视觉模型（如YOLOv8、CNN等）结合，进一步优化模型结构和性能。
消融实验
HiLo与其他注意力机制的比较基于LITv2-S模型，作者将HiLo与其他高效注意力机制进行了比较，包括PVT中的空间缩减注意力（SRA）、Swin中的移位窗口注意力（W-MSA）以及Twins中的局部和全局交替注意力（T-MSA）。实验结果表明，HiLo在减少FLOPs的同时，实现了更好的性能和更快的推理速度。此外，在不同图像分辨率下的基准测试中，HiLo在FLOPs、吞吐量和内存消耗方面均优于其他注意力机制，例如Focal、QuadTree和Performer。 α的影响在HiLo中，α表示分配给低频（Lo-Fi）的注意力头比例。实验表明，在224×224分辨率和窗口大小为2的情况下，较大的α（如0.9）有助于减少FLOPs，因为更多的注意力头被分配给低频分支。当α=0时（即仅保留高频分支），模型性能显著下降，这表明低频信息在自注意力中非常重要。 架构修改的影响
引入3×3深度卷积层（ConvFFN）：在每个FFN中引入3×3深度卷积层后，模型在ImageNet分类和COCO检测任务上的性能显著提升，同时推理速度（FPS）也有所提高。 移除相对位置编码（RPE）：移除RPE后，模型在密集预测任务中的推理速度显著提升，但性能略有下降。 使用HiLo注意力机制：替换为HiLo注意力机制后，模型在FLOPs和推理速度上进一步优化，同时保持了较高的性能。
窗口大小的影响在COCO目标检测任务中，作者尝试了不同的窗口大小（如s=2和s=4）。结果表明，较大的窗口大小（s=4）可以进一步提高模型效率，但性能略有下降。
总结
消融实验表明，HiLo注意力机制在减少计算复杂度和提升模型效率方面表现出色，同时通过调整α和窗口大小，可以在不同任务中实现性能与效率的平衡。此外，引入深度卷积层和移除RPE等架构修改进一步优化了模型的整体性能。
代码实现
import torchimport torch.nn as nnimport torch.nn.functional as Fclass HiLo(nn.Module):def __init__(self, dim, num_heads=8, window_size=2, alpha=0.9):"""HiLo Attention Mechanism.Args:dim (int): 输入特征的通道数。num_heads (int): 注意力头的数量。window_size (int): 局部窗口的大小（用于Hi-Fi）。alpha (float): 高频和低频头的分配比例。"""super(HiLo, self).__init__()self.dim = dimself.num_heads = num_headsself.window_size = window_sizeself.alpha = alphaself.head_dim = dim // num_headsself.scale = self.head_dim ** -0.5# 高频头数量和低频头数量self.hi_heads = int((1 - alpha) * num_heads)self.lo_heads = int(alpha * num_heads)# 线性投影层self.qkv = nn.Linear(dim, dim * 3, bias=False)self.proj = nn.Linear(dim, dim)def forward(self, x):"""前向传播。Args:x (torch.Tensor): 输入特征图，形状为 (B, N, C)。"""B, N, C = x.shapeq, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)q, k, v = q.contiguous(), k.contiguous(), v.contiguous()# 分离Hi-Fi和Lo-Fi的头q_hi, q_lo = q[:, :, :self.hi_heads, :, :], q[:, :, self.hi_heads:, :, :]k_hi, k_lo = k[:, :, :self.hi_heads, :, :], k[:, :, self.hi_heads:, :, :]v_hi, v_lo = v[:, :, :self.hi_heads, :, :], v[:, :, self.hi_heads:, :, :]# Hi-Fi: 局部窗口自注意力x_hi = self._local_attention(q_hi, k_hi, v_hi, self.window_size)# Lo-Fi: 全局注意力x_lo = self._global_attention(q_lo, k_lo, v_lo)# 合并Hi-Fi和Lo-Fi的输出x = torch.cat([x_hi, x_lo], dim=2).permute(0, 2, 1, 3).reshape(B, N, C)x = self.proj(x)return xdef _local_attention(self, q, k, v, window_size):"""局部窗口自注意力。"""B, _, H, N, C = q.shapeq = q.view(B * H, N, C)k = k.view(B * H, N, C)v = v.view(B * H, N, C)# 分割成局部窗口q = self._window_partition(q, window_size)k = self._window_partition(k, window_size)v = self._window_partition(v, window_size)attn = (q @ k.transpose(-2, -1)) * self.scaleattn = attn.softmax(dim=-1)x = (attn @ v).view(B, H, N, C)return xdef _global_attention(self, q, k, v):"""全局注意力。"""B, _, H, N, C = q.shapeq = q.view(B, H, N, C)k = k.view(B, H, N, C)v = v.view(B, H, N, C)# 全局平均池化k = k.mean(dim=2, keepdim=True)v = v.mean(dim=2, keepdim=True)attn = (q @ k.transpose(-2, -1)) * self.scaleattn = attn.softmax(dim=-1)x = (attn @ v).view(B, H, N, C)return xdef _window_partition(self, x, window_size):"""将特征图分割成局部窗口。"""B, N, C = x.shapeH = int(N ** 0.5)x = x.view(B, H, H, C)x = x.view(B, H // window_size, window_size, H // window_size, window_size, C)x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size * window_size, C)return x
本文内容为论文学习收获分享，受限于知识能力，本文对原文的理解可能存在偏差，最终内容以原论文为准。本文信息旨在传播和学术交流，其内容由作者负责，不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题，请及时与我们联系，我们将在第一时间回复并处理。
确保文章为个人原创，未在任何公开渠道发布。若文章已在其他平台发表或即将发表，请明确说明。
建议使用Markdown格式撰写稿件，并以附件形式发送清晰、无版权争议的配图。
【AI前沿速递】尊重作者的署名权，并为每篇被采纳的原创首发稿件提供具有市场竞争力的稿酬。稿酬将根据文章的阅读量和质量进行阶梯式结算。
您可以通过添加我们的小助理微信（aiqysd）进行快速投稿。请在添加时备注“投稿-姓名-学校-研究方向”
长按添加AI前沿速递小助理
