你把一段会议记录发给 AI，请它找出谁答应了什么。接着你把后半场的记录也贴进去。新贴的那段可能引用前半场提过的名字，也可能补齐一项决定的来由。模型处理更长的输入时，除了读入新文字，还多了哪些工作？
先约定一种长度。语言模型会用分词器把文本切成叫作 token 的单元。一个 token 可以是一个词，也可以只是词的一部分，切法取决于分词器。下面说的“变成两倍”，指的是分词后的位置数从 1,000 增到 2,000。两份字符数差一倍的文件，切完后未必也是这个比例。Hugging Face 的分词算法文档
一次比较在比什么？
这里的“比较”是自注意力中的一次打分。先看一层里的一个注意力头，它为序列里的每个位置准备三组数字，分别叫查询、键和值。可以把查询和键想成打分的两张卡片。要更新某个位置的信息，就拿它的查询与各个允许访问的位置的键配对。把两组数字逐项相乘再加起来，得到一个点积。模型缩放这些分数，再把它们转换成权重，用来加权汇总各个位置的值。
例如在“阿宁把钥匙放进书包。她出门时把书包背走了”这句话里，“她”附近的位置可以从前面提到“阿宁”的位置取到信息。这里只是在说明信息能从哪里来，实际分数由模型内部的表示决定。以下算的是有多少组查询与键可以配对，一组配对包含多次数字运算。Transformer 原论文第 3.2 节
计分表的行和列都变长了
做一个教学算例。假设一层自注意力允许任意位置访问整段文字，包括自身。把各个查询放在表格的行上，把各个键放在列上，每一格就对应一次打分。四个位置的表格有四行四列，共 16 格。位置增到八个时，表格有八行八列，共 64 格。
为什么多出来的格子有这么多？把八个位置分成前四个和后四个，整张表就由四个小方块拼起来。原来的位置彼此配对占一块，新增的位置彼此配对占一块。剩下两块分别是原有位置查看新增位置，以及新增位置查看原有位置。查询和键扮演不同角色，这两个方向各占一块。
这个例子换成 1,000 和 2,000 个 token，格子数就从 100 万变成 400 万。每组数字的长度固定时，单看这一头的打分，工作量随位置数的平方增长。Transformer 原论文的表 1 也给出了全连接自注意力的平方项。
如果只能看前文，就数三角形
按顺序生成文本时，用来预测下一步的当前位置只能看到自身和前面已经出现的 token。这条规则由因果掩码落实。在计分表里，第一行只留第一格，第二行留前两格，第八行留八格。四个位置共有 1 + 2 + 3 + 4 = 10 个允许的配对。八个位置共有 36 个，比原来多了 26 个。
一般来说，长度为 n 时，把从 1 到 n 的格数加起来，得到 n × (n + 1) ÷ 2。长度从 1,000 增到 2,000，允许的配对从 500,500 增到 2,001,000，接近四倍。表格只保留一个三角形，但那个三角形的宽和高都在变长。这里数的是能参与计算结果的配对。实现时，可以先算满整张表再盖住一半，也可以跳过整块用不到的区域。
也可以沿用前后两段的分法核对这 36 格。前四个位置内部有 10 格，后四个位置内部也有 10 格。后四个位置还可以各自查看全部四个旧位置，多出 16 格。旧位置看不到新位置，因而这一方向没有格子。即使原来四行完全不变，新增的四行也比原来四行长得多。
四倍的格子，要等四倍久吗？
如果把输入一次性送进模型处理，在密集自注意力下可以用上面的表格估算配对规模。等到模型逐个生成后续 token 时，若启用了键值缓存，先前算好的键和值可以复用。处理一个新位置，只需要用它的查询算新增的一行。也就是说，处理整份输入和处理单个新位置，对应两种不同的计数问题。Hugging Face 的键值缓存说明
计算时还可以少搬一些数据。FlashAttention 的做法是把表格拆成小块，在 GPU 的片上存储中分块运算，同时逐步汇总权重所需的量，避免把整张中间表写进显存。它保留了密集注意力的计算目标。即使要考虑同样的位置配对，换一种安排，临时存储和数据搬运也会变。FlashAttention 原论文第 3 节
还有一种变化是改配对规则。假设某一层里，每个位置最多看包括自身在内的最近 128 个位置，表格就只剩对角线附近的一条带。长度为 2,000 时，每一行最多 128 格，配对总数就不超过 256,000。这个上限只适用于假设中的这一层。隔得更远的位置在这一层少了直接的连接，信息能否经过其他层传过去，还得看模型怎样搭建。
因此，想知道把会议记录加长后要多等多久，需要先问测的是送入全文到收到第一个 token 的时间，还是整段回答生成完的时间。前者和整份输入的处理有关，后者还要看模型生成了多少内容、硬件怎样并行计算。如果目的是找出某项决定，记录里相隔很远的两句能否直接交换信息，则是另一项具体的问题。把这几个问题分开，计分表才告诉我们它能回答的那一部分。
发表回复