Loading...
Loading...
Loading...
今年年初寒假的时候,我曾经给自己的博客网站做过一套搜索和推荐算法。
那套东西的核心实现其实不算复杂,主要是借助 PostgreSQL 的倒排索引、tsvector 和位置向量查询机制,把文章内容拆成一个个可以被检索和匹配的 token,再根据关键词命中、位置关系、相似度等信息算出一个相对合理的推荐结果。它有点像一个很朴素的多路召回系统,关键词负责确定性匹配,位置向量负责补充上下文,相似度再负责把一些没有完全命中的内容拉回来。
但只要一进入这个领域,就不可避免地会遇到“分词”这个问题。当时我选择的是一套静态分词方案,也就是提前准备一批词表,再结合语法规则去逐个切分文本。这种方式在简单场景中当然能用,而且工程实现也很直接。但它的问题也非常明显——一旦句子中出现歧义,分词结果就可能不稳定。
比如同一段文本,这一次可能切出了“因为”,下一次在某个搜索或相似度匹配场景中,又被切成了“因”和“为”两个字。表面上看,这只是分词粒度发生了一点变化,但对于倒排索引、位置向量和相似度计算来说,后面的结果就完全不一样了。轻一点的情况是排序发生变化,严重一点甚至会出现明明应该搜到的内容却搜不到的情况。
一开始我还以为是代码写得有问题,后来查了很久才发现是歧义句切错了,也是那次才发现,分词并不只是字符串处理问题。它牵扯到了语义、上下文和稳定性。如果只靠静态规则,系统很难真正判断一句话里的某个片段到底应该被看作一个整体,还是应该被拆成更细的部分。
于是我当时产生了一个想法:能不能训练出一个小参数模型,让它和原来的静态分析结合起来,同时从语法和语义两个角度辅助分词?
这和 RAG 中多路召回的思路有点接近,关键词搜索负责把确定能命中的内容找出来,语义模型负责补上那些表达方式不同但含义接近的内容,两条链路互相兜底,整体结果就会稳定得多——至少歧义句的效果会好很多。
后来我选择用 BERT 来做这个模型,但结果很遗憾,这个项目失败了。无论我怎么调整数据、训练方式和输入输出设计,模型最终的分词效果始终达不到我的预期,后来我总结了一下原因,可能是参数确实太小了,模型的认知水平有限;当然也可能是数据集不太合理,导致它学不到什么有用的规律。
更关键的是,当时的我其实根本不理解它的内部原理。那个项目几乎完全依赖 AI 帮我写,我只是不断告诉它“我想要什么效果”,然后让它生成代码、改代码、再继续生成。现在回头看,那就是一次非常典型的 100% Vibe Coding。
所以这个项目失败以后,我最不安的地方其实不在训练效果,而在完全没有掌控感。它对我来说像一个黑盒,我知道输入是什么,也知道输出是什么,但中间到底发生了什么,我完全不知道。
这件事后来被我暂时放下了。直到很久以后,我在某次上课摸鱼的时候刷到了几个讲《Attention is all you need》的视频,才又把这条线重新接了回来。
我最开始看这些视频的时候,只是抱着“了解一下”的心态。毕竟 Transformer、LLM、Attention 这些词这几年出现得太频繁了,以至于它们好像已经变成了一种很抽象的时代背景。大家都在说,但真正把它拆开来看的人并不多。
看完几个讲解后,我粗略地读了一下《Attention is all you need》这篇论文。读完以后,我最大的感受就是:他们真是一群天才。
因为全篇没有任何一个数学理论是我没学习过的,也没有任何地方是因为硬性条件导致我无法学习的,只是使用了很多很简单的理论——甚至高中水平就能理解的数学知识天才般地将数学知识与深度学习领域结合,就构建出了Transformer这样一个划时代的架构,不仅如此,整个深度学习领域无处不充斥着这种例子,比如FFN、梯度下降和反向传播也是这样,我以前从来没想过导数在高维视角下还有“梯度”这种意义,也没想过复合函数求导的链式法则居然能如此巧妙地结合到参数学习这个过程中,更没想过原来非线性函数能拟合任何函数,这真的是太聪明了,这其中的智慧让我大受震撼,也是让我第一次感受到了数学的魅力其实并不在于不断创造各种概念,而是将各种现象、模式建模成一个个数学模型,让它能真实地描述一个事情、解决一个问题。
这正是深度学习真正吸引我的地方。很多概念单独看并没有难到无法理解,但它们被组织起来以后,会产生一种很强的表达能力。FFN 本质上就是线性变换加非线性激活;梯度下降本质上就是沿着让损失变小的方向调整参数;反向传播本质上就是复合函数求导的链式法则在神经网络中的展开。我当然学过微积分,也知道链式法则——但从来没想过它们可以在高维参数空间中以“梯度”的形式出现,更没想过它们能被组织成一套完整的参数学习机制。
Attention 更是这样。说到底,它最核心的思想其实非常直观:一个 token 在理解自己时,应该去关注序列中其它 token 对它的影响。它通过 Q、K、V 三组向量,把“我在找什么”、“别人有什么特征”、“别人携带什么信息”这三件事拆开,再通过点积和 softmax 计算出一组权重,最后把所有信息按权重汇总回来。
简单来说,Transformer 就像一套认知系统,它先把输入的自然语言切成一个个具有语义的最小单位(token),然后映射成模型可以理解的、包含实际语义的 token 向量,接着通过 Attention 的处理,让每个 token 都能独立地计算出它与上下文中其它 token 的相似度,并让自己带上对方的部分语义,比如“它”字可能就会在 Attention 计算中与上下文中的某个主语产生联系,如此一来,模型就能理解出各种复杂文本背后的语义、语法关系,从而处理各种自然语言任务。
这个过程一旦想通,Transformer 就不再像一个完全的黑盒,而更像是一套很清晰的计算流程。它当然复杂,但这种复杂是很多简单组件层层叠加出来的,并没有到无法理解的程度。
也正是在这个时候,我突然想起了年初那个失败的分词器项目。它失败的原因有很多,但最让我在意的,始终只有一点——我不理解它。正如我前文所说的,我只是一味地通过自然语言去让AI帮我完成这个项目,而我实际甚至连一行代码都没审查过,因为我当时根本就不懂,也不理解LLM的基本原理,Code Review对我来说没什么意义,这也是为什么我后面无法完全总结出失败的原因。
但现在我至少已经理解了 Transformer 的基本概念,我就在想:能不能自己用 JavaScript 实现一个小模型?哪怕它很小,效果也很一般,但只要它是我自己从头到尾拆出来、写出来、跑起来的,那它对我的意义就完全不一样。
当时,我立刻就将这个想法分享给了GPT:

得到的回复是肯定的——JS可以做到,并且在了解完原理后,真正实现起来确实没有这么困难。
于是就有了这一次Transformer Encoder的实践经历。
Tokenizer就是分词器,模型的输入形式是一系列 Token,简单理解就是一堆具有语义的字词或词缀,要实现将自然语言转为 Token,就需要用到 Tokenizer。不过对于 Tokenizer 的技术选型,确实也是这个项目实现中最纠结的一个部分。
一开始我用的是 char-level tokenizer。它最简单,每个字符都是一个 token,词表也很好构建,代码实现几乎没有门槛。对一个学习项目来说,这个方案很适合起步,因为它能让模型快速跑起来,也方便确认后面的 Embedding、Attention、MLM 和 Backward 链路有没有真的实现。
但 char-level tokenizer 的缺陷当然也很明显。对于英文来说,字符级 token 会让模型把大量精力花在“怎么拼出一个单词”上。它当然也能从字符序列里学到一些规律,比如 tion、ing、ed 这类结构,但它和真实语义之间隔了一层很长的距离——它并没有真的在学习自然语言的语义和语法关系,而是像个单词机一样的在学习单词的拼写结构。
而我希望模型学到的是词与词之间的关系、上下文里的语义约束、句法结构对预测结果的影响,而不是一直在猜下一个字母是什么。所以后面我把 tokenizer 做成了word-lever,接着又不满意,并加了一层抽象——上层训练流程不直接关心当前到底是 char、word,又或者是 BPE(Byte Pair Encoding)这种结构,它只关心 encode、decode、tokenToId、idToToken 这些稳定接口。这样做的好处很实际:后面从 char-level 切到 word-level,再切到 BPE 的时候,不需要把整个训练流程翻一遍。
word-level tokenizer 带来的变化很直接。模型终于可以把 river、training、model 这类单词当成一个整体来处理,训练目标也更接近“根据上下文猜词”。但它也暴露了另一个问题:词表之外的词都会变成 <unk>,也就是未知词,因为模型在训练时没见过它,词表中自然也就没有这个 token 向量,那模型自然就无法理解它的语义。如果验证样本里有很多没见过的词,模型看到的输入就会变成一大片 <unk>,这时候它几乎没办法恢复原文信息。甚至还可能会学到一种很坏的偏好:遇到困难就预测某些高频 token。
BPE 就是在这个阶段接进来的。它不像 word-level 那样把整个单词当成不可拆分的单位,也不像 char-level 那样完全退回字符。它会先从基础符号开始,然后在语料里不断寻找高频相邻片段,把它们合并成新的 token。比如一个词最开始可能被拆成字符,随着 merge 次数增加,ing、tion、ment 这类高频结构就会逐渐变成更大的片段。这样既能降低未知词的比例,又能保留词缀、词根这类有意义的子词结构。

这个实现里,BPE 的构建过程大致是先清洗语料,再统计基础 token,然后反复统计相邻 pair 的频率,选择最高频且超过阈值的 pair 进行合并,直到词表达到目标大小或再也找不到足够高频的 pair。训练完成后,TOKEN_LIST 负责保存最终词表,MERGES 负责保存合并顺序。真正编码文本时,流程会先做基础切分,再按训练好的 merge 顺序合并片段,最后映射成 id。
这里还踩过一个很典型的坑:如果语料太杂,BPE 很容易退化。比如我之前在找数据集的时候,就试过一些语种类型十分复杂的数据集,可以说是主流的各国语言都有了,然后基础字符集合一下子就膨胀的特别大,目标词表容量很快被单字符占满。结果就是整个词表看起来全是一个个孤立字符,几乎没有真正有意义的词缀或片段。
所以后来我单独清洗出了英文部分,只保留 26 个英文字母、数字、常见标点和必要的特殊 token,再配合最低频率过滤,BPE 才重新开始产生合理的 merge。
另一个细节是 whole-word masking。因为 BPE 会把一个单词拆成多个 piece,如果 MLM 随机只 mask 掉其中某一个 piece,模型有时候会从剩下的 pice 里直接猜出来,这就有点“作弊”了;有时候又会因为只预测一个残片而偏离真实的“猜词”语义。whole-word masking 的做法是保留每个原始单词对应的 token span,mask 的时候把整个单词的所有 piece 一起 mask 掉,loss 也对这些 piece 共同计算。这个改动让训练目标更接近 BERT 的 MLM,也更符合一开始想做语义分词的方向。
Tokenizer 输出的是 token id,它并不具备真实的语义,只是 Token 的一种表示方法,真正具备语义,且可被 AI 学习的是这个 token 对应的向量,所以在真正进入 Attention 前,先要查 Embedding 表,把每个 token id 映射成一个 D_MODEL 维的向量。
这张表的形状可以理解成:
embeddingTable: [VOCAB_SIZE, D_MODEL]
tokenIds: [SEQ_LEN]
embeddings: [SEQ_LEN, D_MODEL]每一行都是一个 token 的可训练表示。训练开始时,这些向量只是随机初始化的小数;训练过程中,反向传播会不断调整对应 token 的行。某个 token 在相似上下文中反复出现,它的向量就会逐渐被推向更合适的位置。
但只靠 Embedding 还不够,因为 Attention 本身是并行处理每个 token 的,对顺序并不敏感。它看到的是一组向量之间的相关性,如果不显式加入位置信息,“我爱你”和“你爱我”在它看起来就几乎是一模一样的意思。
而位置编码就是为了解决这个问题的。这个项目用的就是最经典的正弦余弦位置编码,把每个位置 pos 映射成一组不同频率的 sin 和 cos 值,再加到前一个环节得到的向量上。偶数维用 sin,奇数维用 cos,不同维度对应不同频率。低频维度表示更粗的位置变化,高频维度表示更细的位置变化,最后每个位置都会得到一个相对稳定且可区分的向量。

位置编码让“同一个 token 出现在不同位置”这件事变得可见,也让 Attention 后续能在语义相关性之外感知顺序。
Attention 是 Transformer 最核心的一个部分。它解决的问题很朴素:让每个 token 都能关注到与自己语义相似的其它 token,并让彼此互相带上对方的语义。
在 Transformer 被提出之前,业界主要是使用 CNN 和 RNN 等架构来训练模型,但这种架构在部分场景中有几个比较致命的缺陷——模型训练效率低且成本高、处理长上下文时容易丢失信息。
而 Transformer 很好的解决了这些问题,Attention 天然可以并行计算,尤其是在现代 GPU 多核心的支持下,每个 token 都可以独立计算与其它 token 的相似度,并且每次推理都会根据整个上下文综合计算。
这其中,决定了模型能不能真正的理解语义关键就是相似度的计算,业界有一套比较经典的做法,就是计算余弦相似度,说起来倒也比较简单——向量点积计算的过程就是两个向量的模长积乘上它们之间的夹角的余弦值,而方向靠的越近的两个向量,夹角就越小,夹角的取值区间可以简化为[0,π],而在这个区间内,余弦值的变化规律呈现为:角度越大,余弦值越小,因此,夹角越大的两个向量,它们的点积结果就越小,体现在语义空间上就是越不相关,这就是余弦相似度的原理,Attention正是利用了这一原理,用 token 向量计算出 Q、K、V 三个矩阵,然后让它们彼此相乘,体现在向量上就是点积的过程。
当然,在Transformer中,相似度的比较不止是依靠余弦相似度,它是一套更全面的比较机制,被称之为缩放点积,具体的来说,余弦相似度比较的是两个向量之间的夹角,也就是方向,从余弦相似度的角度看来,两个向量的方向越相近,它们的相似度就越高。但这套方法有一个明显的局限性——它没法比较向量的模长,更具体的一点来说,就是它没法比较两个词在语义上的“强度”,它在计算之前会先把向量的模长归一化比较,无论它们原来的模长是多少,最后都会被化成一个单位长度。也就是说,即使一个词表达出了对某个语义的强烈情感,例如强烈的赞同或强烈的反对,在余弦相似度的眼中,这毫无意义,那在部分场景中它天然就会丢失一部分信息的精度。
因此,Transformer选用了缩放点积来进行相似度的计算和比较:
s = |q|·|k|·cosθ / sqrt(dk)在这个计算中,相似度不仅考虑到了角度θ,还考虑到了两个向量的模长关系,最后除以sqrt(dk)是为了防止向量维度过大时计算结果无限膨胀的问题,这样一来,模型就能更全面的理解语义了。

回到正题,Attention 的计算过程可以压缩成几行:
Q = X * Wq
K = X * Wk
V = X * Wv
scores = Q * K^T / sqrt(dk)
weights = softmax(scores)
output = weights * V其中,X 是输入序列,也就是加了位置编码后的 token 表示。Wq、Wk、Wv 是三组可训练参数,输入序列经过这三组参数的计算后会得到 Q、K、V 三个矩阵,Q 可以理解成当前 token 发出的查询,K 是每个 token 暴露出来的匹配特征,V 是真正要被汇总的信息。Q * K^T 会得到一个 [SEQ_LEN, SEQ_LEN] 的矩阵,矩阵中第 i 行第 j 列表示第 i 个 token 对第 j 个 token 的关注程度。

这个过程可以形象的理解为:每个 token 的 Q 逐个的去和每个 token 的 K 比较,看看对方的语义中是不是含有自己想知道的东西,比如下面这段话:
我上午去图书馆借了一本书,下午回家帮妈妈干了一会家务,晚上才开始读它。这个语境中,最后一个“它”字显然是一个代词,代指上下文中的某个主语,那这个时候,有一个天然的问题就产生了:这个代词指代的到底是谁?
这个问题就可以被看作是这个字的其中一个 Q,它想知道自己指的到底是谁,并且根据这个字本身的语义特性,这个指代物必须是一个名词,且大概率会是一个非人物体,所以这个时候,这个 token 就会拿着自己的这个问题,去和上下文中所有的 token 的 K 去匹配。
大部分的非名词 token 身上含有的与“它”字相同的语义非常少,比如“去”字就与这个词可以说是毫无关系,那在缩放点积下,它们的相似度极低,就几乎不会影响到两个 token 本身的含义。
真正能够影响语义的显然就是“图书馆”和“书”了,因为这两个都满足“名词”和“非人物体”的特点,因此经过计算后的相似度就比较高。但最后模型还是会认为“它”字指代的是“书”而不是“图书馆”,为什么呢?
因为 Attention 是并行计算的,“它”字在寻找指代词的时候,“书”字也在寻找相关的其它成分,而当“书”字匹配到后文的“读”字时,会发现它们两个的相似度非常高,因为“读书”这两个词在中文中是非常常见的一个词语搭配,“书”本来就可以被“读”,因此“书”字就会与“读”字去计算相似度,得出的分数极高,接着再用这个分数,去“读”字的 V 矩阵中取了最相关的语义,那显然就是“阅读读物”这一类的意思,这个意思会被附加到“书”字原本的向量中,让它带上“阅读”的语义。
现在回过头来,我们再来看“它”字,这个字紧挨着“读”字,它们俩在语法层面上,一个作谓语,一个作宾语,说明“读”字的作用对象显然就是这个“它”,但是前文里“书”和“读”又扯上关系了,这不仅仅表示“书”带上了“读”的意思,“读”也会带上“书”的意思,同时“它”在语法角度上也会沾有一点“读”的意思,刚好这个时候“它”又在寻找一个具体的指代词,而前文中直接作用于这个词的谓语又和另一个词搭上线了,它们俩都拥有同一个 token 的部分语义,那这个时候,“书”字的相似度很明显就要比“图书馆”要高得多,因为不论怎么想,“图书馆”都不可能与“读”有直接联系——图书馆不可能被“阅读”。因此,“它”指代的物体自然就是“书”。

这样一来二去的,模型在“它”这个字上的注意力自然就被分配到了“它”、“读”等相关的词上,这就表示模型正确的理解了最开始的那个问题:“这个代词指代的到底是谁”。
Attention 就是这样工作的,听起来好像很神秘,实则一步步分析下来之后,其实并没有那么复杂,而且很像我们人类分析语法成分、理解语义结构的过程,只是我们平时几乎不会去刻意的思考“我是怎么理解语义”的这个过程,而 Transformer 则是通过“注意力”这一概念,用数学模型将理解的过程抽象成了一个极简的数学公式,实在是令人惊叹的思维水平。
除以 sqrt(dk) 也是一个很容易被忽略的细节,前文讲解缩放点积时也讲到了。如果不缩放,维度变大以后点积结果会越来越大,后续 softmax 转化为概率分布后的结果会变得非常尖锐,比如某一个位置的值是99999,对应的概率权重无限接近于 1,而其它位置相对的就会非常接近 0,梯度计算也会变得不稳定。因此,这个缩放项的作用就是把 scores 拉回到一个更舒服的范围。
多头注意力是在这个基础上的扩展。它会把 D_MODEL 切成多个 head,每个 head 单独做一套 Attention,最后再拼回来并通过 Wo 做一次输出投影。这样做的意义不只是增加计算量,重点在于让不同 head 可以关注不同类型的关系。有的 head 可能更关心局部搭配,有的 head 可能更关心远距离依赖,还有的 head 可能会学到标点、介词、主谓结构之类的模式。

但这里也踩过一个容易误解的点:在 D_MODEL 不变的情况下,把 head 数从 4 改成 8,并不会让模型容量凭空翻倍。每个 head 的维度会变小,单个 head 的表达空间也会变窄。它可能改善注意力模式的多样性,但不一定直接提高训练速度,也不一定立刻降低 loss。如果想让模型真的有更强的表达能力,通常还要一起考虑 D_MODEL、D_FF、层数、数据规模和学习率。
Attention 的 Backward 是我在这个项目里遇到的第一块真正复杂的反向传播。它比简单线性层绕得多,中间串了 matMul、缩放、softmax、再 matMul,多头版本还要经历 split 和 concat。梯度必须沿着这条链路反着走回来,先从 output = weights * V 得到 dWeights 和 dV,再穿过 softmax 得到 dScores,然后回到 Q 和 K,最后再回到 Wq、Wk、Wv 和输入 X。这部分代码我借助了 AI 完成,但真正把它接进总流程以后,我才对“链式法则在网络里展开”这件事有了更具体的感受。
Softmax 在很多地方都会出现,Attention 里用它把 scores 变成权重,MLM 输出层里用它把 logits 变成 token 概率。它的公式看起来很简单:
softmax(x_i) = exp(x_i) / sum(exp(x_j))但实现时也有两个重要的细节,首先,它通常是按行计算的。Attention 的 scores 是一个二维矩阵,每一行代表“当前 token 对所有 token 的注意力分布”,所以 softmax 必须对每一行单独做归一化。MLM 的 logits 也是类似,如果一次预测多个 mask 位置,那么每个 mask 位置都有自己的一行 vocab 分布。
其次,数值稳定性不能省。exp 对大数非常敏感,如果 scores 里出现比较大的值,很容易溢出。所以实际实现时一般会先减去当前行最大值,再进入 exp。这个操作不会改变 softmax 的最终分布,因为分子分母同时乘了同一个常数,但它能明显降低溢出风险。

Softmax 的反向传播也比线性层绕一点。它不是每个元素独立求导,因为分母里包含整行所有元素,一个位置的变化会影响同一行其它位置的概率。代码里采用的是常见的行级写法:
dScore_i = weight_i * (dWeight_i - sum(dWeight_j * weight_j))这行公式一开始看会有点别扭,但它实际表达的就是 softmax Jacobian 和上游梯度相乘后的简化形式。把它写成代码以后,整个反向过程会清楚很多:先算上游梯度和当前 softmax 权重的点积,再对每个位置做一次修正。
Encoder 里每个子层后面都会接 Add和 Norm。Add 是残差连接,把子层输出和原始输入加起来;Norm 是 LayerNorm,对每个 token 的向量做归一化。
残差连接的价值很实际:神经网络的每一层都在对输入进行加工不同的处理,输出的信息也会一定程度上的弱化输入原本携带的信息。而随着层数增加,如果每一层都完全重写前一层的信息,梯度和表示都很容易在传播中变得不稳定。残差连接相当于给信息留了一条弯道超车的通路,让模型可以在原始表示的基础上做增量更新,让输入与输出彼此相加,非常直接。但同时,它也让多层 Encoder 更容易接起来,因为每一层输入输出都能保持 [SEQ_LEN, D_MODEL] 的形状。
LayerNorm 则负责把每个 token 的向量拉回到比较稳定的分布。它会对一行向量计算均值和方差,再做归一化,最后乘上 gamma、加上 beta:
norm = (x - mean) / sqrt(var + eps)
output = norm * gamma + beta一开始我其实漏掉了 gamma 和 beta。如果只做标准化,模型当然也能跑,但表达能力会被限制,因为归一化以后每一层都被强行拉到固定分布。gamma 和 beta 的作用就是把缩放和平移的自由度还给模型,让它可以自己决定某些维度应该被放大、缩小或偏移。后来我把它们一起补进了正向、反向、参数初始化和参数更新里,这个链路才算完整。

Add & Norm 的 Backward 也比表面看起来更容易出错。残差连接意味着上游梯度会分成两路,一路流向子层输出,一路直接流向原始输入。LayerNorm 又需要同时计算 dX、dGamma 和 dBeta。如果缓存里没有保存归一化前的输入,或者把 normInput 和 addNormOutput 搞混,梯度就会传错位置。这也是我后来越来越重视 forward cache 的原因:反向传播不是凭空算出来的,它依赖正向时留下的中间状态。
之前的步骤都是让信息在线性空间上做各种变换,但现实世界并没有那么“线性”,很多复杂信息都是非线性的,所以如果一直在线性空间上对信息做文章的话,那模型就很难学习到复杂的规律,比如各种语义关系、语法结构等等。
而非线性函数就是用来解决这个问题的,它遵循一个数学特点——非线性函数理论上可以拟合任何函数。
这一块就有点微分的味道了,比如我们都知道电脑 GPU 本质上只能画三角形,无论它如何组合三角形,都无法渲染出一个真正光滑的曲线或曲面,当材质的精度下降时,图形的边界就会出现锯齿状的纹理,那如何解决这个问题呢?最简单的办法当然就是细化精度,让单个三角形的面积不断变小,让更多的三角形彼此组合,从而在宏观的视觉效果中让我们“看不到锯齿”——这其实就是一个非线性拟合的过程:不断用直线段拟合曲线,即“化直为曲”。
而在 Transformer 中,Attention 负责让 token 之间交换信息,FFN (前馈神经网络)则负责对每个 token 的表示做非线性加工。它的结构很简单:
hidden = X * W1 + b1
activated = ReLU(hidden)
output = activated * W2 + b2这里的 W1 通常会把维度从 D_MODEL 扩到 D_FF,再通过 W2 压回 D_MODEL。扩维的意义在于给模型一个更大的中间空间,让它能组合出更复杂的特征。ReLU 是一个激活函数,负责提供非线性关系,否则多个线性层叠在一起仍然等价于一个线性层,表达能力会非常有限。

FFN 和 Attention 的关系也很有意思。Attention 更像是在序列维度上做信息混合,让每个 token 从其它 token 那里取信息;FFN 更像是在特征维度上做加工,让每个 token 自己的表示变得更有表达力。两者一横一纵,叠起来以后,Encoder 才能在上下文和特征两个方向上不断更新表示。
选用ReLU的原因也很简单,从它的公式上就可以看出来了——它的计算过程太简单了,就是一个 Max 函数而已,大于激活值的保留,小于激活值的直接变成零,就靠这样一个极其简单的步骤将线性关系转换为非线性关系,复杂度特别低——甚至交给人来口算都能算的很快,而我们的输入参数是一个个矩阵,ReLU高效的计算速度和计算效果天然适合这一场景。
FFN 的反向传播相对直观一些。output = activated * W2 + b2 先回传到 dW2、db2 和 dActivated,再穿过 ReLU。ReLU 的梯度非常朴素,正数位置梯度保留,非正数位置梯度归零。最后再从 hidden = X * W1 + b1 回到 dW1、db1 和 dX。这个模块写完以后,整个 Encoder 的一大半反向链路就已经能闭合了。
Encoder 本身并不会自然产生“下一个 token”预测,因为它看到的是整个输入序列,更适合做理解类任务。为了让它可以自监督训练,我给它接了一个 MLM 输出头,原理就是随机 mask 掉输入中的某些 token,然后让模型根据上下文猜回原词,类似于一种“完型填空”。
这个流程大致是:先从原始文本里选出一个词,按 tokenizer 编码成 token ids,再把选中的位置替换成 [MASK]。模型前向计算后,会取出 mask 位置对应的 hidden states,通过 vocab head 投影到词表大小:
logits = hiddenAtMask * W_vocab + b_vocab
probabilities = softmax(logits)
loss = crossEntropy(probabilities, targetTokenIds)如果是 whole-word masking,也就是一次性遮住一整个词,而不是一个字母或一个词缀,那一个原词可能对应多个 token,那么 hiddenAtMask 就会变成多行,targetTokenIds 也会是一组 id。loss 会对这些位置分别计算交叉熵,再取平均。这个细节很重要,因为它要求模型对整个被 mask 的词负责,不能只靠某一个残片投机。

交叉熵的直觉很好理解。如果正确 token 的概率越高,loss 就越低;如果模型把概率分给了错误 token,loss 就会上升。对于 softmax + cross entropy 组合,输出层的梯度可以写得很简洁:把预测概率复制一份,然后在正确 token 的位置减 1。这个 probs - oneHot(target) 会作为 dLogits 继续往 vocab head、hidden states 和前面的 Encoder 传播。
MLM 也暴露过一些训练上的问题。早期 char-level tokenizer 下,模型更像是在补字母;word-level 阶段又出现了 <unk> 太多的问题;BPE 接入以后,如果 mask 只落在某个 piece 上,任务目标也会变得不够自然。最后 tokenizer 抽象、BPE 词表、whole-word masking 这几件事一起完成以后,MLM 才更接近我一开始想要的那种“根据上下文恢复语义单位”的效果。
反向传播(Backward)是这个项目里最难的一部分,也是我借助 AI 最多的一部分,因为这部分涉及到的数学理论实在是太多了,所以我选择先实现,再理解原理。
反向传播的概念其实很好理解,就是模型在跑完一遍正向传播(也就是前面一直在说的Transformer Encoder部分)后,拿着自己算出来的结果去对一遍答案,然后计算出自己的答案与正确答案之间有多少误差,也就是损失函数,再反过头来,根据这个损失函数去逐层的检查,看看是哪个参数在计算过程中对最终答案起到的“误导性”最大,也就是去找每个参数各自把模型的认知带偏了多少。
要说原理其实也很容易理解,模型会对每一层参数求导,但因为各层的参数很多时候不是只有一个,而是一个具有多个自变量的函数,所以针对每个参数,就需要求偏导,求出来的这个值,就被称为梯度,也就是函数在这个方向的上升速度,根据梯度就能看出这个参数对最终的损失贡献有多大了。
“反过头来”其实就是反向传播的过程,因为在正向传播中,上一层的输出会直接作为下一层的输入,所以正向传播的整体结构就类似于一种复合函数,一个函数套娃另一个函数,就像一个洋葱,而反向传播就是由外向内的一层层剥开这个洋葱。而因为上面提到了,求梯度本质上就是求导,这个时候,复合函数独特的链式求导法则就派上用场了:复合函数求导,等于外层函数求导加上内层函数——而此时我们需要的正是“单独对外层函数求导”!有了链式法则,我们就可以把某一层的结果单独的拆出来,对它的某个参数计算梯度,真正实现“剥洋葱”的效果。
这是最能体现我在文章开头所说的一个感慨:这群人真的太聪明了。毋庸置疑,链式求导法则连高中生都学过,但是有几个能想到这个法则还有这种意义的?又有几个人能想到它有这一层用途的?反向传播这个词听起来好像很高大上,实际上并没有什么理解门槛,真正厉害的是能够想到为链式法则赋予这一层意义的这个想法。
回到正题,在把洋葱剥开后,接下来要做的就是调参了,模型的自我学习就发生在这一过程中——根据每个参数对损失函数的贡献,即梯度,模型会沿着梯度下降最快的方向调整参数,也就是梯度的负方向。
这就类似于一个人下山,理论上来说最快抵达山底的方法就是不断沿着山坡最陡的方向一路下降——梯度下降正是这一层意思,梯度本身表示的意思是参数往哪个方向改动,损失涨得最快,而沿着梯度的负方向调整参数,就能使得这个参数对损失的贡献减少,其本质上就是沿损失下降最快的方向修正参数,降低整体损失函数的值,让模型不断的自我学习和优化迭代。

总的 backward 流程大概是从 loss 开始,先得到输出层的 dLogits,再反传 vocab head,得到 dW_vocab、db_vocab 和 mask 位置的 hidden 梯度。因为 MLM 只在被 mask 的位置计算 loss,所以这里不能把梯度直接当成整个序列的梯度,还需要 scatter 回原序列。也就是创建一个 [SEQ_LEN, D_MODEL] 的零矩阵,只在 mask 对应的位置填入 dHiddenAtMask。
之后梯度会沿着 Encoder 反着走。每一层都是先穿过最后的 Add & Norm,再穿过 FFN,然后把 FFN 的输入梯度和残差梯度合并,继续穿过前一个 Add & Norm,再进入 Attention。Attention backward 会得到 dWq、dWk、dWv、dWo 和输入梯度。多层 Encoder 的情况下,这个过程会从最后一层循环回第一层。
最后一层传回来的 dX 会继续流向位置编码和 Embedding。位置编码如果是固定的正弦余弦值,本身不需要更新;Embedding 则需要根据 token id 把梯度累加回对应的词表行。这里还有一个容易忽略的点:同一个 token 可能在一句话里出现多次,也可能在一个 batch 里出现多次,所以 embedding 的梯度不能简单覆盖,必须累加。
当然,虽然我能理解 backward 的原理,但具体的代码实现我还没有完全理解透彻,实际实现时很大程度上是靠 AI 帮忙的(感谢 Codex ,给你跪了orz),因为我最开始想快点让整个模型能够被训练,验证我之前的效果,所以距离我真正理解整个模型原理还有很长的一段路要走。
当 backward 能算出梯度以后,模型才真正进入“训练”阶段。参数更新本身很简单,最基础的梯度下降就是:
param = param - learningRate * grad但工程上要做的事情比这一行多很多。模型里有 Embedding、Attention 的四组权重、FFN 的两层权重和 bias、Add & Norm 的 gamma 和 beta、vocab head 的参数,它们都要被初始化、保存、读取和更新。如果某个参数只参与了 forward,没有进入 backward 或 update,它就相当于永远不会学习,所以我后面也补齐了对应的输出和保存功能。
训练脚本后面也经历了一些工程化整理。最开始的浏览器训练是通过 Playwright 跑 WebGL,后来又加了纯 Node 训练链路,再尝试接入 @mni-ml/framework 和 CUDA。为了观察训练状态,我还做了一个控制台面板,实时显示 round、batch、训练损失、验证损失、best loss、elapsed、ETA、backend 和简单 profile。
这部分简直就是噩梦,也就是在这个时候我才真的明白为什么业界实践都是用Python做的了——这些库对JS的兼容性都太差了,要么就是几年前的老项目,要么就是各自版本兼容问题,换了好几次之后才勉强能跑,但内存又时不时会出问题,麻了。
不过这部分并不是这个项目最核心的主题。只是模型能跑以后自然出现的工程问题。真正让我受益最多的,还是把 Encoder 的每个结构都亲手搓了一遍。
模型开始训练以后,损失会变成一个很容易让人焦虑的数字。它下降的时候会让人觉得方向对了,升高的时候又会怀疑是不是哪里写错了。这个项目里也出现过类似情况:训练损失一段时间后反而越来越高,验证损失长期追不上历史最佳损失(best loss),甚至改了层数和 head 数以后效果没有明显提升,下降的很慢。
这些现象不能简单归因于“模型不够大”。模型质量和很多因素有关。数据集太小,模型很容易记住训练样本里的表面模式;数据太杂,tokenizer 和模型都会被噪声干扰;学习率太大,参数可能在最优区域附近反复跳动;batch 太小,梯度方向会抖;batch 太大,如果没有配套调整学习率,也可能让模型更新得不够细。层数、head 数、D_MODEL、D_FF 当然重要,但它们只是容量的一部分,不会自动带来更好的训练结果。
还有一个很容易被忽略的点是状态文件。训练会保存 best loss、round、参数和二进制模型,如果改变了 tokenizer、词表、层数或模型结构,却沿用旧的训练状态,就会出现一些很奇怪的结果。比如看起来已经训练到某个 round,实际上参数结构早就变了;或者 best loss 来自旧模型,新模型永远追不上这个不公平的基准。后来我在训练脚本里专门处理了模型二进制保存和恢复,新的 best validation loss 出现时就覆盖旧模型,这样至少能保证留下来的文件是当前配置下更好的版本。
我后面写了一个简单的预测脚本,用来随机从验证集抽样,自动 mask 输入,然后输出原词、预测 token、概率、loss 等信息。这个脚本很有用,因为单看 loss 有时候并不能直观看出模型到底学到了什么。它的损失可能确实下降了,但预测永远偏爱某个高频 token,出现了即将过拟合的趋势;也可能某些样本预测错了,但 top 概率已经明显比随机猜高很多。把具体样本打出来以后,模型的问题会更容易被看见。
如果只看最终效果,我手搓的这个“ Mini-Transformer ”当然还很弱——它的参数量很小,数据集也不大,训练效率和真正的深度学习框架没法比。它现在更像一个教学性质的 Encoder,而不是一个可以直接拿去做产品的模型。
但它对我的意义并不小,年初那个 BERT 分词器项目让我最难受的地方,是我完全不知道里面发生了什么。现在这个项目刚好反过来。它没有多强,但起码我知道每一步在做什么——Tokenizer 为什么会影响模型的视野,Embedding 为什么不能直接用 id,位置编码为什么要加到输入里,Attention 为什么要有 Q/K/V、为什么要这样计算,softmax 为什么要按行归一化,Add & Norm 为什么能使训练更稳定,FFN 为什么需要非线性,MLM 为什么适合 Encoder,以及 Backward 原理等等,这些东西都从抽象概念变成了具体代码,这是最有价值的。
更重要的是,它让我意识到,理解一个模型并不一定要从特别宏大的角度切入。很多时候,把MVP阶段写出来,反而更容易看清楚结构。代码一旦跑起来,所有抽象名词都会落到向量维度、矩阵乘法、参数更新和 loss 曲线上。这个过程会暴露很多纸面学习时看不到的问题,也会让一些原本觉得“好像懂了”的概念变得更扎实。
当然,我也没有必要把这个项目包装成完全手写、完全独立完成的成果。尤其是 backward 这一块,我确实借助了 AI,而且现在也不能说自己完全掌握了每一行梯度推导。但这并不影响它的学习价值。AI 帮我补上了很多我暂时写不出来的代码,我再通过调试、训练和排错去理解这些代码。这个过程比单纯让 AI 生成一个黑盒项目要好得多,因为我是在不断把黑盒拆开,而不是完全拒绝去理解这个黑盒。
这个项目现在还只走到 Encoder 和 MLM。后面如果继续做,我想我会先把 Decoder 补上,按照《Attention is all you need》原文的路径,将它补全成一个真正的 Transformer 架构,然后不断调整代码和数据集,加大训练量,尽量压低损失,让这次的实践不仅具有学习意义,也能获得一点实在的成就感(
但这些都不急。对我来说,当前最重要的收获已经出现了——我不再只是知道 Transformer 很厉害,也不再只是会调用别人封装好的模型。我至少已经知道,一个小型 Encoder 从文本进入 tokenizer 开始,到 MLM 输出概率,再到 loss 反向传回每一层参数,中间大概经历了什么。
这就是我选择继续迭代下去的原因。它可能很难变成一个真正有用的大模型,但它会持续逼着我理解那些以前只停留在概念层面的东西。从Mini-Attention 到 Mini-Transformer,从 forward 到 backward,从 char-level tokenizer 到 BPE,从能跑到能训练,每一步都不仅仅在堆功能,也是在把我对 Transformer 的理解往前推进一点。