用einops简单明了实现自注意力

1.自注意力(Self-Attention)

Transformer凭借其自注意力机制,有效解决了字符之间或像素之间的长距离依赖,日益称为NLP和CV领域的通用架构。
自注意力机制是Transformer的核心,如何简洁有效实现Self-Attention?这里介绍一种法,使用einops和PyTorch中einsum。自注意力的计算公式如下

2.自注意力计算的详细过程如下图所示


这里假设x的形状(1,4,4),标记(Token)个数为4,每个token转换为长度为4的向量,嵌入(Embedding)的维度为3(dim=3)。

3、详细实现代码

用代码实现上述计算过程

4、测试

运行结果
tensor([[[-0.3127, 0.4551, -0.0695],
[-0.3176, 0.4594, -0.0715],
[-0.3133, 0.4551, -0.0703],
[-0.3116, 0.4531, -0.0702]]], grad_fn=)

einops的使用可参考官网:
https://github.com/arogozhnikov/einops