知名AI研究者深挖谷歌Gemma:参数不止70亿,设计原则很独特

就在几天前,开源大模型领域迎来了重磅新玩家:谷歌推出了全新的开源模型系列「Gemma」。相比 Gemini,Gemma 更加轻量,同时保持免费可用,模型权重也一并开源了,且允许商用。

谷歌发布了包含两种权重规模的模型:Gemma 2B 和 Gemma 7B。尽管体量较小,但 Gemma 已经「在关键基准测试中明显超越了更大的模型」,包括 Llama-2 7B 和 13B,以及风头正劲的 Mistral 7B。与此同时,关于 Gemma 的技术报告也一并放出。

相信大家已经对 Gemma 的相关内容进行了系统研究,本文知名机器学习与 AI 研究者 Sebastian Raschka 向我们介绍了 Gemma 相比于其他 LLM 的一些独特设计原则。

Raschka 首先从模型性能展开,他表示看过技术报告的小伙伴可能都有一个疑问,是什么让 Gemma 表现如此出色?论文中没有明确说明原因,Sebastian Raschka 认为可以从下面两点得出结论:

  • 首先是词汇量大,Gemma 词汇量达到 256000 个单词,相比之下,Llama 的词汇量为 32000 个单词;
  • 其次是训练数据集达 6 万亿 token,作为对比,Llama 仅接受了其中三分之一的训练。

在架构方面,Raschka 列举了 Gemma 与 LLama 2 7B 和 OLMo 7B 的架构概览。

在模型大小上,Raschka 表示 Gemma 2B 有多查询注意力,而 Gemma 7B 没有。另外,与 Llama 2 相比,Gemma 7B 具有相对较大的前馈层,尽管其层数较少(28 VS 32),但 Gemma 中的参数数量却相当大。

Raschka 猜测 Gemma 7B 实际上总共有 93 亿个参数,如果考虑到权重共享(Weight tying)的话,则有 85 亿个参数。权重共享意味着模型在输入嵌入和输出投影层中共享相同的权重,类似于 GPT-2 和 OLMo 1B(OLMO 7B 的训练没有权重共享)。

归一化层

另一个引人注目的细节是以下出自 Gemma 论文中的段落。

归一化位置。谷歌对每个 transformer 子层的输入和输出进行归一化,这与单独归一化输入或输出的标准做法不同。谷歌使用 RMSNorm 作为归一化层。

乍一看,看起来像 Gemma 在每个 transformer 块之后都有一个额外的 RMSNorm 层。但是,通过查看「keras-nlp」项目的官方代码实现,原来 Gemma 仅仅使用了 GPT-2、Llama 2 等其他 LLM 使用的常规预归一化方案,具体如下图所示。

GPT、Llama 2 和其他 LLM 中典型的层归一化位置,Gemma 中没有什么新东西。来源:https://github.com/rasbt/LLMs-from-scratch

GeGLU 激活

Gemma 与其他架构之间的一大区别是它使用了 GeGLU 激活,而 GeGLU 激活是在 2020 年的谷歌论文《GLU Variants Improve Transformer》中提出的。

论文地址:https://arxiv.org/pdf/2002.05202.pdf

GeLU 全称为高斯误差线性单元(Gaussian Error Linear Unit),它是一个激活函数,越来越多地被作为传统 ReLU 的替代方案。GeLU 的流行得益于它有能力引入非线性特征,并允许为负输入值执行梯度传播,这解决了 ReLU 的一大局限,完全阻断了负值。

现在,作为 GeLU 的门线性单元变体,GeGLU 的激活被分割为两部分,分别是 sigmoid 单元和线性映射单元(它与 sigmoid 单元的输出逐元素相乘),具体如下图所示。

GeLU 与 ReLU 激活函数图示比较,来源:https://github.com/rasbt/LLMs-from-scratch

同时,GeGLU 与 Llama 2、Mistral 等其他 LLM 使用的 SwiGLU 激活类似。唯一的区别是 GeGLU 使用的基础激活是 GeLU 而不是 Swish。

下图展示了 GeLU(GPT-2)、SwiGLU(Llama 2)和 GeGLU(Gemma)的伪代码。

需要注意,与使用 GeLU(仅线性)的常规前馈模块相比,使用 SwiGLU 和 GeGLU 的前馈模块各多了一个线性层(分别是 linear_1 和 linear_2)。不过,在 SwiGLU 和 GeGLU 前馈模块中,linear_1 和 linear_2 通常通过将单个线性层分割为两部分而获得,因此不会增加参数规模。

那是否 GeGLU 就比 SwiGLU 强呢?并没有消融实验来证实这一点。Raschka 猜测谷歌选择使用 GeGLU,只是为了让 Gemma 与 Llama 2 略有不同。

举例而言,Gemma 为 RMSNorm 层添加了 + 1 的偏移量,并通过隐藏层维数的开立方根来归一化嵌入。Gemma 论文中没有提及或讨论这些细节,所以它们的重要性也不清楚。

结论

对于开源 LLM 而言,Gemma 做出了非常棒的贡献,展示了 7B 参数规模也能成就强大的模型,并有潜力在真实世界的用例中取代 Llama 2 和 Mistral。

此外,目前 7B 大小规模的开源模型已经有很多了,因此 Gemma 2B 更加有趣,它可以轻松地在单个 GPU 上运行。当然,Gemma 2B 与 2.7B 大小的 phi-2 之间的对比也将会很有趣。

THE END