Transformer注意力层中为什么需要value、key和query三类参数？

我正在学习Transformer模型的基础概念，根据我查阅的论文和教程，注意力层会通过神经网络生成value、key和query三类参数。

以下是我从网上学习到的自注意力层实现代码：

class SelfAttention(nn.Module):
def __init__(self, embed_size, heads):
    super(SelfAttention, self).__init__()
    self.embed_size = embed_size
    self.heads = heads
    self.head_dim = embed_size // heads

    assert (self.head_dim * heads == embed_size),  "Embed size needs to be div by heads"

    self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
    self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
    self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
    self.fc_out = nn.Linear(heads*self.head_dim, embed_size)

def forward(self, values, keys, query, mask):
    N = query.shape[0]
    value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]
    #split embedding into self.heads pieces
    values = values.reshape(N, value_len, self.heads, self.head_dim)
    keys = keys.reshape(N, key_len, self.heads, self.head_dim)
    queries = query.reshape(N, query_len, self.heads, self.head_dim)

    values = self.values(values)
    keys = self.keys(keys)
    queries = self.queries(queries)

    energy = torch.einsum("nqhd, nkhd->nhqk", [queries, keys])

    # queires shape (N, quesry_len, heads, heads_dim)
    # kyes shape : (N, key_len, heads, heads_dim)
    # energy shape: (N, heads, query_len, key_len)
    if mask is not None:
        energy = energy.masked_fill(mask == 0, float("-1e20"))
    attention = torch.softmax(energy / (self.embed_size ** (1/2)), dim=3)
    out = torch.einsum('nhql, nlhd->nqhd', [attention, values]).reshape(N, query_len, self.heads*self.head_dim)
    # attention shape: (N, heads, query_len, key_len)
    # value shpae: (N, value_len, heads, heads_dim)
    # after eisum (N, query_len, heads, head_dim) then flatten the last two dimmsions
    out = self.fc_out(out)
    return out

我目前有一个十分困惑的点：为什么需要key、value、query这三类参数？我可以只使用其中一类吗？或者除了这三类之外，我还可以设置更多参数吗？这三者看起来只是经过三个单层神经网络变换得到的产物。

问题解答 #

我们可以直接从注意力的设计逻辑来拆解三个参数的作用：

• 你完全可以把自注意力计算类比成搜索引擎的检索流程：
  • query对应你输入的搜索关键词，负责发起匹配需求
  • key对应所有待检索内容的索引标签，专门用来和query做匹配度计算
  • value对应待检索内容本身，匹配度越高的value会被赋予更高的权重，最终加权求和得到输出

你说的只用一类参数的情况是可以跑通的，甚至直接让Q=K=V用原始输入也能得到注意力结果，但是这么做会直接损失模型的表达能力：没有独立的Q、K投影，相当于你不能自定义搜索规则，只能用内容本身当关键词做检索，很多复杂的语义关联根本学不到，在实际任务上的效果会差很多。

至于能不能加更多参数？当然可以。目前已经有不少改进型注意力结构做了相关尝试，比如给注意力得分加位置相关的可学习偏置、给value加额外的门控参数等等，只要你新增的参数能帮模型学到更有效的关联规则，同时控制住计算量的增长，都是合理的优化方向。

你观察的没错，标准Transformer里的Q、K、V确实就是三个独立的单层线性变换的输出，这个设计的性价比非常高：三个线性层的总参数量只有3*embed_size²，相比Transformer整体的参数量占比极低，但能让注意力的匹配灵活度提升好几个量级，是经过大量工程验证的最优设计之一。

内容的提问来源于stack exchange，提问作者Xudong