Skip to content

Transformer 深度教程

封面:一条链上只有少数几根粗细不一的长线,一个人拿着放大镜看最粗的那根


目录

章节内容
0. 写在前面这本教程给谁看、怎么读、怎么跑
1. 背景序列建模的三代范式,以及化学人为什么要学它
2. 注意力QKV、softmax 与它的雅可比、dk\sqrt{d_k} 的两种推导、核平滑、完整反向传播
3. 多头注意力定义、凸包限制、多头反向、头数怎么选
4. 位置编码整数/二进制/正弦的编码阶梯、排列等变性定理、正弦加法定理与内积恒等式、RoPE 的证明、化学的坑
5. 残差、LayerNorm 与 FFNFFN 的记忆视角、LayerNorm 雅可比、初始化方差守恒、Pre-LN vs Post-LN、dropout 的期望与方差
6. 完整架构与训练目标三种形态、掩码、交叉熵梯度、自回归生成的正确性
7. 纯 Python 实现不用 NumPy 怎么把矩阵乘法跑快、每个模块的前向/反向、梯度检验
8. 训练动力学warmup 的推导、梯度裁剪、权重共享
9. 化学与材料中的应用表示选择、性质预测、生成、反应、聚合物、3D、可解释性
10. 九个实验每个实验的设置、数字、图与结论
11. 优缺点与边界什么时候该用,什么时候别用
12. 常见坑与调参清单避坑深度备忘录与 5 维调参指南
13. 参考文献47 篇
附录符号表、实验一览、复现指南、手绘示意图清单、分词与 BPE

0. 写在前面

0.1 这本教程站在哪

在整个「生成模型」教程丛书中,每一册都专注于解答一个极其纯粹的底层命题,彼此环环相扣却又各司其职:

册子回答的问题
《循环神经网络》《长短时记忆网络》怎么把序列喂进网络、梯度怎么沿时间传
《自回归模型》一个联合分布怎么拆成一串条件分布
《Transformer》(本册)在每一步条件分布里,网络内部凭什么"看"其他位置
《图神经网络》分子不是字符串时怎么办
《变分自动编码器》《生成对抗网络》《标准流》《扩散模型》《能量模型》另外五条生成路线

如果用一句话来概括本册的核心:Transformer 是定义在序列之上的"全局关系算子"。它本身并不发明新的概率分解形式,其革命性的突破在于构建了一套全新的信息路由机制——它允许处于第 tt 个位置的表征直接越过时空距离去读取第 ss 个位置的信息,既免受距离拉长带来的信息衰减惩罚,又天然支持现代算力集群的高吞吐并行计算。

0.2 你需要什么前置知识

为了完全掌握本册内容,你只需要具备三项最基础的数学工具与直觉:其一是能够理解基础的求和符号 \sum 与矩阵乘法几何意义;其二是熟练掌握微积分中的复合函数求导链式法则;其三是理解"极大似然估计在离散分布下等价于最小化交叉熵损失"这一统计学习公理。除此之外,无论是 softmax 的全雅可比矩阵、LayerNorm 的中间变量偏导,还是旋转位置编码(RoPE)的正交旋转变换,本教程都会从最底层的代数与微积分原理出发,一步步完整推导,不跳过任何关键中间步骤。

0.3 怎么读

面对这本兼具理论深度与实操代码的教程,不同需求的读者可以选择不同的行进轨迹。如果你当前工作节奏极快、急于把握核心物理直觉与核心结论,建议直奔核心机理篇章,精读 §2(注意力机制)、§4(位置编码)、§9(化学与材料前沿应用)以及 §11(架构边界),随后快速查阅 §10 的实验定量结果表。如果你倾向于从代码实操中建构理解,推荐直接克隆仓库并切入 code/ 目录,首先运行 tests_transformer.py(在 0.8 秒内完成全部 59 项严苛的底层测试与逐参数梯度检验),确认数学算子万无一失后,再顺次跑通 §10 中的九个实战实验。而如果你正在撰写学术论文或算法报告,§2、§4 与 §5 中关于梯度动力学与几何不变性的定理可以直接作为理论推导引用;§10 则为你提供了全部九个真实实验的原始测试数字,所有数值均已在 figures/results_*.json 中留存了机器可读的原生记录。

0.4 文件结构

整套教程的代码与资产采用模块化布局,核心算法与领域数据处理保持高度解耦:

text
Transformer教程/
├── Transformer模型教程.md     ← 你正在读的主文件
├── README.md                  ← 一页速览
├── code/                      ← 全部可运行代码
│   ├── transformer.py         ← 核心库(零第三方依赖,约 1300 行)
│   ├── chemdata.py            ← 化学/材料数据与分词器(零第三方依赖)
│   ├── bpe.py                 ← 字节对编码(BPE,零第三方依赖,附录 E)
│   ├── tests_transformer.py   ← 59 项正确性测试(含逐参数梯度检验)
│   ├── demo_*.py              ← 九个主实验 + 两个补充实测
│   ├── make_figures.py        ← 把 results_*.json 画成图
│   ├── models/                ← 训练好的小模型(JSON 权重)
│   └── data/                  ← 数据缓存(DeepChem 原始文件)
├── figures/                   ← 31 张数据图(PNG + SVG)+ results_*.json
└── images/                    ← 19 张手绘示意图 + 生成提示词

0.5 环境

为了践行真正的工匠精神,本教程的核心算法实现(包括 transformer.pychemdata.py 以及所有单元测试)不依赖任何第三方张量计算库,仅凭原生 Python 3.10 及以上版本即可畅行无阻。只有在运行对接工业界基线的数据科学实验时,才需要依赖 DeepChem 及其科学计算生态:

bash
/Users/polyai/anaconda3/envs/deepL/bin/python -m pip list | grep -E "deepchem|rdkit|torch|numpy"
# deepchem  2.8.0
# rdkit     2026.3.6
# torch     2.14.0
# numpy     2.5.3

需要特别说明的是,本册所有实验均在配置完备的 /Users/polyai/anaconda3/envs/deepL 虚拟环境下通过测试(采用 Python 3.12)。除了在 DeepChem 基线对比脚本中不得不调用外部接口外,我们自主构建的 Transformer 体系从头到尾未曾导入一行 NumPy 或 PyTorch。这种刻意的约束绝非卖弄技巧,而是因为唯有彻底脱离封装良好的黑盒张量库,亲手编写内存遍历与反向求导,才能真正刺破算法表象,检验自己对底层数学机理的理解深度。


1. 背景:从"逐字接力"到"全局互看"

1.1 三代序列模型

为了理解序列建模范式的演进,让我们从化学世界中最常见的一条 SMILES 字符串切入:

CC(=O)Oc1ccccc1C(=O)O        (阿司匹林)

当我们希望利用深度学习模型预测该分子的水溶解度、潜在毒性或预测其逆合成反应产物时,模型在底层必须正面回答一个本质问题:第 6 个字符「O」(酯基氧)与第 12 个字符「c」(苯环碳)之间的相互影响与几何/化学约束,究竟应当通过怎样的计算通路被捕捉出来? 回顾深度学习的发展历程,序列建模经历过三次重大范式跃迁。

第一代方案是经典的循环神经网络(RNN/LSTM),其机制本质上是一场逐字接力

ht=ϕ(Wxxt+Whht1+b),y^=g(hT)h_t = \phi(W_x x_t + W_h h_{t-1} + b),\qquad \hat y = g(h_T)

在这一体系下,第 tt 个位置若想获取历史上下文,唯一的通道就是上一时刻传递而来的隐藏状态 ht1h_{t-1}。这种单线推进的结构天然孕育了两个无法调和的致命瓶颈。首先是距离衰减效应,从第 6 个字符跋涉到第 12 个字符,中间必须连续穿过 6 次矩阵乘法与非线性激活函数的洗礼,梯度与语义信号极易在漫长的流水线上指数级衰减或爆炸;即使引入了精巧的门控机制(正如《长短时记忆网络》中所详述的那样),也只能在经验上缓解而无法根除这一物理路径上的损耗。其次是计算无法并行,因为 hth_t 的求解严格以前序状态 ht1h_{t-1} 的完工为前置条件,导致现代算力硬件引以为傲的大规模矩阵并行吞吐能力在面对时序轴时束手无策。

左边是一队人只能把信封往后传,右边是所有人同时互相看

为了突破串行计算的枷锁,第二代方案转向了以时序卷积(TCN/WaveNet)为代表的局部窗口逐层放大路线:

ht(l)=k=0K1Wk(l)ht2lk(l1)h_t^{(l)} = \sum_{k=0}^{K-1} W_k^{(l)} h_{t - 2^l k}^{(l-1)}

时序卷积通过因果卷积核与层间膨胀因子的指数倍增,使网络的感受野随着网络深度 LL 实现 2L2^L 级的指数扩张,从而成功将时间步摊平在 GPU 上实现了全序列的并行前向。然而,这种提速是以连接模式的静态僵化为代价的。卷积核的权重在训练收敛后便完全冻结,它只是一把刻度固定、只看相对位移的标尺,绝不会因为序列当前位置出现了关键的羰基,就主动打破空间窗口去聚焦远端的羟基。

于是,以注意力机制(Attention)为核心的第三代范式应运而生,开启了由内容动态决定关注对象的全新时代:

outputt=sexp(qt,ks/d)sexp(qt,ks/d)权重,由内容算出  vs\text{output}_t = \sum_{s} \underbrace{\frac{\exp(\langle q_t, k_s\rangle/\sqrt{d})}{\sum_{s'}\exp(\langle q_t, k_{s'}\rangle/\sqrt{d})}}_{\text{权重,由内容算出}} \; v_s

在这里,任何两个位置之间的交互权重不再是预先固定的静态参数,而是由 tt 位置的内容向量与 ss 位置的内容向量在运行时通过相似度度量动态推演出来的。这一构造带来了三项颠覆性的理论优势:首先,任意两个位置之间的交互路径被彻底压缩为单步直达,无论物理距离相隔几个字符还是成百上千个 token,信息交互的路径长度恒为 O(1)O(1);其次,整个打分与加权聚合过程全部转化为全局矩阵运算,训练期可以毫无阻碍地一次性吃下整条序列,实现完全的高度并行;最后,模型拥有了**基于内容寻址(content-based addressing)**的自适应能力,关注重点会随着输入分子结构的变幻而实时流转。

1.2 2017 年那篇论文到底解决了什么

在自然语言处理与机器学习的里程碑中,Vaswani 等人在 NeurIPS 2017 发表的奠基之作 Attention Is All You Need,其真正的贡献并非凭空发明注意力机制——早在 2014 年,Bahdanau 等人便已将其作为循环神经网络的辅助对齐部件。Vaswani 等人的魄力在于彻底砸碎旧框架,一次性完成了一场教科书式的三位一体革命:

他们以极大的勇气将循环结构彻底抛弃,使得整个编码器与解码器纯粹由自注意力层与逐位置前馈网络交替堆叠而成;针对抛弃循环之后模型对序列先后次序完全丧失感知(即后文 §4 将严格证明的置换等变性)的致命缺陷,他们创造性地引入了位置编码体系,为无序的特征空间注入了时空几何先验;同时,他们构筑了多头注意力与维度缩放机制,使得模型能够在多个不同的表征子空间中并行捕捉多元语义关系,并在数学上彻底规避了高维内积引发的梯度饱和(详见 §2 与 §3)。

随后的故事已然成为深度学习的黄金历史:以编码器见长的 BERT 开辟了语义表征与属性预测的黄金赛道,以自回归解码器为核心的 GPT 系列确立了生成建模与大语言模型的主流范式,而“计算量、数据量与参数规模越大,泛化性能越强”的 Scaling Law 更是席卷了人工智能的几乎每一个分支。演进至今日,在 AI for Science 与计算化学前沿,无论是小分子理解的 ChemBERTa、长程序列建模的 MolFormer、三维几何对齐的 Uni-Mol,抑或是处理空间物理场的 Equiformer,其底层血脉无一例外均流淌着这套自注意力算子的基因。

1.2.1 出发的地方:Bahdanau 的加性注意力

在被砸碎重组之前,注意力最初只是一个"补丁"。Bahdanau 等人 2014 年的工作(Neural Machine Translation by Jointly Learning to Align and Translate)第一次把它变成了可训练的组件,而它要修的问题正是 §1.1 里那个上下文瓶颈

当时的机器翻译标准架构是"编码器—解码器":编码器(RNN/LSTM)把整句源语言读进去,最终只交出一个定长向量 hTh_T;解码器就从这个向量出发生成译文。信息瓶颈于是变成结构性的——句子越长,句子开头那几个词元的痕迹就越淡,译文的开头往往最先失真。Bahdanau 的改动只有一步:让解码器在第 tt 步回看每一个编码器状态,而不是只依赖 hTh_T

ets=vtanh ⁣(Wsst1+Whhs),αts=eetsseets,ct=sαtshs(1.1)e_{ts} = v^\top \tanh\!\big(W_s s_{t-1} + W_h h_s\big),\qquad \alpha_{ts} = \frac{e^{\,e_{ts}}}{\sum_{s'} e^{\,e_{ts'}}},\qquad c_t = \sum_{s} \alpha_{ts}\, h_s \qquad\textbf{(1.1)}

其中 st1s_{t-1} 是解码器上一步的隐藏状态,hsh_s 是编码器在第 ss 个源词处的隐藏状态,ctc_t 是喂给解码器的上下文向量。对照 §2 将要展开的现代写法,有三处差异必须点明,因为它们解释了后来的一切设计:

  1. 打分器是加性的:一个单隐层 MLP(Wsst1+WhhsW_s s_{t-1}+W_h h_s)加一次线性读出 vv^\top,而不是内积。加性形式的优势是两侧维度可以不同、且在维度很大时不会失控;代价则是每一对 (t,s)(t,s) 都要过一次 MLP,无法退化成一次大矩阵乘法。Transformer 之所以改用点积 qt,ks\langle q_t,k_s\rangle,核心动机就是这一句:只有点积才能让整张打分矩阵 QKQK^\top 一次性算完(见 §2.7 的复杂度实测)。而点积带来的副作用——高维内积方差随 dkd_k 线性增长——又必须用 1/dk1/\sqrt{d_k} 补偿(§2.4)。
  2. 值不是新投影出来的:Bahdanau 直接把编码器隐藏状态 hsh_s 当值用;Transformer 则引入独立的 WVW_V,把"被检索的内容"和"被加权聚合的内容"解耦成两套空间(§2.1)。
  3. 没有"头"的概念:这是单头的交叉注意力;多头要到 2017 年才出现(§3)。

左边:整句被压进一只小罐子,罐里的纸团看不清任何线条;右边:解码的小人边写边回头逐张看同伴举起的卡片

α\alpha 排成矩阵——行是目标语言词、列是源语言词——就得到一张软对齐图:两种语言词序一致时,高权重沿对角线排布;形容词后置之类的语序差异则表现为偏离对角线的连线。这是"注意力可以作为对应关系证据"的最早形态,也是今天在化学反应预测里拿注意力当原子映射的直系祖先(§9.4 与实验九 §10.9 讨论的正是这件事的可靠性边界)。

最后区分两个后面反复出现的名词:

名称Q 来自K/V 来自化学里的例子
交叉注意力(cross-attention)序列 A序列 B英→法翻译;反应物→产物的原子对齐
自注意力(self-attention)同一条序列同一条序列SMILES 内部第 6 个原子看第 12 个原子

本章 §2、§3 讨论的全部是后者;§6.1 的编码器—解码器形态则把两者叠在一起使用(解码器先对自己做因果自注意力,再对编码器输出做交叉注意力)。

1.3 为什么化学/材料方向的人更该学它

对于化学与材料科学领域的科研人员而言,掌握 Transformer 并非赶时髦,而是因为化学与材料体系的研究对象在数字化之后天然表现为高度结构化的序列,并且覆盖了极为广阔的物质尺度:

化学对象序列化写法典型任务
小分子SMILES / SELFIES 字符串溶解度、毒性、logP 预测;分子生成
无机晶体组成式(LiFePO4)、元素序列形成能、带隙、稳定性筛选
聚合物重复单元 SMILES 的拼接玻璃化温度、带隙、介电常数
反应反应物>>产物产物预测、逆合成、条件推荐
生物大分子氨基酸/核苷酸序列蛋白–配体亲和力、酶功能
谱图峰位序列、token 化的质谱/红外结构解析

在上述所有序列表示中,SMILES 语法与 Transformer 的契合度堪称天作之合。作为一种高度精简的一维线性描述符,SMILES 拥有丝毫不亚于自然语言的严苛语法树约束——分支括号必须成对开启与闭合、环闭合数字编号必须跨越千山万水精确配对。这使得利用语言模型处理 SMILES,其底层的概率推断与注意力路由机制与常规文本大模型几乎完全同构(相关细节将在 §9.1 深度拆解)。

然而,化学工作者在拥抱这一工具时,必须时刻在脑海中铭刻一个冷酷的客观事实:SMILES 的书写顺序绝不等同于真实的化学物理顺序。由于分子图遍历遍历起点的任意性,同一个化学分子在计算机中可以拥有多种语法上完全等价的表示:

CC(=O)Oc1ccccc1C(=O)O        (阿司匹林)
O=C(O)c1ccccc1OC(C)=O        (同一个分子!)

这两种写法描述的是完全一致的物理客观实在,但在序列空间里,原子的相对排列与绝对索引却被完全洗牌。这意味着一维位置在化学语境中蕴含着远比自然语言更加微妙、甚至极具欺骗性的物理本质(这一深层矛盾我们将在 §4.7 进行专题批判)。

1.4 作为生成模型,它站在哪里

从统计学习的本质来看,Transformer 本身并非一种独立的概率生成模型,而是一个表达能力极强的参数化骨架。只有当我们把它安插在概率论经典的自回归分解链条之上时,它才化身为一台功能完备的自回归生成机器:

pθ(x1,,xT)=t=1Tpθ(xtx<t),pθ(xtx<t)=softmax(Wht)p_\theta(x_1,\dots,x_T) = \prod_{t=1}^{T} p_\theta(x_t \mid x_{<t}), \qquad p_\theta(x_t \mid x_{<t}) = \mathrm{softmax}\big(W\, h_t\big)

这个等式的前半部分正是《自回归模型》一册所奠定的概率论地基,而 Transformer 所贡献的,正是计算隐藏表征 hth_t 的卓越算法实现——通过因果注意力掩码堆栈,使模型在推演 xtx_t 时能够瞬时吸收 x1,,xt1x_1,\dots,x_{t-1} 的全部历史特征,且不产生任何未来信息的泄漏。将两本册子融会贯通,方能看清现代生成式大模型的完整拼图:概率图模型负责提供严密的联合分布分解定理,而 Transformer 架构则以关系算子的形态提供算力与泛化能力的极致支撑。

本册讲的是这一根线:字符串 -> token -> 向量 -> 关系 -> 概率


2. 注意力:从检索到公式

2.1 检索隐喻:Q、K、V 到底是什么

为了直观理解注意力机制的运转,我们可以暂时抛开抽象的高维张量运算,将序列内部的信息交互想象成一个井然有序的分布式借阅市场。在这个体系中,序列中的每一个位置(每一个 token)在进入注意力层时,都会同时铸造成三种不同身份的向量表征。第一张是写着自己当前迫切需求的信息卡片,即 Query(查询向量);第二枚是向外界公开展示自身属性的主题标签,即 Key(键向量);第三份则是自己实际承载的核心内容宝库,即 Value(值向量)

当处于第 tt 个位置的 token 想要更新自身表征时,它所执行的操作极其自然:手持自己的 Query 卡片,去遍历比对序列中所有位置公开发布的 Key 标签。谁的 Key 标签与当前的 Query 最契合,系统就会赋予谁更高的打分,进而按比例从对应位置的 Value 宝库中借阅更多内容;最终,把借来的多份 Value 线性混合加权,便凝练成第 tt 个位置在当前层的新表征。

一排人举着三种东西:牌子(Q)、标签(K)、书(V)

将上述直觉形式化为数学语言。设输入矩阵为 XRT×dX \in \mathbb{R}^{T \times d},其中第 tt 行对应位置 ttdd 维嵌入向量。我们引入三组可学习的投影参数矩阵及其偏置:

WQ,WK,WVRd×dk,bQ,bK,bVRdkW_Q, W_K, W_V \in \mathbb{R}^{d \times d_k},\qquad b_Q,b_K,b_V \in \mathbb{R}^{d_k}

通过线性映射,三组矩阵被并行生成:

Q=XWQ+1bQ,K=XWK+1bK,V=XWV+1bV(2.1)Q = XW_Q + \mathbf{1}b_Q^\top,\qquad K = XW_K + \mathbf{1}b_K^\top,\qquad V = XW_V + \mathbf{1}b_V^\top \qquad\textbf{(2.1)}

这里的 1\mathbf{1} 代表全 1 列向量,用于将一维偏置广播到序列的每一行。需要特别强调的是,Q、K、V 本质上全部源自输入矩阵 XX 自身的线性投影,这也正是"自注意力"(self-attention)中"自"字的数理本源。而在解码器读取编码器特征的交叉注意力(cross-attention)场景下,QQ 矩阵改由解码器自身状态映射生成,而 KKVV 则源自编码器最终输出的语义流。

2.2 打分:为什么用点积

在完成 Q、K、V 的投影映射后,下一个核心问题是如何量化 Query 与 Key 之间的亲疏程度。在标准注意力机制中,打分矩阵通过如下点积矩阵建立:

S=QKdkRT×T,Sts=qt,ksdk(2.2)S = \frac{QK^\top}{\sqrt{d_k}} \in \mathbb{R}^{T\times T},\qquad S_{ts} = \frac{\langle q_t, k_s\rangle}{\sqrt{d_k}} \qquad\textbf{(2.2)}

为什么在众多相似度度量工具中,不约而同地选择内积作为打分算子?深入其几何与微积分结构,便能发现这一设计的精巧之处。首先是几何直观的天然契合:根据欧氏几何定义 q,k=qkcosθ\langle q,k\rangle = \|q\|\,\|k\|\cos\theta,点积不仅衡量了两个向量在特征空间中的夹角朝向,同时兼顾了它们各自的范数模长,既能刻画语义方向的对齐,又能体现特定表征被激活的客观强度。更妙的在于微积分上的极简代价:由于内积对分量的求导 q,k/q=k\partial \langle q,k\rangle/\partial q = k 具有完全的线性形式,使得反向传播时无需计算任何昂贵的非线性导数,极大地卸下了硬件算力的包袱。最后,内积具备分量可拆解性q,k=iqiki\langle q,k\rangle = \sum_i q_i k_i 允许我们将运算顺畅地摊平到特征维度求和,这恰恰构成了后续线性注意力机制将平方复杂度降为线性复杂度的数学切入点。

2.3 softmax:把打分变成权重

2.3.1 定义与三条基本性质

为了将任意尺度的实数打分矩阵转化为合法的概率权重,我们需要引入逐行作用的 softmax 函数。对于任意实向量 sRns\in\mathbb{R}^{n},其定义为:

softmax(s)j=esjj=1nesj(2.3)\mathrm{softmax}(s)_j = \frac{e^{s_j}}{\sum_{j'=1}^{n} e^{s_{j'}}} \qquad\textbf{(2.3)}

这一算子具有三条极其优美的数学与数值性质。首先是非负性与全概率公理:对任意分量均有 softmax(s)j>0\mathrm{softmax}(s)_j>0 且行和恒等于一 jsoftmax(s)j=1\sum_j \mathrm{softmax}(s)_j=1。这意味着它天然将打分映射为一个合法的离散概率分布,注意力权重 αts\alpha_{ts} 因此获得了极其直观的物理意义——它代表位置 tt 从位置 ss 处借阅信息的置信度比例。其次是平移不变性:对任意标量常数 cc,恒有 softmax(s+c1)=softmax(s)\mathrm{softmax}(s + c\mathbf 1) = \mathrm{softmax}(s),其证明极为平凡,只需在分子分母同时约去公因子 ece^{-c} 即可。然而,这条看似简单的代数性质却是工业级数值稳定实现的核心基石。在浮点计算中,直接对较大的正数求指数极易引发上溢,利用平移不变性,我们可以在进入指数运算前先减去当前行的最大值:

python
mx = max(row)                  # 先减最大值
e  = exp(row[j] - mx)          # 再取指数,指数参数 <= 0,绝不溢出

此时所有指数输入均被严格限制在 (,0](-\infty, 0] 区间内,从根本上杜绝了浮点溢出的风险。第三条性质则刻画了极端温度调控下的极限行为。当我们引入温度超参数 τ>0\tau>0 并定义 p(τ)=softmax(s/τ)p^{(\tau)} = \mathrm{softmax}(s/\tau) 时,系统的熵表现出剧烈的相变特性:当温度趋于零极限界 τ0+\tau \to 0^+ 时,概率分布将完全退化为在最大打分位置取 1 的 one-hot 独裁分布;而当温度趋于无穷大 τ\tau\to\infty 时,各选项差异被抹平,分布平滑收敛为完全均匀的 1n1\frac1n\mathbf 1 分布。这一推导的前半部分可以通过分离出唯一最大值 s\*=maxjsjs^\*=\max_j s_j 来证明:

pj(τ)=e(sjs)/τ1+je(sjs)/τ  τ0+  δj,1+0,p^{(\tau)}_j = \frac{e^{(s_j - s^*)/\tau}}{1+\sum_{j\neq *} e^{(s_j-s^*)/\tau}} \;\xrightarrow[\tau\to0^+]{}\; \frac{\delta_{j,*}}{1+0},

因为对于所有非最大项 jj\neq *,分子指数均以 e(正常数)/τe^{-(\text{正常数})/\tau} 的速度急速衰减至零,后半部分关于均匀分布的收敛则可通过对称的 Taylor 展开严格得到。温度实验(fig_a02_temperature.png)生动展现了这一连续演化过程:针对同一组原始打分,当设定低温 τ=0.1\tau=0.1 时,注意力几近全神贯注地钉在单一 key 上,信息熵近乎为 0;而当升温至 τ=4\tau=4 时,权重几乎被均匀摊平,信息熵逼近理论最大值 ln203.0\ln 20 \approx 3.0

温度控制注意力从"独裁"到"平均"

为什么必须是 softmax,而不是"除以总和"? 在把打分变成权重这一步上,最朴素的做法是线性归一化(简单归一化):linearj=sj/jsj\text{linear}_j = s_j / \sum_{j'} s_{j'}。它同样能让权重和为 1,却会在两个地方直接失效——而这两点恰好是本书后面两条论证的前提:命题 3.1 的凸包限制(§3.2)与注意力即核平滑的视角(§2.5)。

第一,它处理不了负数。取打分 s=[1,2,3,5]s=[1,2,-3,5](和为 55),线性归一化给出

[0.20, 0.40, 0.60, 1.00][\,0.20,\ 0.40,\ \mathbf{-0.60},\ 1.00\,]

第三个分量出现了负的"概率"。这不是无关紧要的数值瑕疵:命题 3.1 的凸包论证完全建立在"系数非负且和为 1"之上,一旦出现负权重,输出 Ot=sαtsvsO_t=\sum_s \alpha_{ts}v_s 便会逃出值向量的凸包,"软检索/软混合"的语义随之瓦解——模型开始对从未见过的值向量组合产生响应。softmax 则因为指数恒正而自动规避:同一个输入给出 [0.017, 0.047, 0.0003, 0.936][0.017,\ 0.047,\ 0.0003,\ 0.936],负分被压到接近零,而不是变成负数。

第二,它缺少"放大差异"的能力。对打分 s=[1,2,3,6]s=[1,2,3,6],线性归一化得到 [0.083,0.167,0.250,0.500][0.083,0.167,0.250,0.500]——最大权重恰好是最小权重的 6 倍,与打分之比一模一样,这个映射只是对尺度做了一次恒等复制,模型无法表达"我几乎确定就是它"。softmax 给出 [0.006,0.017,0.046,0.930][0.006,0.017,0.046,0.930],最大值独占 93%,与最小值的比达到 148 倍。注意力需要的正是这样一个可以远离均匀分布的"推子";而推子的陡峭程度由温度 τ\tau 与缩放因子 dk\sqrt{d_k} 两个旋钮控制(§2.3.1 第三条性质与 §2.4)。

上排:按原水位比例直接舀进大碗(线性归一化);下排:先过一只细漏斗再舀(softmax);右下角是那滴不该出现的负水位

需要澄清的是,这两条论证说明的是"为什么 softmax 是一个自然的选择",而不是"softmax 是唯一正确的选择":任何满足非负、归一、可微的映射都可以当权重函数。softmax 之所以成为默认,是因为它同时具备平移不变性(数值稳定)、处处光滑(梯度好传)、以及"打分即对数几率"的解释——把它的负对数似然对打分求导,得到的恰是 ponehotp-\text{onehot} 这样干净的结果(§6.3.1)。

2.3.2 softmax 的雅可比:反向传播的核心

在构建反向传播引擎时,softmax 的求导是不可绕过的高地。设概率向量 p=softmax(s)p=\mathrm{softmax}(s),我们需要求解雅可比矩阵元素 pi/sj\partial p_i/\partial s_j。分两种情况讨论:当对角线项 i=ji=j 时,对商式应用除法求导法则可得:

pisi=siesiZ=esiZesiesiZ2=pipi2=pi(1pi)\frac{\partial p_i}{\partial s_i} = \frac{\partial}{\partial s_i}\frac{e^{s_i}}{Z} = \frac{e^{s_i}Z - e^{s_i}\cdot e^{s_i}}{Z^2} = p_i - p_i^2 = p_i(1-p_i)

而当非对角线项 iji\neq j 时,分子的导数为零,仅分母提供负向贡献,直接给出:

pisj=esiesjZ2=pipj\frac{\partial p_i}{\partial s_j} = -\frac{e^{s_i}e^{s_j}}{Z^2} = -p_i p_j

将这两类情形合并写为紧凑的矩阵形式,便得到了著名的雅可比表达式:

ps=diag(p)pp(2.4)\frac{\partial p}{\partial s} = \mathrm{diag}(p) - pp^\top \qquad\textbf{(2.4)}

这是一个对称、半正定且秩至多为 n1n-1 的特殊矩阵。之所以秩不满,是因为当我们用全 1 向量或自身概率向量右乘它时,由于 p(diag(p)pp)p=pp(pp)2=0p^\top(\mathrm{diag}(p)-pp^\top)p=p^\top p-(p^\top p)^2=0,概率向量恰好落在其零空间内,这与平移不变性质完全契合。在反向传播过程中,假定上游传递而来的标量损失梯度为 aˉ=L/p\bar a = \partial L/\partial p。如果盲目地显式构造出 n×nn\times n 的雅可比矩阵再执行矩阵乘法,将不可避免地招致 O(n2)O(n^2) 的内存分配与时间开销。但若深入代数求和内部:

sˉj=iaˉipisj=iaˉipi(δijpj)=aˉjpjpjiaˉipi\bar s_j = \sum_{i}\bar a_i \frac{\partial p_i}{\partial s_j} = \sum_i \bar a_i p_i(\delta_{ij}-p_j) = \bar a_j p_j - p_j \sum_i \bar a_i p_i

我们惊喜地发现,整个向量导数可以极其精巧地化简为一个逐元素积与内积的组合:

sˉ=p(aˉaˉ,p1)(2.5)\bar s = p \odot \big(\bar a - \langle \bar a, p\rangle \mathbf 1\big) \qquad\textbf{(2.5)}

其中 \odot 表示逐元素 Hadamard 积,而 aˉ,p=iaˉipi\langle \bar a,p\rangle=\sum_i \bar a_i p_i 仅为一个一维点积标量。这正是 transformer.pysoftmax_backward 函数能够将计算复杂度由平方阶奇迹般降至线性阶 O(n)O(n) 的数学精髓

python
def softmax_backward(A, dA):
    dS = []
    for a, da in zip(A, dA):
        dot = sum(map(mul, a, da))          # ⟨dA, p⟩
        dS.append([ai * (dai - dot) for ai, dai in zip(a, da)])
    return dS

测试脚本 tests_transformer.py::t_softmax_backward_matches_jacobian 将式 (2.5) 的轻量化实现与显式构造矩阵乘法进行了逐位对拍,相对误差稳定保持在 <1012<10^{-12} 的极低水平。

2.4 缩放因子 dk\sqrt{d_k}:两个独立推导

在原初论文中,作者曾以一句略显简略的注记解释引入缩放因子的初衷:“当维度很大时,点积幅值会急剧膨胀,从而将 softmax 推入梯度极小的饱和区域”。若要将这句直觉升华为严格的数学论据,我们需要从统计方差与极值饱和两个独立维度进行双重推导。

推导 A(方差):不缩放时打分方差随 dkd_k 线性增长

考虑基础统计假设:设 qqkk 的各个分量相互独立且同分布,其均值为 0、方差为 1(在神经网络中,这一条件在经过规范的 LayerNorm 处理后能够得到极好满足,且并不严苛要求高斯分布假定):

s=q,k=i=1dkqikis = \langle q,k\rangle = \sum_{i=1}^{d_k} q_i k_i

考察内积标量 ss 的数学期望:

E[s]=iE[qi]E[ki]=0\mathbb E[s] = \sum_i \mathbb E[q_i]\,\mathbb E[k_i] = 0

进而推演其方差,由于各特征维度 ii 之间相互独立,求和算子的方差可直接线性拆解:

Var(s)=iVar(qiki)=i(E[qi2]E[ki2](E[qi]E[ki])2)=i110=dk\mathrm{Var}(s) = \sum_i \mathrm{Var}(q_ik_i) = \sum_i \Big(\mathbb E[q_i^2]\mathbb E[k_i^2] - (\mathbb E[q_i]\mathbb E[k_i])^2\Big) = \sum_i 1\cdot1 - 0 = d_k

由此可见,内积打分的标准差随着特征维度 dkd_k 呈现平方根级的线性膨胀:

std(s)=dkVar ⁣(sdk)=1(2.6)\mathrm{std}(s) = \sqrt{d_k} \qquad\Longrightarrow\qquad \mathrm{Var}\!\Big(\frac{s}{\sqrt{d_k}}\Big) = 1 \qquad\textbf{(2.6)}

如果不进行缩放,当 dk=64d_k=64 时打分的标准差已达 8;而当维度扩展至现代大模型的 dk=512d_k=512 时,标准差飙升至 22.6,使得进入 softmax 的数值动态范围完全失控。在实验基准(results_attention.json)的严谨实测中:当 dk=512d_k=512 时,未缩放打分的方差实测高达 489.07(紧密贴合理论预测的 512),而除以 dk\sqrt{d_k} 缩放后的方差完美回落至 0.955(理论基准为 1)。

打分方差与  成正比;除以  后回到 1

推导 B(饱和):不缩放时注意力退化成 one-hot,梯度消失

仅仅揭示方差的膨胀还不足以说透其危害的核心,我们必须进一步证明它如何直接扼杀梯度。设序列中包含 nn 个候选 key,其打分 s1,,sns_1,\dots,s_n 近似独立且具有标准差 σ=dk\sigma=\sqrt{d_k}。设其中取得的最大值为 s\*=maxjsjs^\*=\max_j s_j,则经过 softmax 归一化后的最大概率权重为:

αmax=11+je(ssj)(2.7)\alpha_{\max} = \frac{1}{1+\sum_{j\neq *} e^{-(s^*-s_j)}} \qquad\textbf{(2.7)}

借助极值统计学理论,当随机变量 sjs_j 服从亚高斯分布时,nn 个独立样本的最大值与次大值之差 Δ=ss(2)\Delta = s^*-s_{(2)} 具有明确的统计尺度估计:

Δ    σ2lnn    Θ(1)(2.8)\Delta \;\approx\; \sigma\sqrt{2\ln n}\;\cdot\;\Theta(1) \qquad\textbf{(2.8)}

虽然相邻次序统计量的平均间距比极值本身略小,但其绝对差距依然与标准差 σ\sigma 保持严格的线性正比关系。将式 (2.8) 代入最大概率表达式 (2.7) 中:

αmax11+(n1)edk2lnnΘ(1)  dk  1\alpha_{\max} \gtrsim \frac{1}{1+(n-1)e^{-\sqrt{d_k}\sqrt{2\ln n}\cdot\Theta(1)}} \;\xrightarrow[d_k\to\infty]{}\; 1

当特征维度 dkd_k 逐渐增大时,指数项被极其猛烈地压缩至零,导致 αmax\alpha_{\max} 迅速逼近于 1。这意味着整个注意力分布完全退化为一个尖锐的 one-hot 向量。而当我们把这一退化的概率向量代回 softmax 雅可比公式 (2.5) 时就会发现,当 pp 退化为基向量 eje_j 时,有:

sˉ=p(aˉaˉj1)=0\bar s = p\odot(\bar a - \bar a_j \mathbf 1) = 0

这表明反向传播传回的梯度被数学上精确地置为零。如果不进行尺度缩放,一旦模型进入较高维度,自注意力模块便会因输出过饱和而主动切断所有前向梯度的流动,使得深度网络陷入彻底无法训练的瘫痪状态——这才是“推入梯度极小区域”背后令人不寒而栗的真实数学机制。

在包含 n=16n=16 个候选 key 的受控实验中,这一推导得到了震撼性的数值印证:

dkd_k未缩放方差未缩放熵缩放后方差缩放后熵均匀分布熵 ln16\ln 16
22.002.1490.9982.3902.773
87.541.3130.9432.3662.773
3229.930.6340.9352.3602.773
128117.320.2680.9172.3692.773
512489.070.1410.9552.3552.773

从上表清晰可见,在未缩放状态下,分布的信息熵由小维度的 2.15 急剧坠落至高维下的 0.14,系统几乎丧失了全部不确定性而沦为孤注一掷的选择;反观缩放之后,信息熵无论维度如何变迁均稳健固守在 2.36 附近,既保留了因输入差异而形成的正常置信度起伏,又始终与理论最大熵(ln162.773\ln 16 \approx 2.773)保持健康的动态距离。

这里需要澄清一个广泛流传的工程误区:除以 dk\sqrt{d_k} 的初衷绝非为了“让注意力分布显得更均匀”,其纯粹的使命是让内积打分的数值尺度彻底脱离对特征维度 dkd_k 的病态依赖。如果你在特定任务中确有平滑分布的先验需求,应当诉诸于显式调节采样温度 τ\tau(见 §2.3.1)。

2.5 注意力其实是核平滑

当我们把归一化后的概率公式 (2.3) 带回加权求和的表达中时,整个自注意力的输出形态展现出了惊人的统计学同构性:

Ot=sαtsvs=sexp ⁣(qt,ks/dk)vssexp ⁣(qt,ks/dk)(2.9)O_t = \sum_s \alpha_{ts} v_s = \frac{\sum_s \exp\!\big(\langle q_t,k_s\rangle/\sqrt{d_k}\big)\,v_s} {\sum_s \exp\!\big(\langle q_t,k_s\rangle/\sqrt{d_k}\big)} \qquad\textbf{(2.9)}

如果熟悉非参数统计学,读者会立刻意识到这正是经典的 Nadaraya–Watson 核回归估计器

f^(x)=sκ(x,xs)yssκ(x,xs),κ(q,k)=exp ⁣(q,k/dk)(2.10)\hat f(x) = \frac{\sum_s \kappa(x, x_s) y_s}{\sum_s \kappa(x,x_s)}, \qquad \kappa(q,k)=\exp\!\big(\langle q,k\rangle/\sqrt{d_k}\big) \qquad\textbf{(2.10)}

借助这一深刻的统计学类比,许多原本晦涩的概念瞬间豁然开朗:自注意力操作在本质上就是在由样本对组成的特征数据库 {(ks,vs)}\{(k_s, v_s)\} 上执行的一次连续核平滑回归。在这套映射中,QQ 构成了目标查询点的连续坐标,KK 标定了已有数据点在度量流形中的参考位置,而 VV 则是这些参考点所承载的函数真实取值。由此,整个注意力层可以被严谨地视为一种完全平滑可微的最近邻检索算子。这一数学隐喻对于后续 §9.7 探讨注意力在化学语境中的可解释性构成了关键认知分水岭:既然它本质上是一个连续加权的软检索器,我们就决不能将其当成具备因果逻辑链条的确定性化学规则去过度解读。

2.6 反向传播:把 L/O\partial L/\partial O 一路推回参数

作为全书在微积分层面的核心推演,本节将展示如何将标量损失的梯度一路沿计算图无损逆推至每一个权重参数。在此统一采用符号约定:ZˉL/Z\bar Z \equiv \partial L/\partial Z 代表标量目标 LL 对张量 ZZ 的偏导梯度,其张量形状与 ZZ 保持严格一致。

面对由上游层级传递而来的梯度张量 OˉRT×dv\bar O \in \mathbb{R}^{T\times d_v},反向传播引擎需要逆向推导并依次求解出中间变量梯度 Sˉ, Qˉ, Kˉ, Vˉ\bar S,\ \bar Q,\ \bar K,\ \bar V 以及最终参数梯度 WQˉ,WKˉ,WVˉ,WOˉ,Xˉ\bar{W_Q},\bar{W_K},\bar{W_V},\bar{W_O},\bar X。这一推演过程犹如一部严丝合缝的逆向机械链条,可以拆解为六个高度协同的阶段:

首先是拆解加权求和矩阵积 O=AVO=AV。根据元素级线性组合定义 Ots=uAtuVusO_{ts} = \sum_{u} A_{tu}V_{us},利用链式求导立即解耦出对注意力权重矩阵与值矩阵的梯度:

Aˉtu=sOˉtsVus,Vˉus=tAtuOˉts\bar A_{tu} = \sum_s \bar O_{ts}V_{us},\qquad \bar V_{us} = \sum_t A_{tu}\bar O_{ts}

将其整理为优雅紧凑的矩阵乘法,便完成了第一段梯度的逆流传递:

Aˉ=OˉV,Vˉ=AOˉ(2.11)\bar A = \bar O V^\top,\qquad \bar V = A^\top \bar O \qquad\textbf{(2.11)}

紧接着是逆向穿越 softmax 归一化层。由于每一行是独立进行概率归一化的,利用此前推导出的线性化雅可比精简公式 (2.5),我们可以按行瞬时求得未归一化打分矩阵的梯度:

Sˉt=At(AˉtAˉt,At1)(2.12)\bar S_{t\cdot} = A_{t\cdot}\odot\Big(\bar A_{t\cdot} - \langle \bar A_{t\cdot}, A_{t\cdot}\rangle\mathbf 1\Big) \qquad\textbf{(2.12)}

第三阶段是逆向穿过维度缩放。鉴于映射 S=(QK)/dkS = (QK^\top)/\sqrt{d_k} 为纯粹的线性缩放算子,其梯度的逆传递只需执行原位的标量除法更新:

SˉSˉ/dk(2.13)\bar S \leftarrow \bar S/\sqrt{d_k} \qquad\textbf{(2.13)}

第四阶段是逆向穿越双线性打分算子 QKQK^\top。展开分量求和 Sts=iQtiKsiS_{ts} = \sum_i Q_{ti}K_{si},通过对各自因式分别求偏导,可知 Q 与 K 的反向梯度由打分梯度矩阵与对侧张量相乘而成:

Qˉti=sSˉtsKsi,Kˉsi=tSˉtsQti\bar Q_{ti} = \sum_s \bar S_{ts}K_{si},\qquad \bar K_{si} = \sum_t \bar S_{ts}Q_{ti}

写成紧凑的张量运算,即为两个对称的标准矩阵乘法:

Qˉ=SˉK,Kˉ=SˉQ(2.14)\bar Q = \bar S K,\qquad \bar K = \bar S^\top Q \qquad\textbf{(2.14)}

第五阶段是逆向穿透输入线性投影层。以查询投影方程 Q=XWQ+1bQQ = XW_Q + \mathbf 1b_Q^\top 为例,依据全微分法则,权重矩阵与偏置向量的参数累加梯度,以及输入端接收的支路梯度分别表示为:

WQˉ=XQˉ,bQˉ=tQˉt,Xˉ(Q)=QˉWQ(2.15)\bar{W_Q} = X^\top \bar Q,\qquad \bar{b_Q}=\sum_t \bar Q_{t\cdot},\qquad \bar X^{(Q)} = \bar Q W_Q^\top \qquad\textbf{(2.15)}

对键投影与值投影的操作同构展开。在自注意力架构中,由于 Q、K、V 源自同一个底层物理输入 XX,因此根据多变量复合函数求导法则,最终注入输入端的总梯度必须将三条分支的逆流严格叠加相加,这也揭示了自注意力中同一个物理表征必须同时承载三重角色的深刻力学实质:

Xˉ=Xˉ(Q)+Xˉ(K)+Xˉ(V)(自注意力)(2.16)\bar X = \bar X^{(Q)}+\bar X^{(K)}+\bar X^{(V)}\quad (\text{自注意力}) \qquad\textbf{(2.16)}

最终,在第六阶段逆向穿越多头输出投影层 Y=O~WO+1bOY = \tilde O W_O^\top + \mathbf 1 b_O^\top 时,其结构与式 (2.15) 完全同构,从而实现整个注意力块参数梯度的完美闭环。

将上述严谨的推导浓缩为直观的工程流水线,它不仅勾勒出信息前向流动的拓扑全貌,更在反向传播中构成了清晰的镜像倒放(在工程实现上,这正是 transformer.py::MultiHeadAttention.backward 所严格遵循的代码组织轨迹):

注意力就是四步流水线:量距离 → 缩短尺子 → 重新分配纸条 → 按比例混合

在前向流水线中,输入张量经历多路线性投影后切分入各个注意力头,通过点积度量语义距离、应用因果掩码遮蔽未来、利用行级 softmax 计算置信权重,最终与 Value 进行加权汇聚并跨头拼合:

X ──Wq/Wk/Wv──► Q,K,V ──切头──► Q_h,K_h,V_h

      S_h = Q_h K_hᵀ / √d_k        ◄────┤   (打分 = 量距离)
      S_h ← S_h + M(掩码)         ◄────┤
      A_h = softmax_rows(S_h)       ◄────┤   (把打分变成权重)
      O_h = A_h V_h                 ◄────┘   (加权求和)

      拼接所有头 ──Wo──► Y

而反向传播过程则严格按照计算图的时间箭头倒转回溯:

text
dY ──Wo──► dOcat ──切头──► dO_h

        dA_h = dO_h V_hᵀ  ◄───┤
        dV_h = A_hᵀ dO_h  ◄───┤

        dS_h = A_h ⊙ (dA_h − ⟨dA_h, A_h⟩)   ← softmax 雅可比

        dS_h ← dS_h / √d_k    ← 缩放

        dQ_h = dS_h K_h ,  dK_h = dS_hᵀ Q_h

             ──Wq/Wk/Wv──► dXq + dXk + dXv = dX

在掩码机制与梯度传播的交织处,蕴含着一个令人拍案叫绝的物理特质:对于任意在注意力矩阵中被因果掩码置为 -\infty 的被屏蔽位置,其 softmax 输出权重数学上严格归零(Ats=0A_{ts}=0)。将此项代回梯度公式 (2.12),可知对应未缩放打分的梯度 Sˉts\bar S_{ts} 同样被精确地置为零。这意味着因果掩码不仅在前向计算中切断了来自未来的信息探视,更在反向传播中在数学底层完全阻断了任何梯度跨越时间维度的回流——掩码在物理意义上保证了“未来不仅不可读,而且在历史中根本不曾存在”。

2.7 复杂度与"纯 Python 怎么跑得动"

从计算资源的角度审视注意力模块的各个操作环节,我们可以清晰地勾勒出其资源消耗图谱:

操作时间复杂度空间复杂度
Q,K,VQ,K,V 投影O(Td2)O(Td^2)O(Td)O(Td)
打分 QKQK^\topO(T2dk)O(T^2 d_k)O(T2)O(T^2)
softmaxO(T2)O(T^2)O(T2)O(T^2)
加权求和 AVAVO(T2dv)O(T^2 d_v)O(Td)O(Td)
FFNO(Tddff)O(Td\,d_{ff})O(Tdff)O(Td_{ff})

在这张复杂度图谱中,最核心的算力与内存瓶颈毫无疑问是与序列长度相关的 O(T2)O(T^2) 二次项。当序列长度处于 T=1024T=1024 的常规尺度时,单层注意力矩阵尚且包含约一百万个浮点元素;但一旦面对超长聚合物链或全长生物大分子(TT 突破百万量级),其单层矩阵元素数量将激增至天文数字般的 101210^{12},令所有硬件内存瞬间告罄。

既然算力账单是 T2T^2,而 TT 完全由分词粒度决定,那么"换一个好 tokenizer"就是与"改架构"同级的手段:在 ESOL 的 1128 条 SMILES 上,把字符级分词换成 512 次合并的 BPE,平均长度从 22.33 降到 4.96,tT2\sum_t T^2 随之降到原来的 6.4%(附录 E 有完整实测与实现)。

在实验一的精密基准测试中,我们对这一增长规律进行了量化验证。在这里需要指出一个极易被浅层理论分析所忽视的工程细节:一个完整的多头注意力前向过程在计算结构上属于混合阶数模型——它由与长度呈线性关系的 O(Td2)O(Td^2) 投影运算,以及与长度呈平方关系的 O(T2d)O(T^2d) 打分和聚合运算叠加而成:

测什么log-log 斜率理论
纯打分矩阵 QKQK^\top1.982
完整多头注意力前向(d=32d=321.431.43 = 混合阶数的结果
t=a+bT2t=a+bT^2 二参数拟合R2=0.9986R^2=0.9986——

当序列较短(例如 T32T\le32)时,由于投影项的开销以及解释器层面的固定调用成本占据绝对主导,导致双对数坐标下的实测拟合斜率仅为 1.43(这也从微观工程层面解释了为什么在短序列小模型上,自注意力并没有表现出令人生畏的计算负担);而一旦序列长度持续延展,平方项的物理规律必然最终全面接管系统的运行耗时。

左:两种复杂度;右:除以 T² 之后仍在下降,说明小 T 由线性项与固定开销主导

针对长程注意力带来的平方复杂度危机,学术界与工业界探索出了三条截然不同的破局路线。第一种是以线性注意力为代表的核化方案,它尝试通过核函数映射将内积指数项近似分解为 ϕ(q)ϕ(k)\phi(q)^\top\phi(k),借助矩阵乘法的结合律优先计算键与值的乘积 ϕ(K)V\phi(K)^\top V,从而将时间与内存复杂度一举压缩至纯粹的线性阶 O(T)O(T);第二种是以 FlashAttention 为代表的硬件协同设计路线,它并不改变任何底层的数学等式,而是利用 GPU 高速 SRAM 缓存执行分块(Tiling)计算与在线归一化重计算,将原本需要驻留显存的中间注意矩阵完全消除,使峰值显存从 O(T2)O(T^2) 降维至 O(T)O(T);第三种则是基于先验的稀疏或局部窗口注意力机制。在化学与材料计算中,这一方案尤为实用,因为长链聚合物的物理性质很大程度上由局部单体微环境决定;然而必须警惕的是,过度局域化的视野会直接切断环闭合数字跨越数十个字符形成的远距离精确依赖。

在厘清了算法复杂度之后,我们有必要回答一个读者心中常有的疑问:在彻底摒弃 NumPy 等成熟 C 扩展库的前提下,本教程的纯 Python 代码究竟凭什么能够承担起真实的化学模型训练?答案在于深入 Python 底层的内存操作:我们规避了性能低下的原生三重 for 循环嵌套,而是将内积计算的核心循环交由 Python 解释器内建的 C 层迭代器处理:

python
def matmul(A, B):
    Bt = transpose(B)
    return [[sum(map(mul, row, col)) for col in Bt] for row in A]

经由基准测试,这种利用内置 mapsum 在 C 语言连续内存块中执行向量内积的方式,在现代苹果 M 系列芯片的单线程环境下能够稳定提供高达 205 Mflops/s 的浮点吞吐量,相比质朴的多层循环提升了整整一个数量级。正是这一工匠级优化,为我们完全依靠原生 Python 训练 d=48d=48 的双层深度网络并完成全套前沿化学实验(见 §10)提供了坚如磐石的算力底座。


3. 多头注意力

3.1 定义

如果仅使用单一组查询与键,模型在给定的表征空间中便只能捕获单一维度的关联模式。然而,现实中的语义关系错综复杂——某个 token 可能需要同时在句法层面上锚定紧邻的修饰词,在长程范围上追踪代词的指代实体,并在功能层面上匹配对应的语法标记。多头注意力(multi-head attention)的构造极其自然地迎合了这一诉求:将整体的 dd 维隐层空间均匀切分为 HH 个互不重叠的子空间,让每一个独立的注意力“头”在各自专属的低维投影中自由寻找关注目标,最后将所有头的观察结果重新拼接并投影回统一的隐层维度:

Q=XWQ,K=XWK,V=XWV(整体算一遍,再切分)headh=Attention(Qh,Kh,Vh),h=1,,HMHA(X)=[head1;;headH]WO+1bO(3.1)\begin{aligned} Q &= XW_Q,\quad K = XW_K,\quad V = XW_V \quad(\text{整体算一遍,再切分})\\ \text{head}_h &= \mathrm{Attention}\big(Q_h, K_h, V_h\big),\qquad h=1,\dots,H\\ \mathrm{MHA}(X) &= \big[\text{head}_1;\dots;\text{head}_H\big]\,W_O + \mathbf 1 b_O^\top \end{aligned} \qquad\textbf{(3.1)}

其中 QhQ_h 对应 QQ 矩阵中属于第 hh 个头的连续切片(每个头的特征维度为 dh=d/Hd_h=d/H),符号 [;][\cdot;\cdot] 表示沿特征轴执行的张量拼接(拼合后的形状恢复为 T×dT\times d),而 WORd×dW_O\in\mathbb R^{d\times d} 则是统领全局的输出投影矩阵。

同一个位置上三副不同形状的眼镜:看邻居、看远处、看配对

在这一架构的设计与落地过程中,有两个至关重要的实现机理需要深刻把握。在计算流水线上,先执行一次完整的高维大矩阵乘法、随后在内存中切片分头,与“为每个头单独配置小矩阵分别相乘”在代数上完全等价;但由于现代硬件对大规模密集矩阵乘法拥有极高的指令流水优化,前者的吞吐吞吐效率远胜于分散的小计算,本教程在 MultiHeadAttention.forward 中采用的正是这种高效的批式投影。更为深邃的一点在于输出矩阵 WOW_O 的拓扑角色:如果我们将其视为由 H×HH\times H 个大小为 dh×dhd_h\times d_h 的子矩阵拼接而成的块矩阵,便会发现每一个小分块本质上承载着“如何将第 hh' 个头的特征流折射并写入第 hh 个输出通道”的线性组合权力。换言之,WOW_O 是整个多头注意力模块中实现跨头信息交流与跨子空间特征融合的唯一法定通道

3.2 多头不是"更宽的单头"

为什么我们需要多头机制,而不是简单地把单头注意力的特征维度 dd 拓宽?为了回答这个问题,我们可以证明一条具有几何美感的重要命题。

命题 3.1(单头的凸包限制) 设一个单头注意力的输出为 O=AVO = AV,其中权重矩阵 AA 的每一行均为严格满足全概率条件的离散分布向量。则对任意给定的序列位置 tt,其输出行向量满足:

Ot=sAtsVsconv{V1,,VT}O_{t\cdot} = \sum_s A_{ts}V_{s\cdot} \in \mathrm{conv}\{V_{1\cdot},\dots,V_{T\cdot}\}

即单头输出的每一行都必然被严格局限在值矩阵 VV 各行向量所张成的**凸包(Convex Hull)**之内。

该命题的证明直接植根于凸分析的基本定义:由于行权重满足非负性 Ats0A_{ts}\ge0 且归一化条件 sAts=1\sum_s A_{ts}=1,由它们加权生成的 OtO_{t\cdot} 在定义上就是集合 {V1,,VT}\{V_{1\cdot},\dots,V_{T\cdot}\} 的一个标准凸组合。\square

凸包限制揭示了单头注意力在表达能力上的根本枷锁:假设当前位置 tt 同时面临两项互相冲突的检索诉求——它既需要在特征维度 1 上重点汲取来自位置 aa 的特征 vav_a,又必须在特征维度 2 上汲取来自位置 bb 的特征 vbv_b。由于单头机制下只存在一个唯一的标量权重向量 AtA_{t\cdot},模型只能在两个目标之间做出妥协与折中,强行寻找折中点。反观多头注意力配合输出投影 WOW_O,其最终输出在数学形式上是 HH源自不同低维子空间凸包的表征在经历矩阵变换后的广义线性组合

Yt=h=1Hs(Ats(h)Vs(h))WO(h)(每个 h 项各自在凸包内)Y_t = \sum_{h=1}^H \sum_{s}\Big(A^{(h)}_{ts}\,V^{(h)}_{s\cdot}\Big)\,W_O^{(h)} \qquad\text{(每个 $h$ 项各自在凸包内)}

一旦跨越了头维度的加和运算,线性组合的系数不再受制于非负与和为 1 的概率公理,最终的输出表征 YtY_t 彻底摆脱了任何单一凸包的几何囚笼,获得了在整个高维向量空间中自由重构的无上自由度。

这一数学结论得到了大量经验实证的强力支撑。Voita 等人(2019)在机器翻译任务上进行严苛的注意力头剪枝实验时发现,在深度网络的较高层级,相当一部分注意力头的分布高度重叠,而真正留存下来的核心头则高度特化,分别承担起“追踪主谓依从”、“跨句法指代对齐”等截然分立的语言学职责。这充分证明多头机制并非简单粗暴的容量线性堆叠,而是真正促成了不同语义维度的分工协作。

这种分工现象在化学任务中同样留下了生动的注脚。在本教程的实验五(§10.5)中,我们对训练在真实化学 SMILES 上的双层小模型进行了逐头解剖:8 个注意力头的信息熵各不相同(分布在 2.08 至 2.36 之间),这表明它们确实捕获了不同的上下文模式;但令人警醒的科学事实是,这些小模型学到的主要是紧邻字符或上一个原子的局部词法依赖,并未自发涌现出文献中大模型所展示的那种远距离“环闭合配对专职头”。这是一个极其真实的负结果,其背后的物理机理将在 §10.5 深入揭晓。

3.3 多头的反向传播

理解多头注意力的反向传播,核心在于看透其并行的拓扑结构:在注意力矩阵与加权求和的计算过程中,各个注意力头在数学上是完全解耦、互不干涉的。它们唯一的耦合点仅仅发生在最前端(共同瓜分输入 XX)与最末端(共同馈入 WOW_O 执行线性映射)。

因此,梯度的回溯过程非常清晰优雅:上游梯度 Oˉ\bar O 首先通过切片被解构为属于各个头的局部梯度 Oˉ1,,OˉH\bar O_1,\dots,\bar O_H;随后,每个头独立并发地执行 §2.6 中详述的完整逆向求导链条;最终,在回退至原始输入 XX 的汇合点时,将来自所有头、涵盖 Q、K、V 三重身份的支路梯度全部累加求和:

Oˉ=拆成 Oˉ1,,OˉH    每头独立做 §2.6 的六步    Xˉ=h(Xˉh(Q)+Xˉh(K)+Xˉh(V))(3.2)\bar O = \text{拆成 } \bar O_1,\dots,\bar O_H \;\Longrightarrow\; \text{每头独立做 §2.6 的六步} \;\Longrightarrow\; \bar X = \sum_h \Big(\bar X_h^{(Q)}+\bar X_h^{(K)}+\bar X_h^{(V)}\Big) \qquad\textbf{(3.2)}

在底层的工程实现中,transformer.py::MultiHeadAttention.backward 函数清晰地映射了这段数理逻辑,其核心循环仅需简洁的六行张量运算:

python
dA  = matmul(dOh, transpose(Vh))          # dA = dO Vᵀ
dVh = matmul(transpose(A), dOh)           # dV = Aᵀ dO
dS  = softmax_backward(A, dA)             # softmax 雅可比
dS  = scale_mat(dS, self.scale)           # 缩放 1/√d_k
dQh = matmul(dS, Kh)                      # dQ = dS K
dKh = matmul(transpose(dS), Qh)           # dK = dSᵀ Q

而当引入旋转位置编码(RoPE)时,只需在完成打分梯度回传后,再追加两行正交逆旋转操作即可(详见 §4.5):

python
dQh = [rope_rotate_inverse(dQh[p], p) for p in range(Tq)]
dKh = [rope_rotate_inverse(dKh[p], p) for p in range(Tk)]

由于旋转变换的变换矩阵具有天然的正交性,其反向传播等价于乘以转置矩阵,在几何上恰好等价于以完全相同的角位移向相反方向执行逆向旋转,代码实现干净纯粹到了极点。

3.4 头数与 dhd_h 怎么选

在工程实践中,选择多少个注意力头以及设定每个头的子空间维度 dhd_h,是一门在表征自由度与单头容量之间寻找帕累托最优的权衡艺术。业界历经千锤百炼沉淀出三条核心设计准则:

首先,将单头维度 dhd_h 锚定在 64 或 128 附近已成为大模型架构的通用准则。无论是经典的 GPT-3 还是开源中坚 LLaMA,其头数均被设定为模型隐层维度与 64(或 128)的商。这一经验法则直接呼应了 §2.4 的数理推导:当 dhd_h 过于膨胀时,内积打分的动态范围极难驾驭,更容易诱发潜在的饱和危机;而若 dhd_h 过于逼仄,单个子空间将被压缩至缺乏必要的表征容量。

其次,头数的扩展必须警惕边际效用递减。在模型总维度 dd 锁定的前提下,盲目增加头数必然以剧烈削减单头维度 dhd_h 为代价,导致各个子空间过于狭窄而丧失复杂的特征表达力。因此在主流工程配置中,头数通常稳定在 8 至 32 之间,以维持多角度观察与单通道深度之间的精妙平衡。

最后,多头注意力与归一化层之间存在着天然的特征冗余。正如剪枝理论所揭示的,许多头在收敛后表现出极高的共线性。为此,现代推理优化架构积极引入了多查询注意力(MQA)与分组查询注意力(GQA),通过让多个甚至全部 Query 头共享同一组 Key 与 Value 权重矩阵,在几乎不折损下游精度的前提下,奇迹般地降低了自回归推理阶段庞大的显存带宽瓶颈。


4. 位置编码

4.1 先证明:没有位置编码时,顺序完全不可见

如果说自注意力的全局交互赋予了它穿透距离的伟力,那么这种机制也带来了一柄隐蔽的达摩克利斯之剑:在纯粹的自注意力层眼中,序列顺序是彻底隐形的。为了将这一认知从经验直觉上升为不可动摇的数理公理,我们需要完整建立如下定理。

定理 4.1(自注意力的置换等变性) 设映射 ff 为一个不包含任何位置信息的自注意力层(即设定 pos_encoding="none"),π\pi 代表作用在序列索引上的任意双射置换,定义置换后的输入矩阵为 (πX)t=Xπ1(t)(\pi X)_t = X_{\pi^{-1}(t)}。则网络层的前向映射满足严格的置换等变性:

f(πX)=πf(X)f(\pi X) = \pi f(X)

证明:让我们沿计算图逐层审视置换操作的流动。首先考察线性投影,因为 Q=XWQQ = XW_Q,矩阵乘法对输入行向量的线性变换完全独立于其所在的行索引,故有 (πX)WQ=π(XWQ)=πQ(\pi X)W_Q = \pi(XW_Q) = \pi Q。根据完全对称的代数结构,键矩阵与值矩阵同样满足 (πX)WK=πK(\pi X)W_K = \pi K 以及 (πX)WV=πV(\pi X)W_V = \pi V

接下来考察打分矩阵在置换作用下的变换规律:

S=(πQ)(πK)/dkS' = (\pi Q)(\pi K)^\top/\sqrt{d_k}

注意到乘积矩阵 (πQ)(πK)(\pi Q)(\pi K)^\top 位于第 (t,s)(t,s) 位置的元素恰好为 qπ1(t),kπ1(s)\langle q_{\pi^{-1}(t)}, k_{\pi^{-1}(s)}\rangle,这意味着打分矩阵在行列两个方向上同时经受了置换算子的作用,即 S=πSπS' = \pi S \pi^\top。由于 softmax 算子是严格沿每一行独立执行的非线性归一化,行的重排与列的置换不会改变相对指数比率,因此归一化后的注意力权重矩阵满足:

A=softmaxrow(πSπ)=πAπA' = \mathrm{softmax}_{\text{row}}(\pi S\pi^\top) = \pi A \pi^\top

进而推导值向量的加权汇聚结果:

O=A(πV)=(πAπ)(πV)O' = A'(\pi V) = \big(\pi A \pi^\top\big) \big(\pi V\big)

由于置换矩阵具有天然的正交性,恒有 ππ=I\pi^\top\pi = I。代入上式,中间项瞬间对消闭合:

O=πA(ππ)V=π(AV)=πOO' = \pi A \big(\pi^\top \pi\big) V = \pi (AV) = \pi O

此后,后续的输出线性投影 WOW_O 逐行作用而不改变置换状态,LayerNorm 逐 token 统计内部均值方差而不涉及跨 token 顺序,FFN 前馈层同样逐位置独立作用。通过数学归纳法将这一性质递推至任意深度 LL 的网络,定理得证。\square

推论 4.2 设模型采用自注意力编码器提取特征,并使用**平均池化(mean pooling)**构造整个序列的全局池化表征 z=1Tthtz = \frac1T\sum_t h_t。则对于任意置换 π\pi,恒有 z(πX)=z(X)z(\pi X) = z(X)换言之,无论如何随意洗牌打乱序列中 token 的物理先后顺序,模型的最终表征与预测输出严格保持纹丝不动。

这绝非停留在纸面上的数学推论,而是可以在计算机中被高精度实测确认的冷酷事实。在实验二(results_positional.json)的严格对比测试中,我们对不同位置编码机制在打乱输入序列后的预测变化幅度进行了基准量化:

位置编码打乱 token 顺序后预测的平均变化相对变化
none(无位置编码)0.00002.9×10162.9\times10^{-16}2.9×10162.9\times10^{-16}
sine(正弦)0.03931.8×1011.8\times10^{-1}
learned(可学习)0.07841.3×1011.3\times10^{-1}
rope(旋转)0.14742.7×1012.7\times10^{-1}

排列实验:只有 none 那一列精确为零

在上述表格中,无位置编码项所测得的 2.9×10162.9\times10^{-16} 绝非“微小的数值扰动”,它在本质上正是 64 位双精度浮点数算术运算中不可避免的机器舍入精度极限——在严密的代数空间中,这一变化量严格为零。

这一本质特性在化学与材料计算中揭示了一对尖锐的矛盾:一方面,在处理小分子时,乙醇分子无论写成 CCO 还是 OCC,在物理客观世界中都是同一客观实体。若盲目引入位置编码,模型会对这两种本该等价的写法给出大相径庭的预测打分;而另一方面,在更深层的物化性质预测中,相对位置信息又构成了识别骨架取向、区分末端官能团与核心母核的绝对前提。正是置换等变性与物理客观性之间的张力,迫使我们必须极其严谨地审视位置信息的注入机制。

4.2 位置的三种候选:整数、二进制、正弦

要告诉模型"谁在前、谁在后",最直接的想法是给每个位置发一个号码,再把号码变成向量加进词嵌入。但号码怎么变成向量?下面三种方案是历史上真实出现过的三级阶梯,我们用同一套可复现的脚本(demo_positional.py::experiment_ladderd=64d=64、位置 0,,2550,\dots,255)量了它们的三个性质:幅度、连续性、粒度

三格里举着三种座位号:越举越沉的号码牌、忽高忽低的方格、始终很轻的波纹纸带

4.2.1 整数编码:最简单,也最快出局

PEt=[t,t,,t]RdPE_t = [\,t,t,\dots,t\,]^\top\in\mathbb R^d,与词嵌入逐元素相加。它的致命伤是量纲

PEt=td(4.1)\|PE_t\| = t\sqrt d \qquad\textbf{(4.1)}

d=64d=64,位置 t=255t=255 时范数已经达到 2040,而一个逐分量标准正态的词嵌入,其范数期望只有 d=8\sqrt d=8(这也是 LayerNorm 输出后的典型尺度)。也就是说,在这个位置上,位置信号比语义信号大 255 倍(实测 2040/8)。这不叫"位置信息更强",而叫语义被淹没:模型在第 255 个位置上无论读到什么 token,输入向量的方向几乎只由位置决定。附带一提,整数编码在位置 0 处天生全零,与"没有位置"无法区分——但这只是小问题,量纲才是判决性证据。

4.2.2 二进制编码:把幅度关进笼子,换来跳变

把位置写成二进制位,每一位占一个维度,取值只有 0/1:

PEt,i=t2imod2,i=0,1,,w1,w=log2T(4.2)PE_{t,i} = \Big\lfloor \frac{t}{2^{i}}\Big\rfloor \bmod 2,\qquad i=0,1,\dots,w-1,\quad w=\lceil\log_2 T\rceil \qquad\textbf{(4.2)}

T=256T=256w=8w=8,于是范数被死死限制在 [0,8]=[0,2.83][0,\sqrt 8]=[0,2.83](实测 0 – 2.83,均值 1.96),量纲问题一举解决。更妙的是,它顺手送出了一个频率阶梯:最低位每 2 步翻转一次,第 1 位每 4 步,第 ii 位每 2i+12^{i+1} 步——实测前 8 位的翻转频率为

1.000, 0.498, 0.247, 0.122, 0.059, 0.027, 0.012, 0.0041.000,\ 0.498,\ 0.247,\ 0.122,\ 0.059,\ 0.027,\ 0.012,\ 0.004

每往左一位就精确减半(顺带一提:二进制编码在 t=0t=0 处是全零向量,与"没有位置"不可区分——这与整数编码共享同一个缺陷,实测范数下界恰为 0)。"低位管细节、高位管全局"这一思想,正是位置编码最核心的设计直觉,它后来被正弦编码完整继承(§4.3)。

但二进制编码在两个地方留下了硬伤,都出在"连续性"上:

  1. 相邻位置的距离是跳变的PEt+1PEt\|PE_{t+1}-PE_t\| 只能取离散值,实测恰好是集合 {1.000, 1.414, 1.732, 2.000, 2.236, 2.449, 2.646, 2.828}\{1.000,\ 1.414,\ 1.732,\ 2.000,\ 2.236,\ 2.449,\ 2.646,\ 2.828\}(即 1,,8\sqrt 1,\dots,\sqrt 8),变异系数 0.31。一旦进位(如 t=6364t=63\to 64),多个比特同时翻面,位置表示在向量空间里"瞬移"。
  2. 每一位关于 tt 都是阶跃函数,几乎处处导数为零。位置本身是离散量,谈不上求导;但在需要把位置当可学习参数输入的场合,梯度只能沿着"某一位恰好翻转"这一条窄路传递,优化器面对的是一片布满悬崖的地形。

值得强调的是:Transformer 中位置编码只参与前向加法(反传到的是嵌入矩阵,而不是位置本身),所以这两条伤都不是致命的。它们真正的意义是解释后来者为什么改用光滑函数

4.2.3 正弦编码:把频率阶梯连续化

正弦位置编码的做法,就是把二进制的每一个比特换成一条连续波:波长从 2π2\pi(局部、快)按几何级数一路拉到 2π1042\pi\cdot 10^4(全局、慢)。它同时保住了二进制编码的"多频率"优点,又消除了它的"跳变"缺点:

性质整数列 {t}\{t\}二进制(w=8w=8正弦(d=64d=64
范数 PEt\lVert PE_t\rVerttdt\sqrt dt=255t=25520400 – 2.83恒为 d/2=5.657\sqrt{d/2}=\textbf{5.657}
相邻位置距离恒为 d=8\sqrt d=8(但整体尺度在漂移)8 个离散值,CV 0.31恒为 1.4718,CV 4.6×10154.6\times10^{-15}
可微性线性可导,但尺度无界阶跃,几乎处处导数为 0处处光滑
可学参数000
外推行为无界增长超过 2w2^w 截断/回绕解析式,可外推(但见 §10.3 的实测)

表中正弦一行的两个"恒为"都不是巧合,它们是 §4.3 将要证明的两条恒等式的直接推论:由内积恒等式立刻得到

PEt+1PEt2=2i=0d/21(1cosωi)(4.3)\|PE_{t+1}-PE_t\|^2 = 2\sum_{i=0}^{d/2-1}\big(1-\cos\omega_i\big) \qquad\textbf{(4.3)}

右端与 tt 完全无关。脚本实测的 1.471848 与这个闭式给出的 1.471848 在双精度下逐位一致。

位置信号的三个候选:整数无界增长、二进制跳变、正弦恒定且光滑

从这三级阶梯可以抽出一条贯穿全章的设计准则:一个理想的位置编码应当同时满足"多频率、有界、连续"。多频率提供多尺度的分辨率(谁都能编码"相邻"与"全局"),有界保证它与词嵌入处于同一量纲(不淹没语义),连续保证梯度能被平稳地传下去。正弦编码是满足这三条的最简解析解;2017 年之后的每一种变体(可学习位置编码、RoPE、ALiBi,见 §4.4–§4.6)都只是在重写这条准则,而从未否证它。

4.3 正弦位置编码:加法定理

每个位置发一张不同波纹的座位号:没有它,模型分不清顺序

在奠基之作中,Vaswani 等人给出了一种纯粹基于解析连续函数的巧妙设计:通过一组不同波长的正弦与余弦波,为每一个离散位置 tt 构造一个确定性的 dd 维嵌入向量:

PEt,2i=sin ⁣(tωi),PEt,2i+1=cos ⁣(tωi),ωi=100002i/d(4.4)PE_{t,2i} = \sin\!\big(t\,\omega_i\big),\qquad PE_{t,2i+1} = \cos\!\big(t\,\omega_i\big),\qquad \omega_i = 10000^{-2i/d} \qquad\textbf{(4.4)}

其中指标 i=0,1,,d/21i=0,1,\dots,d/2-1。这里的频率族构成了一组标准的等比几何级数,其波长 2π/ωi=2π100002i/d2\pi/\omega_i = 2\pi\cdot10000^{2i/d} 极其宽广地从局部的 2π2\pi 一直铺展到宏观的 2π1042\pi\cdot10^4。在这一构造中,低频分量波长极长,犹如二进制编码中的最高位,负责粗粒度地定位大尺度的全局区间;而高频分量剧烈振荡,好比最低有效位,负责精细分辨微小的邻域位移。

4.3.1 加法定理:PEt+kPE_{t+k}PEtPE_t 的固定线性变换

为什么这组看似神秘的三角函数波纹能够让网络感知到相对距离?其奥秘深植于三角函数的和角公式之中。考察每一对正交分量 (2i,2i+1)(2i, 2i+1),根据高中三角恒等式展开:

PEt+k,2i=sin(ωit+ωik)=sin(ωit)cos(ωik)+cos(ωit)sin(ωik)=cos(ωik)PEt,2i+sin(ωik)PEt,2i+1PEt+k,2i+1=cos(ωit+ωik)=cos(ωit)cos(ωik)sin(ωit)sin(ωik)=sin(ωik)PEt,2i+cos(ωik)PEt,2i+1\begin{aligned} PE_{t+k,2i} &= \sin(\omega_i t + \omega_i k) = \sin(\omega_i t)\cos(\omega_i k) + \cos(\omega_i t)\sin(\omega_i k)\\ &= \cos(\omega_i k)\,PE_{t,2i} + \sin(\omega_i k)\,PE_{t,2i+1}\\[2pt] PE_{t+k,2i+1} &= \cos(\omega_i t + \omega_i k) = \cos(\omega_i t)\cos(\omega_i k) - \sin(\omega_i t)\sin(\omega_i k)\\ &= -\sin(\omega_i k)\,PE_{t,2i} + \cos(\omega_i k)\,PE_{t,2i+1} \end{aligned}

令旋转角 θ=ωik\theta=\omega_i k,上述联立方程组可以被紧凑地写为一个经典的二维平面正交旋转矩阵:

(PEt+k,2iPEt+k,2i+1)=(cosθsinθsinθcosθ)R(θ)(PEt,2iPEt,2i+1)\begin{pmatrix} PE_{t+k,2i} \\ PE_{t+k,2i+1}\end{pmatrix} = \underbrace{\begin{pmatrix}\cos\theta & \sin\theta\\ -\sin\theta & \cos\theta\end{pmatrix}}_{R(\theta)} \begin{pmatrix} PE_{t,2i}\\ PE_{t,2i+1}\end{pmatrix}

将所有特征维度成对拼装,便得到了著名的线性变换关系:

PEt+k=RkPEt,Rk=blockdiag(R(ω0k),,R(ωd/21k))(4.5)PE_{t+k} = R_k\,PE_t,\qquad R_k = \mathrm{blockdiag}\big(R(\omega_0k),\dots,R(\omega_{d/2-1}k)\big) \qquad\textbf{(4.5)}

式 (4.5) 蕴含着非同凡响的物理意义:在序列中向后平移 kk 个步长这一操作,映射在位置向量空间中,完全等价于施加一个与当前绝对位置 tt 毫无瓜葛的确定性块对角正交旋转矩阵 RkR_k。这意味着深度神经网络无需为每个绝对位置死记硬背特征,只需通过线性投影层学到一个固定的旋转投影算子,便能在数学底层随时感知“相距 kk 个步长”这一相对几何关系。在实验二的高精度数值验证中,该恒等式的双精度实测最大误差仅为 5.3×10155.3\times10^{-15}

4.3.2 内积恒等式:相似度只依赖距离

除了加法定理之外,正弦编码还蕴藏着另一条更为深刻却往往被教科书忽略的代数性质:对于任意给定的绝对起点 tt,两个位置编码之间的内积在数学上严格独立于起点 tt 本身

PEt,  PEt+k=i=0d/21[sin(ωit)sin(ωi(t+k))+cos(ωit)cos(ωi(t+k))]=i=0d/21cos(ωik)(4.6)\big\langle PE_t,\; PE_{t+k}\big\rangle = \sum_{i=0}^{d/2-1}\Big[\sin(\omega_i t)\sin\big(\omega_i(t+k)\big) +\cos(\omega_i t)\cos\big(\omega_i(t+k)\big)\Big] = \sum_{i=0}^{d/2-1}\cos(\omega_i k) \qquad\textbf{(4.6)}

借助积化和差公式 cos(AB)=cosAcosB+sinAsinB\cos(A-B)=\cos A\cos B+\sin A\sin B,令 A=ωitA=\omega_i tB=ωi(t+k)B=\omega_i(t+k),两项相减后关于绝对位置 tt 的依赖被完美抵消,最终求和项只剩下纯粹关于位移跨度 kk 的余弦叠加。

在实验环境下的数值实测充分证实了这一推导:当设定维度 d=64d=64,位移 kk 遍历 0 到 128,绝对起点 tt 任意取遍 0 到 255 时,不同起点之间余弦相似度的最大离散极差仅为 1.0×10151.0\times10^{-15},而实测均值与理论解析式 1d/2icos(ωik)\frac{1}{d/2}\sum_i\cos(\omega_i k) 之间的最大绝对偏差被严格压制在 3.9×10143.9\times10^{-14} 之内。

左:内积只依赖距离;右:它与 t 的无关性精确到浮点舍入

正弦位置编码热图:每一维是一条不同频率的波

这一数学性质为模型的长度泛化注入了强大的先验:当模型在训练期见识过距离为 k=3k=3 的局部模式后,在测试期即便遭遇长度远超训练集的崭新序列,只要相对距离依然维持为 3,位置编码所提供的内积引导信号在理论上便完全恒定不变。这也是正弦编码在小规模序列上具备一定外推能力的底层根基(虽然如 §10.3 所示,其外推稳定性仍略逊于 RoPE)。

4.4 可学习位置编码

除了构造解析的正弦波形,另一种极其直接的工程路线是引入完全由数据驱动的可学习位置编码(Learned Positional Embeddings)。该方案直接在内存中开辟一个参数矩阵 PRTmax×dP\in\mathbb R^{T_{\max}\times d},将每个离散位置视作普通的可训练参数,在梯度反向传播中随同整个网络一并更新(在早期经典的 BERT 与 GPT-2 中均采用此设计)。

可学习位置编码的优势在于其极致的灵活性与无约束的拟合潜力,在训练数据极其充沛且推理长度严格受限的场景下往往能达到优异的性能。然而,这一设计隐藏着两大致命缺陷:首先是绝对缺乏外推能力,一旦推理阶段序列长度突破了预设的上限(t>Tmaxt>T_{\max}),对应的行向量在参数表中完全未曾定义,网络直接面临无法初始化的窘境;其次是初始化量纲的严重失衡。正如我们在实验二中所测得的冷酷数据,可学习矩阵通常采用微小的正态随机数初始化(其模长范数通常在 0.02 量级,Pt0.16\|P_t\|\approx0.16),相比正弦编码依靠确定性解析式带来的天然范数(PEt=d/25.66\|PE_t\|=\sqrt{d/2}\approx5.66),两者足足相差了 35 倍之巨。这意味着在训练初期,可学习位置信号极其容易被强烈的词嵌入特征全面淹没,必须经历漫长痛苦的优化过程才能将位置的显著性缓慢“长”出来。

4.5 RoPE:把位置放进旋转里

旋转位置编码(Rotary Position Embedding, RoPE,由 Su 等人于 2021 年提出)是当今大语言模型生态(从 LLaMA、Mistral 到 Qwen 系列)毫无争议的统治级方案。其核心哲学发生了一场优雅的范式转移:彻底放弃在输入嵌入层生硬累加位置向量的做法,转而在计算自注意力内积打分时,动态地为 Query 与 Key 向量施加正交平面旋转

具体而言,对于处于第 mm 个物理位置的表征向量,将其拆解为若干个二维平面,在第 ii 对分量 (2i,2i+1)(2i, 2i+1) 构成的子平面上定义正交旋转变换矩阵:

Rm(i)=(cos(mωi)sin(mωi)sin(mωi)cos(mωi)),ωi=100002i/dhR_m^{(i)} = \begin{pmatrix}\cos(m\omega_i) & -\sin(m\omega_i)\\ \sin(m\omega_i) & \cos(m\omega_i)\end{pmatrix}, \qquad \omega_i = 10000^{-2i/d_h}

随后利用该块对角旋转矩阵对查询向量与键向量执行原位旋转映射:q~m=Rmqm\tilde q_m = R_m q_m 以及 k~n=Rnkn\tilde k_n = R_n k_n

定理 4.3(RoPE 的相对内积性) 经由 RoPE 正交旋转映射后的内积打分,在代数上精确地只依赖于两个位置之间的相对位移 nmn-m

q~m,k~n=qmRmRnkn=qmRnmkn(4.7)\langle \tilde q_m,\tilde k_n\rangle = q_m^\top R_m^\top R_n k_n = q_m^\top R_{n-m}k_n \qquad\textbf{(4.7)}

证明:对于每一个独立的 2×22\times2 二维旋转块,由于平面旋转群 SO(2)SO(2) 的阿贝尔性质与正交性,矩阵乘积满足 R(a)R(b)=R(a)R(b)=R(ba)R(a)^\top R(b) = R(-a)R(b) = R(b-a)。鉴于全空间旋转矩阵 RmR_m 是由各个独立的二维块对角拼接而成,其多矩阵乘积逐块封闭进行,最终必然严格得到 RmRn=RnmR_m^\top R_n = R_{n-m}\square

在实验二的严格基准对拍中,我们随机生成了 200 组向量对 (q,k)(q,k) 与位置对 (m,n)(m,n),并分别在相对位移 δ{1,7,33}\delta\in\{1,7,33\} 条件下检验内积的不变性,实测测得的最大相对误差被死死锁定在 6.1×10146.1\times10^{-14} 的极佳机器精度级别。

两条位置恒等式的数值验证:都在双精度机器精度量级

在将 RoPE 落地为纯 Python 代码时,有两个关键的工程洞见必须牢牢掌握。首先,RoPE 的旋转算子严禁施加于 Value 矩阵。因为旋转机制的初衷是为了在度量相似度时注入相对几何关系,而 Value 矩阵承载的是被借阅的内容本身;如果强行对 Value 执行空间旋转,将会彻底扭曲被检索出的语义实体。其次,在反向传播过程中,由于正交矩阵的逆矩阵恒等于其转置矩阵(即反向旋转 qˉ=Rmq~ˉ=Rmq~ˉ\bar q = R_m^\top \bar{\tilde q} = R_{-m}\bar{\tilde q}),代码只需简单调用 rope_rotate_inverse 即可完成梯度的逆流注回,完全无需为固定的旋转矩阵分配任何可学习参数的求导开销。

为什么 RoPE 在工业界能够全面超越正弦编码?其本质差异在于:正弦编码是“加法性”的(将空间座标硬塞进输入特征),相对位移信息在打分层只是“在理论上存在某个线性解”,需要网络在训练中耗费参数容量去艰难拟合;而 RoPE 则是“乘法性”的,它直接在内积算子中以数学恒等式的方式锁死了相对几何关系,赋予了模型无可匹敌的长度外推与语义对齐能力。

4.6 相对位置偏置与 ALiBi

除了对向量本身施加旋转,学术界探索出的第三条典型路线是直接在注意力打分矩阵上叠加一个纯粹与相对距离挂钩的静态标量偏置

StsSts+bts(4.8)S_{ts} \leftarrow S_{ts} + b_{t-s} \qquad\textbf{(4.8)}

其中最具代表性的工作是 Press 等人于 2022 年提出的 ALiBi(Attention with Linear Biases)机制。该机制甚至彻底抛弃了显式的位置嵌入,直接将偏置项设定为与绝对距离呈负线性的惩罚项:bts=ctsb_{t-s} = -c\,|t-s|,并为每个注意力头分配一条预先固定的斜率参数 cc

ALiBi 的最大杀手锏在于其强悍的外推稳定性,因为线性惩罚在理论上不存在长度截断,使得模型即使推理长度翻倍也不会出现数值崩塌。然而,这柄利刃在特定领域同样存在刺眼的短板:它在先验上强行规定了“距离越远、关联度越低”的局域性假定,从而在结构上摧毁了模型捕捉长程精确匹配的能力。在自然语言中这或许可以容忍,但在化学与生物大分子领域,诸如跨越数十个 token 的成环配对、两端遥相呼应的二硫键连结,抑或是远端基团通过空间构象对反应活性中心施加的电子诱导效应,恰恰需要模型彻底无视物理距离去精准捕获远端特征。因此,针对具体的科研任务进行选型时必须权衡:如果任务追求远距离拓扑配对,应坚定拥抱 RoPE 或正弦编码;唯有在仅关注局部短语模式或快速外推时,ALiBi 才是适宜的妥协。

4.7 化学的坑:位置 ≠ 化学距离

在将 Transformer 应用于化学与材料科学时,如果缺乏对物质底层物理规律的敬畏,极其容易陷入由文本一维序列引发的致命误区。这一认知脱节集中体现在三大深层陷阱与三种应对策略之中。

首当其冲的陷阱在于:SMILES 仅仅是分子图在特定算法下的遍历痕迹,绝非真实的分子几何空间。例如在乙醇分子 CCOOCC 的书写中,原本在化学图拓扑中彼此紧密相邻的原子,映射到线性文本上却可能隔着遥远的距离。更极端的例子体现在芳香环系中:分子中空间相邻成环的两个碳原子,在展平为带有开环闭环数字的线性 SMILES 字符串时,其在文本上的物理跨度可能长达数十个字符。如果模型把文本位置死板地当成空间距离去拟合,无异于缘木求求鱼。

第二个陷阱揭示了同一化学分子存在无数种合法写法的双刃剑效应。在分子性质预测中,如果训练集采用了一种遍历起点的写法,而测试集在数据流转中被格式化成了另一种合法的异构书写,原本确定性的物化性质将在模型内部遭遇截然不同的位置分布冲击,引发严重的伪泛化性能滑坡。但反过来,这种冗余性也为我们提供了最廉价、最高效的随机化 SMILES 数据增强手段——在训练期随机打乱遍历起点与分支书写次序,强迫模型在波动的文本位置中提炼不变的化学本质。需要注意的是,执行此类增强的前提是预测目标在物理上具备写法不变性(例如基于规范化母体构建的标签),否则反而会在训练流中注入标签噪声。

第三个陷阱在于必须破除对**“规范化 SMILES(Canonical SMILES)”的化学迷信**。无论是 RDKit 还是 OpenBabel 的规范化算法,其数学目标仅仅是保证“拓扑同构的图在输出端映射为全球唯一的字符串标识符”,其遍历优先级的决定依据是一组严谨但缺乏化学因果含义的图遍历打标规则,绝非按照主链碳链长度或有机官能团的优先级展开。因此,在任何严肃的科研分析中,决不能把字符串层面的第几位直接生硬地解释为某种物理官能团特征(例如得出“位于第 5 个位置的原子对溶解度起决定作用”之类的荒谬结论)。

面对一维文本与高维物质的客观鸿沟,严谨的科技工作者应当建立清晰的应对姿态:当业务任务要求预测结果具备严格的分子书写不变性时,应当在模型顶层采用平均池化并配合高强度的随机化 SMILES 增强;当任务强依赖长程拓扑配对(例如严格闭合分子环)时,应当配备 RoPE 旋转位置编码并配合足够深度的网络层,引导模型在多头中自主学会计算拓扑跨度;而一旦科研命题深入到底层空间构象、非键相互作用或催化过渡态的本质几何时,就必须果断跳出文本序列的狭隘视角,全面转向下文 §9.6 将深入剖析的三维等变几何模型。


5. 残差、LayerNorm 与前馈网络

注意力机制虽然奠定了跨位置信息路由的动态通道,但它在整个 Transformer 架构中仅占半壁江山。另一半由逐位置前馈网络(Position-wise Feed-Forward Network, FFN)、残差跳跃连接以及层归一化(LayerNorm)共同构筑的基石,才真正决定了深层网络能否在反向传播中稳定收敛,并承载起海量的语义记忆。

5.1 逐位置前馈网络(FFN)

一排一模一样的工作台:每个 token 各自过自己的小工厂

在前向流动的序列中,每一个 token 在完成注意力加权汇聚后,都会独立穿过一个两层的全连接多层感知机:

FFN(x)=W2σ(W1x+b1)+b2,W1Rdff×d,  W2Rd×dff(5.1)\mathrm{FFN}(x) = W_2\,\sigma\big(W_1 x + b_1\big) + b_2,\qquad W_1\in\mathbb R^{d_{ff}\times d},\; W_2\in\mathbb R^{d\times d_{ff}} \qquad\textbf{(5.1)}

所谓“逐位置”(position-wise),其深层内涵在于同一个前馈网络完全独立且并发地作用于序列中的每一个 token。在这个运算步骤中,不同位置之间被彻底屏蔽,互不通风报信——这一角色与《图神经网络》中经典的逐节点多层感知机(Node-wise MLP)或卷积神经网络中的 1×11\times1 点态卷积完全同构。

在参数量的分配上,前馈网络展现出了压倒性的统治力。将两个线性投影层与偏置项合并核算,其参数总量为 ddff+dff+dffd+d=2ddff+dff+dd\,d_{ff}+d_{ff}+d_{ff}\,d+d = 2dd_{ff}+d_{ff}+d。在经典标准架构中,中间膨胀维度通常设定为隐层维度的 4 倍(即 dff=4dd_{ff}=4d),此时 FFN 的参数量膨胀至约 8d28d^2,恰好是多头注意力机制参数量(约 4d24d^2)的两倍。这意味着在一个标准 Transformer 中,绝大多数可学习权重实际上沉淀在 FFN 之中(在本教程中,为了确保纯 Python 代码在常规 CPU 上也能实现飞速训练,我们在实验体系中采用了更为轻巧的 dff=2dd_{ff}=2d 设置)。

在激活函数的选择上,现代生成模型(如 GPT 系列)普遍摒弃了在原点处不可导的 ReLU,转而拥抱更为平滑的高斯误差线性单元(GeLU):

GELU(x)=xΦ(x)0.5x(1+tanh(2/π(x+0.044715x3)))(5.2)\mathrm{GELU}(x) = x\,\Phi(x) \approx 0.5x\Big(1+\tanh\big(\sqrt{2/\pi}(x+0.044715x^3)\big)\Big) \qquad\textbf{(5.2)}

为了确保手写反向传播能够以最高精度通过梯度检验,其导数在代码中由解析形式严格给出(见 gelu_grad 函数):

ddxGELU(x)0.5(1+t)+0.5x(1t2)c(1+30.044715x2),t=tanh(c(x+0.044715x3))\frac{d}{dx}\mathrm{GELU}(x) \approx 0.5(1+t) + 0.5x\,(1-t^2)\cdot c\,(1+3\cdot0.044715x^2), \quad t=\tanh\big(c(x+0.044715x^3)\big)

从更深层的可解释性视角审视,Geva 等人(2021)提出了极其精妙的可微键值记忆视角:如果我们把第一层权重 W1W_1 的第 ii 行向量抽象为一个语义键 kik_i,将第二层权重 W2W_2 的第 ii 列向量视为对应的语义值 viv_i,则前馈网络的运算形态可以被完全等价地重写为:

FFN(x)=i=1dffσ(x,ki)激活强度vi(5.3)\mathrm{FFN}(x) = \sum_{i=1}^{d_{ff}} \underbrace{\sigma(\langle x,k_i\rangle)}_{\text{激活强度}}\,v_i \qquad\textbf{(5.3)}

此时,整个 FFN 呈现为一座由成千上万个记忆神经元构成的分布式联想存储器。模型在预训练过程中所汲取的先验模式被分门别类地固化为一组组键值对 (ki,vi)(k_i,v_i);在前向推演时,输入表征与键向量进行模式匹配,一旦触发共鸣便按激活强度将对应的值向量写入特征流。这一认知在化学建模中展现出了极富启发性的应用价值:我们可以将第一层的每一个行向量 kik_i 视作特定的“化学基元检测器”,以此探查模型究竟凭借哪些神经元在瞬时识别羟基、羧酸、酯键或芳香环系——这已成为分子表征可解释性研究中最有力的探针之一。

5.2 残差连接:为什么深网络能训得动

两条平行小路:一条直达(残差),一条经过小工厂,合流处用标尺统一刻度(LayerNorm)

在现代稳健架构中,标准计算单元均采用预归一化(Pre-LN)模式组织残差流:

Xl+1=Xl+Fl(LN(Xl))(5.4)X_{l+1} = X_l + F_l\big(\mathrm{LN}(X_l)\big) \qquad\textbf{(5.4)}

残差连接的存在为深度网络注入了双重物理保障。首先是构筑了梯度的无损高速通路。将残差流沿计算图展开至第 LL 层输出对第 ll 层输入的导数:

XLXl=j=lL1(I+FjXj)=I+jFjXj+j<kFjXjFkXk+(5.5)\frac{\partial X_L}{\partial X_l} = \prod_{j=l}^{L-1}\Big(I + \frac{\partial F_j}{\partial X_j}\Big) = I + \sum_j \frac{\partial F_j}{\partial X_j} + \sum_{j<k}\frac{\partial F_j}{\partial X_j}\frac{\partial F_k}{\partial X_k}+\cdots \qquad\textbf{(5.5)}

展开式括号深处的那个单位矩阵 II 构成了整部深度学习史的奇迹:即便由于深度堆叠导致所有子层雅可比矩阵的乘积项均剧烈衰减趋近于零,在求和级数的最前端永远留存着一个无损直达的单位矩阵 II。这使得反向传播的梯度能够沿着这条开辟的超导通道毫无阻碍地直达最底层输入,彻底粉碎了深度网络中梯度指数消亡的诅咒。

然而,残差通路在化解梯度衰减的同时,也在前向过程中悄然埋下了特征方差线性累积的物理代价。假设每个子层的输出与输入近似独立且具有方差 σF2\sigma_F^2,则在深度为 LL 的网络顶端,残差流的总方差将随深度呈线性膨胀:

Var(XL)=Var(X0)+l=0L1σF2Var(X0)+LσF2(5.6)\mathrm{Var}\big(X_L\big) = \mathrm{Var}(X_0) + \sum_{l=0}^{L-1}\sigma_F^2 \approx \mathrm{Var}(X_0) + L\,\sigma_F^2 \qquad\textbf{(5.6)}

当网络堆叠至 48 层时,顶层特征的标准差将达到第 1 层的近 7 倍之多。这一物理现实直接引出了两项关键的架构契约:其一,层归一化必须严格部署在每一个子层的入口端(即 Pre-LN 范式),在每次特征涌入注意力或前馈网络前强行将其方差拉回 O(1)O(1) 标准尺度;其二,在初始化阶段,残差分支的输出投影应当按照 1/2L1/\sqrt{2L} 进行显式收缩(此为 GPT-2 的经典规范),使得经历 LL 层累加后的总方差在全生命周期内平稳固守在单位量级。

5.3 LayerNorm:前向与完整雅可比

定义

对于序列中任意单个 token 的 DD 维特征向量 xx,层归一化在特征维度上独立统计其均值与方差(这一过程完全独立于样本批次大小以及其他时序位置):

μ=1Di=1Dxi,σ2=1Di=1D(xiμ)2,x^i=xiμσ2+ε,yi=γix^i+βi(5.7)\mu = \frac1D\sum_{i=1}^{D}x_i,\qquad \sigma^2 = \frac1D\sum_{i=1}^{D}(x_i-\mu)^2,\qquad \hat x_i = \frac{x_i-\mu}{\sqrt{\sigma^2+\varepsilon}},\qquad y_i = \gamma_i \hat x_i + \beta_i \qquad\textbf{(5.7)}

其中数值稳定微量常数 ε\varepsilon 通常取 10510^{-5}

前向的雅可比

为了严密推演输入输出之间的微积分映射,我们需要计算全偏导雅可比矩阵 x^i/xj\partial \hat x_i/\partial x_j。首先推演均值与方差对任意分量的敏感度:

μxj=1D,σ2xj=1Di2(xiμ)(δij1D)=2(xjμ)D(5.8)\frac{\partial \mu}{\partial x_j} = \frac1D,\qquad \frac{\partial \sigma^2}{\partial x_j} = \frac1D\sum_i 2(x_i-\mu)\Big(\delta_{ij}-\frac1D\Big) = \frac{2(x_j-\mu)}{D} \qquad\textbf{(5.8)}

第二式利用了零均值正交性质 i(xiμ)=0\sum_i (x_i-\mu)=0。记尺度因子 s=σ2+εs=\sqrt{\sigma^2+\varepsilon},应用多元复合求导法则展开 x^i=(xiμ)s1\hat x_i = (x_i-\mu)s^{-1}

x^ixj=δijs1Ds(xiμ)12s32(xjμ)D\frac{\partial \hat x_i}{\partial x_j} = \frac{\delta_{ij}}{s} - \frac{1}{D s} - (x_i-\mu)\cdot\frac{1}{2}s^{-3}\cdot\frac{2(x_j-\mu)}{D}

整理即得完整的解析雅可比表达式:

  x^ixj=δijs1Ds(xiμ)(xjμ)Ds3  (5.9)\; \frac{\partial \hat x_i}{\partial x_j} = \frac{\delta_{ij}}{s} - \frac{1}{D s} - \frac{(x_i-\mu)(x_j-\mu)}{D\,s^{3}} \; \qquad\textbf{(5.9)}

该矩阵天然具备三条极为漂亮的数理性质:第一,其行和恒为零(jx^i/xj=0\sum_j\partial\hat x_i/\partial x_j = 0),反映出对输入整体平移常数绝不改变归一化输出的物理事实;第二,其与中心化向量 (xiμ)(x_i-\mu) 的内积为零,体现了尺度缩放的不变性;第三,它是一个对称矩阵。

反向的两种写法

在反向求导过程中,面对上游回传的梯度 yˉ\bar y,首先穿透仿射变换层:x^ˉi=yˉiγi\bar{\hat x}_i = \bar y_i\gamma_iγˉi=yˉix^i\bar\gamma_i = \bar y_i\hat x_i 以及 βˉi=yˉi\bar\beta_i=\bar y_i

随后求解输入梯度 xˉ\bar x 时存在两种实现路径。直接套用雅可比矩阵公式 (5.9) 进行求和展开,可以得到直观的第一种形式:

xˉj=ix^ˉix^ixj=x^ˉjs1Dsix^ˉi(xjμ)Ds3ix^ˉi(xiμ)记作 c(5.10)\bar x_j = \sum_i \bar{\hat x}_i \frac{\partial \hat x_i}{\partial x_j} = \frac{\bar{\hat x}_j}{s} - \frac{1}{D s}\sum_i \bar{\hat x}_i - \frac{(x_j-\mu)}{D s^3}\underbrace{\sum_i \bar{\hat x}_i (x_i-\mu)}_{\text{记作 } c} \qquad\textbf{(5.10)}

而在实际工程编码中,更具实用价值的是第二种中间变量递推形式

x^ˉi{σˉ2=ix^ˉi(xiμ)(12)s3μˉ=ix^ˉis+σˉ2(2Di(xiμ))=ix^ˉisxˉj=x^ˉjs+σˉ22(xjμ)D+μˉD(5.11)\bar{\hat x}_i \to \begin{cases} \bar\sigma^2 = \sum_i \bar{\hat x}_i (x_i-\mu)\cdot\big(-\tfrac12\big)s^{-3}\\ \bar\mu = -\dfrac{\sum_i \bar{\hat x}_i}{s} + \bar\sigma^2\cdot\Big(-\dfrac{2}{D}\sum_i(x_i-\mu)\Big) = -\dfrac{\sum_i \bar{\hat x}_i}{s}\\ \bar x_j = \dfrac{\bar{\hat x}_j}{s} + \bar\sigma^2\cdot\dfrac{2(x_j-\mu)}{D} + \dfrac{\bar\mu}{D} \end{cases} \qquad\textbf{(5.11)}

利用 i(xiμ)=0\sum_i(x_i-\mu)=0 的对消特性,式 (5.11) 仅需两次一维点积求和即可闭式求解,以完全线性的时间与空间复杂度 O(D)O(D) 瞬时完成求导,彻底杜绝了分配 D×DD\times D 大矩阵的额外开销。本教程在 LayerNorm.backward 中采用的正是这套精巧的实现。

深入探究为什么自然语言与序列建模全面弃用计算机视觉中盛行的批归一化(BatchNorm),其背后存在三条不可调和的底层矛盾。其一,BatchNorm 依赖样本批次计算统计量,在推理阶段必须借用训练期维护的滑动平均值,而现实中测试文本的长度分布与填充比例随时波动,极易诱发均值漂移;其二,在自回归解码生成中,Batch 维度在单步推理时往往退化为 1,此时计算跨样本统计量在数学上完全丧失意义;其三,LayerNorm 具有逐 token 统计的天然独立性,它对序列长度、批次大小以及填充 token 的存在具有完全的天然免疫力,并且其反向计算能够极其优雅地解耦为一个简单的线性循环。

5.4 初始化:方差守恒

考虑一个标准线性变换层 yj=i=1ninWjixiy_j=\sum_{i=1}^{n_{in}}W_{ji}x_i(暂时略去偏置)。设输入分量 xix_i 彼此独立、均值为零且方差为 σx2\sigma_x^2,权重分量 WjiW_{ji} 同样独立、均值为零且方差为 σW2\sigma_W^2。输出分量的方差展开为:

Var(yj)=iVar(Wjixi)=ninσW2σx2(5.12)\mathrm{Var}(y_j) = \sum_i \mathrm{Var}(W_{ji}x_i) = n_{in}\,\sigma_W^2\,\sigma_x^2 \qquad\textbf{(5.12)}

在前向流动中,若要维持输出方差不发生爆炸或萎缩,必须要求 σW2=1/nin\sigma_W^2 = 1/n_{in}(此即经典的 LeCun 初始化);然而在反向传播中,根据对偶求导关系 xˉi=jyˉjWji\bar x_i=\sum_j \bar y_j W_{ji},为保证反向梯度的方差守恒,又要求 σW2=1/nout\sigma_W^2 = 1/n_{out}

为了兼顾前向激活与反向梯度的双向稳定性,双方妥协后的调和平均给出了著名的 Xavier/Glorot 初始化规范:

σW=2nin+nout(5.13)\sigma_W = \sqrt{\frac{2}{n_{in}+n_{out}}} \qquad\textbf{(5.13)}

本教程中的 xavier_std() 函数正是基于式 (5.13) 构建(当输入输出维度相等时退化为 1/d\sqrt{1/d})。结合前文关于残差流累积的推演,便形成了现代大模型领域牢不可破的三大初始化契约:所有常规线性层权重严格按 Xavier 正态分布采样子样;残差分支的最末输出投影显式乘上缩放系数 1/2L1/\sqrt{2L};层归一化的可学习参数 γ\gamma 恒定初始化为 1、偏置 β\beta 初始化为 0。

5.5 Post-LN 与 Pre-LN:为什么需要 warmup

层归一化在计算图中的安放位置构成了现代架构演进的重要分野:

Post-LN:  Xl+1=LN(Xl+Fl(Xl)),Pre-LN:  Xl+1=Xl+Fl(LN(Xl))(5.14)\text{Post-LN:}\;X_{l+1} = \mathrm{LN}\big(X_l + F_l(X_l)\big), \qquad \text{Pre-LN:}\;X_{l+1} = X_l + F_l\big(\mathrm{LN}(X_l)\big) \qquad\textbf{(5.14)}

在 2017 年的原始论文中,作者采用了将归一化置于残差加法之后的 Post-LN 架构,并因此被迫发明了繁琐的学习率预热(warmup)机制;而当今主流开源大模型则近乎全员转向了将归一化置于各分支入口的 Pre-LN 体系(或其变体 RMSNorm + Pre-LN)。

为什么 Post-LN 必须强制绑定 warmup?这一现象可以通过微积分链条给出清晰的启发式解释:在 Post-LN 架构中,反向传播的梯度链必须连续穿越每一个 LayerNorm 的雅可比算子,而雅可比公式 (5.9) 包含由输入尺度决定的因子 1/s1/s。在网络初始化阶段,残差流的方差随着层数 LL 的加深而线性放大(σL\sigma\sim\sqrt L),导致越是深层的网络,其初始梯度尺度的动态范围越发狂暴而不稳定。如果在训练的第一步就采用激进的大学习率,巨大的更新步长会瞬间将参数轰入病态的局部平原,导致整个训练彻底报废且永不可逆。学习率预热的实质,就是在最初的几百步内强制使用极其微小的步长探索,给予 LayerNorm 充裕的时间建立稳健的统计基准,待系统脱离险境后再全面释放学习率

Xiong 等人(2020)在 On Layer Normalization in the Transformer Architecture 中给出了更为严格的数学证明:在初始化点附近,Post-LN 的期望梯度范数随着网络深度呈现超线性的剧烈增长,而 Pre-LN 的残差分支在入口处便被 LayerNorm 稳稳扼住,使得梯度范数完全独立于深度,从而允许模型在无需任何 warmup 的前提下也能稳定启动。这一理论为科研实践提供了明确的定海神针:在样本量极其有限的化学与材料任务中(例如清理后仅剩 233 条数据的 HOPV 体系),绝对不要在 Post-LN 和复杂的 warmup 调度上虚掷光阴,坚定采用 Pre-LN 才是规避数值崩溃的终极法则。

5.6 Dropout:期望不变,方差上涨

本教程的纯 Python 实现刻意不含 dropout:它是整个网络里唯一会破坏"同样输入必得同样输出"的算子,一旦引入,逐参数有限差分梯度检验(§7.3)与实验复现都要额外处理随机性(§6.5 的配方表对此有明确记录)。但 dropout 在真实训练中绕不过去,因此这里单列一节,用与正文完全相同的算子把它的统计行为量出来(demo_dropout.py,10 万次蒙特卡洛采样,p{0.05,0.1,0.5}p\in\{0.05,0.1,0.5\},结果见图 fig_d01_dropout.png)。

训练时随机关掉几根线、剩下的被轻轻拉长;推理时整张网完好;左下角的天平两端始终平衡

定义:反向缩放的 dropout

设某层激活为 xx,丢弃概率为 pp。取独立同分布的伯努利掩码 miBernoulli(1p)m_i\sim\mathrm{Bernoulli}(1-p)

x^i=mi1pxi,E[mi]=1p    E[x^i]=xi(5.15)\hat x_i = \frac{m_i}{1-p}\,x_i,\qquad \mathbb E[m_i]=1-p \;\Longrightarrow\; \mathbb E[\hat x_i]=x_i \qquad\textbf{(5.15)}

训练期按掩码随机置零,推理期则恒等(不做任何处理)。式 (5.15) 就是"为什么要除以 1p1-p"的全部答案:让期望保持不变。早期实现是反过来在推理期乘 1p1-p,两者在期望上等价,但"反向缩放"让推理路径保持恒等映射,工程上更干净——这就是今天所有框架采用的形式。

三条必须分清的结论

我们取一个真实的 5×55\times5 因果注意力权重矩阵 AA(由 transformer.py 的算子生成,见图左),逐条测量:

(一)它是无偏的。 E[w^]=w\mathbb E[\hat w]=w 精确成立。10 万次采样后,p=0.05p=0.05 时逐元素最大偏差 7.1×1047.1\times10^{-4}(蒙特卡洛理论标准差 σ/N=7.3×104\sigma/\sqrt N=7.3\times10^{-4},即 0.98σ0.98\sigma),p=0.5p=0.56.7×1036.7\times10^{-3}(理论 3.2×1033.2\times10^{-3},即 2.1σ2.1\sigma)——对一个 5×55\times5 矩阵取最大值,(0.98,2.1)σ(0.98,2.1)\sigma 正是纯采样噪声的预期量级,没有任何系统性偏移。

(二)行和不再等于 1。 这是被讲错最多的一点:缩放保持的是期望,不是每一行的归一化。行和的期望确实是 1,但它本身是一个随机变量:

Var(jw^j)=p1pjwj2(5.16)\mathrm{Var}\Big(\sum_j \hat w_j\Big) = \frac{p}{1-p}\sum_j w_j^2 \qquad\textbf{(5.16)}

实测:p=0.1p=0.1 时行和标准差 0.24640.2464(闭式 0.23990.2399),p=0.5p=0.50.73860.7386(闭式 0.71980.7198)。也就是说,dropout 之后的注意力权重不是一个概率分布,而是一个"期望等于原分布"的随机向量。任何"dropout 之后每一行仍然构成合法概率分布"的说法都是错的。

(三)噪声尺度由 p/(1p)\sqrt{p/(1-p)} 决定。 对上下文向量 Oj=iw^ivijO_j=\sum_i \hat w_i v_{ij},各分量相互独立的掩码给出精确的方差:

Var(Oj)=p1piwi2vij2(5.17)\mathrm{Var}\big(O_j\big) = \frac{p}{1-p}\sum_i w_i^2 v_{ij}^2 \qquad\textbf{(5.17)}

10 万次采样与这个闭式的最大相对误差为 1.3×1021.3\times10^{-2}p=0.1p=0.1)。它说明 dropout 给每一维输出注入的噪声幅度,恰好是该维各分量的加权尺度乘以 p/(1p)\sqrt{p/(1-p)}p=0.05p=0.050.2290.229p=0.1p=0.10.3330.333p=0.5p=0.51.000\mathbf{1.000}。当 p=0.5p=0.5 时,噪声已经和信号本身同量级;实测的"相对 O|O| 噪声"甚至因真值的方向相消而更大(p=0.1p=0.1 时 0.385,p=0.5p=0.5 时 1.156)。

位置与顺序:掩码 → softmax → dropout

Transformer 中 dropout 通常出现在四处:词元嵌入与位置编码相加之后、注意力权重上、每个子层输出上(即残差分支内部)、以及 FFN 的中间层。原始论文统一取 p=0.1p=0.1(§6.5 配方表)。注意力上的 dropout 有两条铁律:

  1. 必须在 softmax 之后。只有 softmax 的输出才是行和为 1 的注意力权重;对打分做 dropout 得到的是随机打分,随后 softmax 会把它重新归一化,等于白做一遍。
  2. 绝不解除掩码。因果掩码在 softmax 之前把未来位置压成 -\infty,因此被掩位置恒为 0;dropout 只稀疏化下三角的有效连接。实测:被掩位置在 10 万次采样后的平均绝对值恒为 0.0\mathbf{0.0}。(标准实现会对整张矩阵采样掩码,包括那些本来就是 0 的位置,因此"0 乘任何掩码仍是 0"这一性质天然成立。)

还有一个常被忽略的后果:dropout 可以把一整行清零。第 tt 行有 kk 个有效权重时,整行被清空的概率是 pkp^k,于是序列开头的行比末尾的行脆弱得多——p=0.5p=0.5 时第一行(只有 1 个有效权重)有 50% 的概率变成全零,整张矩阵平均为 19.3%19.3\%(闭式 19.4%19.4\%)。这一行的上下文向量随即变成零向量,与"这个位置什么都没看到"不可区分。这正是 dropout 概率在小模型上不宜取大的另一个理由。

为什么还要用它

两个互补的解释:

  • 统计视角:dropout 是对子网络做 bagging 的廉价近似——每次前向都采样出一个稀疏子网,训练实际上在对这些子网做隐式集成。
  • 注意力视角:用 §2.5 的核平滑语言,dropout 让每一次前向都在一个随机稀疏的核上做核回归。它专门打断"某个 token 固定依赖另一个 token"这条捷径,迫使证据分散到多条连接上——这正是 §3.2 里"多头分散承担"的想法在训练期的随机化版本。

在化学任务里后一条尤其重要:小数据上模型极易记住"某个骨架片段 → 某个标签"的虚假关联(§12.1),dropout 是最便宜的反记忆手段。但代价由式 (5.17) 直接给出:噪声按 p/(1p)\sqrt{p/(1-p)} 增长,而化学任务的信号往往本来就弱(实验七里组成式对形成能的 R2R^2 只有 0.06-0.060.150.15),把 pp 调到 0.5 等于把有效容量压到无法拟合。实践中的经验区间是 p[0.05,0.2]p\in[0.05,0.2],并且推理期必须关闭——这也正是 model.train()model.eval() 需要严格交替的原因。

dropout 是无偏的,但行和与输出的方差都随 p 上涨


6. 完整架构与训练目标

6.1 三种形态

在明确了自注意力、残差与归一化的微观机理后,根据信息交互的方向性约束,Transformer 在宏观上演化为三大经典形态:

形态结构化学里的典型任务
编码器(BERT 式)双向注意力 + 池化 + 预测头溶解度/毒性/带隙预测、谱图分类
解码器(GPT 式)因果注意力 + 语言模型头分子生成、SMILES 补全
编码器–解码器(原始 Transformer)编码器 + 交叉注意力解码器反应产物预测、逆合成、分子翻译

在编码器架构中,不存在任何时间维度的遮蔽,序列中的任意位置能够全局审视全貌,最适宜于全局特征提取与物化性质回归;而在解码器架构中,为了践行自回归生成的因果法则,必须施加下三角因果掩码,严格杜绝向未来的偷窥;至于原始的双端编解码器架构,解码器内部巧妙地堆叠了两层注意力——一层用于审视已生成前缀的因果自注意力,另一层则是横跨双端、负责从编码器语义流中汲取养分的交叉注意力。

两个房间:左边读完把卡片贴到墙上,右边边写边回头抄

在本教程的代码仓库中,我们分别完整实现了这三类架构(对应 TransformerEncoderTransformerLM 以及 EncoderDecoder),并在后续 §10 的实验中逐一派上了用场。

6.2 掩码

6.2.1 因果掩码与 -\infty 的极限论证

斜着拉起的帘子:越往右,能看到的未来越少

因果掩码的核心是构造一个上三角为 -\infty 的掩码矩阵 M{0,}T×TM\in\{0,-\infty\}^{T\times T}

Mts={0,st,s>tM_{ts} = \begin{cases}0, & s\le t\\ -\infty, & s>t\end{cases}

将该矩阵叠加到未缩放打分上(即 S+MS + M)再送入 softmax。为什么引入 -\infty 能够完美等价于在信息流中将未来位置彻底切除?我们可以引入一个带惩罚参数 β>0\beta>0 的有限掩码 Mts(β)=0M^{(\beta)}_{ts}=0sts\le t)或 β-\betas>ts>t),考察其在 β+\beta\to+\infty 极限下的解析表现:

Ats(β)=estseβ1[s>t]stests+eβs>tests  β+  {estsstests,st0,s>t(6.1)A^{(\beta)}_{ts} = \frac{e^{s_{ts}}\,e^{-\beta\mathbf 1[s>t]}} {\sum_{s'\le t}e^{s_{ts'}} + e^{-\beta}\sum_{s'>t}e^{s_{ts'}}} \;\xrightarrow[\beta\to+\infty]{}\; \begin{cases} \dfrac{e^{s_{ts}}}{\sum_{s'\le t}e^{s_{ts'}}}, & s\le t\\[6pt] 0, & s>t\end{cases} \qquad\textbf{(6.1)}

极限推演给出了毫无争议的结论:-\infty 的作用是将未来的 token 在分母的概率归一化求和中物理级地抹除,绝非赋予其一个极其微小的近似权重。反向传播推导更是揭示出掩码的深刻力学本质:代入 softmax 雅可比公式后,因 Ats=0A_{ts}=0 恒成立,对应位置的梯度 Sˉts\bar S_{ts} 精确归零,从微积分的底层绝对杜绝了任何来自未来的梯度回溯。在实验一的 t_causality_no_future_leak 单元测试中,我们通过扰动序列后半程的输入特征,实测验证了前半程输出向量在二进制层面的绝对不变性。

6.2.2 填充掩码(padding mask)

在批处理运算中,由于不同化学分子的 SMILES 长度各异,短序列必须填充特殊的 <pad> 标记以对齐矩阵维度。此时必须引入第二道防线——填充掩码。它的职责是强制禁止任何位置的 Query 去关注填充列。倘若遗漏了这一掩码,即便将填充 token 的输入嵌入强行置零,其经过线性投影后依然会产生非零的 Key 与打分,不仅会非法争夺宝贵的概率权重,更会在后续平均池化中稀释真实官能团的特征表达。

6.2.3 一个著名的 NaN 坑

在严谨的工程落地中,存在一个极其凶险的边界条件:当特定掩码策略导致某一行上的所有位置均被置为 -\infty 时,softmax 的分母求和将变成 0,直接触发除以零异常并引发毁灭性的全局 NaN 崩溃。高水准的工业级实现必须建立显式的安全分支:

python
if mx == -float("inf"):     # 整行被掩
    logZ[i] = 0.0
    continue                # 输出保持全 0,不产生 NaN

单元测试 tests_transformer.py::t_softmax_fully_masked_row_is_safe 正是为守卫这一极端场景而设立的铁律屏障。

6.3 训练目标:交叉熵的梯度

6.3.1 next-token 交叉熵

在自回归序列建模中,给定长为 TT 的离散序列 x1,,xTx_1,\dots,x_T,模型的目标是在每个时间步预测下一个即将出现的 token,其损失函数定义为全序列负对数似然的经验平均:

L=1T1t=1T1logpθ(xt+1xt),pθ(xt+1xt)=softmax(zt)xt+1(6.2)L = -\frac1{T-1}\sum_{t=1}^{T-1}\log p_\theta\big(x_{t+1}\mid x_{\le t}\big), \qquad p_\theta(x_{t+1}\mid x_{\le t}) = \mathrm{softmax}(z_t)_{x_{t+1}} \qquad\textbf{(6.2)}

对该损失函数求导,构成了整个深度学习框架中调用频率最高的核心微积分计算:

Lt=zt,xt+1+logj=1Vezt,j    Ltzt,j=δj,xt+1+ezt,jjezt,j=pt,jδj,xt+1(6.3)L_t = -z_{t,x_{t+1}} + \log\sum_{j=1}^{V}e^{z_{t,j}} \;\Longrightarrow\; \frac{\partial L_t}{\partial z_{t,j}} = -\delta_{j,x_{t+1}} + \frac{e^{z_{t,j}}}{\sum_{j'}e^{z_{t,j'}}} = p_{t,j} - \delta_{j,x_{t+1}} \qquad\textbf{(6.3)}

写成紧凑的向量形态:

zˉt=ptonehot(xt+1)(6.4)\bar z_t = p_t - \mathrm{onehot}(x_{t+1}) \qquad\textbf{(6.4)}

式 (6.4) 散发着无与伦比的物理美感:交叉熵作用于 logits 上的反向梯度,纯粹等价于模型当前预测的概率分布与真实目标 one-hot 分布之间的向量残差。梯度的力学使命清晰明了——将多余的概率质量从预测错误的位置生硬抽离,并等量压注到正确的目标通道中,其抽调的力度与当前的犯错程度保持严格的线性正比。

6.3.2 困惑度、bits 与 nats 的换算

在评估生成质量时,信息论度量提供了超越原始损失的客观标尺:

perplexity=exp(1T1tlogpt)(自然对数),bits/token=nats/tokenln2(6.5)\mathrm{perplexity} = \exp\Big(\frac1{T-1}\sum_t -\log p_t\Big) \quad(\text{自然对数}),\qquad \text{bits/token} = \frac{\text{nats/token}}{\ln 2} \qquad\textbf{(6.5)}

在化学信息学与分子建模文献中,学者们更习惯以二进制的 bits/token 汇报性能,因为它直观地量化了在给定上下文时,每一个分子字符所需的平均信息编码位数。

6.3.3 标签平滑

为了防止模型在训练集上过度自信地将 logits 推向无穷大,标签平滑技术将冷酷的 one-hot 目标软化为一个混合先验:

qj=(1ϵ)δj,t+ϵV    zˉt=ptqt(6.6)q_j = (1-\epsilon)\delta_{j,t} + \frac{\epsilon}{V} \;\Longrightarrow\; \bar z_t = p_t - q_t \qquad\textbf{(6.6)}

这种平滑处理能够显著抑制小样本条件下的过拟合,但天下没有免费的午餐:由于平滑后的目标不再对应纯粹的真实经验分布,模型输出的 logits 将丧失绝对的概率校准意义,因此在需要执行严苛不确定性量化或校准评估的化学场景下必须将其果断关停。

6.4 自回归生成的正确性

通过最大似然估计完成训练后,模型所拟合的条件概率分布能够严丝合缝地支撑自回归生成流程:

pθ(x1:T)=t=1Tpθ(xtx<t)p_\theta(x_{1:T}) = \prod_{t=1}^{T}p_\theta(x_t\mid x_{<t})

在此过程中,采样动力学的控制主要依赖于两大旋钮。其一是采样温度 τ\tau,正如 §2.3.1 中所推导的,低温使分布收敛向确定性的贪心极值,而高温则赋予生成更多随机的探索空间。其二是 Top-k 截断采样,通过在归一化前直接抹杀长尾中概率偏低的候选者,将生成的搜索空间锁定在头部高置信区间。在化学分子生成任务中,Top-k 是极其有力的工程杀手锏:它能够以极小的代价过滤掉位于概率长尾中偶发的括号不匹配等语法噪声。我们在实验四(§10.4)的实测数据显示,在 τ=1.3\tau=1.3 的高探索性环境下,仅仅施加 Top-k=5 的约束,便能将生成分子的 RDKit 语法有效率由孱弱的 11.5% 暴力拉升至 27.0%。

6.5 训练配方

为了保障后续九个真实化学实验能够在零第三方库的纯 Python 环境下顺畅收敛,我们总结并固化了如下整套训练超参数矩阵:

组件本册实验里的取值说明
学习率峰值 2×1034×1032\times10^{-3}\sim4\times10^{-3}小模型用小 lr
调度warmup 150 步后 t1/2\propto t^{-1/2}§8.1
梯度裁剪全局范数 1.0§8.2
权重衰减0(小数据上用 dropout 更安全)AdamW 常用 0.01–0.1
Dropout0(纯 Python 实现不含随机性,便于复现;统计行为见 §5.6)小数据上很关键,常用 0.05–0.2
权重共享SMILES 语言模型开启§8.3
批大小1(逐样本 + 梯度累加)数学等价于批处理

整套配方在数学逻辑与工程落地之间达成了精巧的统一:通过设定 150 步的线性预热与 t1/2t^{-1/2} 的长程退火,我们为网络初期的参数寻优提供了充分的安全缓冲;全局范数 1.0 的梯度裁剪则犹如一具安全阀,彻底封死了任何病态样本引发的梯度爆炸;而采用单样本推演配合梯度累加的批处理模拟策略,不仅在代数上与传统批处理保持绝对等价,更将张量运算的维度死死限制在易于人类追踪与纯 Python 解析的二维平面之内。


7. 纯 Python 实现:公式 ↔ 代码

7.1 设计原则:工匠级底层架构哲学

在构建这套完全脱离第三方重型框架的 Transformer 系统时,我们确立了四项贯穿始终的底层工程原则,力求在数学纯度、代码可读性与运行性能之间取得极致平衡。

首先是张量表征的原生透明性。我们摒弃了工业框架中复杂的 Tensor 对象封装,将矩阵直接映射为 Python 原生嵌套列表 list[list[float]],向量则表现为简单列表 list[float]。系统中不存在黑盒的 shape 属性查询,内存维度的正确性完全依托于前向传播中明确的 assert 断言与全流程的逐参数梯度检验作为终极质量保障。

其次是剥离批次维度的二维纯粹性。代码在单次前向传播中严格只处理一条长为 TT、特征为 dd 的二维序列。传统的多样本批处理在数学上纯粹等价于损失函数的均值加和与梯度的线性叠加,因此我们采用“逐样本前向计算、梯度原地累加后统一更新”的批模拟策略。这一设计不仅在代数优化层面与批处理完全同构,更让所有张量始终保持在直观的二维形态,使得读者在研读代码时能够逐行对照教材上的二维矩阵公式,免受高维张量广播与转置索引的认知干扰。

第三是模块自治与梯度累加闭环。计算图中的每一个神经网络模块(如线性层、层归一化、多头注意力)均内建自洽的 forwardbackward 方法。反向传播方法接受由上游层注入的输出梯度,在内部严密解析出对输入的梯度并向后级传递,同时将当前模块参数的偏导无损累加至自身维护的 self.grad_* 缓冲区中,构建起清晰明了的面向对象微积分引擎。

最后是将瓶颈运算交托 C 语言内联优化。为了打破 Python 解释器在纯脚本执行多重循环时的性能诅咒,我们巧妙利用 Python 底层 C 级迭代器重构了矩阵相乘,利用 sum(map(mul, row, col)) 替代了朴素的三重 for 循环嵌套。经由严格性能实测,该优化将单线程吞吐提升了整整一个数量级,达到了约 205 Mflops/s 的强悍效能,使得完全依托纯 Python 训练双层化学大模型不再是天方夜谭。

7.2 公式 ↔ 代码对照总表

为了帮助读者在严谨的数学符号与具体的代码实现之间建立直观映射,我们将核心公式与代码模块整理如下对照表:

公式代码位置说明
(2.3) softmax 与数值稳定softmax_rows先减行最大值
(2.5) softmax 反向softmax_backwardp(aˉaˉ,p)p\odot(\bar a-\langle\bar a,p\rangle)
(2.11) Aˉ=OˉV, Vˉ=AOˉ\bar A=\bar O V^\top,\ \bar V=A^\top\bar OMultiHeadAttention.backward每个头各自算
(2.12) softmax 雅可比同上softmax_backward
(2.13) 缩放反向同上scale_mat(dS, self.scale)
(2.14) Qˉ=SˉK, Kˉ=SˉQ\bar Q=\bar S K,\ \bar K=\bar S^\top Q同上两个 matmul
(2.15) Wˉ=XQˉ\bar W=X^\top\bar QLinear.backwardmatmul_tn(dy, x)
(2.16) 三路梯度相加MultiHeadAttention.backward自注意力时 dXq+dXk+dXvdX_q+dX_k+dX_v
(4.4) 正弦位置编码sine_cosine_table几何级数频率
(4.7) RoPE 的正反旋转rope_rotate / rope_rotate_inverse正交变换
(5.1)(5.2) FFN 与 GeLUFeedForward / gelu / gelu_grad逐元素激活
(5.4) Pre-LN 残差EncoderBlock.forward / DecoderBlock.forward两个残差支路
(5.7) LayerNorm 前向LayerNorm.forward特征维统计
(5.11) LayerNorm 反向LayerNorm.backward中间量形式
(5.13) Xavier 初始化xavier_std2/(nin+nout)\sqrt{2/(n_{in}+n_{out})}
(6.1) 因果掩码MultiHeadAttention.causal_mask上三角为 True
(6.4) 交叉熵梯度cross_entropy_rowsponehotp-\text{onehot}
Adam 更新Adam.step含偏差修正
全局梯度裁剪clip_grad_norm范数裁剪

7.3 梯度检验:中心差分的数学

在手写反向传播算法时,任何极其细微的符号疏漏或导数转置错误通常不会触发解释器的语法报错,其后果往往极其隐蔽——模型表面上仍在平稳运转,但训练损失却在某一时刻彻底停滞。因此,利用数值微分对解析梯度进行逐参数的精确对拍,是检验系统微积分正确性的唯一金标准。

在数值分析中,中心差分法相比前向单边差分展现出压倒性的精度优势。设目标损失函数 ff 在局部足够阶数可微,我们在待测参数 θ\theta 处分别向两侧施加微小扰动 ε\varepsilon 并做 Taylor 展开:

f(θ+ε)=f(θ)+εf(θ)+ε22f(θ)+ε36f(θ)+O(ε4)f(\theta+\varepsilon) = f(\theta) + \varepsilon f'(\theta) + \frac{\varepsilon^2}{2}f''(\theta) + \frac{\varepsilon^3}{6}f'''(\theta) + O(\varepsilon^4)

f(θε)=f(θ)εf(θ)+ε22f(θ)ε36f(θ)+O(ε4)f(\theta-\varepsilon) = f(\theta) - \varepsilon f'(\theta) + \frac{\varepsilon^2}{2}f''(\theta) - \frac{\varepsilon^3}{6}f'''(\theta) + O(\varepsilon^4)

将两式对称相减,偶数阶导数项完全抵消:

f(θ+ε)f(θε)2ε=f(θ)+ε26f(θ)+O(ε4)(7.1)\frac{f(\theta+\varepsilon)-f(\theta-\varepsilon)}{2\varepsilon} = f'(\theta) + \frac{\varepsilon^2}{6}f'''(\theta) + O(\varepsilon^4) \qquad\textbf{(7.1)}

式 (7.1) 表明,中心差分的数学截断误差是关于步长的平方阶 O(ε2)O(\varepsilon^2),远胜于单边前向差分的 O(ε)O(\varepsilon)。然而在真实的数字计算机中,步长 ε\varepsilon 绝非越小越好。受限于 IEEE 754 双精度浮点的舍入误差(机器精度极限约为 δ1016\delta\approx10^{-16}),两次相近函数值做减法会引发严重的“灾难性有效数字消除”,带来量级约为 δ/ε\delta/\varepsilon 的数值舍入噪声。整个总误差构成了截断误差与舍入误差的权衡博弈:

E(ε)    C1ε2截断+C2δ/ε舍入    ε\*=(C2δ2C1)1/3δ1/36×106(7.2)E(\varepsilon)\;\approx\;\underbrace{C_1\varepsilon^2}_{\text{截断}}+\underbrace{C_2\delta/\varepsilon}_{\text{舍入}} \;\Longrightarrow\; \varepsilon^\* = \Big(\frac{C_2\delta}{2C_1}\Big)^{1/3}\sim \delta^{1/3}\approx 6\times10^{-6} \qquad\textbf{(7.2)}

在工程实践中,为了兼顾极佳的截断精度与数值抗噪鲁棒性,本教程统一设定扰动步长 ε=104\varepsilon=10^{-4}。在涵盖各个模块数千个参数的实测对拍中,解析梯度与数值梯度展现出了惊人的高度重合:

模块参数量最大相对误差最大绝对误差
Linear358.6×10138.6\times10^{-13}9.3×10139.3\times10^{-13}
LayerNorm121.6×10121.6\times10^{-12}2.9×10122.9\times10^{-12}
多头注意力(因果)2888.3×1088.3\times10^{-8}5.1×10105.1\times10^{-10}
FeedForward2972.7×1092.7\times10^{-9}3.4×1093.4\times10^{-9}
TransformerLM(2 层)46723.3×1083.3\times10^{-8}4.9×10104.9\times10^{-10}
EncoderDecoder62204.0×1084.0\times10^{-8}7.1×10107.1\times10^{-10}

逐参数梯度检验:所有模块都在  以下

在制定自动化检验判据时,孤立审视“相对误差”极易被微小的数值陷阱所误导——当参数自身的理论梯度极度微弱(例如本身处于 101610^{-16} 量级)时,哪怕差分计算产生了仅为 101210^{-12} 的无害舍入扰动,换算出的相对误差也可能突增至 10310^{-3}。为了确立客观公允的检验门槛,我们继承了斯坦福 CS231n 经典作业中千锤百炼的联合判据:

通过    (相对误差<105) 或 (绝对误差<108)(7.3)\text{通过} \iff \big(\text{相对误差}<10^{-5}\big)\ \text{或}\ \big(\text{绝对误差}<10^{-8}\big) \qquad\textbf{(7.3)}

只要满足上述两者之一,便可判定该参数的反向传播实现完全正确,这套严格的逻辑在底层全部封装于 grad_ok 检验函数中。

7.4 Adam:偏差修正的推导

Adam 优化器之所以能够横扫深度学习领域,核心在于它同时维护了梯度的一阶有偏矩估计(动量)与二阶有偏矩估计(自适应学习率标尺):

mt=β1mt1+(1β1)gt,vt=β2vt1+(1β2)gt2(7.4)m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t,\qquad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \qquad\textbf{(7.4)}

将一阶动量的递推式从起点展开为显式级数:

mt=(1β1)i=1tβ1tigi(7.5)m_t = (1-\beta_1)\sum_{i=1}^{t}\beta_1^{\,t-i}g_i \qquad\textbf{(7.5)}

假定在局部的优化平原中梯度场近似平稳(即其数学期望满足 E[gi]E[g]\mathbb E[g_i]\approx \mathbb E[g]),对式 (7.5) 两端取期望:

E[mt]=(1β1)(i=1tβ1ti)E[g]=(1β1t)E[g](7.6)\mathbb E[m_t] = (1-\beta_1)\Big(\sum_{i=1}^{t}\beta_1^{\,t-i}\Big)\mathbb E[g] = \big(1-\beta_1^{\,t}\big)\mathbb E[g] \qquad\textbf{(7.6)}

利用等比级数求和公式,我们清晰地看到:由于在初始化时一阶矩被清零赋初值,直接计算出的 mtm_t 是对真实期望的有偏估计,其系统性偏差因子恰好为 (1β1t)(1-\beta_1^t)。为了还原无偏的梯度期望,必须显式除以这一衰减项;对二阶方差矩执行完全同构的推导,便得到了两套经典的无偏修正公式

m^t=mt1β1t,v^t=vt1β2t(7.7)\hat m_t = \frac{m_t}{1-\beta_1^{\,t}},\qquad \hat v_t = \frac{v_t}{1-\beta_2^{\,t}} \qquad\textbf{(7.7)}

随后执行参数更新:

θt=θt1ηm^tv^t+ϵ(7.8)\theta_t = \theta_{t-1} - \eta\,\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} \qquad\textbf{(7.8)}

偏差修正的必要性在训练启动的最初几步表现得淋漓尽致。在步数 t=1t=1 时,若不加修正,未缩放的一阶动量仅为 m1=(1β1)g1=0.1g1m_1 = (1-\beta_1)g_1 = 0.1g_1(当经典动量系数 β1=0.9\beta_1=0.9 时),更新步长被生硬地压低了 10 倍;而在二阶动量处由于常用系数为 β2=0.999\beta_2=0.999,初始未修正的方差被严重缩小了整整 1000 倍,这会导致除法后的有效步长在最初几步发生剧烈失控与尺度畸变。

此外,Adam 还具备优良的全局尺度近似不变性:若将损失函数等比例放大 cc 倍,一阶矩将随之放大 cc 倍,而二阶方差矩放大 c2c^2 倍,最终更新量 cmc2v=mv\frac{cm}{\sqrt{c^2v}} = \frac{m}{\sqrt v} 保持了极其稳健的自适应归一化能力。这种为每个参数量身定制动态有效学习率的特质,极大地减轻了工程调参的负担,尽管在某些极限小数据场景下它可能收敛于平坦度稍逊的局部极值点(Wilson 等,2017)。

7.5 本册三个模型的参数量核算

在搭建深度架构时,对各组件的参数规模建立精准的理论解析模型,是保障算力预算与模型容量匹配的前提。设词表规模为 VV,模型主干维度为 dd,注意力头数为 HH,前馈网络中间维度为 dffd_{ff},堆叠层数为 LL

组件参数量
词嵌入VdVd
位置编码(正弦)00(正弦)/ TmaxdT_{\max}d(可学习)
每层自注意力4(d2+d)4(d^2+d)
每层 FFN2ddff+dff+d2dd_{ff}+d_{ff}+d
每层 LayerNorm(Pre-LN,2 个)4d4d
输出头(不共享权重)dV+VdV+V
末端 LayerNorm2d2d

在单元测试 tests_transformer.py::t_param_count_formula 中,我们以 V=10,d=8,L=2V=10, d=8, L=2 的小型网络为例,实测测得其实际内存参数计数与理论公式计算值完全锁死在 1386 个。而在后续 §10 展开的各项真实化学实验中,我们部署的小型模型参数体量分别核算如下:

模型配置参数量
SMILES 语言模型(实验四)V=28,d=48,H=4,dff=96,L=2V=28,d=48,H=4,d_{ff}=96,L=2,共享权重39 360
ESOL 编码器(实验六)V=28,d=48,H=4,dff=96,L=2V=28,d=48,H=4,d_{ff}=96,L=239 409
USPTO 反应 seq2seq(实验九)V=56,d=48,H=4,L=2V=56,d=48,H=4,L=2 双侧103 160

8. 训练动力学:学习率预热为什么是"必须的"

8.1 原始论文的学习率调度

在经典 Transformer 架构中,Vaswani 等人提出了一种结合线性预热与反平方根退火的学习率调度函数:

lr(t)=dmodel0.5min(t0.5,  twarmup1.5)(8.1)\mathrm{lr}(t) = d_{\text{model}}^{-0.5}\cdot \min\Big(t^{-0.5},\; t\cdot\text{warmup}^{-1.5}\Big) \qquad\textbf{(8.1)}

通过令两段调度函数的代数取值相等,我们可以轻松求得其两阶段的分界转折点:

t1/2=tw3/2    w3/2=t3/2    t=w(8.2)t^{-1/2} = t\cdot w^{-3/2} \;\Longrightarrow\; w^{3/2} = t^{3/2} \;\Longrightarrow\; t = w \qquad\textbf{(8.2)}

求导解析表明,系统学习率的最高峰值恰好精确出现在 t=warmupt=\text{warmup} 步长处,且该极值点的时间坐标完全独立于模型的特征维度 dmodeld_{\text{model}}。整个调度轨迹展现为清晰的两幕剧:在步数小于预热步数 ww 的初期,学习率随时间 tt 严格线性爬坡,执行平稳升温;而一旦越过预热峰值,学习率则平滑切换为按 t1/2t^{-1/2} 缓慢衰减的退火曲线。

这两段设计的背后有着极其严密的统计动力学考量。在前期的线性升温阶段,由于网络刚经历随机初始化,参数空间的方向性极差,特别是未经磨合的层归一化统计量与未校准的注意力分布极不稳定。如果盲目以全速大步长更新,巨大的梯度脉冲极易将自注意力强行推入全选择性饱和的 one-hot 绝境(参见 §2.4),导致所有注意力梯度瞬间归零,使网络陷入不可逆的早衰瘫痪;通过小步预热,能让各个归一化层的均值方差以及各头的注意力视野先平稳着陆。而在后期的反平方根退火阶段,它直接契合了随机优化理论中的经典收敛准则——在梯度噪声主导的深水区,按照 1/t1/\sqrt t 的速率缓慢缩紧搜索步长,能够在最大化压制随机小批量样本噪声的同时,持续汲取真实数据分布的稳定信号。

8.2 梯度裁剪:为什么用全局范数

在反向传播中面对狂暴的梯度震荡,若粗暴地采用逐元素截断 giclip(gi,c,c)g_i\leftarrow \mathrm{clip}(g_i,-c,c),会不可挽回地扭曲整个参数梯度向量的空间几何朝向,使得参数更新偏离原本最陡峭的下降路径。与此形成鲜明对比的是全局范数裁剪(Global Norm Clipping)

ggmin(1,cg2),g2=igi2(8.3)g \leftarrow g\cdot\min\Big(1,\frac{c}{\|g\|_2}\Big),\qquad \|g\|_2=\sqrt{\textstyle\sum_i g_i^2} \qquad\textbf{(8.3)}

该机制先对全局所有参数梯度的总 L2L_2 范数进行度量。只有当整体模长突破阈值 cc 时,才以相同的比例系数将所有维度的分量等比例整体缩放,从而在压制数值爆炸的同时,完美保全了原本宝贵的几何更新方向。在处理化学序列时,由于坏样本干扰、极长链拼接或填充异常,偶尔会出现单个批次梯度范数突增数千倍的异常脉冲,若无全局范数裁剪护航,模型极易在一瞬间被彻底击毁。在本册全部实验中,全局阈值被统一稳妥地固化为 c=1.0c=1.0

8.3 权重共享(tied embeddings)

在自回归分子语言建模中,输入端的词嵌入矩阵 ERV×dE\in\mathbb R^{V\times d} 与输出端的分类预测投影矩阵 WoutW_{\text{out}} 可以进行物理上的深度绑定共享:

logits=HE(8.4)\text{logits} = H E^\top \qquad\textbf{(8.4)}

这种绑定机制不仅能为模型节省整整 VdVd 个参数开销(在小型 SMILES 模型中立减 1344 个参数,而在百亿级大模型中可削减数亿参数),更在表征学习上形成了强烈的归纳偏置:它在几何上强行要求“某特定化学 token 在作为输入上下文时的向量映射”与其“作为自回归生成目标时的分类解码方向”保持空间的一致性。

在反向传播中,权重共享要求我们将源自两条计算通路的梯度流进行严密的叠加闭环:

Eˉ=Eˉ输入侧查表 scatter-add+Hlogitsˉ输出侧(8.5)\bar E = \underbrace{\bar E_{\text{输入侧}}}_{\text{查表 scatter-add}} + \underbrace{H^\top \bar{\text{logits}}}_{\text{输出侧}} \qquad\textbf{(8.5)}

TransformerLM.backward 的具体代码实现中,这种双向汇流表现为先通过矩阵乘法求得输出分类端的投影梯度,随后精准累加进在输入端经过离散索引 scatter-add 形成的嵌入梯度矩阵之中:

python
if self.tie_weights:
    W = self.dec.emb.E.data
    dH = matmul(dlogits, W)          # 输出侧对 H 的梯度
    dE = matmul_tn(dlogits, H)       # 输出侧对 E 的梯度 (V×d)
    grow = self.dec.emb.E.grad       # 输入侧已经 scatter-add 过了
    for i in range(len(grow)):
        for j in range(len(grow[i])):
            grow[i][j] += dE[i][j]   # 两路相加

8.4 训练排障指南与异常症状速查

为了帮助读者在遇到训练异常时能像经验丰富的工程师一样快速定位病灶,我们将训练动力学中最典型的失效模式、背后机理与对策梳理如下:

症状最可能的原因对策
损失一开始就 NaN整行被掩 / 手动除 0 / lr 过大检查掩码;lr 降 10 倍
损失降到某个值后完全不动注意力塌成 one-hot,梯度为 0检查是否漏了 1/dk1/\sqrt{d_k};加 warmup
损失剧烈震荡lr 太大 / 没有裁剪warmup + 全局裁剪 + lr 减半
训练损失 → 0,验证不降过拟合(小数据)减层数/减宽度、dropout、数据增强、换更简单模型(§10.8)
生成全是重复 token解码太贪心温度、top-k/top-p;检查 EOS 是否在训练分布里
生成全是非法 SMILES训练不充分 / 词表切分错误检查 tokenizer(Cl 是否被切成 C+l);加语法约束解码
换长度就崩位置编码外推问题用 RoPE(§10.3)

9. 化学与材料里的 Transformer

同一个分子写成一条链:注意力从一处越过很多字符,直接连到另一处

9.1 先决定"分子怎么写"——六种表示

由于标准的 Transformer 架构在物理底层严格只接受一维离散或连续序列作为输入,因此任何面向化学或材料科学的计算流水线,其首要的核心关卡必然是:如何将高维复杂的分子实体降维投影为一个序列。这是决定整个 AI for Science 任务成败最具决定性的一步,其对模型性能的实质影响往往远超选择几层网络或几个注意力头。

表示例子保留了什么丢了什么适合
SMILES(canonical)CC(=O)Oc1ccccc1C(=O)O原子、键、环、手性的完整信息顺序无化学含义;书写不唯一性质预测、LM 生成
随机化 SMILES同一分子的 100 种写法同上同上(但作为数据增强)抗位置过拟合
SELFIES[C][C][=Branch1][C][=O]...语法永远合法可读性差;长度更长生成模型
分子指纹(Morgan/ECFP)1024 位 0/1子结构统计,固定长度位置、计数、顺序传统 QSAR 强基线
2D 图(原子=节点,键=边)邻接表拓扑精确,置换不变3D 构象GNN、Graphormer
3D 坐标 + 原子类型(x,y,z,element)构象、手性、非键相互作用需要构象搜索等变模型(Uni-Mol、Equiformer)
组成式(材料)LiFePO4Mg2Si元素配比结构、位点材料筛选(CrabNet 等)

从工程实战的深水区总结,有三条血淋淋的经验必须常驻心头。首先,分子指纹配合紧凑的小型多层感知机(MLP)是每一个研究者必须优先建立的对照底线。指纹算法凝聚了计算化学家数十年的领域先验知识,计算极其廉价、数值表现极其稳健,在小规模数据场景下(如后续 §10.8 所展示的 HOPV 体系)经常能够干净利落地击溃繁重庞大的深度 Transformer。如果在基准对决中连这一简单基线都无法超越,问题几乎必然出在不合理的数据泄露评估或者糟糕的表示降维策略上。其次,SMILES 字符的书写顺序带有极强的人为编码痕迹,在输入模型前必须统一进行严格的规范化(Canonicalization),或者反其道而行之,通过大规模随机化重写构建数据增强屏障,以防止网络将注意力虚掷在特定遍历路径的字符伪影之中。最后,当科研命题强依赖三维立体几何时,切忌盲目依赖一维序列模型。立体手性中心、低能构象分布以及氢键或 ππ\pi-\pi 堆积等非共价相互作用,在一维扁平的 SMILES 字符流中几近彻底隐形,必须依赖更高阶的几何算子方能求解。

9.2 性质预测(QSAR / QSPR)

(在这六种表示之下还有一层更细的选择:同一个 SMILES 字符串要切成什么粒度的 token。本教程正文与实验统一采用"原子级"切分——ClBr[nH]、环闭合数字各自成 token;附录 E 给出了它与字符级、BPE 子词级在真实化学语料上的完整对照,包括序列长度、词表大小、tT2\sum_t T^2 账本,以及 BPE 在 SMILES 上"自学"出官能团片段的实证。)

在定量构效关系(QSAR)与构效性质(QSPR)的预测任务中,标准的工业级范式采用编码器主干配合池化聚合与回归预测头:

h1:T=Encoder(x1:T)编码器,z=Pool(h1:T),y^=wz+b(9.1)\underbrace{h_{1:T} = \mathrm{Encoder}(x_{1:T})}_{\text{编码器}},\quad z = \mathrm{Pool}(h_{1:T}),\quad \hat y = w^\top z + b \qquad\textbf{(9.1)}

在将变长的序列隐层状态 h1:Th_{1:T} 坍缩为固定维度的全局分子指征向量 zz 时,业界发展出了三种主流的池化策略。其一是平均池化(Mean Pooling),通过对序列所有位置的特征直接求算术平均,赋予各原子对等的发言权,其对字符排列顺序的抗扰动鲁棒性最高,在小样本场景下表现最为稳健,本教程后续实验均默认采用该策略;其二是继承自 BERT 范式的 [CLS] 分类标记池化,在输入首端预置一个特殊的锚点 token 并专门抽取其最顶层表征,但这一机制通常需要在大规模预训练阶段专门为其注入全局上下文聚拢动力;其三是注意力池化(Attention Pooling),通过引入一个可学习的全局查询向量去动态扫描序列中的全部原子,自适应提取分子中权重最大的关键基团。

在训练目标的设定上,回归任务通常采用均方误差(MSE),但在训练启动前务必对标签目标执行标准分(z-score)归一化;而对于分类任务则标准适配交叉熵。

在此类任务中,最容易引发学术不端或虚假繁荣的深渊是数据划分方式的严谨性。在传统的随机划分(Random Split)下,共享高度相似骨架的同一分子家族衍生物会以极大概率同时跨入训练集与测试集。此时模型根本无需真正掌握复杂的电子效应机理,只需投机取巧地死记硬背这个核心骨架母核,便能在测试集上刷出惊艳的虚高分数。为了彻底粉碎这种骨架记忆作弊,业界推荐采用由 Bemis 与 Murcko(1996)提出的骨架划分(Scaffold Split)——将分子按其拓扑母核严格隔离,强迫测试集面对训练期未曾谋面的崭新化学结构。本教程在实验六中对这两种划分方式在 ESOL 数据集上的真实表现进行了冷酷的量化解剖:

模型随机划分 测试 RMSE骨架划分 测试 RMSE
Transformer(读 SMILES)1.0661.089
DeepChem MLP + Morgan 指纹1.3031.819
均值基线2.1372.292

从上表中可以清晰看到,在骨架划分这面照妖镜下,基于 Morgan 指纹的传统 MLP 的测试 RMSE 由 1.303 暴增至 1.819,测试 R2R^2 由 0.627 坠跌至 0.189。对照其仅为 0.329 的训练集 RMSE(相比测试集误差骤降 5.5 倍),这是典型的骨架过拟合病态反应。反观直接读入原生 SMILES 字符流的 Transformer 模型,其测试 RMSE 仅仅从 1.066 极其轻微地上浮至 1.089,展现出了强大的抗分布偏移鲁棒性。

ESOL:随机划分 vs 骨架划分

9.3 分子生成(SMILES 语言模型)

生成式建模将化学分子设计巧妙地映射为一门字符语言的自回归推演:

词表由原子类型、化学键符号、分支括号以及环闭合数字共同构成(例如在 ESOL 训练集上抽取的紧凑词表仅包含 28 个 token);在输入序列的两端显式包裹开始符 <bos> 与结束符 <eos>;模型以经典的 next-token 交叉熵为优化标靶;而在推理期,则通过逐步条件概率采样持续吐出 token,直至遭遇结束符 <eos> 优雅停机。

在科学评估生成模型的产出品质时,仅仅汇报训练损失是在欺骗自己,业界必须公开发布由三项严苛指标构建的评估闭环

指标定义本册实验四(ESOL,d=48d=48,2 层)
有效率RDKit 能解析的比例τ=0.5\tau=0.572.0%τ=1.0\tau=1.0:32.5%
唯一率去重后 / 生成总数31%–39%
新颖率有效且不在训练集的比例47%–76%

温度与 top-k 的权衡

在调控生成动力学时,核心经验体现在三个关键维度。首先,采样温度 τ\tau 与 Top-k 截断是掌控生成天平最重要的两套砝码,它们在“分子合法性”与“结构多样性”之间进行着永恒的博弈:极低的温度能够确保生成的分子具备极高比例的 RDKit 语法有效性,但代价是输出高度保守单一,新颖率严重下滑;而高温虽能激发出前所未见的新颖化学拓扑,却会带来大面积的分支括号不匹配等语法车祸。本教程在实验四中实测揭示,在 τ=1.3\tau=1.3 的极端环境下引入 Top-k=5 过滤,能够将原本濒临崩溃的 11.5% 有效率有效拉回至 27.0%。其次,转向 SELFIES 编码可以从根本上绕过语法失效的顽疾。SELFIES 基于形式语言理论构建,其语法状态机保证了空间中的任意离散字符串均能无损还原为一个在物理上绝对合法的化学分子,但其代价在于破坏了直接读取成环数字的直观性。最后,在**受控条件生成(Conditional Generation)**的前沿探索中,主流路径通常采用将目标物化性质直接序列化为前缀 token 引导解码,或者在解码步长中引入基于能量或梯度的外部分类器进行引导。

9.4 反应预测与逆合成

左边断链、右边接好的链,中间的小人同时盯着断口两侧

反应产物预测是 Transformer 在化学逻辑推演中最具原始教旨主义色彩的领域:

反应物  编码器–解码器  产物\text{反应物} \;\xrightarrow{\text{编码器–解码器}}\; \text{产物}

Schwaller 等人(2019)发表的里程碑工作 Molecular Transformer,开创性地将有机化学反应预测完全等价于序列到序列(Seq2Seq)的机器翻译工程——将反应物与催化剂的 SMILES 视作源语言,将目标生成产物视为目标语言。

在评价翻译成果时,化学工作者构建了一套超越传统自然语言 BLEU 指标的物理评价准则:其一是 Top-k 命中准确率,考察正确的产物是否被包含在模型采样的前 kk 个候选列表之内;其二是极为苛刻的 完全匹配率(Exact Match),要求生成的字符串在消除书写异构后与真实产物逐字完全一致;其三是 Tanimoto 结构相似度,利用分子指纹度量预测产物与真实目标在子结构层面的偏离程度,从而在科研上区分“差之毫厘的合理副产物”与“荒谬绝伦的胡乱拼凑”。

在此类反应体系中,潜伏着一个极其巨大的数据陷阱:反应模板的严重重复倾斜。在工业级基准 USPTO-50K 中,海量的有机反应共享少数几种高频反应类型(如经典的酰胺缩合或酯化反应)。倘若草率地采用随机划分,测试集将充斥着与训练集骨架几乎雷同的平行反应,使模型的准确率指标被系统性高估 10 到 20 个百分点。因此在严谨的前沿研究中,必须强制按反应核心机理类型或反应物核心骨架展开硬隔离划分。

在后文的实验九中,我们展示了一个基于轻量级模型、仅在 800 条反应上训练的诚实负结果:尽管模型在给定前缀的 Teacher Forcing 条件下取得了看似不错的 token 级预测准确率,但在真正放开自回归解码时,其整串完全匹配率几乎跌落至零——这为我们理解小规模参数模型在宏观离散搜索空间中的局限性提供了无比珍贵的真实样本。

9.5 聚合物与材料基因组

一条被重复单元铺满的长链,两端用一根长箭头连起来

在跨越高分子聚合物与无机固体材料的疆域时,传统面向小分子的图描述符与化学指纹几乎全面失灵,这使得 Transformer 序列化建模成为了开辟新边疆的核心主力:

对象序列化关键挑战
均聚物重复单元 SMILES + 聚合度一条链可能上千个原子,O(T2)O(T^2) 直接爆掉
共聚物 / 嵌段单体序列序列本身是设计变量(序列–性质关系)
交联网络节点度数分布没有唯一"顺序",本质上不是序列
无机晶体组成式需要周期表先验(元素相似性)
MOF / 钙钛矿组成 + 拓扑数据少、分布偏移大

面对上述严酷的物理尺度挑战,材料计算领域摸索出了三套核心应对准则。针对超长聚合物链条,最直接的策略是滑动窗口局部碎片化,将动辄包含数万原子的全长链条截断为固定长度的特征片段(例如仅保留 8 个连续重复单元)进行局部特征编码,在承认“绝大多数物理性质主要取决于局部凝聚态微环境”的前提下规避平方复杂度的算力陷阱;针对无机材料纯文本组成式的表达贫血问题,必须显式注入元素周期律先验知识,将元素的族序、周期、电负性、共价半径等第一性原理特征以附加嵌入的形式绑定在每一个元素 token 之上;同时,科研人员必须时刻保持清醒:无机固体的能带结构、带隙宽度与离子电导率在本质上是由三维晶格对称性与轨道重叠积分决定的,纯粹的化学式文本推演只能学到某种经验性的配方统计,无法洞悉真实的物性根源。

在后文的实验三中,我们专门设计了周期链外推基准任务,定量拷打模型究竟能否自主在序列中学会“锁定恰好相隔固定周期的重复单元”,并检验其向未见长度泛化的理论极限。

9.6 3D 与等变模型

当化学问题涉及真实的物理场(如分子力场推演、催化剂表面吸附构型、蛋白质–配体高亲和力对接)时,最符合物理客观实在的方案是将空间几何对称性严格内嵌于注意力机制的数学定义中,即构筑具有严格三维欧几里得群 SE(3)SE(3) 等变性的几何注意力算子:

f(RX+t)=Rf(X)(三维旋转与平移等变),f(πX)=πf(X)(原子置换等变)(9.2)f(R\cdot X + t) = R\cdot f(X)\quad(\text{三维旋转与平移等变}),\qquad f(\pi\cdot X) = \pi\cdot f(X)\quad(\text{原子置换等变}) \qquad\textbf{(9.2)}

在网络实现上,主流范式是将原本纯粹依赖特征内积的打分算子进行物理维度的重构,显式引入三维坐标欧氏距离与键角二面角诱导的几何空间偏置项:

Sts=qt,ksdh+b(rtrs, 空间角度分量)(9.3)S_{ts} = \frac{\langle q_t, k_s\rangle}{\sqrt{d_h}} + b\big(\|r_t-r_s\|,\ \text{空间角度分量}\big) \qquad\textbf{(9.3)}

这一思想在近年来催生了 AI for Science 领域最耀眼的一批巅峰成果:Zhou 等人(2023)提出的 Uni-Mol 框架,通过在三维构象空间中联合编码原子特征与成对空间距离,在大规模预训练中奠定了多属性预测的新标杆;Liao 与 Smidt(2023)设计的 Equiformer,将高阶不可约球谐张量积与自注意力算子深度融合,在 OC20 催化剂表面吸附能基准中傲视群雄;Ying 等人(2021)提出的 Graphormer,通过在图结构上引入基于最短路径距离与中心性的空间注意力偏置,攻克了分子图表征的难题;而名扬天下的 AlphaFold2 其核心构件 Evoformer,更是通过在氨基酸残基成对表征与多序列比对(MSA)之间交替穿梭双轨注意力,铸就了结构生物学史上的无上丰碑。所有这些改动的底层脉络清晰如一:如果你的化学问题本质上依赖立体构象,一维序列只能作为快速粗筛,三维等变几何模型才是科研的终极战场。

9.7 可解释性:注意力到底能不能当化学证据?

在大量 AI 辅助药物研发与材料设计的宣传论文中,常常能看到将注意力热图生硬渲染在分子结构上、声称“模型关注了该羟基从而证明该基团是活性中心”的论断。作为严谨的科技作家,我们必须在此给出斩钉截铁的学术结论:注意力权重绝不等于特征重要性,更绝不能直接当成化学因果机理的客观物理证据

这一严厉警示立足于三条坚不可摧的数理与实验支柱。首先,注意力矩阵仅仅是前向流动中的一个中间过客。在经过注意力加权求和之后,表征流还要经历输出投影矩阵 WOW_O 的全量混流、前馈网络 FFN 的非线性重构以及残差通路的稀释;在这一连串复杂的非线性重组之下,即便某个原子被分配了高达 0.9 的注意力权重,其承载的语义完全有可能在随后的投影变换中被旋转至零空间而对最终预测贡献为零。其次,高维注意力分布在数学上存在着无穷多组经验等价解。Jain & Wallace(2019)以及 Serrano & Smith(2019)通过严谨的反事实对抗扰动实验证实:可以在人工施加剧烈约束的前提下将原本高度集中的注意力分布完全打散甚至彻底反转,而网络的宏观预测输出却保持纹丝不动。最后,注意力热图缺乏双向对称性与归因因果性,它仅仅反映了在特定子空间中的信息借阅偏好,绝非统计显著的因果决定论。

在后续的实验五(§10.5)中,我们对这一质疑进行了亲自操刀的实验检验:在训练于真实 SMILES 上的语言模型中,全部 8 个头在“成环闭合数字配对”这一决定分子拓扑生死的强化学结构上,其注意力分配质量仅仅录得微弱的 0.040,甚至与纯随机均匀基线(0.030)几无差异;相反,模型压倒性的注意力几乎全数倾泻在紧邻的前序字符上。这如实揭示了小模型本质上只是一个短视的局部文本语法统计器,根本未曾在大脑中构筑起三维环系的化学概念。

为了在科研写作中杜绝伪科学解释,我们建议研究者严格按照下述边界规范审视注意力热图的作用:

用法是否可靠说明
当"假设生成器"看到某个头总看羰基氧,可以提出"模型学到了羰基"的假设,再用消融验证
当"调试工具"检查是否有头在看 padding、看 <eos>、看错位置
当"论文里的机理解释"必须补消融实验(mask 掉那个位置看输出变化)
当"化学家的结构解释"位置是字符串位置,不是化学位置(§4.7)

如果确实需要在学术报告中寻找真正具备因果说服力的可解释性证据,研究者应当采用更加扎实的归因工具链:从局部的积分梯度归因(Integrated Gradients)、到输入特征显式遮蔽(Feature Occlusion)、再到训练专门的隐层探针分类器(Linear Probing),乃至在化学实验室中执行最根本的因果基团消融(把特定官能团真正合成为氢原子替换体再行测定对比)

9.8 已经存在的化学预训练模型(你大概率应该先试它们)

在工业与科研实践中,面对几百至数万条样本的垂直化学任务,从零初始化一个庞大的 Transformer 通常并非最优起点,基于工业级大规模预训练模型进行迁移微调往往能够取得降维打击般的精度收益。以下是活跃在该领域生态位的领军模型代表:

模型表示预训练典型用途
ChemBERTa(2020)SMILES(字符级)在 PubChem 千万级 SMILES 上做 MLM性质预测(MoleculeNet)
MolFormer(2022)SMILES + 线性注意力 + RoPEZINC 十亿级 SMILES性质预测、表征
MolT5(2022)SMILES ↔ 自然语言T5 式文本–分子翻译分子描述、指令式生成
ChemGPT / MolGenSELFIES/SMILES自回归语言模型分子生成
Uni-Mol(2023)3D 构象 + 原子对大规模构象数据3D 性质、对接
Equiformer(2023)3D 图 + 等变注意力OC20 等催化、能量
Graphormer(2021)2D 图 + 结构偏置OGB-LSC图级性质预测

既然工业界已经构筑了如此丰满的巨塔,为什么本教程依然坚持耗费海量心血带领读者从零一行行实现一套原生态的纯 Python 引擎?因为唯有亲手推演过每一个算子的雅可比求导,才能在微调遭遇不收敛的深水区时精准判断梯度是死于注意力饱和还是受阻于层归一化漂移;更因为在极端缺乏标注数据的特种化学场景下,预训练大模型同样可能遭遇维数灾难,此时依托深厚领域先验的小巧模型配合强归纳偏置才是破局的核心杀招;最终,当你日后研读 Uni-Mol 或 AlphaFold 那些名垂青史的顶级论文时,你会豁然开朗地发现,它们所有令人目眩神迷的架构创新,其灵魂全部沉淀在前方 §2 至 §6 那薄薄数页的微积分公式之中。

9.9 化学数据的"规模贫困"

在自然语言大模型的世界里,Scaling Law 的奇迹建立在动辄数以万亿计的海量无标注文本 token 滋养之下。然而当我们把目光转回化学与材料领域时,最核心的基准数据集规模却展现出了触目惊心的数据贫困现实:作为基准测试常客的水溶解度数据集 ESOL 仅包含 1128 条分子;有机光伏带隙 HOPV 在剔除数据坏点后仅存可怜的 233 个样本;即便是名震天下的反应数据库 USPTO-50K 也不过 5 万条记录;而凝聚了无机晶体心血的 Materials Project 形成能库也仅收录约 13.3 万条化合物。

这意味着在涉及高质量实验标签的深水区,化学可用的有监督数据量比语言模型整整匮乏了 4 到 8 个数量级。这一残酷的资源边界,直接宣判了一条铁律的诞生:

数据越是稀缺匮乏,架构就越应当坚定地依赖领域先验(物理定律、空间对称性、拓扑图结构与分子指纹),而绝不能寄希望于用一个参数更臃肿的纯 Transformer 去创造奇迹。

在本教程后续的实验八中,我们将以极为诚实的笔触呈现这一铁律在 186 条 HOPV 极小样本上的灾难性翻车实录:一个仅包含 3 万个参数的单层小型 Transformer,在长达数十个字符的聚合物长链面前甚至无法在训练集上完成有效拟合(训练集 R2R^2 仅为 0.270,测试 R2R^2 仅为 0.205);反观仅用 4 行代码调用的经典 Morgan 指纹搭配微型 MLP,不仅能把训练集背得滴水不漏,其测试集泛化 R2R^2 更是达到了 0.440,将狂妄的自注意力模型彻底挑落马下。


10. 九个实验:每一个数字都是跑出来的

10.0 读这一节的方法

在正式步入这九场贯穿理论推演与工程实测的实验之前,我们必须首先确立三条贯穿始终的治学原则:

第一是百分之百的数据可复现性。文档中呈现的每一个指标与浮点数,都并非从论文中摘录的虚假宣称,而是由实验脚本真实执行后完整落盘在 figures/results_<名字>.json 的原始记录中,并由可视化脚本 make_figures.py 自动化提取渲染而成。

第二是基于第一性原理的机制透光度。本章所有参与评估的 Transformer 模型,全部源自我们在 §7 中亲手逐行编写的纯 Python 算子,参数规模严整收敛在 1.7 万至 15 万之间。我们不追求在庞杂工业框架加持下盲目刷榜,而是力求在最轻量级且透明的代码镜片下,把每一处矩阵运算的动力学机制看透。

第三是对科学真实性抱有绝对坦诚。在全部九个实验中,我们有四个实验如实记录了模型的惨痛失败或理论假说的不成立。在深度学习往往充斥着选择性汇报的当下,深入剖析一次“狂妄的 Transformer 如何被 4 行经典代码挑落马下”,其启发意义远胜于十个粉饰太平的平凡正结果。

全部涉及真实化学分子与晶体材料的实验数据,均严格源自 DeepChem 官方维护的 dc.molnet 标准基准库:

数据来源函数规模用途
ESOL / Delaney(水溶解度)load_delaney1128 条(902/112/114)实验四、五、六
HOPV(有机光伏带隙)load_hopv清理后 233 条实验八
MP 形成能(无机组成)load_mp_formation_energy13.3 万条,取 2 万实验七
USPTO-50K(有机反应)load_uspto 用的同一份 CSV5 万条,取 3200实验九

10.1 实验一:注意力机制的体检

实验一的核心主旨,是将 §2 中推导的每一条抽象数学公式置于数值计算的高倍显微镜下,逐一检验其工程稳定性与数学契约。

我们首先对缩放因子的方差钳制效应进行严苛检验。通过采样 400 组随机的高斯查询与键向量(键数量设定为 n=16n=16),分别测量从未缩放到缩放后的内积方差与香农熵变化:

dkd_k未缩放方差缩放后方差未缩放熵缩放后熵
22.000.9982.1492.390
87.540.9431.3132.366
3229.930.9350.6342.360
128117.320.9170.2682.369
512489.070.9550.1412.355

根据概率论的中心极限定理,未经缩放的点积方差理论上应当严格等于维度 dkd_k(在 512 维时实测为 489.07),而除以 dk\sqrt{d_k} 之后的实测方差被稳稳钳制在 0.955,与理论预期值 1 严丝合缝。然而,右侧的熵指标才是揭示本质的灵魂所在:当缺少缩放时,注意力分布的香农熵随着维度膨胀从 2.15 nats 垂直崩塌至 0.14 nats,意味着最大的注意力权重已被极化推高至 αmax1\alpha_{\max}\approx 1。按照 §2.4 的微分求导法则,此时的导数项 αi(δikαk)\alpha_i(\delta_{ik}-\alpha_k) 精确趋近于零,反向传播的梯度在第一步便宣告死亡。而一旦引入缩放因子,各维度的熵值始终恒定维持在 2.35 nats 附近,注意力分布始终保持着健康的梯度通道。

紧接着我们考察温度参数 τ\tau 对分布平滑性的调控。针对同一组包含 20 个 key 的打分向量,温度参数的变化直接控制着软最大化的尖锐程度:

τ\tau0.10.51.02.04.0
注意力熵(nats)0.0000.4641.6632.5262.854
最大注意力概率1.0000.9010.5460.2460.126

随着温度从极冷态的 0.1 缓释至高温态的 4.0,注意力熵从绝对坍缩的 0.000 nats 线性舒展至 2.854 nats,最大注意力概率也从 1.000 的独占状态稀释至 0.126,完美重现了统计物理中由能量壁垒向最大熵均匀分布演化的热力学图式。

在时间复杂度的实测标定中,我们在 d=32d=32、双头配置的纯 Python 环境下单线程测试了不同序列长度 TT 的运行耗时:

TT8163264128256
纯打分矩阵 QKQK^\top(ms)0.020.070.241.064.2216.21
完整多头注意力前向(ms)0.621.333.008.2425.487.2

对纯打分矩阵 QKQK^\top 的耗时进行双对数回归,其实测斜率达到惊人的 1.98,彻底印证了理论预言的 O(T2)O(T^2) 二次复杂度极限;而包含前后投影项 O(Td2)O(Td^2) 的完整多头注意力前向过程,由于低阶线性运算的摊薄,拟合斜率测得 1.43。

最后是梯度检验与极端数值稳定性的双重护航:针对六大算子模块、四千余个网络参数所进行的双侧有限差分数值梯度检验,其最大相对误差均未超过 10710^{-7};即便我们将输入打分向量人为放大 10310^3 倍施加激进冲击,依靠减去最大值的行平移数值技巧,softmax 的概率求和依然分毫不差地定格在 1.0000000000000001.000000000000000,全流程无任何溢出崩溃。这套严密体检充分证明,§2 所构建的全部微积分公式在工程实现中是完备且坚不可摧的。

10.2 实验二:位置编码

实验二旨在数值上逐一验证 §4 证明的三大几何代数定理,探究序列模型感知拓扑顺序的核心依据。

首先检验正弦位置编码的加法定理(公式 4.2):通过构建旋转块对角阵 RkR_k,实测变换 PEt+k=RkPEtPE_{t+k}=R_kPE_t 的最大数值浮点误差被压制在惊人的 5.3×10155.3\times10^{-15},这表明该代数关系在机器浮点精度内是绝对精确成立的。

紧接着是内积平移不变性恒等式(公式 4.3)的数值勘验:

PEt,PEt+k=i=0d/21cos(ωik)(与绝对坐标 t 完全解耦)\langle PE_t, PE_{t+k}\rangle = \sum_{i=0}^{d/2-1}\cos(\omega_i k) \quad\text{(与绝对坐标 } t \text{ 完全解耦)}

理论推导表明,该内积完全由相对跨度 kk 决定,而与绝对时间基点 tt 彻底无关。我们在不同绝对时间起点 tt 之间实测计算出的内积最大波动差异仅有 1.0×10151.0\times10^{-15},且实测内积与解析三角和之间的最大残差亦仅为 3.9×10143.9\times10^{-14}。这一优雅的代数恒等式,为自注意力机制在连续特征空间中度量相对距离提供了坚如磐石的物理底色。

转向现代大语言模型广泛采用的旋转位置编码(RoPE,公式 4.4),我们在 200 组随机向量与多种跨度位移 δ{1,7,33}\delta\in\{1,7,33\} 构成的广泛测试集中验证其相对性:旋转后的内积偏差 q~m+δ,k~n+δ\langle \tilde q_{m+\delta},\tilde k_{n+\delta}\rangle 与原始内积 q~m,k~n\langle \tilde q_m,\tilde k_n\rangle 之间的最大浮点差异仅为 6.1×10146.1\times10^{-14},再次完美复现了其内在的相对几何不变性。

最后,我们通过排列打乱实验对定理 4.1 所揭示的排列不变性进行数值断言:

位置编码打乱顺序后预测的平均变化相对值
none0.00000.00002.9×10162.9\times10^{-16}
sine0.03930.03931.8×1011.8\times10^{-1}
learned0.07840.07841.3×1011.3\times10^{-1}
rope0.14740.14742.7×1012.7\times10^{-1}

实测数据给出了极富冲击力的结论:在完全剥离位置编码(none)的极端条件下,即便将输入序列顺序彻底打乱,模型的预测输出变化严格锁死在 0.00000.0000(机器底层的噪声波动为 2.9×10162.9\times10^{-16}),这在数值层面上雄辩地证明,缺少位置先验的纯自注意力机制本质上是一台对前后拓扑时序完全致盲的词袋机器;而引入正弦编码(0.03930.0393)、可学习编码(0.07840.0784)以及 RoPE(0.14740.1474)之后,系统对称性被成功打破,模型由此获得了洞察结构时序的关键能力。

10.3 实验三:长度外推(两个任务,四种位置方案)

为了彻底理清不同位置编码方案在面对未见过长度时的泛化行为,我们设计了两个具有鲜明对比意义的合成任务:其一为考察局部周期性的周期链任务,其二为考察长程全局坐标的镜像对称任务。实验统一采用 d=32d=32、4 头、2 层的紧凑解码器(参数量约为 1.8 万),使用 Adam 优化器在长度固定为 16 的样本集上严谨训练,随后将其直接推入跨度在 6 至 30 之间的广泛长度区间进行零样本外推测试。

在周期链任务中,模型若想准确预测下一个 token,其注意力核心必须精准回溯恰好 3 步之前的历史特征,这模拟了化学中重复聚合物单体之间的局部时序依赖:

位置编码长度 6长度 16(训练)长度 20长度 26长度 30
none0.2630.4650.4820.4570.476
sine0.5430.8260.7840.6480.564
learned0.5230.8280.7320.6700.645
rope0.5300.8180.8450.8470.815
随机猜测0.0830.0830.0830.0830.083

任务 A:只有 RoPE 在超长链上站稳

数据揭示出极为清晰的演化轨迹:旋转位置编码(RoPE)表现出了近乎无视长度膨胀的卓越稳定性,在将序列长度推至近乎两倍的 30 步时,准确率依然牢固地立足于 0.815(训练长度为 0.818),这从深层印证了其相对内积恒等式在面对平移变换时的绝对鲁棒性。与之形成鲜明对比的是正弦编码与可学习绝对编码,它们虽然在训练见过的 16 步长度上斩获了高达 0.83 的峰值表现,但一旦迈出训练窗口,准确率便发生明显退化(正弦编码在 30 步时已暴跌至 0.564)。其中可学习编码的外推失效机制尤为具象——在整个反向传播生命周期中,参数仅见过前 16 个绝对坐标,负责标记 16 步以后的参数切片自始至终未曾接收过任何梯度的更新,依然沉沦在初始化时微弱的随机噪声中,自然无法形成泛化。

更有深意的是,完全剥离位置编码的模型居然在该任务上稳稳拿到了 0.48 的准确率,大幅领先于随机盲猜的 0.083。这并非因为位置信息无足轻重,而是揭示了自注意力机制在缺乏先验时自发退化为一种“内容模式归纳头(Induction Head)”:网络只要检索到当前 token 在上文某处出现过,便机械地将其后继出现的元素作为预测输出,这种纯基于内容关联的退化策略虽然能够提供一定的基准表现,但上限极低且完全丧失了对时序的精细控制。

然而,当我们把视线切换到任务 B 的镜像序列对称任务时,局面发生了颠覆性的逆转。镜像序列形如 s1sTsTs1s_1\ldots s_T s_T\ldots s_1,要求网络在序列后半程精准倒序输出前半程内容,这本质上是一个要求极高的长程绝对计数任务:

位置编码长度 6长度 16(训练)长度 20长度 30
none0.2400.1520.1300.104
sine0.0871.0000.0680.092
learned0.1131.0000.1060.131
rope0.2870.1820.1480.119

任务 B:正弦/可学习编码把训练长度"背"到满分,长度一换立刻归零

这个任务深刻地剖析了“绝对位置编码”的核心价值与代价两面性。倒序输出的本质要求模型必须具备精确“数”出当前究竟距离对称中心隔了几个步长的能力,这恰恰是唯有具备绝对时空坐标体系的模型才能做到的事情。因此,正弦与可学习编码在固定的 16 步训练集上展现出了惊人的记忆精度,双双刷出了无可挑剔的 1.000 满分;然而这种完美的代价值得深思:它们实质上不过是拟合了一张仅对“长度 16”有效的高维硬编码查找表,测试长度哪怕仅仅微调至 20 步,准确率便发生崩塌式滑坡骤降至 0.07,甚至低于随机盲猜;更具震撼力的是,仅具备相对距离感知的 RoPE 在该任务上展现出了彻底的无力(准确率仅为 0.18,与随机基线 0.083 处于同一数量级),为我们交出了第一个极为深刻的负结果:相对位置编码绝非万灵之药,面对需要严格全局绝对计数的时序问题,纯相对机制存在天然的表征死角。

两个任务的训练曲线

这两项实验为真实化学场景的架构选型沉淀了坚实的实践指南:在小分子性质预测(QSAR)等推理长度与训练分布高度重合的场景下,采用无需额外参数的正弦编码往往是最具工程效益的解;若面临需要外推到未知聚合度的大分子超长链条,则应毫不犹豫地采用 RoPE 并配合训练期的长度扰动;而一旦面对要求严格依赖绝对拓扑计数的高阶结构任务(例如精确定位共聚物链条第 5 个重复单元上的官能团),则必须重用绝对位置编码,并严格杜绝推理长度突破训练分布的界限。

10.4 实验四:SMILES 语言模型

(本实验使用"原子级"分词:Cl[nH]、环闭合数字各自成 token,词表 28。若想换成字符级或 BPE 子词级再看一遍长度与困惑度的变化,附录 E 提供了可直接套用的实现与全部对照数字。)

在实验四中,我们构建了一个纯原生的自回归字符级 SMILES 语言模型。模型在 ESOL 训练集包含的 902 个小分子上展开自监督学习,词表囊括了包含原子、化学键、闭环数字与括号在内的 28 个基础 token。解码器网络由 d=48d=48、4 头、2 层结构组成(参数量仅为 39 360 个),以 Adam 优化器配合预热退火策略纯 Python 训练 6 个 epoch(耗时约 8 分钟)。

首先观察模型在各测试子集上的语法学习量化表现:

指标训练集验证集测试集
bits/token2.0122.7502.812
困惑度4.036.737.02
token 准确率0.5660.4360.435

对比 28 个类别在均匀分布下的理论盲猜上界 log2284.81\log_2 28 \approx 4.81 bits/token,本模型在测试集上取得的 2.812 bits/token(困惑度 7.02)清晰地证明网络已高度习得了 SMILES 的底层序列语法,能够在极小的数据规模下捕捉字符间的转移概率。

训练损失与验证指标

当进一步将测试集样本按 token 长度切入不同区间分桶观察时,呈现出一种耐人寻味的非单调现象:

token 长度≤2021–3233–48>48
bits/token2.8962.5742.8713.011
token 数20210661770662

按长度分桶的每 token 交叉熵

从数据中可以清晰看到,位于 21 至 32 长度区间的中等分子交叉熵反而取得了最佳的 2.574 bits/token,而短于 20 的小分子与长于 48 的大分子困惑度均发生上升。这一现象有力地澄清了一个技术误区:在序列建模中,单纯的“长度”未必是决定预测难度的根本因子,处于真实化学数据分布两端的“数据稀缺性”(短分子往往结构特异,长分子样本稀少)才是诱发困惑度漂移的核心根源。

在分子生成采样的质检环节,我们对不同温度 τ\tau 与 top-k 约束组合各抽取 200 个生成样本,由 RDKit 严格判定其化学合法性:

温度 τ\tautop-k有效率唯一率新颖率
0.5不限0.7200.3400.529
0.550.7750.3100.581
0.7不限0.5450.3850.468
1.0不限0.3250.3050.623
1.3不限0.1150.1050.762
1.350.2700.1900.579

数据揭示了分子生成领域不可回避的客观妥协规律:温度从 0.5 激进上升至 1.3 时,新颖率虽由 0.529 攀升至 0.762,但代价是分子的合法有效率从 0.720 灾难性暴跌至 0.115;引入 top-k=5 过滤机制能够在高温探索区有效阻断极端低概率字符的误采,在 1.3 高温下将有效率从 0.115 部分挽救回 0.270。

温度与 top-k 的权衡曲线

采样实例

直接审视模型在自由采样(τ=1.0\tau=1.0)下生成的若干代表性缺陷分子:

text
ON2cBrncc1)ccc1
CCOPC(=<unk>)C
CCCCCOO
)c1ccccccOC(=O)c1

诸如括号孤立失配、未知标记 <unk> 突兀渗入、环闭合编号单飞等具体形态,生动地勾勒出“仅具备统计规律却未掌握严格拓扑语法约束”的模型缺陷。这警示所有化学算法研究者:在评估生成模型时绝不可被单调下降的自回归损失指标所蒙蔽,必须强制调用 RDKit 等真实化学软件实施严格的合法性全量体检;若想在工程中彻底杜绝语法失配,采用具备语法自愈能力的 SELFIES 编码(如 §9.3 所述)才是终极之道。

10.5 实验五:注意力头到底在看什么(负结果)

实验五载入实验四训练完成的模型检查点,在 114 条独立测试 SMILES 上逐个统计每一个注意力头的权重在各类化学特定关系上的质量分布:

层-头自己前一个 token上一个原子环闭合配对括号配对随机基线
L0-H02.350.1490.1050.0760.0300.0460.030
L0-H12.340.1830.1130.0950.0220.0340.030
L0-H22.290.1620.1380.0840.0310.0570.030
L0-H32.360.1280.0890.0660.0260.0370.030
L1-H02.220.1680.1080.0970.0170.0130.030
L1-H12.270.1040.1090.1040.0230.0280.030
L1-H22.080.1850.1360.1290.0400.0110.030
L1-H32.110.2460.1470.1330.0240.0270.030

统计结果向我们呈现了极其真实且残酷的科学事实:全网 8 个注意力头的能量几乎完全被自身和紧邻的前驱 token 占据(最高占比达 0.246 与 0.147),这表明网络自发捕捉的仅仅是 SMILES 序列的浅层局部转移规律(例如左括号后通常衔接原子,闭环数字后往往紧跟骨架);而真正具有高级化学拓扑意义的长程“环闭合对偶配对”,其注意力分配权重(表现最佳的头仅为 0.040)与均匀随机基线的 0.030 几乎完全重叠;括号层级配对也仅仅在个别头部略显痕迹(0.057)。

这毫无疑问是一个值得深刻反思的负结果:这个小规模 Transformer 根本没有自发演化出真正意义上的“分子图拓扑注意力”,而仅仅被驯化成了一具记忆字符局部转移的 n-gram 统计机器。

剖析这一机制缺位背后的根源,恰好可以从前沿文献的理论中找到印证:首先是模型容量的硬性制约,双层结构配以每个头仅 12 维的极小表示空间,远未达到大语言模型涌现出“化学语法专门头”的容量阈值;其次是训练样本的极度匮乏,在仅有 902 个分子的微型语料上,捕获长程而精准的闭环跨度所需的泛化信息量巨大,网络选择优先背诵高频局部片段是优化算法在梯度驱动下的自然妥协;最后则是自回归交叉熵目标本身的局限,环闭合标记在长序列中所占的比例微乎其微,模型无需理解空间闭环逻辑便已能将全局损失降至较低水平。这一结论与 §9.7 形成的共识高度统一:“注意力权重显得很符合化学”是一种极其危险的主观投射,未经严格统计推断与因果干预的注意力图谱,绝不能被轻率地作为化学机理的科学解释。

每个头的注意力质量构成

一个真实分子的注意力矩阵(最后一层第 0 头)

10.6 实验六:ESOL 水溶解度(DeepChem 数据 + 两种划分)

实验六聚焦于小分子水溶解度预测(SMILES 序列 \to logS\log S)。模型沿用 d=48d=48、4 头、2 层、39 409 个参数的纯原生 Transformer,目标值先进行标准化的 z-score 转换,以 Adam 优化器(学习率 2×1032\times10^{-3})训练 6 个 epoch。

为了客观标定模型的性能区间,对照组设立了基于经典化学特征的 DeepChem 多层感知机dc.models.MultitaskRegressor(layer_sizes=[128,64], dropouts=0.3),输入 1024 维 Morgan 圆形指纹,训练 200 个 epoch)以及均值基准线。三者均在完全相同的随机划分(Random Split)与骨架划分(Scaffold Split)下接受严苛检验:

划分模型训练 RMSE测试 RMSE测试 R²测试 Pearson
scaffoldTransformer(SMILES)0.7681.0890.7100.844
scaffoldDeepChem MLP + 指纹0.3291.8190.1890.600
scaffold均值基线2.0682.292−0.287——
randomTransformer(SMILES)0.8291.0660.7500.880
randomDeepChem MLP + 指纹0.2651.3030.6270.803
random均值基线2.0922.137−0.004——

ESOL 训练曲线:随机划分 vs 骨架划分

测试集 parity 图

三种模型在两种划分下的测试 RMSE

这一对比极其生动地揭示了化学机器学习中的评估真相。

首先,骨架划分无情地充当了模型真实泛化能力的“试金石”。在传统的随机划分下,DeepChem MLP 展现出了 R2=0.627R^2=0.627、测试 RMSE 1.303 的体面水准;然而一旦切换到按 Bemis-Murcko 骨架严密隔离的分布外测试集时,MLP 的测试 RMSE 暴增至 1.819,R2R^2 悬崖式跌落至可怜的 0.189。审视其训练 RMSE(0.329),它比测试误差整整优异了 5.5 倍,展现了教科书般的结构过拟合。这也向行业敲响了警钟:文献中任何脱离划分方式孤立汇报的 ESOL 指标均不具备横向可比性,唯有经受住骨架划分拷打的系统才具备在未知化学空间探索的资格。

其次,直接输入字符级 SMILES 的原生 Transformer 展现出了反直觉的抗分布偏移韧性。在骨架划分下,其训练 RMSE(0.768)与测试 RMSE(1.089)仅相差 40%,测试 R2R^2 依然坚挺在 0.710,Pearson 相关系数高达 0.844。这背后的归纳偏置机制在于:端到端自注意力无法像指纹模型那样直接检索预先固化的特定子结构模板,它必须被迫从底层字符时序中重构对原子与化学键的全局表征,反而学到了对拓扑骨架扰动更加钝感、更具鲁棒性的本质特征。

最后,均值基线在骨架划分下呈现出的负 R2R^20.287-0.287)直观地证实了测试集与训练集的均值存在显著的系统性宏观偏移,这也再次捍卫了必须将决定系数 R2R^2、相关系数与绝对均方根误差联合汇报的科学诚实性。

10.7 实验七:材料组成式 → 形成能(负结果:表示决定了上限)

实验七将目光投向无机固体材料领域,试图验证仅凭材料化学式字符串能否预测晶体形成能。我们将 Materials Project 数据库的晶体构型压缩为纯文本组成式(如 Mg2SiCu3Sn4U3 等),将元素符号与化学计量数字切分为 token 序列预测单位原子形成能(eV/atom)。训练集、验证集与测试集分别抽取 8000、1000 与 2000 条,网络配置为 d=48d=48、4 头、2 层(共 43 153 个参数),采用预热退火策略训练 6 个 epoch。

实测给出的测试集表现呈现出极具颠覆性的结果:

模型训练 R²测试 RMSE测试 MAE测试 R²
均值基线0.0000.536——0.000
元素分数岭回归(84 维,闭式解,4 行代码)0.1500.4930.3370.154
Transformer(组成式序列,43 153 参数)−0.0390.5530.304−0.063

(需要特别指明,Transformer 的负训练 R2R^2 源自标准化均方误差的直接换算:在经历 6 个完整的 epoch 优化后,其在训练集上的标准化 MSE 依然高达 1.039,甚至劣于永远输出固定训练均值的基线 MSE 1.000——这意味着网络连训练集本身的分布都未能有效拟合。)

训练曲线:验证 RMSE 一直贴着均值基线的 0.536

这一极其深刻的负结果,用冰冷的数据确立了特征表征在机器学习中的决定性地位。

首先,仅仅依托 84 维元素摩尔分数的经典线性岭回归模型,只凭 4 行纯数学闭式解,便在测试集上以 RMSE 0.493、R2=0.154R^2=0.154 干净利落地击败了拥有四万余参数的自注意力网络;而参数臃肿的 Transformer 测试 R2R^2 竟跌至 0.063-0.063,预测质量全面输给均值基线。更关键的洞见在于,拥有解析最优解的岭回归在训练集上也仅能斩获 R2=0.150R^2=0.150。这强有力地宣示:这绝非普通的过拟合或欠拟合工程故障,而是纯组成式化学计量表征对于晶体形成能的可预测性物理上限

在真实的凝聚态物理世界中,无机晶体的形成能本质上是由原子微观空间的三维排布(配位构型、键长分布、空间群对称性及电子轨道重叠态密度)所主导的,而将立体结构生硬压缩为一维化学式,抹杀了绝大部分核心物理变量。在 DeepChem 官方基准中,唯有引入诸如 SineCoulombMatrix 等 3D 几何特征配合树模型,方能将预测 MAE 压缩至具有实用工程价值的 0.116 eV/atom(本实验的组成式模型仅为 0.304)。盲目寄希望于用纯注意力去跨越被彻底抹平的信息鸿沟,无异于缘木求鱼。

不仅如此,针对模型的两项专项机制体检进一步暴露了时序表示与物理真实之间的冲突。一方面,最后一层的全头聚合注意力矩阵显示,自注意力权重仅仅无意义地堆积在极少数常见高频元素上,未曾浮现出任何具有配位化学特征的共现模式;另一方面,对书写顺序的鲁棒性测试表明,若将完全等价的化学式打乱书写次序(例如将 Mg2Si 颠倒为 SiMg2),模型的预测值会产生平均高达 0.0270 eV/atom0.0270\text{ eV/atom} 的波动(占测试集总 RMSE 的 4.9%)。这证明模型对无序的化学计量比产生了错误的绝对位置依赖。这也为材料智能研发提供了一条重要实操法则:若受限于数据条件必须采用一维字符串表征,必须实施严格的元素书写顺序置换随机化增强。

元素–元素注意力矩阵(最后一层,所有头求和)

测试集 parity:组成式模型的点云几乎是扁的

元素越多越难预测;打乱书写顺序会改变预测

10.8 实验八:HOPV 有机光伏(负结果:小数据 + 长序列两头不讨好)

在进入实验八之前,必须强调真实科学数据治理的极端严肃性。HOPV 数据库的原始 CSV 中,electrochemical_gap 核心标签列在全部 350 条数据中竟然潜伏着 117 条静默的 NaN 缺失值。若未经审慎清洗直接投喂,数值梯度会在反向传播首个 step 便发生隐蔽坍塌并污染全网权重。经过编写的 chemdata._finite() 算子严格剔除后,实际可用的合法样本仅存可怜的 233 条(按 186/23/24 切分为训练/验证/测试集)。

面对仅有 186 条训练数据的极端贫困工况,我们特意克制了模型容量,构建了一个仅含单层、2 头、32 维特征空间的微型 Transformer(参数量压缩至仅 3 万左右),最大序列长度限制为 128(样本平均长度约为 79 个 token):

模型训练 RMSE训练 R²验证 RMSE测试 RMSE测试 R²
Transformer(SMILES)0.2940.2700.2170.2930.205
DeepChem MLP + Morgan 指纹0.0250.9950.2480.2460.440
元素分数岭回归0.3330.0640.2610.332−0.021
均值基线0.3450.0000.2420.332−0.022

训练曲线:Transformer 的训练误差一直贴着标准差,说明它在欠拟合

谁在背训练集,谁真的学到了东西

数据展现了深度学习在面对“极小样本与长时序双重挤压”时的典型崩塌实况。

首先,微型 Transformer 陷入了不可自拔的深度欠拟合。其实测训练 RMSE(0.294)紧密贴着训练集本身的方差波动(标准差为 0.345),仅仅勉强解释了 27% 的训练方差,测试集 R2R^2 最终止步于微弱的 0.205。其症结在于双重受制:一方面,186 条样本对于让注意力机制自主学会长程时序语法而言过于贫瘠;另一方面,共轭大分子平均长达 79 个字符的序列跨度严重摊薄了仅有 32 维的嵌入空间,使得简单的平均池化无法承载高阶光电共轭物理信息。

然而,经典方法的表现却形成了刺眼的对比反差:基于预定义化学拓扑构建的 1024 维 Morgan 指纹配合微型 MLP,不仅将训练集记忆至极致的 R2=0.995R^2=0.995(训练 RMSE 仅 0.025),其测试集 R2R^2 更是达到了 0.440(测试 RMSE 0.246),全面压倒自注意力网络。这一结果再次打破了对训练集指标的迷信,更深刻证明了先验归纳偏置的力量:在百级样本的荒漠区,人类化学家凝练数十年的分子环状环境拓扑先验,其信息密度远远超越了数万个漫无目的自适应注意力自由参数。

同时,元素分数岭回归在此处的测试表现彻底坍塌至 R2=0.021R^2=-0.021(等同于随机均值基线),有力地佐证了聚合物光电带隙完全依赖于大 π\pi 键共轭体系的骨架连续性,而与简单的原子计数毫不相干。这为所有从事材料 AI 研发的工程师确立了一条不容违逆的严谨技术路径:当面临几十到几百条微量数据时,正确的选型序列永远是“领域指纹/分子图特征 \to 大规模预训练模型下游微调 \to 最后才审慎考虑从零初始化训练 Transformer”。

10.9 实验九:USPTO-50K 反应产物预测(负结果:小模型的 exact match 很低)

实验九步入有机化学最具挑战性的疆域——反应产物自回归端到端生成。我们从 USPTO-50K 反应数据库中筛选出源序列与目标产物均在 48 字符以内的单产物干净样本共 3200 条,构建基于序列到序列(Seq2Seq)的纯 Python 编码器–解码器架构(d=48d=48、4 头、2 层,总参数量为 103 160 个)。以反应物 SMILES 为源输入,以目标产物为自回归解码目标,仅在 800 条反应小切片上执行 3 个 epoch 的纯原生训练(耗时约 8 分钟)。

评估指标忠实记录了从字符预测准确率到严苛化学全匹配之间的深重鸿沟:

指标数值
参数量103 160
训练损失(3 个 epoch)1.760 → 1.403 → 1.310 nats/token
验证集 token 准确率0.527 → 0.549 → 0.582
测试集 teacher forcing token 准确率0.576
整串 exact match(贪心解码,150 条)0.000
生成产物的 RDKit 有效率0.087
与真实产物的平均 Tanimoto0.135

直接比对贪心解码生成的产物与真实目标分子的文本样例:

text
反应物: Nc1ccnnc1.O=C(Cl)CCl
真实  : O=C(CCl)Nc1ccnnc1
预测  : COC(=O)c1ccc(-c2cccccccccc3)c(Cl)c2)c1     ← 括号不配对,整串作废

这段输出直观展示了化学端到端生成对语法精度的极端敏感性。在 Teacher Forcing 机制辅助下,给定真实正确前缀的条件下,模型的下一个 token 预测准确率达到了看似过半的 57.6%;然而在脱离人工扶持、完全依赖自身逐步生成的自由解码环境下,模型在 150 条样本上的绝对吻合率(Exact Match)被无情定格在绝对的 0.000,生成分子在 RDKit 中的合法化学有效率仅为 8.7%,与真实分子的指纹相似度也低至 0.135。

这种“单字预测良好而全局化学全灭”的巨大断层,源自化学序列生成的独特属性:在一个包含 30 步的产物序列中,全串完全正确要求模型在多达 V30V^{30} 种可能路径的高维迷宫中一步不差地走出一条完美的全局因果链。在 SMILES 复杂的嵌套语法体系中,仅仅多出一个失配的右括号,便会瞬间引发整条分子链的语义雪崩,导致化学引擎直接报错并废弃整串产物。

训练损失与 teacher forcing 准确率

生成质量与实例

面对如此惨淡的指标,我们必须以清醒的科学理性明确该实验的真正定位:业界成熟的 Molecular Transformer 之所以能刷出 80% 以上的 Top-1 精确度,是建立在数以百万计的昂贵反应数据库预训练、极深的网络容量以及结合物理化学过滤器的 Beam Search 集束搜索工程基础之上的。而我们在此处运行实验的目的,绝非为了在微型脚本上虚妄地追求落地,而是为了让读者亲手丈量“零预训练的小模型配以极度匮乏的数据集”在复杂的自回归反应推演中所能触碰的物理边界。

10.10 九个实验的总结

纵览前述九场实验的全景表现,我们将其核心产出与结论类型整齐归纳如下:

#任务我们的最好结果结论类型
注意力体检缩放方差 489→0.96;梯度误差 <107<10^{-7}✅ 公式全部验证
位置编码加法定理 5×10155\times10^{-15};排列不变性 3×10163\times10^{-16}✅ 公式全部验证
长度外推周期链上 RoPE 0.82@长度30(正弦 0.56)✅ 正结果 + ⚠️ RoPE 在镜像任务上失败
SMILES LM2.81 bits/token,有效率 72%✅ 正结果
注意力可解释性环闭合注意力 0.040 vs 随机 0.030❌ 负结果
ESOL 溶解度骨架划分 R² 0.710(指纹 MLP 0.189)✅ 正结果
材料形成能R² −0.063(岭回归 0.154)输给 4 行闭式解❌ 负结果
HOPV 带隙R² 0.205(指纹 MLP 0.440)输给基线❌ 负结果
反应预测exact match 很低(见 JSON)❌ 负结果

四胜四负一警告——这就是褪去学术宣传滤镜后,机器学习在真实化学应用中展现出的硬核全景。如果我们将九场实验沉淀出的核心认知浓缩为一句金科玉律,那便是:

先把表示选对、把基线跑对、把划分切对,再谈模型。


11. 优点、缺点与适用边界

在经历了一连串深刻的实证检验与正负结果的交锋之后,我们得以跳出单纯的工程狂热,以一种冷静而客观的科学视角审视 Transformer 在计算科学中的立身之本与阿喀琉斯之踵。

11.1 优势的几何与系统本质

若将自注意力架构的成功因素按其底层重要性逐层剖析,首当其冲的是其在信息拓扑上实现的路径长度 O(1)O(1) 跃迁。在传统循环神经网络(RNN)中,信息必须沿着序列链条进行逐个离散时间步的串行渗透,长程梯度势必遭受指数级的衰减或爆炸;而 Transformer 将任意两个 token 之间的有效传播跳数硬性坍缩为常数 1。在分子与物质科学中,这一几何特性的价值无可替代:它意味着在复杂的空间回折中,一条长达数十个字符的 SMILES 序列两端的环闭合配对、遥隔数个单体的聚合物两端基团、乃至酶活性中心与远端变构调节取代基之间的跨空间长程耦合,都可以在单层运算内达成瞬时的“一跳直达”,彻底消除了长程信息传递的物理迟滞。

与路径优化并驾齐驱的是训练过程的全局并行性。因为摒弃了隐藏状态 hth_t 依赖于 ht1h_{t-1} 的串行递推约束,整条长达 TT 个 token 的上下文矩阵能够以统一的张量形态完全吞吐进现代并行硬件(GPU/TPU)的张量核心之中,TT 个位置的反向梯度同时汇流更新。正是这种对硬件吞吐能力的极致压榨,构成了过去数年间模型吞吐数以万亿计海量文本或海量分子语料、催生出大语言模型 Scaling Law 的根本工程基石。

更深层次的智能跃升源于基于动态内容的自适应寻址机制。经典的卷积网络(CNN)依赖于空间上固化不可变的滑窗权重核,而自注意力的权重矩阵 A=softmax(QK/dk)A=\mathrm{softmax}(QK^\top/\sqrt{d_k}) 完全是由当前输入的上下文内生计算而成。“在当前的构象环境下,第 3 个碳原子究竟该重点审视远端的羰基氧还是邻近的取代基”,完全由分子本身的电子与几何特征动态裁量,实现了前所未有的上下文感知柔性。

在此之上,Transformer 展现出了近乎恐怖的架构可扩展性与跨模态统一能力。从最初 4 层、几万参数的紧凑微型模型,到横跨 96 层、吞吐上千亿参数的巨型集群,底层核心的注意力数学算子无需进行任何破坏性的重构,即可平稳顺应参数与算力规模的指数级扩张。更为关键的是,它为科学计算建立了一个通用的符号熔炉:无论是线性的化学 SMILES 文本、蛋白质折叠的一维氨基酸序列、无机固体的化学组成式、质谱与核磁共振的连续峰位 token,抑或是实验室日志中的自然语言文本,全部能够被严整统一在相同的自注意力与前馈变换之下,成为现代多模态科学大模型的共同底座;与此同时,其完全显式的注意力权重图谱与线性残差流结构,也为理解高维表示的演化轨迹提供了前所未有的分析切入点。

11.2 固有缺陷与应对工程策略

然而,没有哪一种架构能够在物理世界中享有免受热力学第二定律制约的特权。自注意力的第一个也是最沉重的枷锁,便是其与生俱来的 O(T2)O(T^2) 二次计算与显存复杂度。正如我们在 §2.7 的理论推导与实验一的数值实测中所证实的(实测斜率严整定格在 1.98–2.04),当序列长度 TT 攀升至数千乃至数万时(这在长程高分子聚合物、超大蛋白质复合物以及基因组长序列中是极为常见的常态),显存与算力开销将面临窒息般的爆炸。工程上的标准反击策略包括采用线性注意力(如 Performer 或 Linear Attention)将打分矩阵的结合律翻转、引入滑动局部窗口结合少量跨度全局 token(如 Longformer),或者在生物聚合物建模中采用“先单体局部表征、后高阶宏观组装”的层级式编码范式。

紧随其后的沉重代价是对标注样本的极端饥渴。自注意力机制赋予了模型极高的假设自由度,但也意味着它放弃了所有能够约束搜索空间的强归纳偏置。实验八的惨痛教训警醒我们,在面对仅有数百条小样本化学实验数据时,数十万参数的模型极易在长序列中迷失方向并陷入深度的欠拟合或虚假的过拟合。在真实工业管线中,这要求研发人员坚决杜绝一上来就盲目从零训练巨型模型的冲动,严格遵守以预计算的领域指纹或图神经网络作为 baseline 锚点、优先引入大规模自监督预训练权重进行下游迁移、并辅以强收缩正则化的操作纪律。

更为本质的科学局限,在于纯序列 Transformer 完全缺乏原生的物理与化学空间先验。网络在底层并不知晓微观世界的薛定谔方程、不理解原子轨道的杂化取向、没有立体化学的范德华半径概念、甚至对经典的共价键合规则与八隅体定则一无所知。相比之下,传统的图神经网络(GNN)至少硬性注入了“信息仅沿着化学键在拓扑邻域传递”这一不可撼动的结构契约。为了在纯注意力中弥补这一缺失,现代前沿工作普遍采取在外层人为注入元素周期律先验嵌入、显式添加三维几何欧氏距离偏置,或者直接借助 SELFIES 等在数学上具有 100% 语法鲁棒性的表征体系来提供硬性安全垫。

此外,由于绝对坐标先验的局限,Transformer 在面对超出训练分布的超长时序外推时往往会发生灾难性的表征退化(如实验三中正弦与可学习编码在迈出训练窗口后的全面崩溃),这促使现代大模型全面转向以 RoPE 为代表的相对旋转变换并配合训练期的长度扰动抖动;而在化学分子序列化表达中,由于同一分子在图到序列的遍历中存在海量等价的同构 SMILES 字符串,导致纯自注意力对非规范化表示表现出极高的脆弱敏感性,必须依托标准规范化、大规模字符串随机化增强以及在损失函数中引入置换不变性惩罚来予以抚平;最后,在自回归推理阶段,由于逐字生成必须重复经历 TT 次串行前向传播,若不引入精心优化的键值缓存(KV Cache)消除冗余投影开销,其推理延迟将变得不可承受。

11.3 否定之境:何时坚决弃用 Transformer

在算法工程的决策树中,知晓何时“不该”使用某种工具,其重要性往往远超掌握如何使用它。当一项具体的科学研发任务命中了以下判定准则中的两条或更多时,应当保持清醒的学术克制,果断从候选方案中剔除纯 Transformer 架构:

首当其冲的是数据规模维度的绝对贫困:在有效训练样本量低于 1000 条且缺乏任何通用领域预训练大模型可以借力微调时,从头训练 Transformer 几乎必然沦为被经典算法羞辱的重灾区;

其次是物理空间的维度错配:若目标性质(如分子手性不对称性、药物与靶点受体的空间对接自由能、均相催化反应中的过渡态能垒)主要由微观分子的三维空间立体构型与柔性构象系综决定,强行将其拍扁为一维序列将造成灾难性的信息损耗,此时应当坚定地选择等变图神经网络或三维结构表征;

第三是局部特征的绝对主导性:当任务本质上仅仅取决于某类特异性局部官能团的显式存在与否(例如判定某些剧毒或易燃基团的二元分类、或者单一酸性中心的微观 pKa 测定),耗费全局注意力去通览整个分子无异于杀鸡用牛刀,简单的子结构匹配或指纹模型不仅运算极速而且拥有无可挑剔的因果鲁棒性;

第四是序列长度与计算预算的硬性冲突:当聚合物重复链条或非折叠多肽序列长度轻松突破 5000 个 token,且团队既无工程算力进行分块局部注意力改造、又缺乏层级拆解算法储备时,强上自注意力无异于主动引爆硬件资源;

第五是基准评估流程的科学缺失:如果你在着手敲击 Transformer 代码之前,尚未在相同的严谨划分下跑过一份“经典分子指纹(如 Morgan/ECFP)配合岭回归或微型感知机”的基准对照,那么所有的深度学习尝试都缺乏起码的立论支点;

最后是物理守恒律的刚性契约:如果你需要模型输出具有严苛物理守恒性的物理量(例如保守力场、解析能量梯度张量),而在团队内部又缺乏开发精确满足李代数旋转平移等变性架构的能力,那么一个输出不受连续对称性保护的纯自注意力黑箱将带来极大的安全隐患。

11.4 与其他架构的对照

为了在全景视角下建立对主流神经网络架构的技术定位,我们将 Transformer 与循环网络、时序卷积、图神经网络以及近年来在生成领域异军突起的扩散/流模型展开多维度的横向对齐:

维度RNN/LSTMCNN/TCNGNNTransformer扩散/流模型
长程依赖差(逐步衰减)中(靠层数堆)中(受消息传递深度限制)强(O(1)O(1) 路径)取决于主干
训练并行
置换不变性有(天然)无(要额外处理)视主干
数据需求
化学先验局部性键/拓扑无(可注入)视主干
生成能力自回归自回归自回归自回归(强)一步/少步(强)
典型化学任务早期 QSAR谱图、序列分子性质、力场性质/生成/反应/材料分子/构象生成

现代科学计算的演进历程雄辩地表明,唯架构论的非此即彼往往是一种狭隘的偏见,不同归纳偏置之间的融合共生才是工程前沿的主旋律。当今在分子性质预测、三维结构建模与晶体探索中屡立奇功的顶尖系统,几乎无一例外地采用了杂交融合架构:无论是将图拓扑空间偏置编码进自注意力打分阵的 Graphormer,还是将三维连续欧氏李群不可约表示与多头自注意力深度融合的 Equiformer,抑或是融汇了原子对空间距离张量与 Transformer 主干的 Uni-Mol,它们都在用无可辩驳的实绩证明,唯有将物理图谱的硬性结构偏置与自注意力的全局交互能力融为一体,才能真正铸就通向物理化学真实的坚实阶梯。


12. 常见坑与调参清单:从数据底色到物理现实的五重考验

在真实工程落地与科学计算的前线,构建一套能够稳定运行且真正具备泛化能力的 Transformer 系统,其挑战往往不在于宏观架构的堆叠,而隐藏在海量极易被忽视的工程细节与数学契约之中。本节将理论推演与工程实战中踩过的血泪教训,凝练为五重层层递进的实战深度备忘录。

12.1 数据与表示层面的隐蔽陷阱

在敲击任何一行深度模型代码之前,研发流程的第一道防线必须由双重朴素基准(Baseline)进行坐标校准。面对连续回归任务,首要动作永远是先行计算并冻结一份纯均值基线——它不仅为整个任务提供了最底层的表现底线,更直接决定了决定系数 R2R^2 的物理零点;紧随其后的是必须在相同数据上跑通基于经典分子指纹(如 1024 维 Morgan/ECFP)配合岭回归或微型 MLP 的特征基准。如果在严格的评估下自注意力网络连几行代码构建的指纹基准都无法稳固战胜,那么强行上线庞大模型不仅毫无工程收益,更会掩盖特征工程本身的缺陷。

在数据分布的切分机制上,必须无条件确立骨架划分(Scaffold Split)的法定地位。正如我们在实验六中所揭示的严酷事实:随机划分(Random Split)极易因为骨架相似分子的相互渗透而制造虚假的学术繁荣(指纹基线在随机划分下测试 R2R^2 虚高至 0.627,然而在骨架划分下断崖式崩塌至 0.189;而 Transformer 则从 0.750 微调至 0.710)。在最终的学术论文与技术报告中,必须以骨架划分为主干,随机划分仅能作为对比参考。与划分紧密交织的是数据泄漏的物理排查:在数据预处理阶段,必须将所有分子通过 RDKit 转化为严格的标准规范化字符串(Canonical SMILES),并执行全量去重校验,严防同一分子以不同书写形式在训练集与测试集之间暗度陈仓。

对于预测目标本身,必须深度审查标签分布的离群值与量纲漂移。科学数据往往包含跨越数个数量级的极值,在送入反向传播之前必须施加严格的 z-score 标准化归一化处理(本教程中的全部回归实验均无一例外执行该工序),并在最终反向还原预测值时确保还原公式的代数闭环。

在文本分词(Tokenizer)环节,必须时刻捍卫化学语法的原子性契约。通用的文本分词器极易犯下将二字符元素 Cl 粗暴割裂为碳原子 C 与未知符号 l 的低级错误;在构建专用分词词表时,诸如 ClBrSiSe 以及携带芳香氢原子的整体基团 [nH] 必须作为不可分割的独立单一 token 处理,环闭合数字亦需独立成词。为缓解模型对固定字符书写路径的过度记忆,应当引入随机化 SMILES 增强机制,但前提是必须从化学逻辑上核准所预测的物理标签对分子书写起点的置换具有物理不变性。最后,必须完整统计并通报序列长度的全局分布特征以及因硬性截断参数 max_len 而丢失的信息比例(例如在 ESOL 数据集的 902 条样本中,仅有极低比例的 2 条分子长度超过 64 tokens),绝不可让隐蔽的信息截断破坏实验的可信度。

12.2 架构设计与初始化中的数学契约

步入网络架构的搭设阶段,首要的数学铁律是万不可遗漏注意力矩阵缩放因子 1/dk1/\sqrt{d_k}。正如我们在 §2.4 的求导推导与实验一的数值压测中所反复论证的,一旦在打分矩阵 QKQK^\top 之后遗漏该除数,特征方差将随维度剧烈膨胀,迅速将 Softmax 推入平坦的极值饱和区,造成前向分布熵急剧坍塌、反向梯度瞬间归零的假死现象——整个网络会在“看似平稳运行、实则从未学习”的假象中空耗算力。

在层归一化(LayerNorm)的拓扑拓扑选型上,尤其在科学计算这种样本量天然受限的工况下,应当坚定地采用 Pre-LN 结构。如果因为特定历史原因采用了原始的 Post-LN,则必须在训练调度中强制引入漫长而审慎的学习率线性预热(Warmup)阶段以对冲早期不可控的梯度冲击。

在权重初始化层面,必须严格遵守方差维持与尺度收敛的几何契约。前向网络与注意力内部的通用线性投影层应当统一采用 Xavier/Glorot 正态初始化;为了抑制极深残差分支中方差随层数线性发散的固有倾向,位于残差连接末端的输出线性投影层权重必须额外乘以收缩缩放因子 1/2L1/\sqrt{2L}(其中 LL 为编码器或解码器的总层数);而所有归一化层的仿射变换参数则需严格坚守初值契约:缩放系数固定为 γ=1\gamma=1,平移偏置精确归零 β=0\beta=0

对于超参数的几何配比,应当将每个注意力头的特征维度 dhd_h 维持在 64 左右(根据总维度设定头数 H=d/64H=d/64),避免单个头维度过窄引发的信息表达瓶颈或过宽引发的优化阻尼。最后,在敲定网络深度与层数之前,必须依据 §7.5 给出的严格解析公式将全网参数总量精算清楚——明确 d=48d=48、2 层的紧凑结构仅对应约 3.9 万轻量级参数,而一旦将维度推至 d=256d=256、6 层则瞬间跨入数百万级参数的深水区,确保网络容量与手头可用的样本规模在信息论层面上保持对称。

12.3 训练动力学中的防失控法则

在优化驱动的生命周期中,学习率预热(Warmup)是防范早期发散的法定法则。在网络启动的前 1% 至 5% 训练步数内,必须强制学习率从极小零点线性攀升至预设峰值,给予 Adam 优化器二阶矩估计向量 vtv_t 充足的时间完成对参数真实梯度尺度的平滑标定。与此同时,必须配备全局梯度裁剪机制,将所有参数聚合梯度的全局 L2L_2 范数硬性限制在 1.0 的安全天花板之内,坚决杜绝仅针对局部张量的逐元素裁剪。

为了在早期彻底排查隐蔽的逻辑与梯度漏洞,工程上线前必须执行极小批量过拟合探针测试:随机抽取 16 条微型数据组成的单 batch 强制训练,如果模型在数十个 step 内无法将这 16 条样本的训练损失平滑压缩至接近于零的绝对过拟合状态,则铁证如山地表明网络实现、反向求导或数据流水线中潜伏着硬性 bug,此时应当立刻停机启动有限差分梯度检验(Gradient Checking)全量排查。

在训练状态的动态监控上,切忌单凭经验死盯单一损失标量,必须实施三维联合监测:训练损失、验证集损失、与验证集真实物理指标。仅仅观察训练损失的下降极易掩盖数据还原或尺度缩放中的逻辑错位(即损失平滑走低但验证集 R2R^2 始终锁死负值的灾难情境);建议在双对数坐标系(log-log)下绘制损失衰减曲线的即时斜率,通过收敛速率的几何态势科学研判当前系统究竟处于数据不足的欠拟合瓶颈还是容量失控的过拟合泥潭。

对于自回归文本与反应生成任务,评估视线必须遵循“先单字准确率、后全局完全吻合(Exact Match)”的渐进阶梯。由于同一分子拥有多种合法的序列表示形式,直接追求精确字符匹配对于初生的小模型过于苛刻。最后,在自监督序列建模中,必须格外核验结束标记 <eos> 是否获得了真实的梯度监督——在生成解码的线上测试中,模型无限死循环生成、停不下来的经典惨剧,绝大多数正是由于损失计算阶段通过错误的掩码直接忽视了对 <eos> 预测监督所导致的。

12.4 评估报告的科学诚实性

面对实验产出的数据报告,我们必须建立起坚决杜绝“学术粉饰”与“单指标筛选”的科学操守。

针对连续物理量回归任务,必须完整联袂通报测试集 RMSE、MAE、决定系数 R2R^2 与 Pearson 线性相关系数,并清晰注明物理量的绝对工程单位。单一的 RMSE 极易被数据本身的离散度所混淆,唯有结合 R2R^2 才能清晰暴露模型是否劣于基线均值。

对于自回归分子生成任务,单纯汇报困惑度或损失毫无工程价值,必须联合发布经由真实化学软件验证的合法有效率(Validity)、去重唯一率(Uniqueness)以及超越训练集的新颖率(Novelty),并明确写入用于判定化学有效性的第三方工具(如特定版本的 RDKit)。

对于化学反应生成与逆合成推导任务,应当全面通报具备候选容错性的 Top-k 命中率以及与真实产物骨架之间的 Morgan 指纹 Tanimoto 相似度,绝不能仅仅以严苛的字符全匹配 Exact Match 一叶阻目。面对类别极度失衡的二元毒性或活性筛选分类任务,必须坚决以**精确率-召回率曲线下面积(AUC-PR)**为最终裁判,严密防范受样本偏置欺骗的虚高 AUC-ROC 假象。

在科学统计的严谨性维度上,由于化学小样本数据对初始坐标具有极高的敏感性,单次偶然训练的波动甚至会直接抹平模型之间的架构差异。任何负责任的学术汇报都必须在至少 3 个独立随机种子(Random Seeds)下重复全流程,完整呈现均值与标准差。最后,工程系统必须永久落盘归档全量测试样本的原始标签真实值与网络模型预测值(正如本教程全部实验均完备生成 y_true/y_pred 的 JSON 记录),唯有保留了原始点云,才能支持后续的 Parity 对角图诊断、残差分布分析与深层归因排错。

12.5 化学与物理现实的硬性约束

最后一重考验,来自于从离散符号空间迈向真实物理微观世界的认知跨越。

首先必须在思想认知上彻底斩断**“序列绝对位置等价于化学空间拓扑”的虚妄幻想**。正如我们在 §4.7 与实验五中所反复警示的,SMILES 字符串中的“第 5 个字符”不过是图遍历算法偶然路过的一个离散符号,绝非三维物理空间中原子核之间的真实欧氏间距或化学键路径,盲目从序列位置的注意力跳跃去倒推反应机理属于严重的因果倒置。

在表征分子的精细几何特征时,必须严防手性异构信息的静默湮灭。SMILES 语法中用于标记四面体立体手性构型与顺反异构的 @@@/\ 等标点符号,必须在分词阶段独立成词并赋予专门的嵌入通道,任何将手性标记误当无意义标点直接过滤的做法,都会让模型对至关重要的对映异构体与非对映体彻底致盲。

在真实化合物数据清洗阶段,必须严谨确立针对混合物、盐类与溶剂分子(表现为 A.B 点号连接形式)的处理法则。在反应预测中,通常应当剥离非反应性的金属抗衡离子与结晶溶剂以保持目标单一;而在电解液或络合物性质预测中,这些辅助分子又直接主导着介电环境,必须依据任务的真实物理图式作出明确的技术取舍。

最后,必须警惕材料与化学数据库中隐蔽的“文献族群时间轴泄漏”。在 Materials Project 或大规模专利反应数据库中,往往存在源自同一篇科研论文的大量相似组成的材料或同系物反应。若采用完全无意识的纯随机切分,这些血缘极度亲近的“同门亲属”将同时散落在训练集与测试集两端,造成泛化能力的严重高估。面对此类复杂工程场景,必须强制按照科研论文来源、所属化学家族体系或公开时间跨度实施严密的族群阻断分组划分(Group Split),唯有在物理上彻底阻断血缘关联,所验证出的模型才真正具备走出实验室、迈向真实未知化合物空间探索的科学力量。


13. 参考文献

奠基与架构

  1. Vaswani, A. et al. Attention Is All You Need. NeurIPS 2017.
  2. Bahdanau, D., Cho, K., Bengio, Y. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015.
  3. Devlin, J. et al. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL 2019.
  4. Radford, A. et al. Language Models are Unsupervised Multitask Learners. (GPT-2) 2019.
  5. Brown, T. et al. Language Models are Few-Shot Learners. NeurIPS 2020.
  6. Shaw, P., Uszkoreit, J., Vaswani, A. Self-Attention with Relative Position Representations. NAACL 2018.

位置编码

  1. Su, J. et al. RoFormer: Enhanced Transformer with Rotary Position Embedding. 2021 (arXiv:2104.09864).
  2. Press, O., Smith, N., Lewis, M. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022.
  3. Chen, S. et al. Extending Context Window of Large Language Models via Positional Interpolation. 2023.

归一化、初始化与训练稳定性

  1. Ba, J. L., Kiros, J. R., Hinton, G. E. Layer Normalization. 2016 (arXiv:1607.06450).
  2. Xiong, R. et al. On Layer Normalization in the Transformer Architecture. ICML 2020.
  3. Glorot, X., Bengio, Y. Understanding the Difficulty of Training Deep Feedforward Neural Networks. AISTATS 2010.
  4. He, K. et al. Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. ICCV 2015.
  5. Kingma, D., Ba, J. Adam: A Method for Stochastic Optimization. ICLR 2015.
  6. Wilson, A. C. et al. The Marginal Value of Adaptive Gradient Methods in Machine Learning. NeurIPS 2017.
  7. Zhang, B., Sennrich, R. Root Mean Square Layer Normalization. NeurIPS 2019.
  8. Nguyen, T., Salazar, J. Transformers without Tears: Improving the Normalization of Self-Attention. 2019.

效率与变体

  1. Katharopoulos, A. et al. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020.
  2. Choromanski, K. et al. Rethinking Attention with Performers. ICLR 2021.
  3. Beltagy, I., Peters, M., Cohan, A. Longformer: The Long-Document Transformer. 2020.
  4. Wang, S. et al. Linformer: Self-Attention with Linear Complexity. 2020.
  5. Dao, T. et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
  6. Shazeer, N. Fast Transformer Decoding: One Write-Head is All You Need. 2019.
  7. Ainslie, J. et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023.
  8. Child, R. et al. Generating Long Sequences with Sparse Transformers. 2019.

可解释性与分析

  1. Voita, E. et al. Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. ACL 2019.
  2. Clark, K. et al. What Does BERT Look At? An Analysis of BERT's Attention. BlackboxNLP 2019.
  3. Jain, S., Wallace, B. Attention is not Explanation. NAACL 2019.
  4. Serrano, S., Smith, N. Is Attention Interpretable? ACL 2019.
  5. Geva, M. et al. Transformer Feed-Forward Layers Are Key-Value Memories. EMNLP 2021.
  6. Elhage, N. et al. A Mathematical Framework for Transformer Circuits. Anthropic 2021.
  7. Olsson, C. et al. In-context Learning and Induction Heads. Anthropic 2022.

化学与材料

  1. Schwaller, P. et al. Molecular Transformer: A Model for Uncertainty-Calibrated Chemical Reaction Prediction. ACS Cent. Sci. 2019.
  2. Chithrananda, S., Grand, G., Ramsundar, B. ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction. 2020.
  3. Ross, J. et al. Large-Scale Chemical Language Representations Capture Molecular Structure and Properties. Nature Machine Intelligence 2022.
  4. Zhou, G. et al. Uni-Mol: A Universal 3D Molecular Representation Learning Framework. ICLR 2023.
  5. Liao, Y.-L., Smidt, T. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs. ICLR 2023.
  6. Ying, C. et al. Do Transformers Really Perform Bad for Graph Representation? (Graphormer) NeurIPS 2021.
  7. Wang, A. Y.-T. et al. Predicting Materials Properties without Crystal Structure: Deep Representation Learning from Stoichiometry. (CrabNet) Nature Communications 2021.
  8. Edwards, C. et al. Translation between Molecules and Natural Language. (MolT5) EMNLP 2022.
  9. Krenn, M. et al. Self-Referencing Embedded Strings (SELFIES). Machine Learning: Science and Technology 2020.
  10. Jumper, J. et al. Highly Accurate Protein Structure Prediction with AlphaFold. Nature 2021.
  11. Wu, Z. et al. MoleculeNet: A Benchmark for Molecular Machine Learning. Chemical Science 2018.
  12. Ramsundar, B. et al. Deep Learning for the Life Sciences. O'Reilly 2019(DeepChem 的教科书).
  13. Dunn, A. et al. DeepCoy / Deep Learning for Chemistry — 见 DeepChem 文档与 molnet 数据卡。
  14. Bemis, G. W., Murcko, M. A. The Properties of Known Drugs. 1. Molecular Frameworks. J. Med. Chem. 1996(骨架划分的来源).
  15. Lowe, D. M. Extraction of Chemical Structures and Reactions from the Literature. PhD Thesis, Cambridge 2012(USPTO 数据的来源).

附录 A:符号表

符号含义
TT序列长度(token 数)
dd模型维度(dmodeld_{\text{model}}
HH注意力头数
dhd_h每个头的维度,dh=d/Hd_h=d/H
dffd_{ff}FFN 中间层维度
LL层数
VV词表大小
XRT×dX\in\mathbb R^{T\times d}某一层的输入(每行一个 token)
Q,K,VQ,K,V查询/键/值矩阵
ART×TA\in\mathbb R^{T\times T}注意力权重(行和为 1)
αts\alpha_{ts}位置 tt 对位置 ss 的注意力权重
MM掩码矩阵(因果掩码/填充掩码)
PEtPE_t位置 tt 的位置编码向量
ωi\omega_iii 对分量的角频率
γ,β\gamma,\betaLayerNorm 的缩放与平移
μ,σ2\mu,\sigma^2LayerNorm 的均值与方差(特征维上算)
Zˉ\bar Z损失对张量 ZZ 的梯度 L/Z\partial L/\partial Z
\odot逐元素乘法
,\langle\cdot,\cdot\rangle向量内积
conv\mathrm{conv}凸包
sg()\mathrm{sg}(\cdot)停止梯度(本册未用)

附录 B:九个实验一览

#脚本数据运行时间你要看到什么
demo_attention_basics.py合成3 秒缩放方差、温度、O(T2)O(T^2)、梯度检验
demo_positional.py合成2 秒加法定理、内积恒等式、排列不变性
demo_length_generalization.py合成(周期链/镜像)6 分钟长度外推:谁崩谁稳
demo_smiles_lm.pyESOL(902 条 SMILES)8 分钟SMILES 语言模型的困惑度与采样质量
demo_attention_chem.py实验四的模型3 秒注意力头落在哪种化学关系上
demo_deepchem_esol.pyESOL(DeepChem)10 分钟随机 vs 骨架划分、与指纹基线对比
demo_deepchem_materials.pyMP 形成能(16000 条)20 分钟组成式建模、元素注意力、顺序敏感性
demo_deepchem_hopv.pyHOPV(清理后 233 条)5 分钟小数据 + 长序列:两头不讨好
demo_reaction_seq2seq.pyUSPTO-50K(3200 条)8 分钟反应预测的 token 准确率与 exact match

除九个主实验外,还有一批不训练的补充实测(它们是正文 §4.2、§5.6 与附录 E 里那些数字的来源):

脚本位置运行时间你要看到什么
demo_positional.pyexperiment_ladder§4.2 编码阶梯(含在实验二内)整数编码范数 2040 vs 正弦恒定 5.657;二进制相邻距离 8 个离散值 vs 正弦恒为 1.4718
demo_tokenization.py附录 E5 秒字符级/原子级/BPE 的 token 数与 tT2\sum_t T^2 对比;BPE 在 SMILES 上学到的子词
demo_dropout.py§5.67 秒dropout 的无偏性(偏差 ≤ 采样误差)、行和方差、噪声尺度 p/(1p)\sqrt{p/(1-p)}

附录 C:复现指南

bash
cd code

# 0) 先确认实现无误:59 项测试(含逐参数梯度检验),约 1 秒
/Users/polyai/anaconda3/envs/deepL/bin/python tests_transformer.py

# 1) 实验一:注意力体检(3 秒)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_attention_basics.py

# 2) 实验二:位置编码(2 秒)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_positional.py

# 3) 实验三:长度外推(6 分钟;QUICK=1 约 1 分钟)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_length_generalization.py

# 4) 实验四:SMILES 语言模型(8 分钟;QUICK=1 约 2 分钟)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_smiles_lm.py

# 5) 实验五:注意力头的化学关系(需要实验四的模型,3 秒)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_attention_chem.py

# 6) 实验六:ESOL(DeepChem,10 分钟)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_deepchem_esol.py

# 7) 实验七:材料形成能(20 分钟)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_deepchem_materials.py

# 8) 实验八:HOPV(5 分钟)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_deepchem_hopv.py

# 9) 实验九:反应预测(8 分钟)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_reaction_seq2seq.py

# 10) 补充实测:分词粒度与 BPE(附录 E,5 秒)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_tokenization.py

# 11) 补充实测:注意力 dropout 的期望与方差(§5.6,7 秒)
/Users/polyai/anaconda3/envs/deepL/bin/python demo_dropout.py

# 12) 画全部数据图(约 30 秒)
MPLCONFIGDIR=/tmp/mplcache /Users/polyai/anaconda3/envs/deepL/bin/python make_figures.py

数据第一次运行时自动下载(已缓存在 code/data/),之后完全离线可复现。 每个实验都会把原始数字写进 figures/results_*.json

附录 D:手绘示意图清单

images/ 下的 19 张示意图用 gpt-image-2.5-flare(FlatRouter 第三方 OpenAI 兼容接口) 配合 handraw-style #097 风格生成,参考图仅作画风参考; 每张图的完整提示词在 images/prompts/*.txt,生成脚本是 images/gen.sh所有定量结论都来自 figures/ 下用 Python 画的数据图,不来自示意图。

文件位置画面隐喻
fig00_cover.png封面链条上只有少数几根粗细不一的长线,一个人拿放大镜看最粗的那根
fig01_rnn_vs_transformer.png§1.1左:只能把信封往后传的队伍;右:围成一圈同时互相看
fig14_bahdanau.png§1.2.1左:整句被压进一只小罐子;右:边写边回头逐张看卡片
fig13_limits.png§1.4字符串 → token → 向量 → 关系 → 概率
fig02_qkv.png§2.1一排人举着牌子(Q)、标签(K)、书(V)
fig15_softmax_vs_linear.png§2.3.1按原水位比例直接舀水 vs 先过细漏斗再舀;角落里一滴负水位
fig03_attention_pipeline.png§2.6注意力四步流水线:量距离 → 缩短尺子 → 重新分配纸条 → 按比例混合
fig04_multihead.png§3.1同一位置三副不同形状的眼镜:看邻居、看远处、看配对
fig16_positional_ladder.png§4.2三种座位号:越举越沉的号码牌、忽高忽低的方格、始终很轻的波纹纸带
fig05_positional.png§4.3每个位置发一张不同波纹的座位号
fig07_ffn.png§5.1一排一模一样的工作台:每个 token 各自过自己的小工厂
fig06_residual_layernorm.png§5.2两条平行小路,合流处用标尺统一刻度(LayerNorm)
fig17_dropout.png§5.6训练时随机关掉几根线、其余被轻轻拉长;天平两端始终平衡(期望不变)
fig09_encoder_decoder.png§6.1两个房间:读完把卡片贴上墙 / 边写边回头抄
fig08_mask.png§6.2.1斜着拉起的帘子:越往右,能看到的未来越少
fig10_smiles_chemistry.png§9同一个分子写成一条链,注意力越过很多字符直达另一处
fig12_reaction.png§9.4断链与接好的链,小人同时盯着断口两侧
fig11_polymer.png§9.5被重复单元铺满的长链,两端用一根长箭头连起来
fig18_tokenization.png附录 E同一段字符串被切成不同数量的碎片,碎片越多方格网越大

附录 E:分词与字节对编码(BPE)

正文从"向量"讲起,刻意跳过了向量之前的那一步——文本怎么变成 token。这不是可以省略的细节:分词决定了模型能看见的最小单位,也决定了序列长度 TT,而 TT 直接决定 §2.7 里那张 O(T2)O(T^2) 的注意力账单。本附录补上这一环,并且所有数字都由 demo_tokenization.py 在真实化学语料上跑出来。

E.1 三种粒度,一场关于"词表大小 ↔ 序列长度"的权衡

粒度例子(ClC(Br)C词表序列长度主要缺陷
字符级C l C ( B r ) C极小极长单字符不携带意义;Cl 被拆成 C+l
原子级(正则切分,本教程采用)Cl C ( Br ) C需要化学规则;对未知书写形式脆弱
子词级(BPE/WordPiece)Cl C( Br )C学到的"子词"未必对应化学官能团

三者构成一条不可能三角:词表越小,序列越长;序列越短,词表越大。而序列长度的代价是二次的。

同一段字符串被切成不同数量的碎片:碎片越多,纸带越长、右边的方格网越大

E.2 真实化学语料上的实测

语料:ESOL/Delaney 全部 1128 条 SMILES(BPE 的合并规则只在训练集 902 条上学习,避免泄漏)。

方案合并数词表平均 token/分子P95最长token 总量tT2\sum_t T^2(相对字符级)可解码率
char(字符级)03122.334697251851.000100%
atom(原子级,本教程正文)02821.584597243420.940100%
BPE@32326110.702560120640.257100%
BPE@6464938.972153101180.184100%
BPE@1281281517.35174782860.125100%
BPE@2562562736.08154668610.089100%
BPE@5125125134.96134255910.064100%
BPE@1024830(自动停止)7334.27124248160.053100%
BPE@512 + 词尾标记5125085.13134657890.068100%

分词粒度:序列长度与注意力账单

三个值得记住的结论:

  1. 原子级分词在这份语料上几乎不省 token:平均 21.58 vs 字符级 22.33,只压缩 3.4%(每 token 1.03 个字符)。因为 ESOL 里多字符原子(ClBr[nH])出现得很少——"按化学规则切分"并不自动等于"序列更短"。它真正的价值在于语义(Cl 是一个原子而不是两个符号),这一点在实验四的 SMILES 语言模型里比压缩率更重要。
  2. BPE 的收益非常大且见效极快:32 次合并就把 tT2\sum_t T^2 打到 0.257,512 次合并打到 0.064——同样的算力预算,注意力可以多看约 15.6 倍的分子。原因是注意力的二次成本惩罚长序列,而 BPE 恰好砍掉长尾(最长序列 97 → 42)。
  3. 合并不是越多越好,它会自己停下来:在本语料上,第 830 次合并之后所有剩余符号对的出现频次都低于 2(min_pair_freq=2),于是 @1024@2048 的结果完全相同。词表则从 31 涨到 733,而平均长度只从 4.96 降到 4.27——收益急剧递减,而嵌入层的参数量、低频 token 的梯度质量都在恶化。

一个直接的工程含义:在算力固定时,"换分词器"往往比"调宽模型"更划算。把 TT 缩小到 1/k1/k,注意力部分的算力与显存降到 1/k21/k^2,而词表扩大带来的参数量增长只是线性的(ΔVd\Delta V\cdot d)。

E.3 算法:BPE 的四步

字节对编码(Byte Pair Encoding)原本是 1994 年的文本压缩算法,2015 年被 Sennrich 等人引入神经机器翻译,此后成为 GPT/BERT/LLaMA 系 tokenizer 的基础。它的规则只有四条:

  1. 加词尾标记:把每个词写成 word</w>,让"词尾的 est"与"词中的 est"成为不同的符号;
  2. 拆成字符:初始词表 = 所有出现过的字符 + </w>
  3. 统计并合并:统计全语料中相邻符号对的加权频次,把最高频的一对合并成一个新符号;
  4. 重复 kk 次:每次合并都让词表 +1,直到达到目标词表大小或没有频次 ≥ 阈值的符号对。

词尾标记为什么重要:在语料 {old×7, older×3, finest×9, lowest×4} 上,前三次合并是

(e,s)es (13)    (es,t)est (13)    (est,</w>)est</w> (13)(\text{e},\text{s})\to\text{es}\ (13)\;\Rightarrow\; (\text{es},\text{t})\to\text{est}\ (13)\;\Rightarrow\; (\text{est},\texttt{</w>})\to\texttt{est</w>}\ (13)

(括号里的数字是这次合并让全语料减少的 token 数,见 E.4 的账本。)此后 newest 会被切成 n e w est</w>,而 esteem 会被切成 est e e m </w>——同一个字符串 est 在词尾与词中变成了两个不同的 token。没有词尾标记,分词器就无法表达"后缀 vs 前缀"这一层语法信息。本教程的 tests_transformer.py::t_bpe_end_of_word_separates_prefix_and_suffix 锁定的就是这个性质。

两条精确的账本(都在测试里被断言):

V=初始字符表+k,t=1kgaint=(原始符号数)(最终 token 数)(E.1)|\mathcal V| = |\text{初始字符表}| + k,\qquad \sum_{t=1}^{k} \text{gain}_t = (\text{原始符号数}) - (\text{最终 token 数}) \qquad\textbf{(E.1)}

其中 gaint\text{gain}_t 是第 tt 次合并让语料实际减少的 token 数。注意它必须实打实地前后相减,而不能用"该符号对的出现频次"代替——在 ccc 这类重叠情形下两者并不相等(相邻对 (c,c)ccc 中被数了两次,但一次合并只能消掉一个 token)。这正是本教程把 gain 定义成"实测差值"的原因,也是 t_bpe_gain_ledger_handles_overlapping_pairs 存在的理由。在 ESOL 训练集上,该恒等式成立:178763249=1462717876-3249=14627

还有一个容易被忽略的工程细节:BPE 的合并顺序必须是确定的。当两个符号对频次并列时,必须给出一个与语料顺序无关的裁决规则(本实现取字典序最小者,见 t_bpe_train_is_order_independent),否则同一份数据在不同机器上会训出不同的 tokenizer,实验再也无法复现。

E.4 纯 Python 实现

code/bpe.py 一共 181 行(含文档字符串与辅助方法),核心逻辑不到 40 行,零第三方依赖:

python
def _apply_merge(tokens, pair, merged):
    """把 tokens 中所有相邻的 pair 替换成 merged(从左到右,不重叠)。"""
    out, i, n = [], 0, len(tokens)
    a, b = pair
    while i < n:
        if i + 1 < n and tokens[i] == a and tokens[i + 1] == b:
            out.append(merged); i += 2
        else:
            out.append(tokens[i]); i += 1
    return out

# 训练循环(BPE.train 的核心)
for _ in range(num_merges):
    pc = self._pair_counts(counts)            # 全语料的相邻符号对频次
    if not pc:
        break
    pair = max(sorted(pc), key=lambda p: pc[p])   # 最高频;并列取字典序最小 → 确定性
    if pc[pair] < min_pair_freq:
        break                                 # 只出现 1 次的组合不固化成子词
    merged = pair[0] + pair[1]
    before = sum(len(t) * c for t, c in counts.items())
    new_counts = Counter()
    for t, c in counts.items():
        new_counts[tuple(_apply_merge(list(t), pair, merged))] += c
    after = sum(len(t) * c for t, c in new_counts.items())
    self.merge_log.append((pair, merged, before - after))   # 真实收益
    self.merges.append((pair[0], pair[1], merged))
    counts = new_counts

编码时不必重放整个训练过程:按 rank 从小到大贪心地套用合并规则即可(每次挑 rank 最小的可用符号对合并)。我们逐条对比过这两种做法在全部 902 条训练 SMILES 上的切分结果——差异为 0 条

E.5 BPE 在 SMILES 上学到了什么

这是本附录最有意思的部分:**把化学字符串交给一个纯频率驱动的算法,它会自己发现化学片段。**BPE@512 对几个分子的切分(demo_tokenization.py 实测输出):

分子原子级切分(长度)BPE@512 切分(长度)
阿司匹林 CC(=O)Oc1ccccc1C(=O)O21 个 tokenCC(=O)O / c1ccccc1 / C(=O)O(3 个)
咖啡因 CN1C=NC2=C1C(=O)N(C(=O)N2C)C28CN/1/C=/N/C2/=/C1/C(=O)N/(C(=O)/N/2/C)C(12)
ClC(Br)C6Cl / C( / Br / )C(4)
[nH]1cccc17[nH] / 1 / cccc1(3)

阿司匹林被干净地切成乙酰基 — 苯环 — 羧基三块,而这三个 token 里没有一个是由化学知识定义的。收益最高的子词也很有说服力:

子词收益(训练集上减少的 token 数)化学含义
CC1193乙基/相邻碳链
cc1098芳香环上的相邻碳
(C769分支起点
c1765环闭合标记的开头
O)616氧原子后接分支收尾
(=O)465羰基
c1cc378芳香环片段
l)328⚠️ Cl) 被切碎后的残渣

最后一行是反面教材:字符级 BPE 不知道 Cl 是一个原子,于是把 l) 粘成了 token。这正是"化学感知分词"存在的理由,也提醒我们:BPE 学到的是字符串统计,不是化学知识。更微妙的风险在于环闭合数字——c1ccccc1 被学成一个 token 看起来很漂亮,但 1 这个数字的含义完全取决于图的遍历顺序(§4.7 与 §9.1 讨论过这个坑);把它当作"环"的符号会在同分异构体、不同遍历起点、不同 canonicalization 版本之间立刻失效。

E.6 风险清单与选型建议

四个常见陷阱

  1. 把化学字符串丢给自然语言的 tokenizer:GPT-2 / BERT 的词表是在自然语言语料上学出来的,而 [nH][C@@H]>>、环闭合数字这类 SMILES 专用符号在那些语料里几乎不出现,只能退化成字节或字符序列,序列长度与注意力账单同步膨胀。跨域迁移时这是最常见的静默错误——先打印一条你熟悉的 SMILES 的切分结果,再看它的长度
  2. 困惑度/损失不可比:不同分词器下的 bits/token 天然不同(分母不同),必须换算成 bits/字符 才能比较(换算方式见 §6.3.2)。
  3. 训练集泄漏:合并规则只能在本该"看不见测试集"的那部分数据上学习。本附录的做法是在训练集 902 条上学合并、在全量 1128 条上评测。
  4. 词表外回退:BPE 的鲁棒性来自"最坏情况下退化为字符",代价是遇到 OOV 时序列会突然变长——这也是长尾分子上延迟/显存抖动的一个真实来源。

选型清单

场景建议
小数据、小模型、零依赖复现(本教程)原子级或字符级;简单、可解释、无 OOV 概念
数据充足、追求短序列与吞吐BPE(kk 由词表预算与平均长度曲线的拐点决定,本语料在 256–512 附近)
需要严格的结构合法性SELFIES 等 100% 合法表示,再叠加 BPE
迁移预训练大模型必须沿用它的 tokenizer,不要自己重训(否则权重里的嵌入表全部失效)

E.7 复现

bash
cd code

# 分词粒度实测(约 5 秒)→ figures/results_tokenization.json + 图 fig_t01
python3 demo_tokenization.py

# BPE 相关测试(含词尾标记、重叠账本、顺序无关性等 8 项)
python3 -c "import tests_transformer as t; t.main()"

results_tokenization.json 里保留了每个方案的全部原始数字,figures/fig_t01_tokenization.png 是它的图,code/bpe.py 是可以直接拿走用的实现。

《原子智能》· 纸质出版预备版 · PolyAI Team 著