#

一款无注意的变形金刚

本文用一种新的高效运算取代了自我注意力层，该运算的存储复杂度为 $O (T d)$ ，其中 $T$ 是序列长度， $d$ 是嵌入的维度。

本文介绍了 AFT 以及 AFT-Local 和 AFT-conv。这里我们实现了 aft-Local，它关注自回归模型中的 cloby 代币。

免注意变形金刚

A@@

FT（类似于 MHA）首先将嵌 $X$ 入转换为具有学习权重的查询 $Q = X W^{Q}$ 、键 $K = X W^{K}$ 和值 $V = X W^{V}$ 张量。每个位置的输出都是通过以下运算计算 $t \in [1, T]$ 的。

$Y_{t} = σ (Q_{t}) ⊙ \frac{\sum _{t^{'} = 1}^{T} e x p ( K _{t^{'}} + w _{t, t^{'}} ) ⊙ V _{t^{'}}}{\sum _{t^{'} = 1}^{T} e x p ( K _{t^{'}} + w _{t, t^{'}} )}$

，其中 $⊙$ 是元素乘积， $σ$ 是非线性（sigmoid）， $w \in R^{T \times T}$ 是成对位置偏差的学习矩阵。

这意味着我们取值的加权平均值并将其乘以查询。这样就无需计算 MHA 所需的 $T \times T$ 注意力矩阵，从而降低了内存需求。

AFT 本地

AFT Local 仅在本地应用学习的配对位置偏差：

w_{t, t^{'}}^{'} = {w_{t, t^{'}}, 0, for ∣ t - t^{'} ∣ < s otherwi s e

，其中 $s \leq T$ 是本地窗口大小。

尽管 $w_{t, t^{'}}^{'}$ 不在本地窗口 $0$ 之外，但 AFT 操作仍使用来自其他区域的键值对。这与本地转换器不同，本地窗口之外的嵌入完全不可见。

以下是 AFT Local 模型的训练代码。

59from typing import Optional
60
61import torch
62from torch import nn
63
64from labml_helpers.module import Module

#

AFT 本地操作

$Y_{t} = σ (Q_{t}) ⊙ \frac{\sum _{t^{'} = 1}^{T} e x p ( K _{t^{'}} + w _{t, t^{'}} ) ⊙ V _{t^{'}}}{\sum _{t^{'} = 1}^{T} e x p ( K _{t^{'}} + w _{t, t^{'}} )}$

在哪里，

w_{t, t^{'}}^{'} = {w_{t, t^{'}}, 0, for ∣ t - t^{'} ∣ < s otherwi s e

67class AFTLocal(Module):

#

d_model 是query 、key 和value 向量中的要素数。
seq_len 是 $T$
local_window_size 是本地窗口大小 $s$
bias 是是否为 $K$ 和的变换 $Q$ 设置偏置参数 $V$ 。

86    def __init__(self, d_model: int, seq_len: int, local_window_size: int, bias: bool = True):

#

94        super().__init__()

#

本地窗口大小 $s$

97        self.local_window_size = local_window_size

#

这些变换query 、key 和value 向量。

99        self.query = nn.Linear(d_model, d_model, bias=bias)
100        self.key = nn.Linear(d_model, d_model, bias=bias)
101        self.value = nn.Linear(d_model, d_model, bias=bias)

#

成对位置偏差 $w \in R^{T \times T}$

103        self.pos_bias = nn.Parameter(torch.zeros(seq_len, seq_len), requires_grad=True)

#

面具用于 $w_{t, t^{'}}$

105        self.local_mask = nn.Parameter(self.create_local_mask(seq_len, local_window_size), requires_grad=False)

#

激活 $σ$

107        self.activation = nn.Sigmoid()

#

输出层

109        self.output = nn.Linear(d_model, d_model)

#

创建局部蒙版

这会为以下对象创建遮罩

m_{t, t^{'}} = {1, 0, for ∣ t - t^{'} ∣ < s otherwi s e

111    @staticmethod
112    def create_local_mask(seq_len, local_window_size):

#

初始化为一

128        local_mask = torch.ones(seq_len, seq_len, dtype=torch.bool)

#

做 $t^{'} - t \geq s$ 零

130        local_mask = torch.tril(local_mask, local_window_size - 1)

#

做 $t - t^{'} \geq s$ 零

132        local_mask = torch.triu(local_mask, -(local_window_size - 1))

#

135        return local_mask

#

query key 和value 是存储查询、键和值的令牌嵌入集合的张量。它们有形状[seq_len, batch_size, d_model] 。

mask 有形状[seq_len, seq_len, batch_size] 并mask[i, j, b] 指示是否为批量查询b ，位置处的查询i 有权访问位置处的键值j 。

137    def forward(self, *,
138                query: torch.Tensor,
139                key: torch.Tensor,
140                value: torch.Tensor,
141                mask: Optional[torch.Tensor] = None):

#

query ，key 并且value 有形状[seq_len, batch_size, d_model]

153        seq_len, _, _ = query.shape
154
155        if mask is not None:

#

mask 有形状[seq_len_q, seq_len_k, batch_size] ，其中第一个维度是查询维度。如果查询维度等于 $1$ 它将被广播。

159            assert mask.shape[0] == 1 or mask.shape[0] == query.shape[0]
160            assert mask.shape[1] == key.shape[0]
161            assert mask.shape[2] == 1 or mask.shape[2] == query.shape[1]

#

转换查询、键和值嵌入

164        query = self.query(query)
165        key = self.key(key)
166        value = self.value(value)

#

得到

w_{t, t^{'}}^{'} = {w_{t, t^{'}}, 0, for ∣ t - t^{'} ∣ < s otherwi s e

使用口罩

179        pos_bias = self.pos_bias[:seq_len, :seq_len] * self.local_mask[:seq_len, :seq_len]
180        pos_bias = pos_bias.unsqueeze(-1)
181        pos_bias.masked_fill_(~mask, float('-inf'))

#

Y_{t} = σ (Q_{t}) ⊙ \frac{\sum _{t^{'} = 1}^{T} exp ( K _{t^{'}} + w _{t, t^{'}} ) ⊙ V _{t^{'}}}{\sum _{t^{'} = 1}^{T} exp ( K _{t^{'}} + w _{t, t^{'}} )} = σ (Q_{t}) ⊙ \frac{\sum _{t^{'} = 1}^{T} e x p ( w _{t, t^{'}} ) ⊙ e x p ( K _{t^{'}} ) ⊙ V _{t^{'}}}{\sum _{t^{'} = 1}^{T} e x p ( w _{t, t^{'}} ) ⊙ e x p ( K _{t^{'}} )}

我们 $e x p (K_{t^{'}})$ 分别计算 $e x p (w_{t, t^{'}})$ ， $e x p (K_{t^{'}}) ⊙ V_{t^{'}}$ 然后进行矩阵乘法。为了清楚起见，我们使用 einsum。

#

我们在计算指数 $m a x_{t^{'}} (w_{t, t^{'}})$ 之前减去 $m a x_{t^{'}} (K_{t^{'}})$ 和，以稳定softmax的计算。

$x_{i}$ if 大 $e x p (x_{i})$ 变大，计算 $\frac{\sum e x p ( x _{i} ) y _{i}}{\sum e x p ( x _{i} )}$ 变得不稳定。在计算分子和分母的指数之前减去一个常数将抵消。并且可以帮助稳定计算。所以我们减去 $max (x_{i})$ 以稳定计算。

203        max_key = key.max(dim=0, keepdims=True)[0]
204        max_pos_bias = pos_bias.max(dim=1,  keepdims=True)[0]

#

$exp (K_{t^{'}} - m a x_{t^{'}} (K_{t^{'}}))$

207        exp_key = torch.exp(key - max_key)

#

$exp (w_{t, t^{'}} - m a x_{t^{'}} (w_{t, t^{'}}))$

209        exp_pos_bias = torch.exp(pos_bias - max_pos_bias)

#

分子部分 $\sum_{t^{'} = 1}^{T} e x p (w_{t, t^{'}}) ⊙ e x p (K_{t^{'}}) ⊙ V_{t^{'}}$

212        num = torch.einsum('ijb,jbd->ibd', exp_pos_bias, exp_key * value)

#

分母部分 $\sum_{t^{'} = 1}^{T} e x p (w_{t, t^{'}}) ⊙ e x p (K_{t^{'}})$

214        den = torch.einsum('ijb,jbd->ibd', exp_pos_bias, exp_key)

#

输出 $Y_{t} = σ (Q_{t}) ⊙ \frac{\sum _{t^{'} = 1}^{T} e x p ( w _{t, t^{'}} ) ⊙ e x p ( K _{t^{'}} ) ⊙ V _{t^{'}}}{\sum _{t^{'} = 1}^{T} e x p ( w _{t, t^{'}} ) ⊙ e x p ( K _{t^{'}} )}$

219        y = self.activation(query) * num / den

#

输出层

222        return self.output(y)

#

测试局部掩码

225def _test_local_mask():

#

229    from labml.logger import inspect
230    inspect(AFTLocal.create_local_mask(10, 4))

#

234if __name__ == '__main__':
235    _test_local_mask()