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