python 实现结合门控单元的自注意力机制

结合门控单元的自注意力机制通常被称为“自注意力门控”（Self-Attention with Gates，SAG）。下面是一个基于 Python 的实现：

import torch
import torch.nn as nn

class SAG(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(SAG, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        
        self.query = nn.Linear(input_dim, hidden_dim)
        self.key = nn.Linear(input_dim, hidden_dim)
        self.value = nn.Linear(input_dim, hidden_dim)
        
        self.gate = nn.Linear(hidden_dim, hidden_dim)
        
    def forward(self, x):
        Q = self.query(x)  # query
        K = self.key(x)  # key
        V = self.value(x)  # value
        
        # 计算注意力权重
        attn_weights = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.hidden_dim, dtype=torch.float32))
        attn_weights = torch.softmax(attn_weights, dim=-1)
        
        # 计算加权和
        attn_output = torch.matmul(attn_weights, V)
        
        # 计算门控
        gate_output = torch.sigmoid(self.gate(attn_output))
        output = gate_output * attn_output + (1 - gate_output) * x
        
        return output

这个实现中，我们首先定义了一个 `SAG` 类，它继承自 PyTorch 的 `nn.Module` 类。在 `__init__` 方法中，我们定义了四个全连接层，分别用于计算查询（query）、键（key）、值（value）以及门控（gate）。在 `forward` 方法中，我们首先分别通过 `query`、`key` 和 `value` 层计算出查询、键和值，然后通过矩阵乘法计算注意力权重，再将权重与值相乘得到加权和。接着，我们通过门控函数计算门控输出，最后将门控输出与加权和相乘并与原始输入相加得到最终输出。

需要注意的是，在计算注意力权重时，我们在最后一个维度上进行 softmax，即每个时间步的权重之和为 1。在计算加权和时，默认是将注意力权重与值相乘，这被称为“乘性注意力”（multiplicative attention）。如果要使用加性注意力，可以将注意力权重与值先拼接起来再通过一个全连接层计算。