具身智能论文学习1：PaLM: Scaling Language Modeling with Pathways

第一部分基本信息正式发表于Journal of Machine Learning ResearchJMLRVolume 242023年全文约113页。第二部分核心思想文本token→Embedding→Decoder-only Transformer→下一个文本tokenPaLM 是一个decoder-only、自回归语言模型。论文在引言中说明它采用从左到右的语言建模目标根据前面已经出现的 token预测下一个 token。第个 token第个 token 之前的所有 token模型给出的概率分布。其中举例假设前面输入The robot picks up the模型需要预测下一个 token。它可能给出候选 token概率cup0.52box0.18object0.11table0.04其他 token0.15那么模型选择或者采样出cup后把它接到前面The robot picks up the cup这就是“自回归”预测一个 token → 放回输入序列 → 再预测下一个 token → 不断重复在训练数据中正确的下一个 token 是已知的。但它不是 PaLM 最核心的论文贡献。因为 GPT、GPT-2、GPT-3 本质上也都是token → decoder-only Transformer → next token第三部分核心要解决的问题存在的问题当时一个核心问题是如果继续扩大语言模型的参数规模、训练数据和计算规模通用语言理解和推理能力还能不能持续提升是否会出现小模型不具备的新能力解决方法PaLM 因此将模型规模扩大到540B 参数通过Pathways在大规模 TPU v4 集群上训练并系统评估 scaling 后的语言、推理、多语言、代码等能力。第三部分核心创新点细节1核心架构创新①SwiGLU 激活函数采用形式相比 ReLU、GeLU 等激活函数显著提升模型质量尽管需增加一次矩阵乘法但在算力等价实验中表现更优②并行层设计将标准串行结构改为并行结构使 MLP 与注意力机制的输入矩阵乘法可融合大规模训练速度提升约15%。2训练的损失函数假设序列是The robot picks up the cup.训练时可以形成多个预测任务输入The 目标robot 输入The robot 目标picks 输入The robot picks 目标up 输入The robot picks up 目标the 输入The robot picks up the 目标cup模型希望正确 token 的概率越来越高因此通常最小化负对数似然3模型规模大小4文字转变变成 token过程Transformer 不能直接接收字符串The robot picks up the cup.第一步需要使用 tokenizer 将文字拆成 token。PaLM 使用的是SentencePiece tokenizer词表大小为也就是词表中共有约25.6万个可选择的 token。词表支持多语言、代码、空格以及UTF-8字节论文还说明数字会被拆成单独的数字 token例如123.5被拆成1 | 2 | 3 | . | 5一个句子可能被分成The | robot | picks | up | the | cup | .实际 SentencePiece 的分割可能更细例如把低频词拆成子词。也就是说token 不一定等于完整单词它可能是单词、子词、符号、空格形式或字节5token转变序号形式token 还只是编号不能直接进入 Transformer每个 token 在词表中都有一个整数 ID。例如仅作示意tokentoken IDThe315robot8402picks12617up207cup5138于是文字序列变成[315, 8402, 12617, 207, 5138]但这些数字只是索引。8402并不表示“robot 的数值大小是8402”也不代表它比cup5138更大。它们只是词表中的编号。Transformer需要的是连续向量因此还要经过 embedding 层。6token 进入 embedding 空间Embedding 层嵌入层的作用就是把计算机无法直接理解的离散编号转换成神经网络能够计算的连续向量。Embedding 层本质上是一个可以训练的巨大矩阵(|V|)词表大小 (d)每个 token 的向量维度维度越高⇒模型有更大的内部表示空间。假设token ID(robot)8402模型就从 embedding 矩阵中取出第8402行E[8402]得到一个长度为的向量例如 PaLM 8B 中4096这个里面的​就是特征向量概念上可能是真实模型中这个向量有数千到上万个维度。所以完整过程是文字 ↓ SentencePiece 分词 ↓ token ↓ token ID ↓ 查询 embedding 矩阵 ↓ 连续向量 ↓ Transformer写成数学形式就是其中就是 token embedding 映射。7embedding 向量学会的信息一开始embedding 向量只是随机初始化的参数。经过大量训练后与相似语言环境有关的 token 会形成有用的内部表示。例如模型可能学到robot与machine、arm有一定语义联系cup与bowl、container有一定联系pick up与grasp、lift有一定联系语法角色、上下文用法和多语言关系。但不要简单理解为embedding 的某一个维度就明确表示“是不是机器人”。语义通常分散在整个高维空间中并经过多层 Transformer 继续加工。8引入位置信息如果只有 token embeddingrobot picks cup和cup picks robot包含的 token 相同但语义完全不同。因此模型还必须知道位置和顺序。PaLM 使用RoPERotary Position Embedding旋转位置编码而不是普通的绝对位置向量。RoPE主要在注意力计算中将位置信息作用于 query 和 key使模型能够区分谁在前面谁在后面token之间相距多远。核心就是token内容向量位置信息→有顺序的序列表示9预测下一个tokentoken embedding 进入多层 decoder-only Transformer 后当前位置得到隐藏状态这个隐藏状态包含模型对前文的综合理解。是位置t的输入向量经过多层 Transformer 后得到的隐藏表示不过它不只包含的信息还融合了前面所有位置的信息即接下来模型计算词表中每个 token 的分数其中是输出投影矩阵它负责把 Transformer 输出的隐藏向量转换成词表中每个 token 的预测分数例如robot→token ID→4096维向量输出时模型又需要把4096维的 hth_tht​ 映射回25.6万个词表分数再经过 softmax得到预测的每个单词的概率最终得到cup 0.52 box 0.18 object 0.11 ...完整流程符号含义第 t 个离散 tokentoken 的初始 embedding 向量Transformer 处理后的隐藏向量输出投影矩阵不是编码器所有词表 token 的 logitssoftmax将 logits 转成概率10训练的数据11Pathways核心架构“PaLM 为什么能够扩展到 540B 参数”实现这一目标的关键TPU 可以先理解成Google 专门为机器学习/深度学习设计的一种加速芯片。pod指的是一大群通过高速网络连在一起的 TPU 芯片组成的“超级计算机集群”Pathways 是 Google 推出的新型机器学习系统支持跨数千个加速器芯片的高效训练其核心设计包括Pod 级双向数据并行单个 Python 客户端将训练批次拆分到两个 Pod每个 Pod 独立执行前向/反向计算随后交换梯度并并行更新参数确保参数一致性。异步调度与分片数据流通过 Pod 级调度器的异步调度掩盖任务分发延迟采用分片数据流模型降低数据传输成本。优化的跨 Pod 梯度传输将梯度数据拆分为小块通过多路径路由避免网络拥堵实现97% 的弱扩展效率理论 2 倍吞吐量实际达到1.95 倍。Figure 2 展示了 PaLM 如何借助 Pathways 跨多个 TPU v4 Pod 进行分布式训练。图中不同颜色表示不同 TPU Pod模型参数和训练数据被分布到大量 TPU 上并通过高速数据中心网络进行通信。PaLM 540B 的训练使用了6144 个 TPU v4 芯片。因此Pathways 并不是 PaLM 神经网络中的一个新模块而是支撑超大规模模型训练的分布式计算基础设施。它解决的核心问题是当模型规模增长到数千亿参数后如何让成千上万个计算核心高效协同工作。12Pathways 的训练效率为了衡量如此大规模的分布式训练是否真的高效论文进一步比较了不同大型语言模型的Model FLOPs UtilizationMFU模型计算利用率。从 Table 3 可以看到PaLM 540B 使用6144 个 TPU v4进行训练模型 FLOPs 利用率达到 46.2%高于表中 GPT-3、Gopher 和 Megatron-Turing NLG 等同期大模型。这说明 Pathways 不只是简单地“堆更多 TPU”而是能够较高效地利用大规模计算资源为训练 540B 参数的 dense Transformer 提供了工程基础。13模型变大真的有用吗——Scaling 带来的能力提升Figure 3 展示了 PaLM 在 BIG-bench 大规模任务集合上的 Scaling 现象。横轴表示模型参数规模纵轴表示多个任务上的归一化平均性能。可以看到随着 PaLM 从较小规模逐渐扩大到540B 参数整体任务表现明显提高而且 1-shot、5-shot 通常进一步优于 0-shot。这说明扩大模型规模带来的不仅是更好的语言流畅度还提高了模型在大量不同任务上的few-shot learning、知识利用和复杂推理能力。这也是 PaLM 最重要的结论之一大规模预训练 模型 Scaling 可以形成更强的通用语言能力。第四部分缺陷落到具身智能PaLM 提供语言知识与推理底座但仍是 Text→Text缺少视觉、affordance 和 action因此自然引出 SayCan / PaLM-E / RT-2。第五部分参考文献https://blog.csdn.net/m0_65010824/article/details/155193911