NLP学习记录八：理解注意力机制-CSDN博客

本文链接：https://blog.csdn.net/scongx/article/details/141575298

一、问题假设

1.1 题目

考虑如下回归问题：

给定回归方程：

现有成对的“输入-输出”数据集{(x1, y1), ...... , (xn, yn)}，输入x与输出y之间满足：

方程中， ε表示噪声，服从均值为0和标准差为0.5的正态分布。

请设计一个估计器f，使得估计器的输入x和输出的y_hat=f(x)尽可能地契合回归方程表示的等式关系。

1.2 基础代码

utils.py：

attentionPooling.py：

导入依赖包：

import torch
from torch import nn
from matplotlib import pyplot as plt
from utils import plot, show

定义回归方程：

def f(x):
    return 2 * torch.sin(x) + x**0.8

生成训练集，对应于问题描述的“输入-输出”数据集{(x1, y1), ...... , (xn, yn)}：

if __name__=='__main__':
    # 训练样本的大小
    n_train = 50

    # torch.rand(n_train)表示生成一个形状为(n_train,)的张量，其中包含了从0到1之间均匀分布的随机数
    # 乘5操作使得张量中随机生成的每个元素都乘5
    # torch.sort表示由小到大排序
    x_train, _ = torch.sort(torch.rand(n_train) * 5)

    # 训练样本的输出，前一项表示回归方程输出，后一项表示均值为0、方差为0.5的噪声
    y_train = f(x_train) + torch.normal(0.0, 0.5, (n_train,))

生成测试集：

# main:
    # 测试样本，在[0, 6)之间生成等差数列，公差为0.1
    x_test = torch.arange(0, 6, 0.1)

    # 测试样本的真实输出
    y_truth = f(x_test)

    # 测试样本的大小
    n_test = len(x_test)

另外再定义一个用于绘制样本输出结果的函数：

def plot_kernel_reg(y_hat, x_test, y_truth, x_train, y_train):
    # 绘制测试样本的真实输出、输入估计器f后的输出
    plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'], xlim=[0, 5], ylim=[-1, 5])

    # 绘制训练集（离散点状分布）
    plt.plot(x_train, y_train, 'o', alpha=0.5)
    show()

二、解决方法

下面由浅入深地介绍设计该估计器的三种方法。

2.1 平均池化

2.1.1 数学建模

使用平均池化时，f的输入输出满足下式关系：

2.1.2 代码实现

代码上我们可以定义如下函数描述平均池化：

def averagePooling(y_train, n):
    return y_train.mean().repeat(n)

测试平均池化的输出结果：

# main:
    y_hat = averagePooling(y_train, n_test)
    plot_kernel_reg(y_hat, x_test, y_truth, x_train, y_train)

容易看出，平均池化最简单，但其输入输出关系与我们想要的结果相差甚远，所以不采用这种方法。

2.2 不带参数的注意力池化

2.2.1 数学建模

思考一下，既然训练集{(x1, y1), ...... , (xn, yn)}的分布是契合目标回归方程的，那么我们可以先计算估计器的输入x与训练集中的{x1, ......, xn}每一项的相似度，假如x与xi特别相似，那对应的输出y也应该越接近于yi才对。于是我们可以将估计器f的输入输出关系建模成如下形式：

上式中，α表示注意力权重，α的值由x与xi的相似度决定，x与xi越接近，给到yi的权重就越大。训练集中所有yi加权求和后的结果便是我们想要的预测值。

我们的需求是x与xi的距离越小，α(x, xi)就越大，那么应该如何设计α(x, xi)呢？。

上式是高斯核（Gaussian kernel）的定义，可以看到，x与xi的相差越小，d(x, xi)就越大，正好满足我们的需求。

由于α(x, xi)是一个权重，为了把α(x, xi)限制在0到1之间，我们可以使用分类问题上常用的softmax函数：

最终估计器f的建模结果为：

使用高斯核和softmax函数来计算权重只是本例的一种做法，事实上，上式可以写作更通用的形式：

其中K是核（kernel），上式所描述的估计器被称为 Nadaraya-Watson核回归（Nadaraya-Watson kernel regression），我们只需要对NW核回归有一个概念就行，此处不过多赘述。

2.2.2 代码实现

不带参数的注意力池化函数：

def weightlessAttentionPooling(x_test, x_train, y_train, n_test, n_train):
    # X_repeat的尺寸：(n_test, n_train)
    X_repeat = x_test.repeat(n_train).reshape((-1, n_test)).T

    # 使用softmax计算权重，dim=1会取一行中的所有元素进行归一化运算
    attention_weights = nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)

    return torch.matmul(attention_weights, y_train)

测试不带参数的注意力池化的输出结果：

# main:
    y_hat = weightlessAttentionPooling(x_test, x_train, y_train, n_test, n_train)
    plot_kernel_reg(y_hat, x_test, y_truth, x_train, y_train)

从图中的紫色曲线可以看出，估计器的预测结果已经比较接近真值了，但精度仍然有提升空间。

2.3 带参数的注意力池化

2.3.1 数学建模

在不带参数的注意力池化章节中，我们对估计器f的建模结果是这样的：

而带参数的注意力池化也很简单，只需要在x与xi的距离后增加一个学习的参数w即可：

2.3.2 代码实现

定义带参数的注意力估计器模型：

class NWKernelRegression(nn.Module):
    def __init__(self, n=1, **kwargs):
        super(NWKernelRegression, self).__init__(**kwargs)
        self.w = nn.Parameter(torch.rand((n, ), requires_grad=True))

    def forward(self, queries, keys, values):
        # repeat_interleave(n)表示对原张量每个元素重复n次，重复n次是为了方便后续和n个“键”进行相减运算
        # 最终queries的尺寸为：(查询个数, “键－值”对个数)
        queries = queries.repeat_interleave(keys.shape[1]).reshape((-1, keys.shape[1]))

        # attention_weights的尺寸为：(查询个数, “键－值”对个数)
        self.attention_weights = nn.functional.softmax(-((queries - keys) * self.w)**2 / 2, dim=1)

        # values的尺寸为：(查询个数, “键－值”对个数)
        # self.attention_weights.unsqueeze(1)的尺寸为：(查询个数, 1, “键－值”对个数)
        # values.unsqueeze(-1)的尺寸为：(查询个数, “键－值”对个数, 1)
        # torch.bmm表示矩阵批量乘法，相乘所得矩阵的尺寸为：(查询个数, 1)
        return torch.bmm(self.attention_weights.unsqueeze(1), values.unsqueeze(-1)).reshape(-1)

训练函数：

def train(net, x_train, y_train):
    X_tile = x_train.repeat((n_train, 1))
    Y_tile = y_train.repeat((n_train, 1))

    # torch.eye用于生成对角线全1，其余部分全0的二维张量
    # keys和values的尺寸：(n_train, n_train-1)
    # 这么处理是为了删除“键-值”中和输入样本相同的元素，不然训练时输入样本每次都能直接从“键-值”找到和自己一模一样的值，会影响模型的预测效果
    # 由于每次迭代可查找的“键-值”只有n_train-1个，所以模型参数w也只能设置为n_train-1了
    keys = X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))
    values = Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape((n_train, -1))

    # 均方误差损失
    loss = nn.MSELoss(reduction='none')

    optimizer = torch.optim.SGD(net.parameters(), lr=0.5)
    animator = Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])

    for epoch in range(100):
        optimizer.zero_grad()
        y_hat = net(x_train, keys, values)
        l = loss(y_hat, y_train)
        l.sum().backward()
        optimizer.step()

        print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
        animator.add(epoch + 1, float(l.sum()))

    show()

测试预测结果：

# main:
    # 带参数的注意力池化，n_train-1的目的在train函数中说明
    net = NWKernelRegression(n_train-1)

    train(net, x_train, y_train)

    x_train = x_train[:-1]
    y_train = y_train[:-1]

    keys = x_train.repeat((n_test, 1))
    values = y_train.repeat((n_test, 1))

    y_hat = net(x_test, keys, values).unsqueeze(1).detach()
    plot_kernel_reg(y_hat, x_test, y_truth, x_train, y_train)