Transformer是推断还是记忆？初始化大小很重要

2024-05-17 12:05

©作者 | 张众望，许志钦，GPT-4o（负责幽默部分）

单位 | 上海交大深度学习基础理论团队

Transformer 架构在当前大语言模型中的地位，堪比麦当劳里的牛肉饼——谁不爱呢？这位“全能选手”能解数学题、写诗作赋，简直是 AI 界的“跨界网红”。不过，要揭开它才华横溢的秘密，可比哄孩子吃饭还费劲。

直接研究真实的大语言模型，难度堪比在菜市场里向大妈们讲解狭义相对论——你会被讨价还价的嘈杂声淹没。所以，咱们换个思路，采用“锚点-关键项”这招“偷天换日”，既省时又高效。接下来，让我们一探 Transformer 学习复合函数的究竟：它是靠推理，还是单纯拼记忆？

假设我们有四种基本运算,它们可以组合出 16 种“变身形态”（16 种复合运算）。如果我们只拿其中 15 种训练模型，你猜：它能掌握剩下那个“遗珠”吗？

答案是肯定的！

我们给模型输入一个 Token 序列，其中有一对从 {1, 2, 3, 4} 中选出的“天选锚点”，每个代表一种加减法运算。还有一个从 20 到 100 中挑选的“关键项”，它在经过“天选锚点”的洗礼后，幻化成序列的目标值。我们令四个锚点分别代表 “+5”、“+1”、“-2”、“-8” 四种魔法。举个栗子,如果关键项是 80，锚点是（1, 2），那目标值就是（80+5）+1=86。

如果我们只用 15 个锚点组合训练，模型竟然能掌握第 16 个的套路。这就像隔着手机屏幕学会做大餐，是不是很酷？那模型到底是怎么做到的呢？为了方便讨论，我们就假设缺席训练的是（4, 3）这对“神仙眷侣”。

其实,模型的学习机制有两种可能。一是走“福尔摩斯”路线，通过演绎推理出每个单锚点的能力，再将 “4”、“3” 的技能组合，找到推理解，破解（4, 3）的奥秘。二是走“复制粘贴”路线，发现（a, b）和（b, a）这两类锚点组合的效果一样（因为加减法可交换嘛），于是把（3, 4）的规律“拿来主义”，安到（4, 3）头上，找到对称解。Transformer 究竟更爱哪一套呢？

破案的方法很简单,我们把（3, 4）的效果从 “-10” 改成“-6”，看（4, 3）是否也跟着变。如果（4, 3）的结果仍是 “-10”，说明模型是靠推理；如果变成了 “-6”，就说明模型其实在背公式。

实验发现，以默认初始化为分水岭，大的初始化推动模型多用背诵，学到对称解，小的初始化则有利于推理。这就像健身房的两类人，一类死记硬背每个动作，一类则掌握基本原理，灵活运用。

这其实与模型复杂度有关。对于推断机制，模型仅需要记住 4 种单 anchor 运算，而对于记忆机制，模型需要学习所有非对称关系的运算规律（如（1，1），（1，2）（1，3），（1，2），（2，2）等共 10 种）。因此学习推断机制所需要的模型复杂度更低。对于小初始化，模型初始复杂度低，模型参数会发生凝聚现象。简单来说，就是同层的神经元会趋同。例如矩阵的列，在小初始化时，有明显地相似度（低秩），但大初始化下并没有这种现象。

训练开始时，参数聚集在少数方向，随着训练深入，聚集的方向逐渐增多，模型复杂度水涨船高。当模型发现学 4 种映射就够拟合训练集，学习就停止了，所以它更爱简单的推理解。而在大初始化下，复杂度够高，模型很快记住其中的对称性。如果初始化再大一点,它干脆把每个输入序列到输出的映射都记住，反而学不到算术规律，连见过的锚点组合都不能保证完美复现。

通过这项研究，我们发现，Transformer 模型的初始化大小决定了它是像福尔摩斯一样通过推理解谜，还是像我奶奶一样通过记忆菜谱来做饭。小初始化让模型像侦探一样，只需要记住几个关键的线索（运算规则），就能推理出所有结果。而大初始化则像孙悟空，把所有知识吃下去的方式记下来。

展望未来，科学家们可能会像调配宇宙飞船的发动机一样，精确调节 Transformer 的初始化参数，以便它们在各种复杂任务中都能表现出色。想象一下，将来有一天，Transformer 不仅可以帮你证定理、写论文，还能帮你规划度假行程，甚至给你家的猫咪制定健身计划。

不过，要实现这些目标，我们还需要深入探讨不同初始化对模型复杂度和泛化能力的影响，就像农民研究如何让奶牛在听音乐时产奶更多一样。这不仅需要大量实验，还需要一点点运气和许多好奇心。

参考文献

[1] Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, Zhi-Qin John Xu*, Initialization is Critical to Whether Transformers Fit Composite Functions by Inference or Memorizing, arxiv 2405.05409 (2024)

[2] Zhongwang Zhang#, Zhiwei Wang#, Junjie Yao, Zhangchen Zhou, Xiaolong Li, Weinan E, Zhi-Qin John Xu*, Anchor function: a type of benchmark functions for studying language models, arxiv 2401.08309 (2024)

更多阅读