Implementing Embedding Gemma from scratch in PyTorch

# 在PyTorch中从零实现Embedding Gemma Embedding Gemma 是 Google DeepMind 于2025年11月发布的一款高效嵌入模型，以3.08亿参数量超越了参数量两倍的其他模型，并能以4位精度量化后几乎无性能损失，足以在手机上运行。本文将依据论文方法，详解如何使用 PyTorch 从零构建该模型的核心架构与训练流程。 ## 嵌入模型的作用与应用场景 嵌入模型将文本映射为固定长度的数字向量（例如768维），使得语义相似的文本在向量空间中彼此靠近。这一特性支撑了多种应用： - **语义搜索**：将查询和文档向量化，通过向量距离找到最相关的文档，这是谷歌搜索和聊天机器人检索增强生成的基础。 - **聚类**：无需标签，通过聚类嵌入向量即可对海量数据（如客户评论）进行主题分组。 - **去重**：通过计算向量间距离，识别近似重复的文档。 - **分类**：在嵌入向量上训练简单的线性分类器即可实现高效分类。 Embedding Gemma 的工程突破在于，它在极小的参数规模下实现了顶尖性能，并优化了移动端部署能力。 ## 模型训练的四个阶段 Embedding Gemma 的训练流程仿照运动员培养路径，分为四个递进阶段： ### 第一阶段：编码器-解码器预训练 以 Gemma 3 通用语言模型为起点，使用 UL2 目标（混合去噪任务）将其适配为编码器-解码器架构。此阶段迫使编码器发展双向注意力机制，使其能同时利用上下文双向信息，为理解词义消歧打下基础。 ### 第二阶段：预微调 将编码器作为嵌入模型，在大规模、多样化且略带噪声的数据集（涵盖100多种语言的网页、问答对、代码等）上训练。此阶段仅使用大批量数据，尚无困难负例，旨在学习广泛的语义表示基础。 ### 第三阶段：微调 使用规模较小但质量更高的任务特定数据进行训练，数据中包含通过贝叶斯优化得到的多个困难负例。此阶段针对不同领域进行专门化训练，使模型达到性能巅峰。 ### 第四阶段：模型融合 将所有微调检查点的权重进行平均，得到一个最终模型。此方法融合了各次专门训练中的最佳特性，提升了模型的泛化能力。 训练过程中同时运行三个损失函数： - **对比损失**：拉近正确配对向量的距离，推远错误配对。 - **分散正则化器**：促使嵌入向量均匀利用整个向量空间。 - **嵌入蒸馏**：让模型从大型教师模型（Gemini嵌入模型）中学习。 ## 核心架构的五大关键改进 Embedding Gemma 的编码器基于标准双向Transformer，但引入了五个关键改进，使其更高效、更强大。 ### 1. 无因果掩码的双向注意力 与标准解码器只允许位置`i`看到位置0至`i-1`不同，双向注意力允许每个词元同时关注左右两侧的所有上下文。这对于消除词义歧义至关重要（例如“bank”在“河流”和“银行”语境中的不同含义）。论文消融研究证明，此初始化方式在MTEB多语言基准上比仅用解码器初始化的模型高出整整一个百分点。 ### 2. 使用RMSNorm替代LayerNorm RMSNorm 去除了LayerNorm中的均值中心化步骤，仅进行缩放：`x / sqrt(mean(x^2)) * 学习到的缩放因子`。它更简单、稍快，在大模型上性能相当。 ### 3. 使用旋转位置编码替代正弦位置编码 RoPE 通过根据位置旋转查询和键向量来编码相对位置信息，而非添加绝对位置偏移。其关键数学性质是：两个词元的注意力分数仅取决于它们的相对距离`i-j`，而非绝对位置。这使模型具有更好的长度泛化能力。 ### 4. 分组查询注意力 采用12个查询头与4个键值头的分组结构（每三个查询头共享一组键值头）。这在推理时显著节省了键值缓存内存，且不会实质损害模型学习能力。 ### 5. GeGLU激活函数 在前馈网络中，GeGLU 使用双分支门控机制：一个分支通过GeLU激活生成0到1之间的门控值，另一个分支承载信息，两者逐元素相乘。这赋予了模型对每个维度信息流更精细的控制能力。 ## 从序列到向量：池化与多尺度表示 ### 平均池化与掩码 编码器输出序列中每个词元的向量，而嵌入模型需要整个输入的单一向量表示。通过对所有词元向量进行平均池化可以实现这一点。关键在于需要使用掩码来忽略填充词元，防止它们稀释真实信号。论文发现，简单的带掩码平均池化优于带学习权重的注意力池化，因其泛化性更好。 ### 投影头 池化后的768维向量会通过一个两层投影头（中间维度为3072，即模型维度的4倍）进行重塑。最终通过L2归一化将向量映射到单位球面上，使得余弦相似度计算简化为点积，极大提升了大规模向量检索的效率。 ### 多尺度表示学习 Embedding Gemma 采用了一种名为“多尺度表示学习”的巧妙策略：同一个模型会同时生成四个不同维度的嵌入向量（128、256、512和768维）。这些嵌入呈套娃式嵌套关系，每个维度的嵌入都是完整可用的，允许用户在模型容量和计算开销之间灵活权衡。在训练期间，损失函数会同时作用于所有尺度的嵌入表示，确保每个尺度都学习到高质量的语义信息。 > **来源**：本文章内容整理自 YouTube 视频《Implementing Embedding Gemma from scratch in PyTorch》（频道：Winter_Mistake_3185）。视频链接：https://youtu.be/rxOUmSJTQBk