动态卷积（轻量级卷积）替代多头自注意力

动态卷积，它比自注意力更简单、更有效。我们仅基于当前时间步长预测单独的卷积核，以确定上下文元素的重要性。这种方法所需的操作数量随输入长度呈线性增长，而自注意力是二次的。在大规模机器翻译、语言建模和抽象摘要上的实验表明，动态卷积比强自注意模型有更好的改进。

1. 引言

RNN通过在每个时间步更新一个隐藏状态来整合上下文信息，CNN通过多层总结固定大小的上下文，而自关注则直接总结所有上下文。注意力为上下文元素分配注意权重，该权重定义了上下文表示的加权和。源-目标注意从另一个序列(如机器翻译)中总结信息，而自注意力则对当前序列进行操作。自注意力被表述为基于内容的，其中通过比较当前时间步与上下文中的所有元素来计算注意权重。在这种不受限制的背景大小上计算比较的能力被视为自注意力的一个关键特征。

由于输入长度的二次复杂度，无限上下文大小在计算上非常具有挑战性。此外，在实践中，长序列需要引入层次结构。

本文中介绍了一种轻量级卷积，它们是深度可分离的，经过softmax归一化，并在通道维度上共享权重。与标准的不可分离卷积相比，轻量级卷积的权重数量减少了几个数量级。不同于自注意力机制，轻量级卷积在处理上下文元素时，无论当前时间步长如何，都重复使用相同的权重。

动态卷积在轻量级卷积的基础上进行改进，通过在每个时间步长预测不同的卷积核来实现。与自注意力机制需要考虑整个上下文不同，动态卷积的卷积核仅是当前时间步长的函数。动态卷积类似于局部连接层，后者在每个位置的权重会发生变化，但不同之处在于动态卷积的权重是由模型动态生成的，而不是训练后固定的。与基于位置的注意力相似，后者不需要访问上下文来确定注意力权重，但我们并不直接考虑前一个时间步的注意力权重。

实验表明，轻量级卷积在性能上能够与强大的自注意力机制结果相媲美，而动态卷积的表现甚至更好。

2. 背景

序列到序列学习：序列到序列学习是一种将源序列映射到目标序列的方法，通常通过两个独立的网络实现，这在机器翻译等任务中尤为常见。在这个过程中，编码器（Encoder）网络负责计算源序列（如英文句子）的表示，而解码器（Decoder）网络则基于编码器的输出，通过自回归的方式生成目标序列。这种框架允许模型在处理诸如语言翻译、文本摘要等任务时，能够捕捉到源序列和目标序列之间的复杂关系。

自注意力机制：自注意力机制由Vaswani等人提出，并在Transformer模型中得到了广泛应用。它通过对输入X（ $X\in R^{n\times d}$ ，其中n是时间步的数量，d是输入/输出维度）进行三个投影操作，得到键（K）、查询（Q）和值（V）三种表示。自注意力机制的核心在于它能够同时关注输入序列中的不同位置，通过计算键和查询之间的点积，并对结果进行缩放和softmax归一化，从而得到每个位置的注意力权重。最后，这些权重被用于计算值的加权和，从而得到考虑了序列中所有位置信息的输出表示。自注意力机制通过多个头（Heads）的并行处理，能够捕捉到输入序列中的多种不同特征，并提高了模型处理长序列的能力。

深度卷积：在每个通道上独立地执行卷积。参数的数量可以从 $d^2k$ 减少到 $dk$ ，其中 $k$ 是核宽度。对元素 i 和输出维度 c 进行权值为 $W\in R^{d\times k}$ 的深度卷积的输出 $O\in R^{n\times d}$ 定义为:

$O_{i,c}=\text{DepthwiseConv}(X,W_{c,:},i,c)=\sum_{j=1}^{k}W_{c,j}\cdot X_{(i+j-\lceil\frac{k+1}{2}\rceil),c}$

对于d个输入通道，我们只需要d * k个参数（每个通道一个卷积核，每个卷积核宽度为k），而不是传统卷积中的d^2 * k个参数

3. 轻量级卷积（LightConv）

LightConv，这是一种深度卷积，它共享某些输出通道，其权重在时间维度上使用softmax进行归一化。与自注意力相比，LightConv有一个固定的上下文窗口，它用一组不随时间变化的权重来确定上下文元素的重要性。

LightConv为序列和输出通道 c 中的第 i 个元素计算以下内容:

$\text{LightConv}(X,W_{\lceil\frac{cH}{d}\rceil,:},i,c)=\text{DepthwiseConv}(X,\text{softmax}(W_{\lceil\frac{cH}{d}\rceil,:}),i,c)$

$W_{\lceil\frac{cH}{d}\rceil,:}$ $\lceil \frac{cH}{d} \rceil$ ：表示将通道数 c 和特征图高度 H 按比例 d 缩减后向上取整得到的索引。

,:表示选择该索引处的所有列，即取出某个特定位置的整行权重。

i：当前时间步或图像块的索引

c：当前通道索引

3.1 权重共享

将每个后续的 $\frac{d}{H}$ 通道（即共享的通道数）的参数捆绑在一起，这将参数的数量减少了 $\frac{d}{H}$ 倍数。例如，对于d = 1024和k = 7，一个正则卷积需要7,340,032 ( $d^2\times k$ )个权重，一个深度可分离卷积有7,168个权重( $d\times k$ )，而对于H = 16的权重共享，只有112个( $H\times k$ )个权重。参数数量的大量减少对于在当前硬件上实现动态卷积至关重要。Wang & Ji提出在所有通道之间共享权重（即H=1）。

使用softmax操作将权重 $W\in R^{H\times K}$ 在卷积核宽度 k 上归一化:

$\mathrm{softmax}(W)_{h,j}=\frac{\exp W_{h,j}}{\sum_{j'=1}^k\exp W_{h,j'}}$

图2b显示了LightConv的模块架构：首先应用从d维到2d维的输入投影映射，然后是门控线性单元(GLU)，以及实际的轻量级卷积。GLU通过应用s型单元使用一半的输入作为门，然后计算与其他输入的点积。将大小为 $W^O\in R^{d\times d}$ 的输出投影应用于LightConv的输出。

DropConnect是一种有效的正则化方法（Wan et al, 2013）。具体来说，DropConnect以概率p随机丢弃softmax(W)中每个归一化权重的条目，并在训练过程中将剩余的权重除以1-p以保持权重的总体比例。这种操作实质上是在每个通道内去除部分时间信息，从而防止模型过度拟合训练数据。通过减少模型对特定时间信息的依赖，DropConnect增强了模型的泛化能力，使其能够在未见过的数据上表现更好。

在尝试实现LightConv时，现有的CUDA卷积原语并不适合用于处理短序列，性能表现不佳。因此，我们采用了一种更快的解决方案来处理短序列。首先，我们将归一化的权重W（形状为H×k）复制并扩展到一个大小为 BH×n×n 的带状矩阵中，其中B是批次大小，n是序列长度（或卷积的输出维度）。接着，将输入数据重新整形并转置为 BH×n×d_H 的形状，其中d_H是输入通道数或特征维度。然后，使用批量矩阵乘法来计算输出。

4. 动态卷积

动态卷积具有随时间变化的核，作为单个时间步长的学习函数。标准卷积的动态版本对于当前的gpu来说是不切实际的，因为它们需要大量的内存。通过建立LightConv来解决这个问题，大大减少了参数的数量。

DynamicConv采用与LightConv相同的形式，但使用时间步长相关的内核，该内核使用函数 $f:\mathbb{R}^d\to\mathbb{R}^{H\times k}$ : $\mathbf{DynamicConv}(X,i,c)=\mathbf{LightConv}(X,f(X_i)_{h,:},i,c)$

具体来说，DynamicConv 在时间步 i 和通道 c 的情况下等同于使用 $f(X_i)$ 生成的权重的 LightConv。用一个简单的线性模块来建模函数 f，该模块具有学习到的权重 $W^Q\in\mathbb{R}^{H\times k\times d}$ 。具体而言，函数 $f(X_i)$ 的计算方式是 $f(X_i)=\sum_{c=1}^{d}W_{h,j,c}^{Q}X_{i,c}$ ，即通过线性组合输入特征 $X_i$ 生成卷积核。

与自注意力类似，DynamicConv会随时间改变分配给上下文元素的权重。然而，DynamicConv的权重并不依赖于整个上下文，它们只是当前时间步长的函数。自注意力需要在句子长度上进行二次运算来计算注意权值，而DynamicConv的动态核计算在序列长度上呈线性缩放。

5. 模型架构

在序列到序列学习（sequence to sequence learning）任务中使用的一种编码器-解码器（encoder-decoder）架构。具体来说，这个架构在很大程度上遵循了Transformer模型的设计，同时在一些模块中引入了LightConv和DynamicConv来替代传统的自注意力模块。下面是对这个架构的详细解释：

1. 基本架构

编码器-解码器架构：该模型使用了编码器-解码器架构，用于序列到序列学习任务。

Transformer Big：该架构的自注意力基础模型是fairseq重新实现的Transformer Big。

2. 编码器和解码器

编码器和解码器网络：每个编码器和解码器网络都有N个块（blocks）。

N的值：通常设定为7（为了大致匹配Transformer Big的参数数量）。

编码器块

第一个子块：可以是自注意力模块、LightConv模块或DynamicConv模块。

第二个子块：前馈神经网络模块（Feed-forward module）。

形式为： $ReLU(W^1X+b_1)W^2+b_2\text{ where }W^1\in\mathbb{R}^{d\times d_{ff}},W^2\in\mathbb{R}^{d_{ff}\times d}$

参数：d = 1024， $d_{ff}$ = 4096。

残差连接和层归一化：每个子块都包裹在残差连接和层归一化中。

解码器块

结构与编码器块类似，但有一个额外的源-目标注意力子块，位于自注意力和前馈模块之间。

源-目标注意力：等同于自注意力模块，但值和值键是对每个源词的编码器输出的投影。

3. 嵌入与位置编码

词嵌入：词嵌入维度为d。

位置编码：使用正弦位置编码来表示序列中每个词的绝对位置。

4. 词汇分布计算

词汇分布：通过线性层（权重为 $W^V\in R^{d\times V}$ ）和softmax归一化，将解码器输出转化为词汇V的分布。

5. LightConv和DynamicConv的设置

替换自注意力模块：LightConv和DynamicConv与Transformer Big相同，只是将自注意力模块替换为固定或动态卷积。

参数更少：每个块的参数更少，因此增加了编码器块的数量（N = 7）。

卷积核大小：编码器和解码器的卷积核大小分别设为3, 7, 15, 31；解码器中只有前三层使用31的卷积核大小。

H的值：一般设置为16。

6. 代码示例

6.1 轻量级卷积

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

class LightConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(LightConv, self).__init__()
        self.depthwise_conv = nn.Conv1d(in_channels, in_channels, kernel_size, groups=in_channels, padding=kernel_size//2)
        self.pointwise_conv = nn.Conv1d(in_channels, out_channels, 1)
        
    def forward(self, x):
        # 对深度卷积的权重进行softmax归一化
        weight = F.softmax(self.depthwise_conv.weight, dim=-1)
        x = F.conv1d(x, weight, bias=self.depthwise_conv.bias, padding=self.depthwise_conv.padding, groups=self.depthwise_conv.groups)
        
        # 进行逐点卷积
        x = self.pointwise_conv(x)
        
        return x

# 参数设置
in_channels = 3
out_channels = 5
kernel_size = 3

# 创建 LightConv 模块
light_conv = LightConv(in_channels, out_channels, kernel_size)

# 创建一个示例输入
input_tensor = torch.randn(2, in_channels, 10)
output_tensor = light_conv(input_tensor)

print("LightConv output shape:", output_tensor.shape)

6.2 动态卷积

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

class DynamicConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(DynamicConv, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        
        # 深度卷积层
        self.depthwise_conv = nn.Conv1d(in_channels, in_channels, kernel_size, groups=in_channels, padding=kernel_size//2)
        
        # 用于动态生成卷积核权重的线性层
        self.dynamic_weight_generator = nn.Linear(in_channels, in_channels * kernel_size)
        
        # 点卷积层
        self.pointwise_conv = nn.Conv1d(in_channels, out_channels, 1)
        
    def forward(self, x):
        # 动态生成权重
        batch_size, channels, length = x.size()
        dynamic_weight = self.dynamic_weight_generator(x.transpose(1, 2))
        dynamic_weight = dynamic_weight.view(batch_size, channels, self.kernel_size, length)
        dynamic_weight = F.softmax(dynamic_weight, dim=2)
        
        # 使用动态权重进行深度卷积
        x = x.transpose(1, 2)
        output = torch.zeros_like(x)
        for i in range(length):
            output[:, :, i] = F.conv1d(x[:, :, max(0, i - self.kernel_size // 2):i + self.kernel_size // 2 + 1],
                                       dynamic_weight[:, :, :, i].transpose(1, 2), groups=channels)
        x = output.transpose(1, 2)
        
        # 进行逐点卷积
        x = self.pointwise_conv(x)
        
        return x

# 参数设置
in_channels = 3
out_channels = 5
kernel_size = 3

# 创建 DynamicConv 模块
dynamic_conv = DynamicConv(in_channels, out_channels, kernel_size)

# 创建一个示例输入
input_tensor = torch.randn(2, in_channels, 10)
output_tensor = dynamic_conv(input_tensor)

print("DynamicConv output shape:", output_tensor.shape)