注意力机制 attention Transformer 笔记

动手学深度学习

这里写自定义目录标题

注意力
加性注意力
缩放点积注意力
多头注意力
自注意力
自注意力+缩放点积注意力：案例
Transformer

注意力

注意力汇聚的输出为值的加权和

查询的长度为q，键的长度为k，值的长度为v。
${\bf{q}} \in {^{1 \times q}},{{\bf{k}}} \in {^{1 \times k}},{{\bf{v}}} \in {\mathbb{R}^{1 \times v}}$
n个查询和m个键-值对
${\bf{Q}} \in {^{n \times q}},{\bf{K}} \in {^{m \times k}},{\bf{V}} \in {\mathbb{R}^{m \times v}}$
${\bf{a}}\left( {{\bf{Q}},{\bf{K}}} \right) \in {\mathbb{R}^{n \times m}}$ 是注意力评分函数
${\boldsymbol{\alpha}} \left( {{\bf{Q}},{\bf{K}}} \right) = {\rm{softmax}}\left( {{\bf{a}}\left( {{\bf{Q}},{\bf{K}}} \right)} \right) = \frac{{\exp \left( {{\bf{a}}\left( {{\bf{Q}},{\bf{K}}} \right)} \right)}}{{\sum\limits_{j = 1}^m {\exp \left( {{\bf{a}}\left( {{\bf{Q}},{\bf{K}}} \right)} \right)} }} \in {\mathbb{R}^{n \times m}}$ 是注意力权重
$f({\bf{Q}},{\bf{K}},{\bf{V}}) = {\boldsymbol{\alpha}} {\left( {{\bf{Q}},{\bf{K}}} \right)^ \top }{\bf{V}} \in {\mathbb{R}^{n \times v}}$ 是注意力汇聚函数

加性注意力

${\bf{q}} \in {\mathbb {R}^{1 \times q}},{\bf{k}} \in {\mathbb {R}^{1 \times k}}$
${{\bf{W}}_q} \in {{\mathbb R}^{h \times q}},{{\bf{W}}_k} \in {{\mathbb R}^{h \times k}},{{\bf{w}}_v} \in {{\mathbb R}^{h \times 1}}$
$a({\bf{q}},{\bf{k}}) = {\bf{w}}_v^ \top {\rm{tanh}}({{\bf{W}}_q}{{\bf{q}}^ \top } + {{\bf{W}}_k}{{\bf{k}}^ \top }) \in \mathbb {R}$ 是注意力评分函数

缩放点积注意力

${\bf{q}} \in \mathbb{R}{^{1 \times d}},{\bf{k}} \in \mathbb{R}{^{1 \times d}},{\bf{v}} \in {{\mathbb R}^{1 \times v}}$
$a\left( {{\bf{q}},{\bf{k}}} \right) = \frac{1}{{\sqrt d }}{\bf{q}}{{\bf{k}}^ \top } \in \mathbb{R}$ 是注意力评分函数
$f({\bf{q}},{\bf{k}},{\bf{v}}) = \alpha {\left( {{\bf{q}},{\bf{k}}} \right)^ \top }{\bf{v}} = {\rm{softmax}}\left( {\frac{1}{{\sqrt d }}{\bf{q}}{{\bf{k}}^ \top }} \right){\bf{v}} \in {{\mathbb R}^{1 \times v}}$ 是注意力汇聚函数

n个查询和m个键-值对
$\mathbf Q\in\mathbb R^{n\times d}, \mathbf K\in\mathbb R^{m\times d}, \mathbf V\in\mathbb R^{m\times v}$
${\bf{a}}\left( {{\bf{Q}},{\bf{K}}} \right) = \frac{1}{{\sqrt d }}{\bf{Q}}{{\bf{K}}^ \top } \in {\mathbb{R}^{n \times m}}$ 是注意力评分函数
$f({\bf{Q}},{\bf{K}},{\bf{V}}) = {\boldsymbol{\alpha}} {\left( {{\bf{Q}},{\bf{K}}} \right)^ \top }{\bf{V}} ={\rm{softmax}}\left( {\frac{1}{{\sqrt d }}{\bf{Q}}{{\bf{K}}^ \top }} \right){\bf{V}} \in {\mathbb{R}^{n \times v}}$ 是注意力汇聚函数

Attention Is All You Need

多头注意力

${\bf{q}} \in {{\mathbb R}^{1 \times {d_q}}},{\bf{k}} \in {{\mathbb R}^{1 \times {d_k}}},{\bf{v}} \in {{\mathbb R}^{1 \times {d_v}}}$
${\bf{W}}_i^{(q)} \in {{\mathbb R}^{{p_q} \times {d_q}}},{\bf{W}}_i^{(k)} \in {{\mathbb R}^{{p_k} \times {d_k}}},{\bf{W}}_i^{(v)} \in {{\mathbb R}^{{p_v} \times {d_v}}}$
${\bf{qW}}{_i^{(q)\top} } \in {{\mathbb R}^{1 \times {p_q}}},{\bf{kW}}{_i^{(k)\top} } \in {{\mathbb R}^{1 \times {p_k}}},{\bf{vW}}{_i^{(v)\top} } \in {{\mathbb R}^{1 \times {p_v}}}$
${{\bf{h}}_i} = f\left( {{\bf{qW}}{{_i^{(q)\top}} },{\bf{kW}}{{_i^{(k)\top}} },{\bf{vW}}{{_i^{(v)\top}} }} \right) \in {{\mathbb R}^{1 \times {p_v}}}$ 是注意力头

多个注意力头连结然后线性变换
${{\bf{W}}_o} \in {{\mathbb R}^{{p_o} \times h{p_v}}}$
${{\bf{W}}_o}\left[ {\begin{array}{c} {{{\bf{h}}_1^ \top}}\\ \vdots \\ {{{\bf{h}}_h^ \top}} \end{array}} \right] \in {{\mathbb R}^{{p_o}}}$

$p_q h = p_k h = p_v h = p_o$
多头注意力：多个头连结然后线性变换
多头注意力：多个注意力头连结然后线性变换

自注意力

${{\bf{x}}_i} \in {{\mathbb R}^{1 \times d}},{\bf{X}} = \left[ {\begin{array}{c} {{{\bf{x}}_1}}\\ \cdots \\ {{{\bf{x}}_n}} \end{array}} \right] \in {{\mathbb R}^{n \times d}}$
${\bf{Q}} = {\bf{X}},{\bf{K}} = {\bf{X}},{\bf{V}} = {\bf{X}}$
$f({\bf{Q}},{\bf{K}},{\bf{V}}) = {\boldsymbol{\alpha}} {\left( {{\bf{Q}},{\bf{K}}} \right)^ \top }{\bf{V}} ={\rm{softmax}}\left( {\frac{1}{{\sqrt d }}{\bf{Q}}{{\bf{K}}^ \top }} \right){\bf{V}} \in {\mathbb{R}^{n \times d}}$
${{\bf{y}}_i} = f\left( {{{\bf{x}}_i},\left( {{{\bf{x}}_1},{{\bf{x}}_1}} \right), \ldots ,\left( {{{\bf{x}}_n},{{\bf{x}}_n}} \right)} \right) \in {{\mathbb R}^d}$

n个查询和m个键-值对
${\bf{Q}} = {\rm{tanh}}\left( {{{\bf{W}}_q}{\bf{X}}} \right) \in {{\mathbb R}^{n \times d}}$
${\bf{K}} = {\rm{tanh}}\left( {{{\bf{W}}_k}{\bf{X}}} \right) \in {{\mathbb R}^{m \times d}}$
${\bf{V}} = {\rm{tanh}}\left( {{{\bf{W}}_v}{\bf{X}}} \right) \in {{\mathbb R}^{m \times v}}$

自注意力+缩放点积注意力：案例

J. Xu, F. Zhong, and Y. Wang, “Learning multi-agent coordination for enhancing target coverage in directional sensor networks,” in Proc. Neural Information Processing Systems (NeurIPS), Vancouver, BC, Canada, Dec. 2020, pp. 1–16.
https://github.com/XuJing1022/HiT-MAC/blob/main/perception.py

类比多头注意力中
${d_q} = {d_k} = {d_v} = {d_{in}}$
${p_q} = {p_k} = {p_v} = {d_{att}}$

${{\bf{x}}_i} \in {{\mathbb R}^{1 \times d_{in}}},{\bf{X}} = \left[ {\begin{array}{c} {{{\bf{x}}_1}}\\ \cdots \\ {{{\bf{x}}_{nm}}} \end{array}} \right] \in {{\mathbb R}^{nm \times d_{in}}}$
${\bf{W}} \in {{\mathbb R}^{d_{att}\times d_{in}}}$
${\bf{Q}} = {\rm{tanh}}\left( {{{\bf{W}}_q}{\bf{X}}^\top} \right)^\top \in {{\mathbb R}^{nm \times d_{att}}}$
${\bf{K}} = {\rm{tanh}}\left( {{{\bf{W}}_k}{\bf{X}}^\top} \right)^\top \in {{\mathbb R}^{nm \times d_{att}}}$
${\bf{V}} = {\rm{tanh}}\left( {{{\bf{W}}_v}{\bf{X}}^\top} \right)^\top \in {{\mathbb R}^{nm \times d_{att}}}$
$f({\bf{Q}},{\bf{K}},{\bf{V}}) = {\boldsymbol{\alpha}} {\left( {{\bf{Q}},{\bf{K}}} \right)^ \top }{\bf{V}} ={\rm{softmax}}\left( {\frac{1}{{\sqrt d }}{\bf{Q}}{{\bf{K}}^ \top }} \right){\bf{V}} \in {{\mathbb R}^{nm \times d_{att}}}$

class AttentionLayer(torch.nn.Module):
    def __init__(self, feature_dim, weight_dim, device):
        super(AttentionLayer, self).__init__()
        self.in_dim = feature_dim
        self.device = device
        self.Q = xavier_init(nn.Linear(self.in_dim, weight_dim))
        self.K = xavier_init(nn.Linear(self.in_dim, weight_dim))
        self.V = xavier_init(nn.Linear(self.in_dim, weight_dim))
        self.feature_dim = weight_dim
    def forward(self, x):
        # param x: [num_agent, num_target, in_dim]
        # return z: [num_agent, num_target, weight_dim]
        # z = softmax(Q,K)*V
        q = torch.tanh(self.Q(x))  # [batch_size, sequence_len, weight_dim]
        k = torch.tanh(self.K(x))  # [batch_size, sequence_len, weight_dim]
        v = torch.tanh(self.V(x))  # [batch_size, sequence_len, weight_dim]
        z = torch.bmm(F.softmax(torch.bmm(q, k.permute(0, 2, 1)), dim=2), v)  # [batch_size, sequence_len, weight_dim]
        global_feature = z.sum(dim=1)
        return z, global_feature