大型语言模型真的需要那么多层吗?
研究发现,70%的注意力头和20%的前馈网络可以被移除而几乎不影响上下文学习性能,这表明大型语言模型可能存在训练不足。
研究背景
大型语言模型(LLMs)虽然已存在一段时间,但随着ChatGPT的出现,今年真正引起了公众的广泛关注。LLMs通常在海量数据上进行预训练;最近的变体还通过强化学习进行指令微调并融入人类反馈。
这些LLMs展示的一个迷人能力是上下文学习,即模型只需遵循少量(有时甚至零个)与新输入一起提供的示例就能学会执行任务。在这种学习范式下,当预训练数据量固定时,较大的LLMs被证明比较小的模型更能执行各种任务。
在计算语言学协会(ACL)年会上发表的一篇论文中,我们从架构可解释性的角度研究了模型规模对上下文学习的重要性。我们特别提出了一个问题:执行上下文学习真的需要所有LLM组件吗?
研究方法
我们以OPT-66B模型为案例进行研究,这是一个拥有660亿参数的LLM,由某机构去年开源,作为GPT-3的开放复刻版(在我们研究时是最大的公开可用的仅解码器LLM)。我们发现,模型的很大一部分可以被丢弃而不影响性能,这表明OPT-66B以及很可能其他知名LLMs都存在训练不足。
我们相信这些发现有助于通过识别(或更普遍地提供识别方法)可能需要更好训练的架构元素来构建更强大的LLMs。
LLM构建模块
现代LLMs使用Transformer架构,该架构依赖于注意力机制:模型学习预测在预测当前token时应关注序列中的哪些先前token。
具体来说,LLMs使用多头注意力,意味着它们并行应用多个注意力机制或头。OPT-66B有64层,每层有72个注意力头。多头注意力的输出在每一层通过一个单独的前馈网络(FFN)。
分析结果
我们分析OPT-66B的第一种方法是为每个注意力头和FFN分配一个分数,指示它们对给定任务的重要性。基于这些分数,我们对模型进行剪枝。
我们发现重要的注意力头主要聚集在模型的中间层,而重要的FFN主要在后几层。当移除高达70%(约157亿参数)的注意力头时,在14个不同自然语言处理(NLP)数据集/任务上执行零样本/少样本上下文学习的能力几乎保持不变。
用于上下文学习的重要(和不重要)注意力头在不同任务和样本量之间似乎也存在重叠。这表明一个通用的与任务无关的注意力头子集负责上下文学习。我们还发现,移除高达20%的FFN(约85亿参数)对零样本/少样本上下文学习的性能影响极小。
我们的第二种分析技术是量化OPT-66B中所有注意力头执行与上下文学习相关的两种任务无关原始操作的能力。这些原始操作是前缀匹配和复制:显式搜索上下文中当前token的先前出现并复制其后继token(其后缀)。
专门用于这两种操作的头最初由机器学习研究公司Anthropic发现,并称为归纳头。我们发现OPT-66B中的一小部分头对这两种原始操作都有显著的分数。我们还发现这些头与之前识别的特定任务重要头存在(不同程度)重叠。这表明归纳头能够执行与上下文学习相关的更复杂行为,如潜在概念匹配,但它们并非唯一具有这种能力的头。
主要结论
我们最重要的观察是,只有注意力头和FFN的核心核心似乎对上下文学习重要,这表明OPT-66B以及很可能其他知名LLMs都存在训练不足。这也强化了最近的研究,这些研究质疑在扩大模型规模时保持预训练数据量固定的有效性,表明所见预训练数据量必须与模型本身同步扩展以达到最佳性能。
看看自我们研究发布以来更新的LLMs变体(如那些经过指令微调的变体)在此类分析中的表现将会很有趣。