Transformer学习二

2018年那个深夜，我蹲在实验室里对着满屏乱码发呆。一条30个字的中文句子，经过我精心调教的LSTM，硬是给翻成了“一棵树在河边跳舞”。你说气不气？

跑了整整14个小时的epoch，三天才能看一眼验证指标。这感觉，就像你写封信寄出去，三个礼拜才收到回信——黄花菜都凉了。

后来呢？我换成了Transformer base。同样的数据量，8小时跑完，BLEU直接往上蹿了4.2个点。

你肯定要问：凭啥？

看，这就是Self-Attention最绝的地方。RNN天生是串行的——你想算第5个词，必须等前4个算完。就像排队做核酸，一个人做完才能轮到下一个。效率能高才怪。

但Self-Attention没有这个锁。一句话里的所有token，训练时一起做矩阵乘法，GPU直接开挂。信息传播路径从序列长度N降到了一步直达，每个词都能直接“看见”所有其他词，信息衰减？不存在的。

是不是以为这就完美了？别高兴太早。

我第一次手写Self-Attention的时候，踩了一个最大的雷，说出来你都不信——忘记除以√d_k。

那叫一个惨烈。loss曲线从0.6直接跳到12，然后变成NaN，梯度爆炸到整个模型直接瘫痪。就跟开车只管油门不管刹车，最后只有一个结果。

查了一整天，问题出在哪？当d_k=512时，点积值的方差大约是512。你想想，softmax一口吞进去这么大的数，梯度全落到饱和区，不死才怪。

现在你明白为什么原版公式里那个除法这么重要了吧？除以√512≈22.6，方差瞬间回到1，训练稳如老狗。

说到这儿，我给你个实在的建议：写attention函数的时候，先把那个缩放常量写死在代码里。就算你的d_k很小，除一下也没坏处。我现在已经把它做成模板，每次新建项目直接拷：

scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

记住，这个d_k别硬编码，从query的最后一维取。我试过动态获取和硬编码性能没差别，但能避免你复制代码时忘记改参数。

再说多头注意力。

很多人觉得多头就是多复制几份QKV，并行算。我以前也是这么想的，结果在一个长文本摘要任务上，把head数从8调到64（d_model保持512），你猜怎么着？显存直接爆了。40G的A100都扛不住，测试集ROUGE还降了1.1个点。

为什么呢？你看，每个头分到的通道数是d_k = d_model / h。头越多，每个头的维度就越小。64个头每个只有8维，你想想8维的空间里能装下什么复杂信息？还不如少头数、高维度来得有效。

我现在一般这么配：

• d_model=256用4个头
• d_model=512用8或16个头
• d_model=1024用16个头

这几个搭配，我在文本分类、机器翻译、情感分析上反复试过，最稳。如果你显存够大，也可以试试32个头，但d_model得跟着涨到1024以上。

还有一个容易踩的坑：MultiheadAttention最后的线性投影层。有人觉得这层可有可无，直接省略，结果模型能力直接断层。如果你用PyTorch的nn.MultiheadAttention，这层它内置好了，放心用。但如果是自己手写，一定记得最后加一个nn.Linear(d_model, d_model)。

说到这儿，有人肯定要反驳了：“Self-Attention确实好，但复杂度O(n²)，序列一长不得炸啊？”

你说得对。我做过一个法律文书匹配的任务，句子平均长度3000，full attention根本跑不了，一个batch就吃了30G显存。

怎么办？有这几个招：先上GQA（Grouped Query Attention）——让多个query头共享一组key-value头，减少KV缓存大小。我试过4个query头对应2个KV头，速度提升了40%，效果几乎不掉。其次，对于中等长度（512到1024），混合精度能省一半显存，再配合梯度检查点还能再省30%，训练时间多花5%但非常划算。还有一个常被忽视的地方——位置编码。Self-Attention自己没有顺序感，你必须给它加上。我正余弦编码和可学习编码都试过：分类任务可学习稍好一点点（+0.3%准确率），但生成任务上正余弦更稳定。最关键的是，当序列长度超出训练时的最大长度时，可学习编码直接崩，正余弦可以外推。如果你做对话系统，我强烈推荐正余弦编码。

有人会说：“那干脆用稀疏注意力或者线性注意力。”我告诉你，如果序列长度不超过4096，full attention配这些优化完全能打；超过4096再考虑稀疏方案。我自己在768长度下永远只用full attention，干净高效。

说到头来，Self-Attention确实是Transformer能横扫序列建模的核心，但它不是个开箱即用的零件。你直接复制论文公式，很可能训练不收敛或效果不达标。

我建议你亲自动手做一次实验。从一个最简单的文本分类任务开始，手写一个不带mask的多头自注意力。先分别测试有没有除以√d_k，观察loss曲线有什么差别。再修改头数h=1,4,8,16，对比验证集指标和显存变化。最后换用不同的位置编码，看看外推能力如何。这四件事做完，你对Self-Attention的理解，会比看十篇笔记都深刻。

最后送你一句话：技术文章喜欢吹“Attention is All You Need”，但真正做工程的人都懂——

Attention只是入场券，把细节调对，才是你的“All You Need”。