大语言模型真的需要这么多层吗? [译]

研究表明,移除 70% 的注意力头和 20% 的前馈网络对上下文学习影响甚微,这暗示大语言模型或许训练得不够充分。

大语言模型(LLMs)已经存在一段时间,但随着 ChatGPT 的出现,今年它们真正进入了公众视野。LLMs 通常会在庞大的数据集上进行预训练,并且最近的版本还通过应用强化学习进行调整,以便更好地遵循指令并整合人类反馈。

这些 LLMs 展示出的一项令人着迷的能力是上下文学习,也就是说,模型能够仅凭几个(有时甚至不需要)好的示例,就学会如何执行新任务。遵循这种学习模式,当预训练数据量固定时,更大的 LLMs 在执行多种任务方面也表现出比小型模型更强的能力。

在我们即将在今年计算语言学协会(ACL)会议上展示的一篇论文中,我们从架构的可解释性角度出发,探讨了模型规模对于上下文学习的重要性。我们特别探讨了这个问题:所有 LLM 组件真的对于执行上下文学习都是必需的吗?

我们的研究以 OPT-66B 模型为案例,这是 Meta 去年开源的一款拥有 660 亿参数的大语言模型(LLM),旨在模仿 GPT-3,并且在我们研究时期是最大的公开解码器仅限大语言模型。我们的发现显示,模型的很大一部分可以被剔除而不影响其性能,这暗示 OPT-66B 乃至其他著名大语言模型可能存在训练不足的情况。

我们相信,识别并改善那些需要更精细训练的架构要素,可以帮助我们构建更强大、更有效的大语言模型。

大语言模型的基础构件

现代大语言模型(LLM)主要基于 Transformer 架构,核心在于一种称为注意力机制的技术。这种机制让模型能够学习在预测当前 Token 时,应该重点关注之前序列中的哪些 Token。

具体来说,LLM 采用多头注意力机制,这意味着模型同时运用多个不同的注意力机制,我们称之为“头”。例如,OPT-66B 模型就包含 64 层,每层有 72 个这样的头。这些多头注意力的输出会在每一层通过一个独立的前馈网络(FFN)进行处理。

我们分析 OPT-66B 的方法之一是对每个注意力头和 FFN 进行打分,以此判断它们在特定任务中的重要性。基于这些评分,我们进一步对模型进行了精简。

我们的研究发现,对于模型来说,较为重要的注意力头多集中于中间层,而重要的 FFN 则多位于后期层。令人惊讶的是,即便我们删去了高达 70% 的注意力头(在 OPT-66B 中约占 15.7B 个参数),模型在处理 14 种不同自然语言处理(NLP)任务时的零样本/少样本学习能力几乎未受影响。

这张热图展示了 OPT-66B 模型在 14 个自然语言处理 (NLP) 任务中,用于五次少样本学习的注意力头的综合重要性评分,覆盖了模型的每一个层级。
这张热图展示了 OPT-66B 模型在 14 个自然语言处理 (NLP) 任务中,用于五次少样本学习的注意力头的综合重要性评分,覆盖了模型的每一个层级。

研究发现,在不同任务和次数的学习过程中,对上下文学习至关重要(或不那么重要)的注意力头有相似之处。这说明存在一批既通用又与特定任务无关的注意力头,它们在上下文学习中发挥着核心作用。此外,我们还观察到,OPT-66B 模型中多达 20% 的前馈神经网络 (FFNs,大约有 85 亿参数) 可以在对零样本/少样本上下文学习的性能影响极小的情况下被移除。

我们采用的第二种分析方法是,量化 OPT-66B 模型中所有注意力头在执行两种与上下文学习相关的任务无关的基础操作的能力。这两种操作分别是前缀匹配复制:即在上下文中查找当前 Token 之前出现的位置,并复制其后续的 Token(即后缀)。

前缀匹配和复制操作。
前缀匹配和复制操作。

这两种操作专用的头部最初由机器学习研究公司 Anthropic 发现,被称为归纳头部。我们的研究表明,OPT-66B 模型中只有少数头部能够在这两种基础操作上都表现出色。而且,这些头部在不同程度上与我们之前识别的对特定任务重要的头部有所重叠。这意味着归纳头部不仅能够进行更为复杂的上下文学习相关行为,如隐性概念匹配,还表明它们并非是唯一具备这类能力的头部。

我们的一个关键发现是,只有少数的注意力头 (attention heads) 和前馈神经网络 (FFNs) 对于上下文学习显得至关重要,这揭示了 OPT-66B 乃至其他知名的大语言模型 (LLMs) 可能存在训练不充分的问题。这一点进一步印证了近期的一些研究成果,它们对在扩大模型规模时保持预训练数据量不变的有效性提出了质疑,认为要想取得最佳性能,必须让预训练数据的数量与模型的规模同步增长。对于我们的研究发布后出现的新版大语言模型,比如那些经过调整以更好地遵循指令的模型,了解它们在此类分析中的表现将会非常有意义。