前言

假设你在读这样一句话:

The pink elephant tried to get into the car but it was too []

读到这里,你几乎会下意识地猜到,下一个词很可能是 big

但你有没有想过,你是怎么做出这个判断的?你的大脑其实在瞬间做了一次“信息筛选”。它不会平均对待句子里的每一个词,而是自动挑出那些真正重要的线索。比如:

  • elephant 很重要,因为大象通常比车大。

  • get into the car 也很关键,因为这说明问题和尺寸有关。

因此这句话在你脑里应该“长成”这样:

The pink elephant tried to get into the car but it was too []

如果句子变成

The pink elephant tried to get into the swimming pool but it was too []

那你可能就不会想到 big,而是想到 scared

与此同时,有些信息却几乎被完全忽略了:

  • 大象是不是 pink 这件事,对预测下一个词没有什么帮助。

  • 像 the, but 这样的词,虽然让句子语法完整,但对我们判断最后的形容词并不重要。

换句话说,当我们理解一句话时,大脑其实一直在做一件事情:

在所有词中寻找最重要的那几个,并把注意力集中在它们身上。

这正是现代语言模型中一个核心机制的思想:注意力(Attention)。其作用就是让模型能够在一整段文本里决定:

  • 哪些词最值得关注

  • 哪些信息应该被忽略

  • 每个词对当前预测应该贡献多少

这样,当模型需要预测下一个词时,它就可以像人类一样,从相关的上下文中提取最有用的信息,而不会被无关细节干扰。那么问题来了:

模型到底是怎么决定该关注哪些词的?

为了回答这个问题,我们搭配着大神 Andrej Karpathy 的 microgpt 的代码给大家讲透其原理。


1. 简介 (Introduction)

Two hundred lines, zero dependencies.

Microgpt 是一个艺术品,一个只有 200 行纯 Python 代码、无任何依赖的单文件脚本。这份代码包含了完整的算法核心内容,包括:

  • 数据集(dataset)- a dataset to learn from

  • 分词器(tokenizer)- a tokenizer to read with

  • 自动求导(autograd)- autograd to compute the gradients

  • 变形金刚转换器(transformer)- a transformer to think

  • 训练循环(training)- a training loop to practice

  • 推理循环(inference)- an inference loop to speak

该代码已经无法再进一步简化,microgpt 是他多个项目(micrograd | makemore | nanogpt 等)的集大成之作,也是他十年来不断尝试将大语言模型(large lauguage model, LLM)简化到最本质形态的一次总结,优美至极 。它甚至可以完美地排成三栏显示:

本文的思路是以 Python 代码为基础,讲透实现 gpt 需要的核心概念和技术。由于该 Python 代码不依赖任何包,最后用 codex 很轻易的转换成 Kotlin 并且一次跑通,而且提速了 12 倍多。完整的 Python 和 Kotlin 代码见附录。


2. 数据集 (Dataset)


LLM 的数据通常是一连串文本数据流,这些数据可以选择性地被划分为多个文档。在生产级应用中,每个文档通常是一整个互联网网页;但在 microgpt 中,我们使用了一个更简单的示例:32,033 个名字,每行一个。


if not os.path.exists('input.txt'):
    import urllib.request
    names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'
    urllib.request.urlretrieve(names_url, 'input.txt')
docs = [line.strip() for line in open('input.txt') if line.strip()]
random.shuffle(docs)
print(f"num docs: {len(docs)}")  
# num docs: 32033

上面这段代码做了 3 件事:1)如果本地没有数据文件就下载;2)读入并清洗成文档列表; 3)打乱顺序并打印数据数量。

2.1 数据下载 (download)

if not os.path.exists('input.txt'):
    import urllib.request
    names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'
    urllib.request.urlretrieve(names_url, 'input.txt')

代码解释如下:

  • os.path.exists('input.txt'):检查当前目录下是否已经有这个文件

  • if not ...:如果没有,才执行下载(避免每次运行都重新下载)

  • import urllib.request:Python 自带的下载工具库

  • names_url:指向 GitHub 上的 names.txt 原始文本文件链接

  • urllib.request.urlretrieve(names_url, 'input.txt')

    • 从 names_url 下载内容

    • 保存到本地文件 input.txt

所以第一次跑会下载;以后再跑发现本地已有 input.txt 就跳过下载。


2.2 数据处理 (process)

docs = [line.strip() for line in open('input.txt') if line.strip()]

这一行是列表推导式,等价于:

  • 打开 input.txt,逐行读取(每次 line 是一行字符串,通常结尾带 \n

  • line.strip()

    • 去掉行首行尾的空白字符(空格、\n\t 等)

    • 例如 "emma\n".strip() → "emma"

  • if line.strip()

    • 这是过滤条件

    • strip() 后如果是空字符串 "" 就会被过滤掉(跳过空行)

结果:docs 是一个列表,里面每个元素都是一个名字字符串,比如 ["emma", "olivia", "liam", ...]

上面数据太大可能不好理解,下面用一个小数据集来说明。

注意 Steven 上一行是空白行,ETH 后三行是空白行,BTC 前面有一个空白字符,ETH 后面也有个空白字符(从截屏中看不出来,下面代码可以看出来)

print([line for line in open('input1.txt')])
# ['\n', 'Steven\n', '\n', 'Wang\n', 'Deep Quant\n', '\n', ' BTC\n', 'ETH \n', '\n', '\n', '\n']

print([line.strip() for line in open('input1.txt')])
# ['', 'Steven', '', 'Wang', 'Deep Quant', '', 'BTC', 'ETH', '', '', '']

print([line.strip() for line in open('input1.txt') if line.strip()])
# ['Steven', 'Wang', 'Deep Quant', 'BTC', 'ETH']

Line 1:直接打印每行,发现 ETH 那行是 'ETH \n'

Line 4:用 strip() 将每一行的首尾的空白字符都去掉,中间的空白字符保留,比如 'Deep Quant'

Line 7:用 if line.strip() 将空白字符的行都去掉,最后只剩有字符的行。

2.3 数据打乱 (shuffle)

random.shuffle(docs)
print(f"num docs: {len(docs)}")
  • random.shuffle 会打乱 docs 的顺序,让训练更稳一些

  • len(docs):名字(文档)的条数

  • f-string:把数量插进字符串里输出,会打印类似:num docs: 32033


3. 分词器 (Tokenizer)


和人不同,计算机处理的是数字而不是字符,因此需要一种方法,把文本转换为一串整数形式的 token id,并且还能再转换回来。


3.1. 一般 Token


在生产环境中使用的分词器(如 GPT-4 使用的 tiktoken)为了效率,会基于字符片段进行编码;但最简单的分词器做法是:为数据集中每个唯一字符分配一个整数编号。

uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1
BOS = len(uchars) # token id for a special Beginning of Sequence (BOS) token
vocab_size = len(uchars) + 1 # total number of unique tokens, +1 is for BOS
print(f"vocab size: {vocab_size}") 
# 27

在上面的代码中,我们收集了整个数据集中所有的唯一字符(这里只是小写字母 a–z),对它们进行排序,然后根据它们在列表中的索引为每个字母分配一个 id。需要注意的是,这些整数值本身没有任何语义意义;每个 token 只是一个独立的离散符号。

2. 特殊 Token


此外,我们还创建了一个额外的特殊 token,叫做 BOS(Beginning of Sequence),它充当一个分隔符:告诉模型“一个新的文档在这里开始/结束”。在后续训练时,每个文档都会在两侧加上 BOS,对于 emma 一次可表示为:


[BOS, e, m, m, a, BOS]

模型会学习到第一个 BOS 表示一个新名字的开始,第二个 BOS 表示这个名字的结束。

因此,最终的词汇表大小是 27:

  • 26 个小写字母(a–z)

  • 再加上 1 个 BOS 特殊 token



4. 自动求导 (Autograd)


训练神经网络需要梯度(gradients),对于模型中的每一个参数,我们想知道“如果把这个数值稍微调大一点,损失会变大还是变小?”整个计算图(computation graph)有很多输入,但只有一个标量输出。


反向传播(backpropagation)从这个单一的输出节点开始,沿着计算图向后传播,计算损失相对于每一个输入的梯度。它依赖于微积分中的链式法则(chain rule)。在生产环境中,像 PyTorch 这样的库会自动帮我们完成这一切。而在这里,我们从零开始自己实现它,用一个名为 Value 的类来完成。

4.1 初始化 (Initialization)

首先 __shots__ 告诉 Python Value 类只会有 'data', 'grad', '_children', '_local_grads' 这 4 个属性,不要用默认的 dict,更省内存也更快一点,对 microgpt 这种会创建大量节点的计算图很有用。

# Python optimization for memory usage
__slots__ = ('data', 'grad', '_children', '_local_grads') 

def __init__(self, data: float, children: tuple = (), local_grads: tuple = ()):
    self.data = data # scalar value of this node calculated during forward pass
    self.grad = 0                   # derivative of the loss w.r.t. this node, calculated in backward pass
    self._children = children # children of this node in the computation graph
    self._local_grads = local_grads # local derivative of this node w.r.t. its children

每个 Value 是计算图里的一个“节点”,需要记住四件事:

  1. data:前向计算出来的数值标量

  2. grad:反向传播时,loss 对当前节点的导数 ∂L/∂v,初始化为 0

  3. _children:这个节点由哪些子节点算出来的,也就是它的输入

  4. _local_grads:当前节点对每个子节点的局部导数,举例:z = x * y,dz/dx = y 和 dz/dy = x 这些就是 _local_grads

简单来说,一个 Value 对象封装了一个标量数值.data,记住自己的_children以及该运算的局部导数_local_grads,不停的更新grad。对于每个运算,它记载了两个方向传播的信息:

  • 前向传播(forward pass):接收一些输入,产生一个输出

  • 反向传播(backpropogation):记下该输出相对于每个输入是如何变化的,这个变化叫局部梯度(local gradient)

这正是自动求导(autograd)所需要的全部信息。

4.2. 运算符 (operator)


4.2.1. Value 类


4.2.1.1. 二元运算符 (Binary Operator)


先看两个操作数参与的运算需要的二元运算符。


加法运算符:__add__(self, other) = self.__add__(other)

def __add__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    return Value(self.data + other.data, (self, other), (1, 1))

这里 self 是 Value 类,但 other 可能是 float 或 Value 类,如果是 float 那就先包装成 Value。最后返回一个新的 Value 作为结果节点:

  • data = self.data + other.data(前向)

  • _children = (self, other)(图结构:这个结果由谁算出来)

  • _local_grads = (11)(局部梯度:d(a+b)/da = 1,d(a+b)/db = 1

乘法运算符:__mul__(self, other) = self.__mul__(other)

def __mul__(self, other):
    other = other if isinstance(other, Value) else Value(other)
    return Value(self.data * other.data, (self, other), (other.data, self.data))

同理,这里 self 是 Value 类,但 other 可能是 float 或 Value 类,如果是 float 那就先包装成 Value。最后返回一个新的 Value 作为结果节点:

  • data = self.data * other.data(前向)

  • _children = (self, other)(图结构:这个结果由谁算出来)

  • _local_grads = (other.dataself.data)(局部梯度:d(ab)/da = b,d(ab)/db = a

指数运算符:__pow__(self, other) = self.__pow__(other)

def __pow__(self, other): 
    return Value(self.data**other, (self,), (other * self.data**(other-1),))

这里 self 是 Value 类,但 other 只能是 float 类,因此不需要封装成 Value 类。最后返回一个新的 Value 作为结果节点:

  • data = self.data ** other(前向)

  • _children = (self, )(图结构:这个结果由谁算出来,注意只有一个)

  • _local_grads = (other * self.data**(other-1),)(局部梯度:d(a^n)/da = n*a^{n-1}

4.2.1.2. 一元运算符 (Unary Operator)

再看一个操作数参与的运算需要的一元运算符。

def log(self): return Value(math.log(self.data), (self,), (1/self.data,))
def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))
def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))

对于一元运算符来说,它们的 _children 属性都是 (self, )

  • log(a) 的导数是 1/a

    • data = math.log(self.data)

    • _local_grads = (1/self.data,)

  • exp(a) 的导数是 e^a

    • data = math.exp(self.data)

    • _local_grads = (math.exp(self.data),)

  • relu(a) = max(0, a) 的导数是 1 (a > 0) 否则 0,

    • data = max(0, self.data)

    • _local_grads = (float(self.data > 0),)

4.2.1.3. 总结

上面代码的运算总结在下表。

在 Value 类里可以搭积木一样任意添加运算符,需要注意的是在 autograd 里:

  • Binary operator 节点通常有两个 children

  • Unary operator 节点只有一个 child

  • 每个 operator 需要:

    • 在 data 属性中存计算的值

    • 在 _children 属性中存 children 或者 child

    • 在 _local_grads 属性存每个 child 的局部梯度值

最后抽象一下:

Binary operator:    data = f(self, other)    _children = (self, other)    _local_grads = (df/dself, df/dother)
Unary operator:    data = f(self)    _children = (self,)    _local_grads = (df/dself,)


4.2.2. Float 类


下面的运算符就是基于 float 类的操作,比如 __neg__ 就是求相反数,函数签名中带 r 的就是 other 放在运算符的左边,self 放在运算符的左边。


def __neg__(self): return self * -1
def __radd__(self, other): return self + other
def __sub__(self, other): return self + (-other)
def __rsub__(self, other): return other + (-self)
def __rmul__(self, other): return self * other
def __truediv__(self, other): return self * other**-1
def __rtruediv__(self, other): return other * self**-1

注意这里的 self 和 other 都是 float 类而不是 Value 类,因此只返回运算结果,不返回 _children 和 _local_grads


4.3. 反向传播 (backpropgation)


给定一个神经网络模型:

  • 正向传播就是把参数 p 当已知量从头算到尾得到输出值 L(p)。

  • 反向传播就是把参数 p 当未知量从尾推到头得到偏导数 ∂L/∂p。

此外反向传播能一次性计算出所有的偏导数,因此不仅能算,主要是算得快。

4.3.1. 抽象理论


方法 backward() 会按照拓扑排序的逆序遍历整个计算图(从损失值 L 开始,最终到模型参数),在每一步应用链式法则。如果损失是 L,某个节点 v 有一个子节点 c,且局部梯度为 ∂v/∂c,那么:


上式本质上用一种直观的方式“乘两个数”。一个类比是“如果一辆汽车的速度是自行车的 2 倍,而自行车的速度是行人的 4 倍,那么汽车的速度就是行人的 2 × 4 = 8 倍。”链式法则也是同样的思想:沿着路径把变化率相乘。

下面的 backward()函数是最核心也不好懂的,我会事无巨细的讲清楚。

反向传播需要保证的是当算一个节点的梯度时,它的“上一个梯度”已经准备好了,如下图所示。

所以要按计算依赖顺序排好:先子后父的后序遍历,得到 topo

build_topo 是用深度优先遍历(Depth First Search, DFS)整个图

  • visited 防止重复访问(图会分叉,节点可能被多个路径引用)

  • topo.append(v) 放在递归最后 → 保证 children 先进入 topo

visited 的意义:同一个节点可能被多个地方引用(图会“分叉再汇合”),防止重复加入 topo。

接下来 self.grad = 1 设置在 loss 节点开始,因为 ∂L/∂L=1,也就是说损失对自身的变化率显然是 1。从这里开始,链式法则会沿着所有路径,把局部梯度一路乘回到模型参数。

从 topo 中后序遍历每个子节点和局部梯度,借助父节点的梯度来更新子节点的梯度。注意这里使用的是 +=是累加,而不是直接赋值。当一个值在计算图中的多个地方被使用时(也就是说,计算图发生了分支),梯度会沿着每一条分支分别向后传播,因此必须把来自各个分支的梯度相加。这源于多变量链式法则:如果某个变量 c 通过多条路径对最终的损失 L 产生影响,那么它的总导数就是各条路径贡献之和。

当 backward() 执行完成后,计算图中的每一个 Value 对象都会有一个 .grad 属性,其中存储的是 ∂L/∂v,也就是说,如果稍微“推动”一下这个值,最终的损失 L 会如何变化。

4.3.2. 具体实例

下面是一个具体例子。注意变量 a 被使用了两次(计算图产生分支),因此它的梯度是两条路径贡献的总和。

a = Value(2.0)
b = Value(3.0)
c = 3*a*b # c = 3*2*3 = 18
L = 4*a + 2*c # L = 4*2 + 2*18 = 44
L.backward()

这里 c = 3ab, L = 2c + 4a = 2*(3ab) + 4a,那么

  • dL/dL = 1

  • dL/dc = 2

  • dL/da = 2*3b + 4 = 22

  • dL/db = 2*3a = 12

  • dc/dc = 1

  • dc/da = 3b = 9

  • dc/db = 3a = 6

为了将整个 backward()过程讲透,我在 Value 类中加入一些额外的信息(用高亮区分新加的内容)。

改动的地方说明如下:

  • 属性多加了 name 纯粹为了打印好看,方便知道这是哪个节点。

  • 加了 __repr__ 打印格式:

    • nm:>13s 右对齐 13 宽度,让输出整齐

    • datagrad 保留两位小数

    • 输出会像:Value(      var a: data= 2.00, grad= 0.00)

  • Value 类运算符都加了 name = "函数名",为了打印好看

输出 DEBUG 信息的 backward()的函数如下:

当 debug 时会打印:

  • 当前访问节点 v

  • 每个 child 收到的 delta

  • child.grad 从 before 到 after 的变化

写一个函数专门演示:

运行得到下面结果。一开始 grad 都是初始化成 0,最后反向传播后得到 dL/da = 22, dL/db = 12, dL/dc = 2, dL/dL = 1,和 3.4.1 里计算出来的结果吻合。

[START] forward values:
Value( var a: data= 2.00, grad= 0.00)
Value( var b: data= 3.00, grad= 0.00)
Value( c = 3ab: data=18.00, grad= 0.00)
Value(L = 2c + 4a: data=44.00, grad= 0.00)

[END] final grads:
Value( var a: data= 2.00, grad=22.00)
Value( var b: data= 3.00, grad=12.00)
Value( c = 3ab: data=18.00, grad= 2.00)
Value(L = 2c + 4a: data=44.00, grad= 1.00)

下面的输出是在 L.backward(debug=True) 里打印的“反向传播访问顺序”。反传会按这个顺序从 loss 往回走,依次把梯度分发给 children。

[DEBUG] backwardorder (loss -> ... -> leaf):
00: Value(L = 2c + 4a: data=44.00, grad= 1.00)
01: Value( mul: data= 8.00, grad= 0.00)
02: Value( const: data= 4.00, grad= 0.00)
03: Value( mul: data=36.00, grad= 0.00)
04: Value( const: data= 2.00, grad= 0.00)
05: Value( c = 3ab: data=18.00, grad= 0.00)
06: Value( var b: data= 3.00, grad= 0.00)
07: Value( mul: data= 6.00, grad= 0.00)
08: Value( const: data= 3.00, grad= 0.00)
09: Value( var a: data= 2.00, grad= 0.00)

按打印的序号来解释上面的产出:

  • 00 Value(L = 2c + 4a: data=44.00, grad= 1.00)

    • 这是最终 loss 节点 L

    • grad=1.00 是因为在 backward 里做了 self.grad = 1.0(seed)

    • 反传从这里开始

  • 01 Value(mul: data= 8.00, grad= 0.00)

    • 这是 t3 = 4 * a,data=8 对应 4*2

  • 02 Value(const: data= 4.00, grad= 0.00)

    • 这是常数 4(它也被包装成 Value(4.0)

    • 注意常数节点也会有 grad(只是一般不关心)

  • 03 Value(mul: data=36.00, grad= 0.00)

    • 这是 t2 = 2 * c,data=36 对应 2*18

  • 04 Value(const: data= 2.00, grad= 0.00)

    • 这是常数 2

  • 05 Value(c = 3ab: data=18.00, grad= 0.00)

    • 这是给 c 手动命名的节点,data=18 对应 3*2*3

  •  06 Value(var b: data= 3.00, grad= 0.00)

    • 这是叶子变量 b

  • 07 Value(mul: data= 6.00, grad= 0.00)

    • 这是中间节点 t1 = 3 * a,data=6 对应 3*2

  • 08 Value(const: data= 3.00, grad= 0.00)

    • 这是常数 3

  • 09 Value(var a: data= 2.00, grad= 0.00)

    • 这是叶子变量 a

更多细节在 for child, local_grad in zip(v._children, v._local_grads) 的循环中打出,我们能更清楚的看出父节点到子节点的局部梯度 local_grad=xxx 和累计梯度变化 grad before -> after

最后总结出这个具体实例的计算图:


5. 生成式预训练转化器 (GPT)


生成式预训练转化器翻译过来 Generative Pre-trained Transformer,也是大家耳熟能详的 GPT 的全称,GPT 里的 T,即 transformer 就是一个能够在整段文本中动态寻找重要信息,并不断组合这些信息来理解语言的模型。

5.1. 模型 (model)


5.1.1. 概览 (overview)


和几乎所有 LLM 一样,GPT 里的核心模型是 Transformer Block。它本身只是“一层如何计算”的问题;而我们通常所说的 Transformer 模型或大模型架构,则是在此基础上,围绕 Block 的堆叠方式、信息流动方式以及训练与推理范式所形成的整体设计。它做的事情其实很简单:

  1. 分词:用 tokensizer 技术将文本拆分成 token(本文忽略其细节)

  2. 编码:用 embedding 算法把 token 转成向量(小节 5.2)

  3. 转换:每一个 transformer接收向量,进行一次完整的上下文建模,输出维度相同的新向量

    1. 自我注意:在 attention 模块里让每个词去“查看”句子里的其他词,看看哪些词对自己最有帮助,根据这些重要程度,把信息组合起来

    2. 加强记忆:在 feed-forward 网络里,将自注意力捕获到的上下文信息加工抽象强化,为后续的语言建模输出提供“高阶特征”

  4. 重复:把 transformer 加量成 transformer 模块,将步骤 3 做很多次,让理解越来越深

  5. 预测:训练好模型后它就能预测下一个词

来来来,一图胜千言。

在 microgpt 中,为了让代码更简单,我们只用一个转换器模块

5.1.2. 参数 (parameter)


模型参数就是模型所包含的知识。一开始是随机初始化的,然后在训练过程中通过迭代不断优化。每个参数具体起什么作用,要等到我们在下面定义模型结构时才会更清楚;不过现在,我们只需要先把它们初始化出来。

n_layer = 1                 # depth of the transformer neural network (number of layers)
n_embd = 16                 # width of the network (embedding dimension)
block_size = 16             # maximum context length of the attention window (note: the longest name is 15 characters)
n_head = 4                  # number of attention heads
head_dim = n_embd // n_head # derived dimension of each head

上面参数的解释如下:

  • n_layer:1 个 transformer block

  • n_embd:词向量的维度,每个 token 变成一个 16 维向量

  • block_size:最大上下文长度,即模型在预测下一个 token 时,只能看到前面的 16 个 token

  • n_head:attention block 的数量,即做了 4 次 attention,每个 head 会关注不同的信息

  • head_dim = n_embd // n_head:每个 attention head 的维度。

在 microgpt 里,多头会把词向量分成几块,原始向量是 16 维,

[ x1 x2 x3 x4 | x5 x6 x7 x8 | x9 x10 x11 x12 | x13 x14 x15 x16 ]

每个 head 处理其中一块,head1 处理 [x1 x2 x3 x4],head2 处理 [x5 x6 x7 x8],head3 处理 [x9 x10 x11 x12],head4 处理 [x13 x14 x15 x16]。这样模型可以从不同子空间理解信息。

由于该 Python 代码不依赖任何库,因此用 list of list 来代表矩阵,行数为 nout,列数为 nin。矩阵里每个参数都会被初始化为一个从高斯分布(Gaussian distribution)中采样得到的小随机数。

matrix = lambda nout, nin, std=0.08: 
         [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]

定义 state_dict 把这些参数组织成带名字的矩阵。

state_dict = {
                'wte': matrix(vocab_size, n_embd), 
                'wpe': matrix(block_size, n_embd), 
                'lm_head': matrix(vocab_size, n_embd)
             }
for i in range(n_layer):
    state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)

state_dict 包括:

  • embedding 权重

    • wte:token 编码矩阵,大小是 vocab_size * n_embd

    • wpe:position 编码矩阵,大小是 block_size * n_embd

  • attention 权重

    • layer{i}.attn_wq:第 i 层的 query 矩阵,大小是 n_embd * n_embd

    • layer{i}.attn_wk:第 i 层的 key 矩阵,大小是 n_embd * n_embd

    • layer{i}.attn_wv:第 i 层的 value 矩阵,大小是 n_embd * n_embd

    • layer{i}.attn_wo:第 i 层的 输出矩阵,大小是 n_embd * n_embd

  • MLP 权重

    • layer{i}.mlp_fc1:第 i 层的升维矩阵 ,大小是 4*n_embd * n_embd

    • layer{i}.mlp_fc2:第 i 层的降维矩阵 ,大小是 n_embd * 4*n_embd

  • 最终的输出投影层 lm_head,大小是 vocab_size * n_embd

把所有参数展平成一个列表 params,这样优化器之后就可以遍历它们并进行更新。

params = [p for mat in state_dict.values() for row in mat for p in row] # flatten params into a single list[Value]
print(f"num params: {len(params)}") 
# num params: 4192

在这个非常小的模型里,总共有 4,192 个参数,计算过程如下:

5.1.3. 架构 (architecture)


模型架构本质上是一个无状态函数(stateless function),它接收一个 token、一个 position、模型参数以及来自之前位置缓存的 key/value,然后输出 logits,表示模型认为序列中下一个token 最可能是什么。我们基本遵循 GPT-2 的架构,但做了一些小的简化:

  • 归一化操作用 RMSNorm 代替 LayerNorm

  • 不使用偏置项(bias)

  • 激活函数用 ReLU 代替 GeLU

接下来,我们先定义三个小的辅助函数:

linear 是一个矩阵-向量乘法。它接收一个向量 x 和一个权重矩阵 w,然后对 w 的每一行与 x 做一次点积。

softmax 会把一组原始分数转换成一个概率分布。这些分数的取值范围可以从 −∞ 到 +∞,而经过 softmax 后,所有值都会落在 [0, 1] 区间内,并且总和为 1。我们会先减去最大值来提高数值稳定性(在数学上不会改变结果,但可以防止 exp 计算时发生溢出)。

rmsnorm 全称是 Root Mean Square Normalization,会对一个向量进行缩放,使其数值的均方根为 1。这样可以防止激活值在网络传播过程中不断变大或变小,从而让训练更加稳定。

整个模型的函数如下:

该函数会在某个时间位置 pos_id 处理一个 token(token_id),并利用之前迭代中得到的一些上下文信息。这些信息以 keys 和 values 的激活值形式缓存起来,被称为 KV Cache。下面每个小节按步骤说明 gpt 整个过程。

5.2. 编码 (embedding)


神经网络无法直接处理像 5 这样的原始 token id,它只能处理向量。因此我们为每一个 token 学习出一个向量,通过词编码和位置编码的方式。

5.2.1. 词编码 (word embedding)


词编码的核心思想是把文本转换成数字。这种转换是必要的,因为许多计算机无法直接处理普通的文本字符串、要求输入必须是连续数值的向量。因此,词编码技术被用来把词汇表中的单词或短语映射为对应的实数向量。除了能够被学习算法处理之外,这种向量表示还具有两个重要而有利的特性:

  • 降维(dimensionality reduction)它是一种更加高效的表示方式。

  • 上下文相似性(contextual similarity)它是一种表达能力更强的表示方式。

直觉来说,在向量空间里,相似的词离得近,不相似的词离得远。

看下面词向量空间,你会把苹果放在哪里?

不出意外的话,苹果会放在离水果近的区域。

通常词向量的维度很大,512, 1024, 2048 很常见,这里为了便于可视化就选择 2 维向量来代表词。给一个简单的例子“苹果真好用”,分成“苹果”、“真”和“好用”三个词,转成三个词向量:

用 tn来表示词编码的结果是因为词的英文是 word 或 token,因为字母 w 通常表示权重,因此用 t 来表示。

5.2.2. 位置编码 (position embedding)


模型在看到一句话时,其实是同时看到所有词的。如果没有位置信息,下面两句话在模型眼里几乎是一样的:


    The dog looked at the boy and

    The boy looked at the dog and 


但这两句话的意思完全不同,如果要预测下一个词的话,很有可能


    The dog looked at the boy and [barked]

    The boy looked at the dog and [smiled]


因此词在句子里的顺序很重要,解决方法很简单,给每个词加上一个“位置标签”,每个词最终表示不是只有词的含义,而是词的含义加上位置标签。


用 p_n 来表示位置编码的结果是因为位置的英文是 position,最终编码后的输入用字母 x 来表示。

5.2.3. 代码分析


前三行代码做的事是:

  • 用 token_id在词编码表 state_dict['wte']里查找对应的一行

  • 用 pos_id在位置编码表 state_dict['wpe'] 里查找对应的一行

  • 得到两个向量后,将它们相加

这样模型在理解了这个 token 词义的同时,也知道了它在句子里的位置。最后一行对输入做了归一化,使得训练时更加稳定。

总结来说,上面代码处理下图红框里的流程:编码。

编码完成之后,向量随后被送入多个转换器模块,每一个转换器模块主要由两部分组成:

  1. 自注意力模块 (self-attention block)

  2. 前馈神经网络 (feed-forward neural net)

它们共同构成了一个稳定的计算单元。下面两节分别介绍自注意力模块和前馈神经网络。

GPT 里面的核心是 Transformer,而 Transformer 的设计就是在通信(通过 attention block)和计算(通过 Feed-forward net)之间交替进行。


5.3. 自注意力模块 (self-attention block)


5.3.1. 用途介绍


自注意力的核心目标,是让模型能够在上下文中动态聚焦关键信息,从而解决长距离依赖和歧义理解的问题。

比如看到下面句子时:

The dog chased the llama because it

当模型处理到  it  这个 token 时,需要判断它指代的是 dog 还是 llama。如果上下文语义表明更可能指向 llama,那么在处理  it 时,自注意力机制就会将 llama 相关的信息更强地融入当前 token 的表示中。

从计算流程上看,自注意力可以拆解为两个阶段:

  • 相关性打分 (relevance scoring):判断当前 token 与上下文中其他 token 的相关程度

  • 信息融合 (information aggregation):将高相关性的 token 信息按权重融合进当前表示

现在的问题时在自注意力模块里是怎么实现相关性打分和信息融合而使得模型“拥有”注意力的,这里要引出模块里的三巨头了:query,key 和 value,我称它们为 QKV 三巨头。

我相信大家看到上面的公式已经懵逼了,我要开始讲故事了。

5.3.2. QKV 三巨头


想象你走进一座巨大的图书馆,想研究一个问题 (spoil alert:三巨头的 query 来了):

为什么大象进不了一辆小汽车?

图书馆里有大量书籍,每本书都记录着不同的信息。但你不会自己一本一本去翻。你请了 4 位不同领域的图书管理员帮你一起查资料,比如有

  • 一位负责动物学

  • 一位负责机械结构

  • 一位负责物理尺寸

  • 一位负责交通工具设计

他们会同时开始查资料。每个管理员都会做同样三件事情:

  1. 看你的问题(query)

你的问题是:为什么大象进不了汽车?这就是 query —— 当前需要的信息。

  1. 扫描所有书的标签(key)

图书馆的每本书都有标签,例如:动物、汽车、颜色、尺寸、历史等。这些标签就是 key

管理员会判断:哪些书和这个问题最相关?

  1. 阅读书里的内容(value)

当管理员找到相关的书后,他会阅读书里的内容,例如:

  • 大象平均高度

  • 汽车车门尺寸

  • 车辆内部空间

这些具体的信息就是 value


为什么要多个管理员?不同管理员会关注不同的信息:

管理员

关注点

动物学专家

大象的体型

工程专家

汽车的结构

物理专家

空间和尺寸

交通专家

车辆的设计

他们各自提取信息,然后把结果汇总。这就是多头注意力 (multi-head attention)。

自注意力模块最终会:

  1. 让每个词提出一个 query

  2. 在所有词的 Key 中寻找匹配

  3. 根据匹配度分配 attention weights

  4. 从对应的 value 中提取信息

  5. 汇总多个 attention head 的结果

这样模型就能在整句话中找到最重要的信息。

一句话总结:

Attention 就像在图书馆查资料:Query 是你的问题,Key 是书的标签,Value 是书里的内容,而 Multi-Head Attention 就像多位图书管理员同时帮你找答案。

5.3.3. 相关性打分 (relevance scoring)


首先自注意力不是直接作用在输入向量上,而是先通过三组线性投影矩阵:


  • 查询矩阵 (query matrix) WQ

  • 键矩阵 (key matrix) WK

  • 值矩阵 (value matrix) WV

将每个 token 映射为三种不同用途的向量:

  • Q:当前 token 想从上下文中获取什么信息

  • K:当前 token 能为其他 token 提供什么样的匹配线索

  • V:当该 token 被关注时,真正参与信息融合的内容

给一个具体例子,query 是“苹果真好用”,包含 3 个 token,每个 token 可编码成一个 2 维向量,因此词编码 X 是一个 3×2 的矩阵,而查询矩阵、键矩阵和值矩阵的大小都是 2×2。通过矩阵乘法生成的 Q, K, V 也是 3×2 的矩阵。具体操作如下图所示:

在相关性打分阶段,首先模型快速匹配出 query 和 key,等价的数学操作就是对每个 query (矩阵 Q 为载体) 与所有 key (矩阵 K 为载体) 做点积计算。这样得到一个 3×3 的矩阵 QKT,如下图所示:

矩阵 QK称为相似度矩阵,里面元素的值越大,表示相关性越强。但该矩阵是一个未缩放相似度矩阵,如果词向量维度很大,那么点积数值也很大。为了保持不同维度向量的量纲一致,将相似度矩阵每个元素除以 sqrt{dk},得到缩放相似度矩阵 QKT/sqrt{dk}。

最后对缩放相似度矩阵的每一行做 softmax,得到一个权重矩阵 S,使其每个元素都大于等于 0 ,每一行的和都是 1,妥妥的权重矩阵。权重矩阵可以理解为:当前 token 在更新自身表示时,应当从其他 token 吸收多少信息。

相关性打分就是用 query 去找最匹配的 key。

5.3.4. 信息融合 (information aggregation)


在信息融合阶段,用注意力权重对 value 做加权求和,等价的数学操作就是把权重 (矩阵 S 为载体) 乘上 value (矩阵 V 为载体)。这样得到一个 3×2 的矩阵 Z =SV,称为 attention score,注意它的大小和一开始的输入 X 的大小是一样的。如下图所示:


相关性打分和信息融合就是用 query 去找最匹配的 key,然后把对应的 value 拿回来。

5.3.5. 多头 (multi-head)


多头注意力的核心思想是:

让模型在多个不同的子空间中,并行关注不同类型的上下文关系。

每个子空间通过独立的 Q/K/V 投影矩阵形成,每个注意力头只在自己对应的子空间计算注意力权重。

如果按照前面介绍的方法进行 self-attention 计算,只是使用不同的权重矩阵重复执行 4 次,那么我们最终会得到 4 个不同的 Z 矩阵。

这就带来了一个问题:下节介绍的前馈网络并不会处理 4 个矩阵,它只期望接收一个矩阵,因此,我们需要一种方法,把这 4 个 attention head 的输出合并成一个矩阵。解决方法很简单:

  1. 先把这 4 个矩阵拼接 (concatenate) 在一起

  2. 再将拼接后的结果乘以一个额外的权重矩阵 WO

这样就能把多头的结果重新变换成单个输出矩阵,并传递给下一层。

5.3.6. 代码分析


上面代码从第 4 行到第 24 行是在实现自注意力模块,其中行 6,7,24 是在做归一化和残差连接,这个在小节 5.5 会介绍。核心代码在:

  • 行 8-10 是用三个投影矩阵将 token 词向量转成 q, k, v

  • 行 11-12 是在做 KV cache

  • 行 14-22 是在做多头操作,其中行 19-20 实现相关性打分,行 21-22 实现信息融合

总结来说,上面代码处理下图红框里的流程:自注意力模块。


5.4. 前馈网络 (feed-forward net)


5.4.1. 用途介绍


如果说自注意力负责让 token 看见上下文,那么前馈网络 (feed-forward net, FFN) 的作用是:

在此基础上,对每个 token 的表示进行非线性变换和增强。

换句话说,FFN 并不负责建模 token 之间的关系,而是负责 “加工”已经包含上下文信息的 token 表示。举个例子,输入是:

“the Shawshank”

自注意力会让模型意识到 Shawshank 在当前上下文中具有强语义指向性;在大量训练数据中,Shawshank 往往和 redemption 一起出现。在这种情况下,FFN 会将已经捕获到的上下文关系,转化为更有利于预测下一个 token 的内部特征表示。最终,这些特征会被送入语言建模头,用于生成下一个 token(例如 redemption)。

在结构上,FFN 通常采用:

  • 第一层:将输入向量的维度扩展到更高的维度

  • 中间部分:通过非线性激活函数对信息进行变换

  • 第一层:再将向量维度压缩回原来的大小

这种“先扩展、再压缩”的结构,并不是为了改变序列长度,而是为了在更高维的空间中,对信息进行重组和抽象。

5.4.2. 代码分析


在 microgpt 里,FFN是一个两层的全连接神经网络:

  1. 第 10 行将 embedding 维度扩展到 4 倍

  2. 第 11 行经过 ReLU 激活

  3. 第 12 行再投影回原来的维度

总结来说,上面代码处理下图红框里的流程:前馈网络。


5.5. 归一与残差连接 (normalization & residual connection)


5.5.1. 用途介绍


输入
 ↓
Attention 或 FFN
 ↓
Residual Connection(把原输入加回来)
 ↓
Normalization(调整数值范围)
 ↓
输出

这里:

  • 残差连接 (residual connection):把原来的输入加回来,防止信息丢失,也让深层网络更容易训练。

  • 归一化(normalization):把向量的数值范围重新调整,让训练更加稳定。

5.5.2. 代码分析


在计算之前,我们进行归一化,体现在注意力模块 (第 7 行) 和前馈网络 (第 13 行) ;在计算之后,我们进行残差连接,注意力模块 (第 6 和 9 行) 和前馈网络 (第 12 和 15 行) 的输出都会加回到输入上。

5.6. 语言建模头 (LM head)


最终转换器模块的输出会通过 lm_head 投影到 vocabulary 大小,从而得到每个 token 的 logit。在 microgpt 的例子里,vocabulary 大小 = 27,因此输出是 27 个数。输出的 logit 越大,模型认为对应 token 越可能是下一个 token。


对应的代码如下:




6. 训练 (Training)


首先设定 Adam 优化器的参数和训练循环次数 num_steps

learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
m = [0.0] * len(params) # first moment buffer
v = [0.0] * len(params) # second moment buffer
num_steps = 1000 # number of training steps

在每次循环中训练循环会不断重复以下步骤:


  1. 选择文档分词

  2. 前向计算损失

  3. 反向计算梯度

  4. 更新模型参数

完整代码如下:

下面四节分别解释训练的四个步骤。

6.1. 选择文档分词

每个训练步骤都会选取一个文档,并在两端加上 BOS  标记。例如名字 “emma” 会变成:

[BOS, e, m, m, a, BOS]

模型的任务是根据之前的 tokens 来预测下一个 token。

6.2. 前向计算损失

我们将 tokens 一个一个地送入模型,同时逐步构建 KV cache。在每个位置,模型会输出 27 个 logits,然后通过 softmax 将它们转换为概率,其损失值为 :

    -ln(ptarget)

也就是正确 token 的负对数概率。这被称为交叉熵损失(cross-entropy loss)。直观来说,损失衡量的是模型预测错误的程度,也就是模型对真实结果有多“惊讶”。

  • 如果模型给正确 token 的概率是 1,模型完全不惊讶,loss = 0

  • 如果模型给正确 token 的概率接近 0,模型非常惊讶,loss → +∞

最后我们会把整个文档中每个位置的损失取平均,得到一个单一的标量损失。

6.3. 反向计算梯

调用一次 loss.backward() 就会运行反向传播,沿着整个计算图从 loss 开始,一路传回:

    loss → softmax → model → parameters

之后,每个参数的 .grad 就表示:如果改变这个参数,loss 会如何变化。

6.4. 更新模型参数

我们可以用最简单的梯度下降 p.data -= lr * p.grad,但这里用更先进的 Adam 优化方法。对于每个参数,它维护两个运行平均:

  • m:最近梯度的平均值(类似动量,像滚动的小球)

  • v:最近梯度平方的平均值(用于自适应调整学习率)

另外 m_hat 和 v_hat 是偏差校正(bias correction)因为 m 和 v 从 0 开始,需要一段“热身期”。训练过程中学习率会线性衰减(linear decay)。每次更新完参数之后,我们都会把 p.grad = 0,为下一步训练做好准备。

经过 1000 次训练步骤:loss 从 3.366 下降到 2.6497。

step 1 / 1000 | loss 3.3660
step 2 / 1000 | loss 3.4243
step 3 / 1000 | loss 3.1778
step 4 / 1000 | loss 3.0664
step 5 / 1000 | loss 3.2209
step 6 / 1000 | loss 2.9452
step 7 / 1000 | loss 3.2894
step 8 / 1000 | loss 3.3245
step 9 / 1000 | loss 2.8990
step 10 / 1000 | loss 3.2229
...
step 998 / 1000 | loss 2.4764
step 999 / 1000 | loss 2.4730
step 1000 / 1000 | loss 2.6497

为什么最开始是 3.3?因为随机猜测 27 个 token:−ln(1/27) ≈ 3.3。所以:

  • 3.3 ≈ 模型几乎随机

  • 2.6497 ≈ 模型已经学到了一些模式

Loss 越低越好,而理论最低值是 0,即完全正确预测。因此虽然还有提升空间,但模型已经明显学会了名字的统计规律。


7. 推理 (Inference)

训练完成后,我们就可以从模型中采样(sample)新的名字了。此时模型参数已经固定,我们只需要在循环中不断运行前向计算,并把每一步生成的 token 作为下一步的输入再喂回模型。

每个采样过程都从 BOS token 开始,它告诉模型:“开始生成一个新名字”。模型会输出 27 个 logits,我们把它们转换成概率分布,然后根据这些概率随机采样一个 token。这个 token 会被作为下一步的输入继续喂回模型,如此重复,直到模型再次生成 BOS(表示“我结束了”),或者达到最大序列长度为止。


temperature 参数用于控制随机性。在应用 softmax 之前用 temperature 去除以 logits:

  • 当 temperature = 1.0 时,采样直接按照模型学到的概率分布进行

  • 当 temperature 较高时,概率分布会变得更平坦,生成结果更加多样,但可能会降低连贯性

  • 当 temperature 较低时,概率分布会变得更尖锐,模型更保守,更倾向于选择概率最高的 token

  • 当 temperature 接近 0 时,模型几乎总是选择最可能的那个 token,称为贪心解码(greedy decoding)

生成的名字如下:

--- inference (new, hallucinated names) ---
sample 1: kamon
sample 2: ann
sample 3: karai
sample 4: jaire
sample 5: vialan
sample 6: karia
sample 7: yeran
sample 8: anna
sample 9: areli
sample 10: kaina
sample 11: konna
sample 12: keylen
sample 13: liole
sample 14: alerin
sample 15: earan
sample 16: lenne
sample 17: kana
sample 18: lara
sample 19: alela
sample 20: anton


8. 总结 (Conclusion)


microgpt 已经包含了 GPT 的算法精髓:

token → embedding → attention → MLP → lm_head → next token

下面把 microgpt vs Real LLM(生产级模型 GPT-4 / ChatGPT) 的差异整理成一个对比表。核心思想是:算法本质一样,但规模和工程复杂度完全不同。

真实 LLM 只是更大、更快、工程更复杂,但核心算法仍然是预测下一个 token。


附录 (Appendix)

列出完整的 Python 和 Kotlin 代码,运行一次分别大概耗时 203s 和 16s,Kotlin 相对于 Python 提速了 12.68 倍。


Python 代码


import os # os.path.exists
import math # math.log, math.exp
import time # time.perf_counter
import random # random.seed, random.choices, random.gauss, random.shuffle
random.seed(42) # Let there be order among chaos

def process_input(file_name: str = "input.txt"):
    ifnot os.path.exists(file_name):
        import urllib.request
        names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'
        urllib.request.urlretrieve(names_url, file_name)
    docs = [line.strip() for line in open(file_name) if line.strip()]
    random.shuffle(docs)

    return docs

# Let there be Autograd to recursively apply the chain rule through a computation graph
class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads') # Python optimization for memory usage

    def __init__(self, data, children=(), local_grads=()):
        self.data = data # scalar value of this node calculated during forward pass
        self.grad = 0                   # derivative of the loss w.r.t. this node, calculated in backward pass
        self._children = children # children of this node in the computation graph
        self._local_grads = local_grads # local derivative of this node w.r.t. its children

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data, (self, other), (other.data, self.data))

    def __pow__(self, other):return Value(self.data**other, (self,), (other * self.data**(other-1),))
    def log(self):return Value(math.log(self.data), (self,), (1/self.data,))
    def exp(self):return Value(math.exp(self.data), (self,), (math.exp(self.data),))
    def relu(self):return Value(max(0, self.data), (self,), (float(self.data > 0),))
    def __neg__(self):return self * -1
    def __radd__(self, other):return self + other
    def __sub__(self, other):return self + (-other)
    def __rsub__(self, other):return other + (-self)
    def __rmul__(self, other):return self * other
    def __truediv__(self, other):return self * other**-1
    def __rtruediv__(self, other):return other * self**-1

    def backward(self):
        topo = []
        visited = set()

        def build_topo(v):
            if v notin visited:
                visited.add(v)
                for child in v._children:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1
        for v in reversed(topo):
            for child, local_grad in zip(v._children, v._local_grads):
                child.grad += local_grad * v.grad

def linear(x: list[Value], w: list[list[Value]]) -> list[Value]:
    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

def softmax(logits: list[Value]) -> list[Value]:
    max_val = max(val.data for val in logits)
    exps = [(val - max_val).exp() for val in logits]
    total = sum(exps)
    return [e / total for e in exps]

def rmsnorm(x: list[Value]) -> list[Value]:
    ms = sum(xi * xi for xi in x) / len(x)
    scale = (ms + 1e-5) ** -0.5
    return [xi * scale for xi in x]


# Define the model architecture: a function mapping tokens and parameters to logits over what comes next
# Follow GPT-2, blessed among the GPTs, with minor differences: layernorm -> rmsnorm, no biases, GeLU -> ReLU

def gpt(token_id, pos_id, keys, values):
    tok_emb = state_dict['wte'][token_id] # token embedding
    pos_emb = state_dict['wpe'][pos_id] # position embedding
    x = [t + p for t, p in zip(tok_emb, pos_emb)] # joint token and position embedding
    x = rmsnorm(x) # note: not redundant due to backward pass via the residual connection

    for li in range(n_layer):
        # 1) Multi-head Attention block
        x_residual = x
        x = rmsnorm(x)
        q = linear(x, state_dict[f'layer{li}.attn_wq'])
        k = linear(x, state_dict[f'layer{li}.attn_wk'])
        v = linear(x, state_dict[f'layer{li}.attn_wv'])
        keys[li].append(k)
        values[li].append(v)
        x_attn = []
        for h in range(n_head):
            hs = h * head_dim
            q_h = q[hs:hs+head_dim]
            k_h = [ki[hs:hs+head_dim] for ki in keys[li]]
            v_h = [vi[hs:hs+head_dim] for vi in values[li]]
            attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5for t in range(len(k_h))]
            attn_weights = softmax(attn_logits)
            head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]
            x_attn.extend(head_out)
        x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])
        x = [a + b for a, b in zip(x, x_residual)]

        # 2) MLP block
        x_residual = x
        x = rmsnorm(x)
        x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
        x = [xi.relu() for xi in x]
        x = linear(x, state_dict[f'layer{li}.mlp_fc2'])
        x = [a + b for a, b in zip(x, x_residual)]

    logits = linear(x, state_dict['lm_head'])
    return logits


if __name__ == "__main__":
    start_time = time.perf_counter()
    try:
        docs = process_input('input.txt')
        print(f"num docs: {len(docs)}")

        # Let there be a Tokenizer to translate strings to sequences of integers ("tokens") and back
        uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1
        BOS = len(uchars) # token id for a special Beginning of Sequence (BOS) token
        vocab_size = len(uchars) + 1# total number of unique tokens, +1 is for BOS
        print(f"vocab size: {vocab_size}")

        # Initialize the parameters, to store the knowledge of the model
        n_layer = 1# depth of the transformer neural network (number of layers)
        n_embd = 16# width of the network (embedding dimension)
        block_size = 16# maximum context length of the attention window (note: the longest name is 15 characters)
        n_head = 4# number of attention heads
        head_dim = n_embd // n_head # derived dimension of each head

        matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
        state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd),
                      'lm_head': matrix(vocab_size, n_embd)}
        for i in range(n_layer):
            state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)
            state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)
            state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)
            state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)
            state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)
            state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)
        params = [p for mat in state_dict.values() for row in mat for p in row] # flatten params into a single list[Value]
        print(f"num params: {len(params)}")

        # Let there be Adam, the blessed optimizer and its buffers
        learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
        m = [0.0] * len(params) # first moment buffer
        v = [0.0] * len(params) # second moment buffer

        # Repeat in sequence
        num_steps = 1000# number of training steps
        for step in range(num_steps):

            # Take single document, tokenize it, surround it with BOS special token on both sides
            doc = docs[step % len(docs)]
            tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]
            n = min(block_size, len(tokens) - 1)

            # Forward the token sequence through the model, building up the computation graph all the way to the loss
            keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
            losses = []
            for pos_id in range(n):
                token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
                logits = gpt(token_id, pos_id, keys, values)
                probs = softmax(logits)
                loss_t = -probs[target_id].log()
                losses.append(loss_t)
            loss = (1 / n) * sum(losses) # final average loss over the document sequence. May yours be low.

            # Backward the loss, calculating the gradients with respect to all model parameters
            loss.backward()

            # Adam optimizer update: update the model parameters based on the corresponding gradients
            lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay
            for i, p in enumerate(params):
                m[i] = beta1 * m[i] + (1 - beta1) * p.grad
                v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
                m_hat = m[i] / (1 - beta1 ** (step + 1))
                v_hat = v[i] / (1 - beta2 ** (step + 1))
                p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
                p.grad = 0

            print(f"step {step + 1:4d} / {num_steps:4d} | loss {loss.data:.4f}")

        # Inference: may the model babble back to us
        temperature = 0.5# in (0, 1], control the "creativity" of generated text, low to high
        print("\n--- inference (new, hallucinated names) ---")
        for sample_idx in range(20):
            keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
            token_id = BOS
            sample = []
            for pos_id in range(block_size):
                logits = gpt(token_id, pos_id, keys, values)
                probs = softmax([l / temperature for l in logits])
                token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
                if token_id == BOS:
                    break
                sample.append(uchars[token_id])
            print(f"sample {sample_idx + 1:2d}: {''.join(sample)}")
    finally:
        elapsed = time.perf_counter() - start_time
        print(f"\nTotal runtime: {elapsed:.3f}s")


Kotlin 代码


import java.io.File
import java.net.URL
import java.util.Collections
import java.util.Locale
import kotlin.math.exp
import kotlin.math.ln
import kotlin.math.max
import kotlin.math.pow
import kotlin.math.sqrt
import kotlin.system.measureTimeMillis

typealias Vector = MutableList<Value>
typealias Matrix = MutableList<Vector>
typealias Cache = MutableList<Matrix>

private val rng = java.util.Random(42L)
private lateinit var stateDict: MutableMap<String, Matrix>
private var nLayer = 0
private var nHead = 0
private var headDim = 0

fun processInput(fileName: String = "input.txt"): MutableList<String> {
    val file = File(fileName)
    if (!file.exists()) {
        val namesUrl = "https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt"
        URL(namesUrl).openStream().use { input ->
            file.outputStream().use { output ->
                input.copyTo(output)
            }
        }
    }

    val docs = file.readLines()
        .map { it.trim() }
        .filter { it.isNotEmpty() }
        .toMutableList()
    Collections.shuffle(docs, rng)
    return docs
}

class Value(
    var data: Double,
    private val children: List<Value> = emptyList(),
    private val localGrads: List<Double> = emptyList()
) {
    var grad: Double = 0.0

    operator fun plus(other: Value): Value =
        Value(data + other.data, listOf(this, other), listOf(1.0, 1.0))

    operator fun plus(other: Double): Value = this + Value(other)

    operator fun times(other: Value): Value =
        Value(data * other.data, listOf(this, other), listOf(other.data, data))

    operator fun times(other: Double): Value = this * Value(other)

    fun pow(exponent: Double): Value =
        Value(
            data.pow(exponent),
            listOf(this),
            listOf(exponent * data.pow(exponent - 1.0))
        )

    fun logValue(): Value = Value(ln(data), listOf(this), listOf(1.0 / data))

    fun expValue(): Value {
        val expValue = exp(data)
        return Value(expValue, listOf(this), listOf(expValue))
    }

    fun relu(): Value =
        Value(max(0.0, data), listOf(this), listOf(if (data > 0.0) 1.0else0.0))

    operator fun unaryMinus(): Value = this * -1.0
    operator fun minus(other: Value): Value = this + (-other)
    operator fun minus(other: Double): Value = this + (-other)
    operator fun div(other: Value): Value = this * other.pow(-1.0)
    operator fun div(other: Double): Value = this * (1.0 / other)

    fun backward() {
        val topo = mutableListOf<Value>()
        val visited = mutableSetOf<Value>()

        fun buildTopo(v: Value) {
            if (visited.add(v)) {
                v.children.forEach(::buildTopo)
                topo.add(v)
            }
        }

        buildTopo(this)
        grad = 1.0
        for (v in topo.asReversed()) {
            for (i in v.children.indices) {
                val child = v.children[i]
                val localGrad = v.localGrads[i]
                child.grad += localGrad * v.grad
            }
        }
    }
}

private operator fun Double.plus(other: Value): Value = Value(this) + other
private operator fun Double.minus(other: Value): Value = Value(this) - other
private operator fun Double.times(other: Value): Value = Value(this) * other
private operator fun Double.div(other: Value): Value = Value(this) / other

fun sumValues(values: Iterable<Value>): Value =
    values.fold(Value(0.0)) { acc, value -> acc + value }

fun linear(x: List<Value>, w: List<List<Value>>): Vector =
    w.map { row -> sumValues(row.zip(x).map { (wi, xi) -> wi * xi }) }.toMutableList()

fun softmax(logits: List<Value>): Vector {
    val maxVal = logits.maxOf { it.data }
    val exps = logits.map { (it - maxVal).expValue() }
    val total = sumValues(exps)
    return exps.map { it / total }.toMutableList()
}

fun rmsnorm(x: List<Value>): Vector {
    val ms = sumValues(x.map { it * it }) / x.size.toDouble()
    val scale = (ms + 1e-5).pow(-0.5)
    return x.map { it * scale }.toMutableList()
}

fun gpt(tokenId: Int, posId: Int, keys: Cache, values: Cache): Vector {
    val tokEmb = stateDict["wte"]!![tokenId]
    val posEmb = stateDict["wpe"]!![posId]
    var x = tokEmb.zip(posEmb).map { (t, p) -> t + p }.toMutableList()
    x = rmsnorm(x)

    for (li in 0 until nLayer) {
        var xResidual = x
        x = rmsnorm(x)
        val q = linear(x, stateDict["layer${li}.attn_wq"]!!)
        val k = linear(x, stateDict["layer${li}.attn_wk"]!!)
        val v = linear(x, stateDict["layer${li}.attn_wv"]!!)
        keys[li].add(k)
        values[li].add(v)

        val xAttn = mutableListOf<Value>()
        for (h in 0 until nHead) {
            val hs = h * headDim
            val qH = q.subList(hs, hs + headDim)
            val kH = keys[li].map { it.subList(hs, hs + headDim) }
            val vH = values[li].map { it.subList(hs, hs + headDim) }

            val attnLogits = kH.indices.map { t ->
                sumValues((0 until headDim).map { j -> qH[j] * kH[t][j] }) / sqrt(headDim.toDouble())
            }
            val attnWeights = softmax(attnLogits)
            val headOut = (0 until headDim).map { j ->
                sumValues(vH.indices.map { t -> attnWeights[t] * vH[t][j] })
            }
            xAttn.addAll(headOut)
        }

        x = linear(xAttn, stateDict["layer${li}.attn_wo"]!!)
        x = x.zip(xResidual).map { (a, b) -> a + b }.toMutableList()

        xResidual = x
        x = rmsnorm(x)
        x = linear(x, stateDict["layer${li}.mlp_fc1"]!!)
        x = x.map { it.relu() }.toMutableList()
        x = linear(x, stateDict["layer${li}.mlp_fc2"]!!)
        x = x.zip(xResidual).map { (a, b) -> a + b }.toMutableList()
    }

    return linear(x, stateDict["lm_head"]!!)
}

fun matrix(nout: Int, nin: Int, std: Double = 0.08): Matrix =
    MutableList(nout) {
        MutableList(nin) {
            Value(rng.nextGaussian() * std)
        }
    }

fun sampleIndex(weights: List<Double>): Int {
    val total = weights.sum()
    var target = rng.nextDouble() * total
    for ((index, weight) in weights.withIndex()) {
        target -= weight
        if (target <= 0.0) {
            return index
        }
    }
    return weights.lastIndex
}

fun main() {
    val time = measureTimeMillis {
        val docs = processInput("input.txt")
        println("num docs: ${docs.size}")

        val uchars = docs.joinToString("").toSet().toMutableList().sorted()
        val bos = uchars.size
        val vocabSize = uchars.size + 1
        println("vocab size: $vocabSize")

        nLayer = 1
        val nEmbd = 16
        val blockSize = 16
        nHead = 4
        headDim = nEmbd / nHead

        stateDict = mutableMapOf(
            "wte"to matrix(vocabSize, nEmbd),
            "wpe" to matrix(blockSize, nEmbd),
            "lm_head" to matrix(vocabSize, nEmbd)
        )
        for (i in 0 until nLayer) {
            stateDict["layer${i}.attn_wq"] = matrix(nEmbd, nEmbd)
            stateDict["layer${i}.attn_wk"] = matrix(nEmbd, nEmbd)
            stateDict["layer${i}.attn_wv"] = matrix(nEmbd, nEmbd)
            stateDict["layer${i}.attn_wo"] = matrix(nEmbd, nEmbd)
            stateDict["layer${i}.mlp_fc1"] = matrix(4 * nEmbd, nEmbd)
            stateDict["layer${i}.mlp_fc2"] = matrix(nEmbd, 4 * nEmbd)
        }
        val params = stateDict.values.flatMap { mat -> mat.flatMap { row -> row } }
        println("num params: ${params.size}")

        val learningRate = 0.01
        val beta1 = 0.85
        val beta2 = 0.99
        val epsAdam = 1e-8
        val m = MutableList(params.size) { 0.0 }
        val v = MutableList(params.size) { 0.0 }

        val numSteps = 1000
        val charToIndex = uchars.withIndex().associate { it.value to it.index }

        for (step in 0 until numSteps) {
            val doc = docs[step % docs.size]
            val tokens = mutableListOf(bos)
            tokens.addAll(doc.map { charToIndex.getValue(it) })
            tokens.add(bos)
            val n = minOf(blockSize, tokens.size - 1)

            val keys: Cache = MutableList(nLayer) { mutableListOf<Vector>() }
            val values: Cache = MutableList(nLayer) { mutableListOf<Vector>() }
            val losses = mutableListOf<Value>()

            for (posId in 0 until n) {
                val tokenId = tokens[posId]
                val targetId = tokens[posId + 1]
                val logits = gpt(tokenId, posId, keys, values)
                val probs = softmax(logits)
                val lossT = -probs[targetId].logValue()
                losses.add(lossT)
            }
            val loss = sumValues(losses) * (1.0 / n.toDouble())

            loss.backward()

            val lrT = learningRate * (1.0 - step.toDouble() / numSteps.toDouble())
            for (i in params.indices) {
                val p = params[i]
                m[i] = beta1 * m[i] + (1.0 - beta1) * p.grad
                v[i] = beta2 * v[i] + (1.0 - beta2) * p.grad.pow(2.0)
                val mHat = m[i] / (1.0 - beta1.pow(step + 1.0))
                val vHat = v[i] / (1.0 - beta2.pow(step + 1.0))
                p.data -= lrT * mHat / (sqrt(vHat) + epsAdam)
                p.grad = 0.0
            }

            println(String.format(Locale.US, "step %4d / %4d | loss %.4f", step + 1, numSteps, loss.data))
        }

        val temperature = 0.5
        println("\n--- inference (new, hallucinated names) ---")
        for (sampleIdx in 0 until 20) {
            val keys: Cache = MutableList(nLayer) { mutableListOf<Vector>() }
            val values: Cache = MutableList(nLayer) { mutableListOf<Vector>() }
            var tokenId = bos
            val sample = StringBuilder()

            for (posId in 0 until blockSize) {
                val logits = gpt(tokenId, posId, keys, values)
                val probs = softmax(logits.map { it / temperature })
                tokenId = sampleIndex(probs.map { it.data })
                if (tokenId == bos) {
                    break
                }
                sample.append(uchars[tokenId])
            }

            println(String.format(Locale.US, "sample %2d: %s", sampleIdx + 1, sample.toString()))
        }
    }

    println("Total runtime: ${time/1000.0} s")
}


Logo

欢迎加入 MCP 技术社区!与志同道合者携手前行,一同解锁 MCP 技术的无限可能!

更多推荐