了解Transformers是如何“思考”的

Transformer 模型是 AI 系统的基础。已经有了数不清的关于 "Transformer 如何工作" 的核心结构图表。

但是这些图表没有提供任何直观的计算该模型的框架表示。当研究者对于 Transformer 如何工作抱有兴趣时，直观的获取他运行的机制变得十分有用。

Thinking Like Transformers 这篇论文中提出了 transformer 类的计算框架，这个框架直接计算和模仿 Transformer 计算。使用 RASP 编程语言，使每个程序编译成一个特殊的 Transformer。

在这篇博客中，我用 python 复现了 RASP 的变体 (RASPy)。该语言大致与原始版本相当，但是多了一些我认为很有趣的变化。通过这些语言，作者 Gail Weiss 的工作，提供了一套具有挑战性的有趣且正确的方式可以帮助了解其工作原理。

!pip install git+https://github.com/srush/RASPy

在说起语言本身前，让我们先看一个例子，看看用 Transformers 编码是什么样的。这是一些计算翻转的代码，即反向输入序列。代码本身用两个 Transformer 层应用 attention 和数学计算到达这个结果。

def flip():
    length = (key(1) == query(1)).value(1)
    flip = (key(length - indices - 1) == query(indices)).value(tokens)
    return flip
flip()

本文内容目录

  • 部分一：Transformers 作为代码
  • 部分二：用 Transformers 编写程序

Transformers 作为代码

我们的目标是定义一套计算形式来最小化 Transformers 的表达。我们将通过类比，描述每个语言构造及其在 Transformers 中的对应。(正式语言规范请在本文底部查看论文全文链接)。

这个语言的核心单元是将一个序列转换成相同长度的另一个序列的序列操作。我后面将其称之为 transforms。

输入

在一个 Transformer 中，基本层是一个模型的前馈输入。这个输入通常包含原始的 token 和位置信息。

在代码中，tokens 的特征表示最简单的 transform，它返回经过模型的 tokens，默认输入序列是 "hello":

tokens

如果我们想要改变 transform 里的输入，我们使用输入方法进行传值。

tokens.input([5, 2, 4, 5, 2, 2])

作为 Transformers，我们不能直接接受这些序列的位置。但是为了模拟位置嵌入，我们可以获取位置的索引:

indices
sop = indices
sop.input("goodbye")

前馈网络

经过输入层后，我们到达了前馈网络层。在 Transformer 中，这一步可以对于序列的每一个元素独立的应用数学运算。

在代码中，我们通过在 transforms 上计算表示这一步。在每一个序列的元素中都会进行独立的数学运算。

tokens == "l"

结果是一个新的 transform，一旦重构新的输入就会按照重构方式计算:

model = tokens * 2 - 1
model.input([1, 2, 3, 5, 2])

该运算可以组合多个 Transforms，举个例子，以上述的 token 和 indices 为例，这里可以类别 Transformer 可以跟踪多个片段信息:

model = tokens - 5 + indices
model.input([1, 2, 3, 5, 2])
(tokens == "l") | (indices == 1)

我们提供了一些辅助函数让写 transforms 变得更简单，举例来说，where 提供了一个类似 if 功能的结构。

where((tokens == "h") | (tokens == "l"), tokens, "q")

map 使我们可以定义自己的操作，例如一个字符串以 int 转换。（用户应谨慎使用可以使用的简单神经网络计算的操作）

atoi = tokens.map(lambda x: ord(x) - ord('0'))
atoi.input("31234")

函数 (functions) 可以容易的描述这些 transforms 的级联。举例来说，下面是应用了 where 和 atoi 和加 2 的操作

def atoi(seq=tokens):
    return seq.map(lambda x: ord(x) - ord('0')) 

op = (atoi(where(tokens == "-", "0", tokens)) + 2)
op.input("02-13")

注意力筛选器

到开始应用注意力机制事情就变得开始有趣起来了。这将允许序列间的不同元素进行信息交换。

我们开始定义 key 和 query 的概念，Keys 和 Queries 可以直接从上面的 transforms 创建。举个例子，如果我们想要定义一个 key 我们称作 key。

key(tokens)

对于 query 也一样

query(tokens)

标量可以作为 key 或 query 使用，他们会广播到基础序列的长度。

query(1)

我们创建了筛选器来应用 key 和 query 之间的操作。这对应于一个二进制矩阵，指示每个 query 要关注哪个 key。与 Transformers 不同，这个注意力矩阵未加入权重。

eq = (key(tokens) == query(tokens))
eq

一些例子：

  • 选择器的匹配位置偏移 1:
offset = (key(indices) == query(indices - 1))
offset
  • key 早于 query 的选择器:
before = key(indices) < query(indices)
before
  • key 晚于 query 的选择器:
after = key(indices) > query(indices)
after

选择器可以通过布尔操作合并。比如，这个选择器将 before 和 eq 做合并，我们通过在矩阵中包含一对键和值来显示这一点。

before & eq

使用注意力机制

给一个注意力选择器，我们可以提供一个序列值做聚合操作。我们通过累加那些选择器选过的真值做聚合。

(请注意：在原始论文中，他们使用一个平均聚合操作并且展示了一个巧妙的结构，其中平均聚合能够代表总和计算。RASPy 默认情况下使用累加来使其简单化并避免碎片化。实际上，这意味着 raspy 可能低估了所需要的层数。基于平均值的模型可能需要这个层数的两倍)

注意聚合操作使我们能够计算直方图之类的功能。

(key(tokens) == query(tokens)).value(1)

视觉上我们遵循图表结构，Query 在左边，Key 在上边，Value 在下面，输出在右边

一些注意力机制操作甚至不需要用到输入 token 。举例来说，去计算序列长度，我们创建一个 " select all " 的注意力筛选器并且给他赋值。

length = (key(1) == query(1)).value(1)
length = length.name("length")
length

这里有更多复杂的例子，下面将一步一步展示。(这有点像做采访一样)

我们想要计算一个序列的相邻值的和，首先我们向前截断:

WINDOW=3
s1 = (key(indices) >= query(indices - WINDOW + 1))  
s1

然后我们向后截断:

s2 = (key(indices) <= query(indices))
s2

两者相交:

sel = s1 & s2
sel

最终聚合:

sum2 = sel.value(tokens) 
sum2.input([1,3,2,2,2])

这里有个可以计算累计求和的例子，我们这里引入一个给 transform 命名的能力来帮助你调试。

def cumsum(seq=tokens):
    x = (before | (key(indices) == query(indices))).value(seq)
    return x.name("cumsum")
cumsum().input([3, 1, -2, 3, 1])

层

这个语言支持编译更加复杂的 transforms。他同时通过跟踪每一个运算操作计算层。

这里有个 2 层 transform 的例子，第一个对应于计算长度，第二个对应于累积总和。

x = cumsum(length - indices)
x.input([3, 2, 3, 5])

用 transformers 进行编程

使用这个函数库，我们可以编写完成一个复杂任务，Gail Weiss 给过我一个极其挑战的问题来打破这个步骤，我们可以加载一个添加任意长度数字的 Transformer 吗？

例如: 给一个字符串 "19492+23919", 我们可以加载正确的输出吗？

如果你想自己尝试，我们提供了一个版本你可以自己试试:
https://colab.research.google.com/github/srush/raspy/blob/main/Blog.ipynb

挑战一 : 选择一个给定的索引

加载一个在索引 i 处全元素都有值的序列

def index(i, seq=tokens):
    x = (key(indices) == query(i)).value(seq)
    return x.name("index")
index(1)

挑战二 ：转换

通过 i 位置将所有 token 移动到右侧。

def shift(i=1, default="_", seq=tokens):
    x = (key(indices) == query(indices-i)).value(seq, default)
    return x.name("shift")
shift(2)

挑战三 ：最小化

计算序列的最小值。(这一步开始变得困难，我们版本用了 2 层注意力机制)

def minimum(seq=tokens):
    sel1 = before & (key(seq) == query(seq))
    sel2 = key(seq) < query(seq)
    less = (sel1 | sel2).value(1)
    x = (key(less) == query(0)).value(seq)
    return x.name("min")
minimum()([5,3,2,5,2])

挑战四：第一索引

计算有 token q 的第一索引 (2 层)

def first(q, seq=tokens):
    return minimum(where(seq == q, indices, 99))
first("l")

挑战五 ：右对齐

右对齐一个填充序列。例："ralign().inputs('xyz___') ='—xyz'" (2 层)

def ralign(default="-", sop=tokens):
    c = (key(sop) == query("_")).value(1)
    x = (key(indices + c) == query(indices)).value(sop, default)
    return x.name("ralign")
ralign()("xyz__")

挑战六：分离

把一个序列在 token "v" 处分离成两部分然后右对齐 (2 层):

def split(v, i, sop=tokens):

    mid = (key(sop) == query(v)).value(indices)
    if i == 0:
        x = ralign("0", where(indices < mid, sop, "_"))
        return x
    else:
        x = where(indices > mid, sop, "0")
        return x
split("+", 1)("xyz+zyr")
split("+", 0)("xyz+zyr")

挑战七：滑动

将特殊 token "<" 替换为最接近的 "<" value (2 层):

def slide(match, seq=tokens):
    x = cumsum(match) 
    y = ((key(x) == query(x + 1)) & (key(match) == query(True))).value(seq)
    seq =  where(match, seq, y)
    return seq.name("slide")
slide(tokens != "<").input("xxxh<<<l")

挑战八：增加

你要执行两个数字的添加。这是步骤。

add().input("683+345")
  1. 分成两部分。转制成整形。加入
❝

“683+345” => [0, 0, 0, 9, 12, 8]

  1. 计算携带条款。三种可能性：1 个携带，0 不携带，< 也许有携带。
❝

[0, 0, 0, 9, 12, 8] => “00<100”

  1. 滑动进位系数
❝

“00<100” => 001100"

  1. 完成加法

这些都是 1 行代码。完整的系统是 6 个注意力机制。（尽管 Gail 说，如果你足够细心则可以在 5 个中完成！）。

def add(sop=tokens):
    # 0) Parse and add
    x = atoi(split("+", 0, sop)) + atoi(split("+", 1, sop))
    # 1) Check for carries 
    carry = shift(-1, "0", where(x > 9, "1", where(x == 9, "<", "0")))
    # 2) In parallel, slide carries to their column                                         
    carries = atoi(slide(carry != "<", carry))
    # 3) Add in carries.                                                                                  
    return (x + carries) % 10
add()("683+345")
683 + 345
1028

完美搞定！

参考资料 & 文内链接:

  • 如果你对这个主题感兴趣想了解更多，请查看论文：Thinking Like Transformers
    https://arxiv.org/pdf/2106.06981.pdf
  • 以及了解更多 RASP 语言
    https://github.com/tech-srl/RASP
  • 如果你对「形式语言和神经网络」(FLaNN) 感兴趣或者有认识感兴趣的人，欢迎邀请他们加入我们的线上社区！
    https://flann.super.site
  • 本篇博文，包含库、Notebook 和博文的内容在
    https://github.com/srush/raspy
  • 本博客文章由 Sasha Rush 和 Gail Weiss 共同编写:Sasha Rush 博客地址
    https://rush-nlp.com/Gail Weiss 博客地址
    https://gailweiss.github.io/
❝

英文原文: Thinking Like Transformers: 
https://srush.github.io/raspy/

中文译者: innovation64 (李洋)


END



年度增速最快语言：Rust



这里有最新开源资讯、软件更新、技术干货等内容

点这里 ↓↓↓ 记得 关注✔ 标星⭐ 哦

相关推荐

  • Manticore Search 6正式发布，C++高性能搜索引擎
  • SeaTunnel，Apache基金会中第一个诞生自中国的数据集成平台项目
  • RISC-V有望与X86、ARM形成三足鼎立之势
  • 包含脏话的开源代码质量更高，“粗口大师”Linus：YES!
  • 台网红“谷阿莫”落选时力决委；尹锡悦称应在不阻碍韩中交流前提下考虑防疫措施；安倍枪击案搜查工作将结束 | 每日大新闻
  • 西班牙女权有多疯狂？警察都害怕女黑社会
  • 我也来玩一玩OpenAI ChatGPT，人工智能时代终究会来临 - 第459篇
  • 论文写作实操手册，yyds!
  • .Net 问世 | 历史上的今天
  • 微软“砍掉”元宇宙部门，刚成立四个月！网友：ChatGPT 会是下一个“元宇宙”吗？
  • Android 14 的首个开发者预览版，推出四大安全增强功能
  • 赶时间发布 ChatGPT 竞品后，Google 遭员工反讽：太仓促、拙劣、非 Google Style
  • 监控Python内存使用情况和代码执行时间！
  • 又一个接私活神器！
  • 他是中国第一美男，70岁依然风华绝代，却无父无母无名无姓，一生孤苦
  • 她走遍70个国家，拍下全球最美500张面孔：看到中国这张，我震惊了...
  • Java 近期新闻：JDK 20、新的 JEP 草案、JobRunr 6.0、GraalVM 22.3.1
  • 在面试中如何巧妙的展现架构能力？附200道面试真题+100例经典架构案例拆解 | 极客时间
  • 2023 年前端十大 Web 发展趋势
  • 告别SVN，Git成“独苗”：GitHub 在 13 年后宣布淘汰Subversion支持