《SeTformer Is What You Need for Vision and Language》

会议：AAAI

年份：2024

论文：DDAE: Towards Deep Dynamic Vision BERT Pretraining - AMinerhttps://www.aminer.cn/pub/6602613613fb2c6cf6c387c2/ddae-towards-deep-dynamic-vision-bert-pretraining

摘要

这篇论文介绍了一种新型的变换器模型，名为SeTformer，它针对视觉和语言任务进行了优化。SeTformer的核心创新是将传统的点积自注意力（DPSA）替换为一种基于自我最优传输（Self-optimal Transport，简称SeT）的机制。这种替换旨在解决传统变换器在处理长序列时面临的计算效率低下的问题，同时保持或提升性能。SeT基于两个关键的softmax属性：维持非负的注意力矩阵和使用非线性重加权机制来强调输入序列中的重要token。通过引入核成本函数来实现最优传输，SeTformer有效地满足了这些属性。实验结果表明，SeTformer在多个任务上都取得了令人印象深刻的性能，包括在ImageNet-1K上的分类准确率、目标检测和语义分割等。

拟解决的问题

传统变换器（如Vision Transformer, ViT）在处理长序列数据时，由于softmax操作导致的二次时间和内存复杂度，使得模型训练和推理变得非常昂贵。这限制了变换器在处理长文档或高分辨率图像等任务时的应用。

创新之处

1. 自我最优传输（SeT）：提出了一种新的自注意力机制，使用最优传输（OT）和核方法来计算输入特征之间的对齐分数，代替了传统的点积操作。
2. 核特征映射：通过将输入特征映射到再生核希尔伯特空间（RKHS），保持了注意力矩阵的非负性质。
3. 效率与性能的平衡：在减少参数和计算量的同时，SeTformer在多个视觉和语言任务上取得了与现有最先进方法相媲美或更好的结果。

方法

 传统的自注意力方法：输入序列的每个元素（token）都会计算与其他所有元素的关系，通常使用点积（Dot Product）来衡量元素间的相似度。然后，通过softmax函数对这些相似度进行归一化，得到注意力权重，这些权重用于加权求和，形成每个元素的输出表示。

尽管自注意力在捕捉序列内长距离依赖方面非常有效，但其计算复杂度随着序列长度的增加而呈二次方增长，这在处理长序列时（如高分辨率图像或长文本）变得非常昂贵。

 为了解决以上问题，提出了自我最优传输（SeT），这是一种新颖的自注意力替代方案。SeT利用最优传输（OT）理论来计算输入特征之间的对齐分数，而不是传统的点积操作。

SeT的关键步骤：

1. 特征映射到RKHS：将输入特征向量映射到再生核希尔伯特空间（RKHS），使用正定核函数保证映射后的向量是非负的。
2. OT对齐：使用最优传输（OT）来对齐输入特征和参考特征。OT通过计算最小成本的传输计划，来找到输入特征与参考特征之间的最佳匹配。
3. 加权聚合：通过OT得到的对齐分数作为权重，对输入特征进行加权聚合，形成输出特征。

 特征映射到RKHS

1. 选择或定义核函数：首先，选择一个合适的正定核函数 𝐾(𝑥,𝑥′)，它能够将原始空间 𝑋 中的数据点映射到一个高维特征空间 𝐹。常用的核函数包括高斯核、多项式核、拉普拉斯核等。

2. 特征映射：通过核函数定义一个特征映射 𝑢:𝑋→𝐹，其中 𝐹 是一个希尔伯特空间。在RKHS中，任意两个点 𝑥 和 𝑥′ 之间的核函数值 𝐾(𝑥,𝑥′) 可以表示为这两个点在 𝐹 空间中映射的内积，即 𝐾(𝑥,𝑥′)=〈𝑢(𝑥),𝑢(𝑥′)〉𝐹。

由于核函数是正定的，它保证了映射后的向量之间的内积是非负的，这与softmax操作的非负性相一致。

OT对齐

OT是数学中的一个概念，它描述了将一个概率分布转换为另一个概率分布的最优方式，这可以类比为将一堆“土”（代表概率质量）从一个位置移动到另一个位置，最小化运输成本。

用于在特征空间中对齐输入和参考特征。在SeTformer中，OT对齐用于计算输入特征与一组参考特征之间的相似性或对齐分数，这些分数随后用于加权聚合，形成注意力机制的输出。

OT对齐的关键步骤：

1. 用Kantorovich形式表示，通过最小化成本函数加上熵正则化项来求解：

C 是成本矩阵，𝑇 是传输计划矩阵，𝜇 和 𝜈 是分别与 𝑋 和 𝑌 相关的分布，𝜖 是正则化参数，𝐻(𝑇) 是传输计划的熵，计算公式如下：

2. 使用Sinkhorn算法来求解上述优化问题，该算法通过迭代过程不断调整传输计划 𝑇 以逼近最优解。OT 根据各个元素/标记在输入中的重要性为不同的元素/标记分配不同的权重，类似于 softmax 注意力中的重新加权方案。换句话说，T(x, x′) 中的每个值表示 x 与 x′ 的元素对齐的权重或重要性。

Self-optimal Transport (SeT)

• 将输入输入特征 𝑋 和参考特征 𝑌 映射到RKHS，这些参考特征用于与输入特征进行对齐
• 使用OT来计算输入特征 𝑋 和参考特征 𝑌 之间的对齐分数，通过最小化传输成本来找到最优的传输计划 𝑇。
• 根据OT得到的传输计划 𝑇，进而形成对齐矩阵 𝐴：，其中U是输入特征在RKHS中的表示
• 将输入特征 𝑥 通过传输计划 𝐴 与参考特征 𝑦 的加权聚合

输入特征：通过相关网络提取的图像特征

参考特征：参考特征是一组预先计算好的向量，它们代表了输入数据中的关键特征或模式。可以通过聚类算法（如K-means）从训练数据中生成的。这种方法可以捕捉到数据中的代表性特征，并将它们用作参考特征集。参考特征在模型中起到了锚点的作用，帮助模型通过OT对齐来聚合输入特征。

总的来说，在SeTformer中，SeT用于替代传统的自注意力机制，通过以下步骤：

• 输入特征通过卷积层进行下采样和嵌入。
• 使用SeT计算输入特征与参考特征之间的对齐分数。
• 根据对齐分数进行特征的加权聚合，生成新的表示。
• 这些表示被送入后续的网络层，如多层感知机（MLP）和归一化层。

结论

SeTformer展示了在视觉和语言任务中的高效性和有效性。通过一系列实验验证了其在图像分类、目标检测和语义分割等任务上的优越性能。此外，SeTformer在语言建模任务上也取得了良好的结果，证明了其在不同领域的泛化能力。论文指出，SeTformer为探索更好的基于内容的交互提供了新的视角，有助于提升视觉识别模型的性能。