Patches Are All You Need?
Lyndon Wong

论文地址
代码

摘要

  • 背景:ViT的性能提升是因为Patches还是因为自注意力?

  • 工作:

    • 提出了ConvMixer,一个和ViT、MLP-Mixer相似的网络。

    • 相似在什么地方呢?:ConvMixer直接以patch分割作为输入,分离了空间维度(spatial)和通道维度(channel)的混合,并在整个网络中保持相同的大小和分辨率

    • 不同之处呢?:ConvMixer只使用标准卷积来进行Mixing步骤

  • 结果:

    • ConvMixer:简单(仅仅几行代码!),优于ViT、MLP-Mixer和ResNet等一系列结构

介绍

Transformer(以下简称TF)被应用到视觉任务上,但是由于自注意力机制的平方级别的计算量,无法对图片直接进行像素维的计算,所以,一个妥协(compromise)的方式就是:把图片分割为一个个patch,然后对patch进行线性嵌入为向量后,直接送入TF进行处理。

但是在本次工作中,我们发现,ViT强大的性能可能更多地来自这个基于patch的表示,而不是TF结构本身!我们由此提出了一个简单的网络,名为“ConvMixer”,因为它和MLP-Mixer十分相像:

  • 直接对patch进行操作

  • 在所有层都保持了相同的分辨率和尺寸表示

  • 它不在连续的层之间进行下采样

  • 它分离了信息的“通道维混合”和“空间维混合”

但是不同指出在于:

  • ConvMixer只使用标准卷积

当然,我们的这些结果只是类似于一个“快照”(snapshot),但是我们相信它提供了一个“基于patch但是卷积”的网络的基线,这样的更加强大的网络可能在未来出现。

一个简单的模型:ConvMixer

模型结构

image-20211106190201558

模型结构很简单(几乎都不需要解释把= =),前面主要是利用一个kernel size = stride的卷积核进行了patch分割,激活然后BN。后面ConvMixer Layer采用了可分离卷积,其中depthwise采用了残差连接结构。最后Global Average Pooling跟一个FC形成head,其实作者也认为这个head如果改改设计应该还能提升。

设计动机

​ 设计思路延续了MLP-Mixer的思想,其中depthwise和pointwise分别来进行空间(spatial)和通道(channel)维度的混合(mixing)。MLP的一个关键的设计思想是利用self-attention来混合远距离空间上的位置(distant spatial location),也就是说它可以拥有任意大小的感受野。所以,我们采用非同寻常的大卷积核来达到类似的效果。

​ 诚然,MLPs和self-attention理论上具有更大的灵活性,允许超大的感受野和内容认知行为(content-aware behavior),但是卷积的归纳偏置使得它天生适合视觉任务并因此具有很高的数据有效性。通过使用这种标准操作,我们还可以一窥与传统的金字塔形、渐进式下采样的卷积网络设计相比,patch表示的根本作用是什么。

实验

  • CIFAR-10:仅用0.7M参数达到了96%准确率
  • 训练设置:
    • 数据集:ImageNet-1k,无额外数据
    • 使用了一些timm框架下的数据处理方法
    • 优化器AdamW
    • 由于计算资源限制,没有对超参数进行调参
  • 结果

image-20211106192402089

源码阅读

由于代码很简单直接看一下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import torch.nn as nn

class Residual(nn.Module):
def __init__(self, fn):
super().__init__()
self.fn = fn

def forward(self, x):
return self.fn(x) + x

def ConvMixer(dim, depth, kernel_size=9, patch_size=7, n_classes=1000):
return nn.Sequential(
nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size),
nn.GELU(),
nn.BatchNorm2d(dim),
*[nn.Sequential(
Residual(nn.Sequential(
nn.Conv2d(dim, dim, kernel_size, groups=dim, padding="same"),
nn.GELU(),
nn.BatchNorm2d(dim)
)),
nn.Conv2d(dim, dim, kernel_size=1),
nn.GELU(),
nn.BatchNorm2d(dim)
) for i in range(depth)],
nn.AdaptiveAvgPool2d((1,1)),
nn.Flatten(),
nn.Linear(dim, n_classes)
)

其中

1
2
nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size)
# patch分割
1
2
nn.Conv2d(dim, dim, kernel_size, groups=dim, padding="same"),
#深度维卷积,输入/输出维数不变,信息只在spatial维度上交流
1
2
nn.Conv2d(dim, dim, kernel_size=1)
#通道维卷积,1*1卷积核,信息在channel维度上交流

读后感

  • spatial和channel维度的信息交流分开交替进行。(猜测:这样做主要是降低了网络的学习负担?,梯度的来源比较“纯净”?)
  • 目前看来,关键不是自注意力还是卷积,而是网络是否保证了全部layers都有很大的感受野(patch分割之后,大卷积核也能产生很大的感受野)
 评论