【官方双语】直观解释注意力机制，Transformer的核心 | 【深度学习第6章】

UP主: 3Blue1Brown · 时长: 26:10 · 🔗 B站原视频

发布: 2024-05-04 · 收录: 2024-09-04

标签: Transformer · 注意力机制 · 深度学习 · 大语言模型 · 3Blue1Brown

模型要做什么：预测下一个词

The goal of the model that you and I are studying is to take in a piece of text and predict what word comes next.

输入文本会被切成一些小片段，叫 tokens。它们通常是词或者词的一部分。为了让例子更好想象，我们先假装 token 总是一个完整的词。

Transformer 的第一步，是给每个 token 配一个高维向量，也就是它的 embedding（词嵌入）。

词嵌入的核心直觉：高维空间里的“方向”对应语义

我最希望你记住的点是：在这个高维 embedding 空间里，“方向”可以对应语义含义。

上一章我们看过一个例子：某个方向对应“性别”，在这个空间里加上一个特定的“步子”，就能从阳性名词的 embedding 走到对应阴性名词的 embedding。这只是一个例子。你可以想象还有很多别的方向，对应各种语义关系。

Attention 的作用：把“泛化的词义”更新成“上下文里的具体词义”

你脑子里可以这么想：在 embedding 空间里，可能有很多不同的方向，编码了同一个词的不同含义（比如 mole 这个词）。一个训练得好的 attention 模块，会计算“需要往原始 embedding 上加什么”，把它推到某个更具体的方向上。

举个例子：如果某个词前面有 large，那么这个词的向量可能会被更新，让它更相关于“高的、大的”这些概念；如果前面还有 miniature，那它又应该继续更新，让它不再和“大”强相关。

更一般地说，Attention 不只是“把一个词变得更具体”。它允许模型把一个 embedding 里编码的信息，搬运到另一个 embedding 里，甚至两者在序列里相距很远；而且这种信息也不一定只是一两个词的简单信息，可能更丰富。

训练时的本质：一堆可学习矩阵在数据上被调出来的行为

Where the matrices are full of tunable weights, things that the model will learn based on data.

我这里用“形容词更新名词”的例子，只是为了说明你可以想象的一类行为。和很多深度学习一样，真实行为更难直接读出来，因为它来自大量参数的微调，目的是最小化某个损失函数。

Query 和 Key：用“提问-匹配”来算相关性

你可以把它理解成：每个词都要问一个问题，比如“我是不是一个形容词，并且我在这个位置上要影响后面的名词？”这类“问题”会被编码成一个向量（一个数字列表），叫 query。

接着，你还会用另一个矩阵，把每个 embedding 映射成另一组向量，叫 keys。

Conceptually, you want to think of the keys as potentially answering the queries.

query 矩阵和 key 矩阵都充满可学习参数，它们把原来的 embedding 映射到同一个更小的维度空间。直觉上：当某个 key 和某个 query 在方向上对得很齐（更对齐），就表示“更匹配”。

点积网格与 Softmax：把匹配程度变成权重

To measure how well each key matches each query, you compute a dot product between each possible key-query pair.

我喜欢把它想象成一个网格：每个格子是一个点积结果，点越大表示匹配越强，表示 key 和 query 对齐得越好。

一个技术细节：为了数值稳定，通常会把这些点积除以 key/query 空间维度的平方根。然后外面包一层 softmax，并且是按列（column by column）去做的。

至于公式里的 v（value）项，马上就讲。

训练时为什么要“同时预测很多位置”：并行让训练更高效

还有个我刚才略过的细节：训练时，当你把一段文本喂给模型，并根据它给“真实下一个词”的概率高不高来奖励或惩罚并调整权重时，让训练更高效的一种做法是：让它同时预测每个位置的“下一个 token”。

比如你关注一个短语时，它不只预测整段末尾的下一个词，也会预测 creature 后面是什么、the 后面是什么之类的。

Value：真正用来“改 embedding”的那部分信息

Computing this pattern lets the model deduce which words are relevant to which other words. Now you need to actually update the embeddings.

接下来要做的就是更新 embedding，让相关的词把信息传给它相关的另一些词。

最直接的做法：用第三个矩阵，叫 value matrix。你把某个词的 embedding（比如 Fluffy）乘上 value matrix，得到一个 value vector。然后把这个 value vector 加到另一个词的 embedding 上（比如 creature），从而让 creature 的 embedding 变成更贴合上下文的“Fluffy creature”。

在网格图的那一列里，你会把与 creature 强相关的那些词（比如 fluffy、blue）的 value vectors 按权重加起来，其他不相关的 value vectors 被压到接近 0。把这一列加权求和得到一个改变量（我会把它标成 delta e），再加回原始 embedding。

结果是：embedding 会变得更精细、更上下文相关，比如变成“a fluffy blue creature”这种更丰富的含义。

而且这不是只对一个词做，你会对所有位置都做同样的加权求和，于是 attention block 输出一整串“更精细的 embeddings”。

Zooming out, this whole process is what you would describe as one head of attention.

一个 Attention Head 的参数规模：Key/Query/Value 三类矩阵

As I've described things so far, this process is parametrized by three distinct matrices: key, query, value.

我也想顺便延续上一章的“参数记账”。以 GPT-3 的数字为例：

• key 和 query 矩阵：列数 12,288（embedding 维度），行数 128（更小的 key/query 空间维度）。每个大约 150 万参数量级。
• value 如果按最直觉的说法，会像是一个 12,288×12,288 的大方阵，因为输入输出都在大 embedding 空间里。

但实际实现里会把 value 映射拆成两步（相当于低秩分解）：先 down 到小维度，再 up 回大维度。用线性代数的话说，这是把整体 value 映射约束成一个 low-rank transformation。

按这种实现方式，这里总共会出现四个矩阵（Q、K、V-down、V-up），它们尺寸相近。加起来，一个 attention head 大约 630 万参数。

Self-Attention 和 Cross-Attention 的区别（顺带一提）

顺带说一下：我讲的这些通常叫 self-attention，用来和 cross-attention 区分。

cross-attention 常见在处理两种不同数据源时，比如一边是某种语言的文本，另一边是另一种语言的生成翻译；或者音频输入和正在生成的转写。它看起来几乎一样，差别是 key 和 query 来自不同的数据集：keys 可能来自一种语言，queries 来自另一种语言，于是 attention pattern 描述“哪边的哪些词对应哪边的哪些词”。

在这种设置里通常不需要 mask，因为没有“后面的 token 影响前面的 token”这种因果约束。

Multi-Headed Attention：并行跑很多个 Head，学到很多种“上下文改变含义”的方式

Everything we've described is a single head of attention, and a full attention block inside a transformer consists of multi-headed attention.

GPT-3 每个 block 里有 96 个 attention heads。意思就是：

• 有 96 组不同的 key/query 矩阵，产生 96 种不同的 attention patterns。
• 每个 head 也有自己的 value 映射，产生 96 组 value vectors 序列。
• 对于上下文里的每个位置，每个 head 都会提出一个“应该加到该位置 embedding 上的改变量”。
• 把所有 head 的改变量加起来，再加回原始 embedding，就得到 multi-headed attention 的输出：一串更精细的 embeddings。

整体直觉是：并行跑很多头，给模型足够容量，去学习很多种“上下文如何改变含义”的方式。每个 head 可能关注不同类型的关系（比如指代、修饰、句法结构等等）。

多层堆叠后：更抽象的语义可能被编码出来

There are many more chances for this more nuanced embedding to be influenced by its more nuanced surroundings.

网络越往下走，每个 embedding 从其他 embeddings 里吸收的意义越来越多，而这些来源 embedding 本身也越来越“细腻”。希望最终模型能编码更高层、更抽象的东西，不只是描述词和语法结构，还包括情感、文本风格（比如是不是诗），以及和文本相关的更抽象事实等等。

回到参数记账：GPT-3 有 96 层，所以 attention head 相关的 key/query/value 参数总量还要再乘 96，总共接近 580 亿个参数都花在 attention heads 上。

为什么 Transformer 适合做大规模：并行化带来的优势

这也意味着：你可以用 GPU 在短时间内并行做海量计算。

过去十几二十年深度学习一个很重要的经验是：光靠 scale（规模）就能带来很大的质变提升。所以，像 Transformer 这种非常适合并行化的架构有很大优势。

如果你想继续深入，我在视频描述里放了很多链接。特别是 Andrej Karpathy 或 Chris Olah 的内容通常都非常棒。我的朋友 Vivek 也做了几期更偏“动机和历史”的视频，The Art of the Problem 频道的 Brit Cruz 也有一支讲大语言模型历史的很不错的视频。