自回归的大语言模型在各种自然语言处理任务上表现优秀，但是，他们需要大量的计算资源消耗。

大规模的自回归模型在解码时，如果要解码K个token，就需要串行K次，因为自回归模型是由前面的分布决定后面的分布的。除了这些大规模的自回归模型之外，还有更高效的、小型的模型（比如同系列的参数比较多的模型，和参数比较少的模型），这些模型可以更快地完成解码

Leviathan等人观察到，大型模型进行推理通常不是在算术运算上受到瓶颈，而是在内存带宽和通信上，因此可能会有额外的计算资源可用（大部分时间GPU都在干等）。这些额外的计算资源可以用来运行更高效的模型用于解码，从而在不损失生成质量的前提下实现对大语言模型的加速。

这种思想与分支预测比较类似——分支预测通过猜测哪一个指令电路会被执行，并且先一步执行这个电路，如果预测成功，可以节省相当的时间，如果预测失败，则退回预测之前的状态。Speculative Decoding（SD）的思想也是这样，利用GPU的空闲时间（其实也可以是CPU的空闲时间）运行一个更高效的小型模型（这个模型可能会犯错误，就像分支预测器一样）去猜测可能的解码结果是什么。如果猜对了，就省了计算资源，否则，也不会比原来更差——因为这个更高效的模型本身就是跑在被浪费的计算资源上的。

方法

SD利用一个规模较小的模型作为草稿，预测大模型的输出token，随后，大模型对输出token进行验证，从而达到减少大模型forward，节省并充分利用算力。

形式化表述

约定：

• $M_p$: Target Model（大模型），我们要加速的模型。其概率分布为 $p(x|x_{ \lt t})$。
• $M_q$: Draft Model（小模型），计算速度快但精度较低的模型。其概率分布为 $q(x|x_{ \lt t})$。
• $x_{ \lt t}$: 当前已生成的上下文序列。
• $K$: 每次投机生成的步数（Lookahead steps），即小模型一次生成多少个草稿 token。

SD通过让小模型串行生成多个token（但很快），随后让大模型并行进行验证（但很慢），在保证输出分布与$p(x|x_{ \lt t})$相同的情况下有着更快的生成速度。
为了理解SD，我们需要先了解拒绝采样原理。

拒绝采样原理

给定目标分布$p(x)$，我们希望从$p(x)$中取样$x \sim p(x)$，但是，$p(x)$由于计算成本过高等特点，难以直接采样，因此，我们改为从提议分布$q(x)$中采样，只要存在$M>0$，使得对于任意的$x$，都有：

$$ Mq(x) \ge p(x) $$

即，$q(x)$在扩大若干倍后可以完全在$p(x)$的上方，我们就可以利用$q(x)$去辅助$p(x)$采样。事实上，在SD场景下，我们一般设$M=1$。

我们需要证明，$\forall x_i \sim q(x)$，我们遵循以下步骤：

1. 从 $U(0, 1)$ 中采样随机数 $u$
2. 计算接受概率 $\alpha = \frac{p(x_i)}{Mq(x_i)}$
3. 若 $ u \le \alpha$，接受$x_i$
4. 否则，考虑修正分布 $\max(0, p(x) - q(x))$，我们从其中采样$x_i \sim (p(x) - q(x))$

这样得到的$x_i$满足 $x_i \sim p(x_i)$。

需要注意的是，这种表述不是标准的拒绝采样，我们称之为不退回的拒绝采样，对于标准的拒绝采样，它会在第4步退出，转而从原分布中重新采样，而我们则从修正分布$\max(0, p(x) - q(x))$中采样，因为如果$x_i$被拒绝，说明样本落在$q(x)>p(x)$的区域，我们必须要从这一部分采样弥补，这与从原分布中采样是等价的。

我们现在对这一命题给出证明。

设最终输出的样本为 $X$，$X$要么源自直接被接受的$q(x)$中采样的样本，这表明：

$$ P_1(x) = q(x) \min (1, \frac{p(x)}{q(x)}) $$

，要么源自从修正分布中采样的样本，此时，拒绝发生的总概率为 $P_{rej} = 1 - \sum_x \min(q(x), p(x))$，这表示采样点恰好落在 $q(x)$ 和 $p(x)$ 之间。此时，修正分布为

$$ p'(x) = \frac{\max(0, p(x) - q(x))}{\sum_x \max(0, p(x) - q(x)} $$

这一步可能稍微有些难理解，如果接受，那么$x_i$一定在目标分布$p(x)$中，否则，它必定在$p(x)$和$q(x)$之间，这个空隙就是

$$ \text{Gap}(x) = p(x) - \min(p(x), q(x)) $$

而 $p(x) - \min(p(x), q(x))=\max(0, p(x) - q(x)$，这表明：

$$ \text{Gap}(x) = \max(0, p(x) - q(x))) $$

由于此时是在采样为 $x$ 的时候被拒绝，根据条件概率公式：

$$ P(x|\text{拒绝}) = \frac{P(x; \text{拒绝})}{P(\text{拒绝})} $$

而 $P(x; \text{拒绝})$ 是从 $p(x)$ 和 $q(x)$ 的空隙中采样得到的，因此 $P(x; \text{拒绝}) = \max(0, p(x) - q(x))$。

而拒绝的情况就是对所有拒绝发生时的概率求和，这表明

$$ P(\text{拒绝}) = \sum_x \max(0, p(x) - q(x)) $$

这个总和实际上就是 p 和 q 两个分布之间的 $L_1$ 距离的一半。

综上，我们要么从 $P_1(x)$ 中采样，要么是 $p'(x)$，从而：

$$ P(x) = P_1(x) P(\text{接受}) + p'(x) P(\text{拒绝}) = q(x) \min (1, \frac{p(x)}{q(x)}) P(\text{接受}) + \frac{\max(0, p(x) - q(x))}{\sum_x \max(0, p(x) - q(x))} P(\text{拒绝}) $$，

如果 $p(x) \ge q(x)$，

$$ P(x) = q(x) + (p(x) - q(x)) = p(x) $$

，否则

$$ P(x) = q(x)\frac{p(x)}{q(x)} = p(x) $$

。

这样就证明了，通过这种方式从提议分布$q(x)$中采样的$x$满足$x \sim p(x)$。

Speculative Decoding 算法

我们已经证明了拒绝采样可以保障分布一致性，这样，我们可以正式提出拒绝采样算法：

输入：上下文 $x_{\lt t}$，投机步数 $\gamma$，目标模型 $M_p$ ，草稿模型 $M_q$

1. 当未达到起草长度 $\gamma$，且生成 token 不是 EOS 时：
2. 起草 token：
   • 对于 $i=1,...,\gamma$：
     • 采样 $x_{t+i-1} \sim q(x | x < t)$
     • 记录该位置的概率分布 $q_i$
   • 得到草稿序列 $\tilde{x} = [x_t, x_{t+1}, ..., x_{t + \gamma - 1}]$
3. 验证起草
   • 将 $\tilde{x}$ 送入 $M_p$ 做一次 forward，得到概率分布 $[p_1, ..., p_\gamma, p_{\gamma+1}]
   • 由 $p_i$ 和 $q_i$ 做拒绝采样：
     • 接受：加入接受序列中
     • 拒绝：从修正分布中采样，停止起草
   • 结束验证
   • 根据目标模型算出的分布采样额外token $x_{t+\gamma}$
4. 同步上下文并回滚KVCache

这样，我们就实现了Speculative Decoding。

效率提升

平均 token 生成数

设每次生成的接受率 $\beta_{x \lt t}$ 代表对于前缀 $x \lt t$ 的接受率，假设这是一个独立同分布的变量，我们设 $\alpha = \mathbb{E}(\beta_{x \lt t}$ 为平均接受率。

那么，单次算法运行所生成的token数量 N 服从一个带上限的几何分布，其成功概率（即拒绝概率）为 1−α，上限为 γ+1，设 $N$ 为单次迭代生成的总token数。其生成的token总数的期望值为：

$$\mathbb{E}(N)=\frac{1-\alpha^{\gamma+1}}{1-\alpha}$$

证明过程如下：
设 $I_i$ 为第 $i$ 个位置产生token的指示随机变量，第 1 个token总是会产生：$P(I_1=1) = 1$，第 2 个token产生的条件是第 1 个被接受：$P(I_2=1) = \alpha$，第 $i$ 个token产生的条件是前 $i-1$ 个均被接受：$P(I_i=1) = \alpha^{i-1}$
因此，总期望为：
$$E(N) = \sum_{i=1}^{\gamma+1} P(I_i=1) = \sum_{i=1}^{\gamma+1} \alpha^{i-1}$$
上述各项构成一个首项为 1，公比为 $\alpha$，项数为 $\gamma + 1$ 的等比数列：
$$E(N) = 1 + \alpha + \alpha^2 + \dots + \alpha^\gamma$$
根据等比数列求和公式 $S_n = \frac{a_1(1-q^n)}{1-q}$：
$$E(N) = \frac{1 \cdot (1 - \alpha^{\gamma+1})}{1 - \alpha}$$

证明完毕。

根据不同的 $\gamma$ 值，Speculative Decoding 预期的 token 数量与 $\alpha$ 的函数关系如上图。这给出了一个简单的推论，接受率总是越高越好。

接受率与模型差异的关系

前面的事情以后再来探索吧

The post Speculative Decoding及其原理（一） first appeared on 南极滑稽的博客.

他们在Github上公开了MetaRAG的源代码。

原文摘要

检索增强生成技术（RAG）因其在生成事实内容方面的作用而成为自然语言处理的核心。虽然传统方法采用单次检索，但最近的方法已转向多跳推理任务的多次检索。但是，这些策略受预定义的推理步骤的约束，可能会导致响应生成不准确。本文介绍了MetaRAG，这是一种将RAG过程与元认知（Metacognition，对认知的认知）相结合的方法。在认知心理学中，元认知允许实体进行自我反思，并批判性地评估其认知过程。通过集成元认知，MetaRAG 使模型能够监控（monitoring）、评估（evaluating）和规划（Planning）其响应策略，从而增强其内省推理能力。通过三步元认知调节管道，该模型可以识别初始认知反应中的不足之处并加以修复。实证评估表明，MetaRAG的性能明显优于现有方法。

方法简析

概要

之前的各种方法虽然从各种地方改进了RAG的检索效果，但是这些方法中的检索步骤都是固定的。与此不同，我们人类可以反思并改进我们的思考方式，这种能力被称为元认知（Metacognition）。作者们认为，让RaLM(Retrieval-Augmented Language Model，i.e. 采用RAG的LLM)模拟人类的元认知能力，可以提高其在复杂任务上的表现。

如上图，人类的元认知可以分为两个部分：知识（Knowledge）和元认知知识（Metacognitive Knowledge），其中，知识表明我们对客观事实、方法等的认知，而元认知知识则是我们对认知过程的管理和调控，我们可以由元认知知识了解事物运行背后的逻辑，并由此找到自我提升的手段。向我们的输入会先由我们的知识部分得出一个响应，随后根据具体情况决定是否进入元认知部分进行分析，最终输出一个不错的响应。

仿照上文中人类的知识和元认知知识，作者将他们设计的LLM分为外部知识部分和元认知知识部分，给LLM一个输入，它会先像普通的RaLM一样从外部知识库中检索，并生成一个响应，随后由根据具体情况决定是否要进入元认知知识部分进行分析，最终输出一个还不错的响应。他们称这种方法为MetaRAG。

具体的，他们的方法可以归结为下图：

首先，他们将MetaRAG分作认知空间(Cognition Space)以及元认知空间(Metacognition Space)，对应之前提到过的External Knowledge和Cognitive Knowledge。认知空间执行正常的RAG步骤。

元认知空间由Monitoring, Evaluating和Planning部分组成，其中，Monitoring部分负责评估回答质量，决定是否进入元认知空间，如果不进入，则按正常的RAG步骤输出，如果决定进入，则进入Evaluating部分，这一部分负责判断现有的内外知识是否足够解决问题，以及是由于犯了哪些错误导致回答质量不佳，随后的Planning阶段根据Evaluating阶段所得出的结论判断应该采取哪些操作提升回答的质量。

细节

Monitoring部分

如上文所说，Monitoring是负责评估回答的质量是否足够好的。对此，他们选择使用专家模型生成相同问题的回答，并通过对比LLM与专家模型的回答的相似度来决定答案是否足够好的。具体来说，他们将模型生成的嵌入与专家模型生成的嵌入作比较，如果相似性小于k，则认为这个答案是不好的。

（为什么他们不直接输出专家模型的结果呢？我也不知道，但是这个论文已经被A会接受了，或许大佬们认为这样做是没问题的？）

Evaluating部分

Evaluating是Monitoring认为答案不够好后进入的部分，主要聚焦于两点：

（a） 内部和外部知识来源都足以解决所提出的问题吗？

（b） QA LLM 的推理过程是否容易受到多跳 QA 中经常遇到的常见问题的影响？

为了解决这两个问题，他们又引入了两个概念：过程知识和陈述性知识。程序知识体现了对面对特定任务所必需的方法和策略的把握，而陈述性知识则锚定在特定事实或基于内容的信息中，涵盖与解决问题相关的事实和概念。他们在prompt中让模型充当evaluator-critic系统，认为这样做效果比让模型从问题的回答者去评估生成的答案要好。

在过程性知识的评估过程中，他们使用模型充当的evaluator-critic系统评估模型的内部知识，用NLI（自然语言推理）模型去衡量模型的外部知识，NLI模型可被抽象为形如 f⁢(premise,hypothesis) 的函数，在premise可以推出hypothesis时返回1，否则返回0，同时，他们也给模型做了一个prompt：

Please act as an evaluator-critic system, determine if you can provide a reliable answer to the {question} based on your own knowledge?

判断完后，他们可以将知识上的不足分为四类，我们稍后会谈到。

在陈述性知识的评估过程中，他们将推理错误分为三类：不完全推理、答案冗余、歧义理解，并调用evaluator-critic系统去判断错误可能是什么，prompt如下：

Please act as an evaluator-critic system, assess whether the {response} based on {references} for the {question} contains any {error types}?

Planning部分

对错误有了认知后，就可以对症下药，尝试去解决错误了。现在填上文挖的坑，知识不足可以根据哪方面的知识来源不足被分为四类：

• 无知识
• 仅外部有
• 仅内部有
• 内部和外部都有

严格来说，最后一个不算知识不足，不过也属于可能的一种情形。无论如何，判断出来知识是哪方面不足后就可以对症下药了，对症生成prompt后，模型就可以从相应的地方继续搜索，直到问题解决或放弃为止。

知识不足不是模型回答错误的全部原因，另一种可能导致不准确回答的情况是内部和外部知识之间存在差异，即知识冲突，解决方案即直接放弃某个部分的知识，仅依赖单一知识来源就可以做出比较好的回答，不过他们这么做并没有考虑模型内部知识与外部知识的质量的差异，如果能考虑到评估知识库的质量，决定放弃某一个知识库可能会更好。

哪怕知识已经足够，模型可能还会在推理过程中犯错误，即错误推理，如果推理是用CoT等多步推理方法进行的，LLM会被要求重新检查推理过程中哪里可能出错，并对症重新进行推理。这里看起来似乎可以和ToT或者MCTS这样的方法结合一下，说不定会更好用呢？

评估部分

总体来说，

1. MetaRAG 始终优于两个数据集的所有基线方法。与 Reflexion 相比，Reflexion 在其推理过程中还集成了自我批评机制，MetaRAG 在所有指标上都表现出了实质性的改进。这表明使用元认知策略比仅仅依靠自我批评更有益。
2. 配备自我批评机制的模型与没有自我批评机制的模型相比，表现出优异的性能。

他们做了一个消融实验，证明剥离元认知空间中的任何一个组件都会造成效果的极大下降。

他们还研究了相似性阈值k，不过这些说起来也比较繁琐，在这里就不说了。

这篇理论上八月就看的差不多了，不过当时一直忙着没写完，现在才有时间写完发出来，不得不说我真的是……

The post 元认知增强检索大预言模型：MetaRAG论文导读 first appeared on 南极滑稽的博客.

参考文献：Saurav Kadavath, Tom Conerly, Amanda Askell, Tom Henighan, Dawn Drain, Ethan Perez, Nicholas Schiefer, Zac Hatfield-Dodds, Nova DasSarma, Eli Tran-Johnson, et al.Language models (mostly) know what they know.arXiv preprint arXiv:2207.05221, 2022.

如果你关注过大模型推理，可能会对推理框架有所了解，其中比较有代表性的大概是CoT（尤其是CoT with Self-Consistency）、ToT、MCTS了，但是这些推理框架想要更深入的思考，就要消耗更多的token，也就是消耗更多的资源、时间和钱了，我们有没有办法在思考深度和token消耗数做一个权衡，让大模型在token消耗数控制在我们满意的范围内，同时取得还不错的推理准确率呢？

我们人类可以评估问题的难度，动态调整思维的深度，例如，对1+1=2，我们并不需要深入思考，但是如果是高考的导数题，我相信大多数人都要认真的思考一番，显然，对1+1=2和高考导数题有相同的思考深度是很不现实的——但有些思考框架确实是这么做的。

直觉上，我们可以给大模型定义一个指标，这个指标衡量了大模型对做出这个题目是否有信心。

Saurav K. 等人就做了这方面的工作，他们认为，大模型（大多）知道自己知道什么，为此，他们考察了两个指标——P(True)，“我回答的答案正确吗”，与P(IK)，“我知道这个问题的答案吗？”，这与上文的引入可能有点不太相关，不过让我们继续往下走。

上图中，"Fraction of Problems" 指的是模型在特定任务中正确回答问题的比例，作者从模型的回答中sample了一些答案，其中有正确的也有错误的，由于Incorrect Samples（蓝色）本身就是不对的，我们希望它们分布在P(True)较低的地方，也就是模型应该认为这个回答是不对的；同上，由于Correct Samples（绿色）本身就是对的，我们希望它们分布在P(True)较高的地方。因此，我们希望蓝色分布在图的左侧，绿色分布在图的右侧。

作者采取了四种方式来衡量（文中称为校准/calibration）模型知不知道自己知道什么，分别为选择题、最后一项为“None of the above”的选择题、温度调整后的RLHF、判断题，下文为其阐述。

选择题能够让模型好的衡量他们自己

他们采取了如下的格式询问模型：

Question: Who was the first president of the United States?

Choices:

(A) Barack Obama

(B) George Washington

(C) Michael Jackson

Answer:

其中，他们只通过选择题的序号(A)、(B)、(C)识别答案，在选择题的情境下，模型展现了良好的“自我认知”能力（得到了很好的“校准”）

理论上，模型作出回答的频率应该与概率一致，因此在理想情况下，频率关于概率的分布应该是一条斜率为1的过原点的射线（在这条射线上为Under Confident，在这条平面下为Over Confident），而模型在此的表现也的确与其很接近。

作者认为，“在选择答案选项之前，模型能够明确地看到答案选项，这一点至关重要；没有这个，我们就不会期望得到校准的回答，因为任何给定答案选项的可能释义和专业化之间存在歧义和退化”，换句话说，正是因为有了选项的存在，模型才能在不存在“歧义和退化”的情况下评估自身的回答。

他们还发现了另外一种情况，只要把最后一个选项改为“None of the above”，就可以大大干扰模型对自己的认知，在这种情况下，模型总是毫无根据地反感选项“None of the above”，因此会造成准确度的大大的下降。

判断题可以让模型好的衡量它们自己

让模型单独判断某一个部分正不正确就可以排除选项对模型的自我认知造成的干扰（例如“None of the above”），尽管模型判断的能力也与模型的参数大小有关，但模型达到52B（黄线）后，尽管在概率很低的地方Under Confident及在概率高的地方Over Confident，但在大部分区域都很接近理想值。

上述让模型判断的都是人类给定的选项，如果让模型自己判断自己的答案正不正确，情况又会有所不同。简而言之，模型判断自己生成的答案更加困难，也更加接近模型的能力边界，而模型判断人类给出的选项更加简单，也更接近模型的能力舒适区。

尽管RLHF会干扰模型的自我判断能力，但温度调整后干扰会消失

强化学习的微调可能会让模型的预测崩溃，呈现图中蓝色的情况，但调整以后（红色）就与理想情况非常接近。

（*我不太懂强化学习，因此在这部分没办法做太多的阐述）

训练模型以预测它们是否能够正确回答问题

这里似乎来到了重头戏——我们如何让模型知道自己知不知道，从而调整推理的思维深度呢？作者希望训练模型来预测它们是否知道任何给定自由形式问题的答案，这与我们开头提出的情况高度相关。

• Value Head：将 P（IK） 训练为添加到模型中的附加值 'head' 的 logit（独立于语言建模的 logit）。这种方法的一个优点是我们可以轻松地在一般token位置探测 P（IK）。
• 自然语言：训练 P（IK），让模型按字面意思回答“你凭什么信心回答这个问题？”，并输出 0%、10%、20%、 ⋯ 100% 等答案。

自然语言方法效果不佳，因此作者采用了Value Head方法。

具体训练出的模型写不动了，留给以后再写（

The post 模型对自己也有认知——它们知道自己“知道”什么(1) first appeared on 南极滑稽的博客.

Zhihao Zhang1†, Alan Zhu1, Lijie Yang1, Yihua Xu2, Lanting Li1, Hitchaya Mangpo Phothilimthana3, Zhihao Jia1‡

1Carnegie Mellon University, School of Computer Science
2 University of California, Berkeley 3 Google DeepMind
† zhihaoz3@andrew.cmu.edu,
‡ zhihao@cmu.edu

注：本文相当部分的内容是我根据自己的理解写的，如有错误，敬请谅解，也欢迎指正与讨论。

原文摘要：

检索增强语言模型 （RaLM） 已经证明了通过将非参数知识库与参数语言模型相结合来解决知识密集型自然语言处理 （NLP） 任务的潜力。RaLM 不是对完全参数化的模型进行微调，而是作用在低成本地适应最新数据和更好的源归因机制方面。在各种 RaLM 方法中，由于检索器和语言模型之间的交互更频繁，迭代 RaLM 提供了更好的生成质量。尽管有这些好处，但由于频繁的检索步骤，迭代 RaLM 通常会遇到高开销。为此，我们提出了 RaLMSpec，这是一个受推测启发的框架，它为迭代 RaLM 提供通用加速，同时通过推测检索和批量验证保留相同的模型输出。通过进一步结合预取、最优推测步幅调度器和异步验证，RaLMSpec 可以自动充分利用加速潜力。对于朴素迭代的RaLM服务，在4个下游QA数据集上对三种语言模型的广泛评估表明，当检索器为精确密集检索器、近似密集检索器和稀疏检索器时，RaLMSpec的加速比分别为1.75-2.39×、1.04-1.39×和1.31-1.77×与基线相比。对于 KNN-LM 服务，当检索器是精确的密集检索器和近似的密集检索器时，与baseline相比，RaLMSpec 可以实现高达 7.59× 和 2.45× 的加速比。

方法简析：

一般的RaLM从知识库中查找与用户输入相似度最高的文档（它们都会被转化成向量进行相似性比较），但是由于知识库中的文档数量很多，查找所需的时间长，所以迭代RaLM会因为冗长的查找流程而遇到堪忧的性能问题。

对此，作者等人想到，可以推测一部分可能会用到的文献，提前存放在本地缓存中（这就是Speculative的地方），在本地缓存中查找文档与在知识库中查找文档的方法相似，但由于本地缓存的文档数远远小于知识库中的文档数，因此，查找会变得非常快速（这就是Speculative Retrieval）。

具体来说，RaLM有一个本地缓存，如果这个本地缓存是空的，就进行一次Initial Query，查找一些文档放入其中。随后查找的时候就只从这个缓存里面查找了。

但是本地缓存毕竟不是知识库（除非它有着和知识库几乎一样的文档，但是这样就没意义了），所以这样子查找会出错。作者们想到了一个法子，可以先查s个文档，然后再批量对这些文档进行验证（也就是从知识库批量对比这些文档是不是最相似的那些，这就是Batched Verification），如果发现某s个文档中的第i个文档出错了，就回滚到第i个文档，然后代之以正确的文档并将之加入到本地缓存里，随后再用新的文档生成新的prompt。这里的s被称为推测步长，而且s要取一个比较好的值，如果s太短，比如我们让s=1，那每从本地缓存里查出来一个文档，就要从知识库里查这个文档对不对，不管用不用这个方法，每查一次文档都要上一次知识库，那这岂不是白用这个方法了（笑）；如果s太长，假如s=114514，但s从第二个开始就错了，那后面的114512个文档就白查了。至于怎么取，后文有所介绍。

此外，检索步骤和验证步骤可以异步进行，这样效率会更高。

现在我们来说一说s是怎么实现的，作者们称这个方法为Optimal Speculation Stride Scheduler(OS³)，它能自动地确定最好的s。原理是这样的：

至于这个数学期望是怎么算出来的，就超出我的理解范围了（叹气）。

但是事情到这里还没有完，我们之前定义了三个参数，现在我们要讨论怎么估计它们了。对于a和b，可以直接依靠RaLM最近的表现来预测。但对于γ(X)，我们需要用到对数最大似然估计。作者们所用的式子是这样的：

无论如何，我们已经有办法确定最优步长s了。下面就是RaLMSpec的流程：

他们先进行3个Speculative Retrieval(i.e. q0, q1, q2, 所查询到的文档为A, B, A)，随后对这些文档进行批量验证，在这里，q3所对应的A出错了（其实应该是C），所以将之修改为C，并将C存入到本地缓存中，然后重新进行LM Decoding环节。更精确的说，这段过程的伪代码是这样的：

1: Input: Input tokens X = {x0, x1,... , xt−1}, external corpus C, language model f(·)
2: Output: RaLM generated outputs
3: Initialize local cache Q = {}, speculation stride s, model generation stride k
4: q = encode(X), Q.insert(C.retrieve(q)) ▷ cache prefetching
5: while EOS not in X do
6: for i = 1 to s do
7: qi = encode(X),ˆdi = Q.retrieve(qi) ▷ speculative retrieval
8: Xˆi = f(X, ˆdi, k) ▷ model generation step that generates k new tokens
9: X = [X, Xˆi]
10: end for
11: d1,... , ds = C.retrieve(q1,... , qs) ▷ batched verification
12: m = arg miniˆdi != di
13: if m ≤ s then ▷ do correction if needed
14: Roll X back to the m-th speculation step
15: Xˆ = f(X, di, k)
16: X = [X, Xˆ]
17: end if
18: end while

大体上来说，这段代码主要是在：初始化本地缓存，并进行初始预取。在主循环中，进行推测性检索和生成新的token。每经过s次推测后，进行批量验证。如有需要，进行校正，确保生成结果的准确性。直到生成结束符 (EOS)，生成过程结束。

评估：

实验设置如下：

•LLMs: GPT2-medium, OPT-1.3B, LLaMA-2-7B. 对于KNN-LM服务，采用了16 layers, decoder-only, 247M的LLM

•Datasets:

•四个 QA 数据集：Wiki-QA、Web Questions、Natural Question, Trivia-QA;

•外部知识库: Wikipedia Corpus;

•KNN-LM评估: Wikitext-003.

•Retrievers:

•Exact Dense Retrievers(EDR): Dense Passage Retriever(DPR);

•Approximate Dense Retriever(ADR): DPR-HNSW;

•Sparse Retriever(SR): BM25 Retriever.

•Baseline:

•Naïve Iterative RaLM: 使用RaLMSeq作为基准，在每个LLM生成4个token后执行检索;

•KNN-LM:对生成的每个token执行检索.

评估结果：

1. RaLMSpec对EDR提升效果最显著，OS³对ADR，SR的提升作用最大

2. EDR使用较大步长效果较好，但ADR则需要使用OS³。

3. EDR相较ADR和SR表现较好，在1.70x~1.85x.

The post RaLMSpec论文导读：使用推测方法加速RAG first appeared on 南极滑稽的博客.

*本文使用ChatGPT翻译。

原文标题：PaLM: Scaling Language Modeling with Pathways

作者：Aakanksha Chowdhery∗ Sharan Narang∗ Jacob Devlin∗
Maarten Bosma Gaurav Mishra Adam Roberts Paul Barham
Hyung Won Chung Charles Sutton Sebastian Gehrmann Parker Schuh Kensen Shi
Sasha Tsvyashchenko Joshua Maynez Abhishek Rao† Parker Barnes Yi Tay
Noam Shazeer‡ Vinodkumar Prabhakaran Emily Reif Nan Du Ben Hutchinson
Reiner Pope James Bradbury Jacob Austin Michael Isard Guy Gur-Ari
Pengcheng Yin Toju Duke Anselm Levskaya Sanjay Ghemawat Sunipa Dev
Henryk Michalewski Xavier Garcia Vedant Misra Kevin Robinson Liam Fedus
Denny Zhou Daphne Ippolito David Luan‡ Hyeontaek Lim Barret Zoph
Alexander Spiridonov Ryan Sepassi David Dohan Shivani Agrawal Mark Omernick
Andrew M. Dai Thanumalayan Sankaranarayana Pillai Marie Pellat Aitor Lewkowycz
Erica Moreira Rewon Child Oleksandr Polozov† Katherine Lee Zongwei Zhou
Xuezhi Wang Brennan Saeta Mark Diaz Orhan Firat Michele Catasta† Jason Wei
Kathy Meier-Hellstern Douglas Eck Jeff Dean Slav Petrov Noah Fiedel

1 引言

在过去几年中，针对语言理解和生成训练的非常大型神经网络在各种任务中取得了令人惊讶的成果。其中许多模型，如 BERT（Devlin 等，2019）和 T5（Raffel 等，2020），使用填充（“掩码语言模型”或“跨度破坏”）的预训练目标，在大型文本语料库上训练了仅有编码器或编码器-解码器结构，然后通过微调来适应一个或多个特定任务。尽管这些模型在成千上万的自然语言任务中几乎达到了普遍的最先进水平，但缺点是它们需要大量特定任务的训练示例来微调模型。此外，至少部分模型参数必须更新以适应任务，这增加了模型微调和部署的复杂性。

GPT-3（Brown 等，2020）证明了极其大型的自回归语言模型（LMs）可以用于少样本预测，其中模型只接收自然语言任务描述（可选）以及一些示例，演示任务应如何完成。这类模型使用仅有解码器的架构和标准的从左到右的语言建模目标在大型文本语料库上训练，其目标是在给定示例中的前面令牌的情况下预测下一个令牌。已经证明少样本评估可以在不需要大规模任务特定数据收集或模型参数更新的情况下实现非常强大的结果。

自 GPT-3 以来，已经开发了许多其他大型自回归语言模型，继续推动最新技术的进步。这些后 GPT-3 模型中最强大的包括 GLaM（Du 等，2021）、Gopher（Rae 等，2021）、Chinchilla（Hoffmann 等，2022）、Megatron-Turing NLG（Smith 等，2022）和 LaMDA（Thoppilan 等，2022），它们在发布时在大量任务中实现了少样本最先进的结果。与 GPT-3 一样，这些模型都是 Transformer 架构（Vaswani 等，2017）的变体。这些模型的改进主要来自以下一种或多种方法：（1）在深度和宽度上扩展模型的规模；（2）增加模型训练的令牌数量；（3）训练更干净的来自更多来源的数据集；以及（4）通过稀疏激活模块增加模型容量而不增加计算成本。

在这项工作中，我们延续了语言建模改进的规模化路线，并在 7800 亿标记的高质量文本上训练了一个拥有 5400 亿参数、密集激活的自回归 Transformer 模型。这是通过使用 Pathways（Barham 等，2022）实现的，Pathways 是一个新的 ML 系统，可以在数千个加速器芯片上实现非常大型神经网络的高效训练，包括跨多个张量处理单元（TPU）v4 Pod 的芯片。这个新模型，称为 Pathways 语言模型（PaLM），在数百个自然语言、代码和数学推理任务中取得了最先进的少样本结果。我们在其中一些任务中取得了突破性的表现。

这项工作的主要要点如下：

• 高效的扩展 - 我们展示了 Pathways（Barham 等，2022）的首次大规模应用，这是一种新的 ML 系统，可以以高效的方式跨数千个或数万个加速器芯片训练单个模型。利用 Pathways，我们在 6144 个 TPU v4 芯片上以前所未有的效率水平训练了一个 5400 亿参数的语言模型。大多数先前的大型语言模型要么在单个 TPU 系统上训练（Du 等，2021；Thoppilan 等，2022），要么使用管道并行性（Huang 等，2019）在 GPU 集群（Smith 等，2022）或多个 TPU v3 Pod（Rae 等，2021）上进行扩展，最大规模为 4096 个 TPU v3 芯片。在第 4 节中，我们描述了如何能够将 PaLM 540B 的无管道训练扩展到 6144 个芯片，跨越两个 TPU v4 Pod，并在模型 FLOPs 利用率方面实现了非常高的效率，达到了46.2%（相对于理论最大吞吐量的观察吞吐量）以及57.8%的硬件 FLOPs 利用率。

• 规模化带来的持续改进 - 在第 6 节中，我们评估了 PaLM 在数百个自然语言、代码和数学推理任务上的表现，并在绝大多数基准测试中取得了最先进的结果，通常差距很大。这个关键的证明表明，大型 LM 的扩展改进既没有达到平稳状态，也没有达到饱和点。例如，在表 4 中，我们提出了对 29 个最广泛评估的英语语言理解基准测试中的 28 个新的少样本最先进结果，与 GLaM（Du 等，2021）、GPT-3（Brown 等，2020）、Megatron-Turing NLG（Smith 等，2022）、Gopher（Rae 等，2021）、Chinchilla（Hoffmann 等，2022）和 LaMDA（Thoppilan 等，2022）每个任务的最佳结果进行比较。

• 突破性能力 - 我们在一些困难任务中展示了语言理解和生成的突破性能力。具体来说，第 6.3 节展示了对一系列推理任务的评估，这些任务需要多步数学或常识推理来产生正确答案。先前的最先进结果使用了任务特定的微调、领域特定的架构和任务特定的验证器来获得强大的结果。在这项工作中，我们证明了当模型规模与思维链提示（Wei 等，2022b）相结合时，简单的少样本评估可以在广泛的推理任务上超过或与微调状态下的最先进结果相匹配。在第 6.2 节中，我们还突出展示了在 BIG-bench（BIG-bench collaboration，2021）上的突破性能力，这是一个最近发布的包含 150 多个新的语言理解和生成任务的套件，其中许多任务即使对于人类来说也非常难以正确回答。在图 1 和第 9 节中，我们展示了 PaLM 在解释和说明复杂推理链方面的探索能力。

• 不连续的改进 - 为了更好地了解规模化行为，我们在三种不同的参数规模下提供了结果：8B、62B 和 540B。通常，从 62B 扩展到 540B 的效果与从 8B 扩展到 62B 的效果相似，这与神经网络扩展中常见的“幂律”规则一致（Kaplan 等，2020）。然而，对于某些任务，我们观察到不连续的改进，即从 62B 扩展到 540B 导致准确度与从 8B 扩展到 62B 相比大幅跃升。这种行为在第 6.2 节的大约 25% 的 BIG-bench 任务中观察到。这表明当模型达到足够的规模时，大型 LM 的新能力可能会出现，并且这些能力将持续出现超过先前研究的规模。

• 多语言理解 - 先前对大型语言模型的研究在多语言领域进行了有限的评估。在这项工作中，我们对多语言基准进行了更全面的评估，包括机器翻译（第 6.5 节）、摘要（第 6.6 节）和问答（第 6.7 节），涉及多种语言。即使在训练语料库中非英语数据的比例相对较小（约为 22%），540B 模型的少样本评估结果也能弥合与以前的微调最先进在非英语摘要任务中的差距，并在翻译任务中胜过以前的最先进水平。进一步的工作是必要的，以了解增加多语言数据比例对英语和多语言任务的影响。

• 偏见和毒性 - 我们还评估了模型在分布偏见和毒性方面的表现，得出了几个见解（第 10 节）。首先，在性别和职业偏见方面，我们发现 Winogender 共指任务的准确性随着模型规模的增加而提高，而 PaLM 540B 在 1 次射击和少样本设置下创造了一个新的最先进结果。其次，在种族/宗教/性别提示延续的共现分析中，我们发现模型可能错误地肯定了刻板印象，例如，将穆斯林与恐怖主义、极端主义和暴力联系起来。这种行为在模型规模上保持一致。最后，在提示延续任务的毒性分析中，相对于 8B 模型，62B 和 540B 模型的总体毒性水平略高。然而，模型生成的延续的毒性与提示文本的毒性高度相关，而人类生成的延续则没有强烈的毒性相关性。这表明模型更受提示样式的影响，而不像人类生成的文本那样。在未来的工作中，我们计划将这些基准扩展到非英语语言，并更全面地考虑潜在风险。

2 模型架构

PaLM 使用了标准的 Transformer 模型架构（Vaswani 等，2017），采用了仅有解码器（Decorder-Only）的设置（即，每个时间步只能注意到自身和过去的时间步），具有以下修改：

• SwiGLU 激活函数 - 我们使用 SwiGLU 激活函数（Swish(xW) · xV），因为已经证明与标准的 ReLU、GeLU 或 Swish 激活函数相比，它们可以显著提高质量（Shazeer，2020）。需要注意的是，这在 MLP 中需要三次矩阵乘法，而不是两次，但 Shazeer（2020）在等价计算实验中展示了质量的提高（即，标准的 ReLU 变体具有比例上更大的维度）。

• 并行层 - 我们在每个 Transformer 块中使用“并行”形式（Wang & Komatsuzaki，2021），而不是标准的“串行”形式。具体来说，标准形式可以写成：

y = x + MLP(LayerNorm(x + Attention(LayerNorm(x)))

而并行形式可以写成：

y = x + MLP(LayerNorm(x)) + Attention(LayerNorm(x))

并行形式在大规模下训练速度大约快了 15%，因为 MLP 和 Attention 输入矩阵乘法可以合并。消融实验表明在 8B 规模下存在一定程度的质量下降，但在 62B 规模下没有质量下降，因此我们推断并行层的影响在 540B 规模下应该是质量中性的。

• 多查询注意力 - 标准的 Transformer 公式使用 k 个注意力头，其中每个时间步的输入向量被线性投影到形状为 [k, h] 的“查询”、“键”和“值”张量中，其中 h 是注意力头的大小。在这里，键/值投影被共享到每个头部，即“键”和“值”被投影到 [1, h] 的形状，但“查询”仍然被投影到形状 [k, h]。我们发现这对模型质量和训练速度没有影响（Shazeer，2019），但在自回归解码时间上节省了显著的成本。这是因为标准的多头注意力在自回归解码期间在加速器硬件上效率低下，因为键/值张量在例子之间不共享，并且每次只解码一个令牌。

• RoPE 嵌入 - 我们使用 RoPE 嵌入（Su 等，2021），而不是绝对或相对位置嵌入，因为 RoPE 嵌入已经被证明在长序列长度上具有更好的性能。

• 共享输入-输出嵌入 - 我们共享输入和输出嵌入矩阵，这在过去的工作中经常（但不是普遍）这样做。

• 无偏见 - 在任何密集核或层归一化中都没有使用偏见。我们发现这样做可以提高大型模型的训练稳定性。

• 词汇表 - 我们使用了一个具有 256k 个标记的 SentencePiece（Kudo＆Richardson，2018a）词汇表，这个选择是为了支持训练语料库中众多语言的需求，而不会产生过度的标记化。该词汇表是从训练数据中生成的，我们发现这样做可以提高训练效率。该词汇表完全无损和可逆，这意味着词汇表中的空白符完全保留（对于代码特别重要），并且超出词汇表的 Unicode 字符被拆分为 UTF-8 字节，每个字节都有一个词汇表标记。数字总是被拆分为单个数字标记（例如，“123.5 → 1 2 3 . 5”）。

2.1 模型规模超参数（Model Scale Hyperparameters）

在这项工作中，我们比较了三种不同的模型规模：540B 参数、62B 参数和8B 参数。每个标记的 FLOPs 数量大约等于参数数量，因为这些模型是标准的密集 Transformer。这些模型是使用表格 1 中的超参数构建的。这三个模型在相同的数据和词汇表下进行了相同（除了批量大小）的训练。训练过程在第 3 和第 5 节中有更详细的描述。

2.2 模型卡片

PaLM 的模型卡片（Mitchell等人，2019）提供在附录 E 中。这提供了对模型架构、训练设置、训练数据和预期用途的高层次摘要。

3 训练数据集

PaLM 的预训练数据集包括一个高质量的语料库，共包含 7800 亿个标记，代表了广泛的自然语言使用情况。该数据集是由经过筛选的网页、书籍、维基百科、新闻文章、源代码和社交媒体对话的混合物组成。这个数据集基于用于训练 LaMDA（Thoppilan等人，2022）和 GLaM（Du等人，2021）的数据集。我们对所有三个模型都训练了完全相同的数据集（对所有模型进行了相同的洗牌），并选择了混合比例，以避免在任何子组件中重复数据。

除了自然语言数据之外，预训练数据集还包含代码。预训练数据集中的源代码来自 GitHub 上的开源代码库。我们通过存储库中包含的许可证对文件进行了筛选； copyleft 许可证被排除在外。我们通过文件名扩展名对文件进行过滤，以限制到 24 种常见编程语言之一，包括 Java、HTML、Javascript、Python、PHP、C#、XML、C++ 和 C，从而获得了 196GB 的源代码。此外，我们根据文件之间的 Levenshtein 距离删除重复项，因为在源代码存储库中已知存在大量重复文件（Lopes等人，2017；Allamanis，2019）。

表 2 列出了用于创建最终 PaLM 数据集混合物的各种数据源的比例。我们在第 8 节中检查数据污染，并展示了我们的训练数据集与评估数据之间的重叠分析。附录 D 包含了一份数据表（Gebru等人，2021），其中包含更多信息，包括语言比例的详细分析。

4 训练的基础设施

我们的训练和评估代码基于 JAX（Bradbury等人，2018）和 T5X（Roberts等人，2022），所有模型都在 TPU v4 Pods（Jouppi等人，2020）上进行训练。PaLM 540B 是通过两个通过数据中心网络（DCN）连接的 TPU v4 Pod（Xu等人，2021）上进行训练的，使用了模型和数据并行的组合。每个 Pod 上连接了 3072 个 TPU v4 芯片，连接了 768 个主机。这个系统是迄今为止描述的最大 TPU 配置，使我们能够有效地将训练扩展到 6144 个芯片，而无需使用任何管道并行性（Huang等人，2019）。

在可比较规模的端到端模型训练的先前报告中，有两种方法。LaMDA（Thoppilan等人，2022）和 GLaM（Du等人，2021）分别在单个 TPU 系统上进行训练，而没有利用管道并行性或 DCN。Megatron-Turing NLG 530B（Smith等人，2022）使用模型、数据和管道并行性的组合，在 2240 个 A100 GPU 上进行训练，而 Gopher（Rae等人，2021）则在四个通过 DCN 连接的 TPU v3 Pod（每个具有 1024 个 TPU v3 芯片）上进行训练，并在 Pod 之间使用流水线。

流水线通常与 DCN 一起使用（Smith等人，2022），因为它具有较低的带宽要求，并且提供了超出模型和数据并行性所能承受的最大有效规模的额外并行化。流水线通常将训练批次分成“微批次”，但它也具有重要的缺点。首先，在流水线“气泡”的步骤时间开销中，许多设备在填充和排空管道时完全闲置，这是前向和后向传递的开始和结束。其次，由于要为每个微批次中的小批次重新加载权重，它要求更高的内存带宽。在某些情况下，它还涉及增加的软件复杂性。我们能够使用以下策略有效地将 PaLM 540B 的训练扩展到 6144 个芯片而无需使用流水线。

每个 TPU v4 Pod 包含模型参数的完整副本，每个权重张量都使用 12 路模型并行和 256 路完全分片数据并行（称为“2D finalized”方法）进行划分（Xu等人，2021）。在前向传递期间，权重在数据并行轴上进行全收集，并且从每个层保存一个完全分片的激活张量。在反向传递期间，其余的激活被重新生成，因为与备选重新计算选择相比，这会在更大的批次大小时提高训练吞吐量。

我们使用 Pathways 系统（Barham等人，2022）将训练扩展到超出单个 TPU v4 Pod 的范围。PaLM 540B 利用 Pathways 的客户端-服务器架构，在 Pod 级别实现双向数据并行。一个单独的 Python 客户端将一半的训练批次分发给每个 Pod，每个 Pod 使用标准的 Pod 内数据和模型并行计算前向和反向计算以并行计算梯度。然后，Pod 传输（计算其一半批次上的）梯度与远程 Pod，最后，每个 Pod 积累本地和远程梯度，并并行应用参数更新，以获得下一个时间步的位相同的参数。

图2显示了Pathways系统如何执行双向Pod级数据并行。单个Python客户端构建了一个分片数据流程序（在图2左侧显示），该程序在每个包含TPU Pod的远程服务器上启动JAX/XLA工作。该程序包含一个用于Pod内前向+后向计算的组件A（包括Pod内梯度归约）、用于跨Pod梯度传输的转移子图，以及一个用于优化器更新的组件B（包括本地和远程梯度的求和）。Pathways程序在每个Pod上执行组件A，然后将输出梯度传输到另一个Pod，最后在每个Pod上执行组件B。Pathways系统设计具有几个特点，使其能够将程序执行扩展到数千个加速器芯片——首先，它通过Pod级调度器（在图2右侧显示）的异步组调度掩盖了从单个Python客户端向远程服务器分发JAX/XLA工作的延迟，其次，通过分片数据流执行模型摊销了管理数据传输的成本（有关详细信息，请参阅Barham等人（2022年））。

双向Pod级数据并行的一个有趣方面是在连接到两个Pod的总共1536个主机上的6144个TPU v4芯片的规模下实现高训练吞吐量的挑战。请注意，跨Pod梯度传输仅需要在两个Pod上对应的主机之间进行1：1的传输，因为每个核心只需要其模型分片参数的远程梯度。此外，两个Pod之间的主机通过谷歌数据中心网络连接。由于传输直到每个核心完成梯度计算才开始（如图2所示），这导致了一个非常突发的工作负载，其中所有主机在数据中心网络链路上同时传输其梯度。特别地，每对主机在每个训练步骤中交换大约1.3GB的梯度，总共产生了81 Tbps的聚合突发。这种工作负载的突发特性引入了挑战，我们通过精心设计Pathways网络堆栈来解决这些挑战，以实现最佳的数据中心网络链路利用率。例如，为了减轻拥塞的影响，梯度传输的数据被分解为较小的块，并通过多个较小的流路由到多样化的数据中心网络链路。通过这些优化，我们实现了相对于训练时单个Pod的约1.95倍的训练吞吐量（相对于单个Pod，我们将批处理大小在两个Pod之间翻了一番，相当于理论上的完美弱可伸缩性的97%）。与理论上的2倍吞吐量相比，性能差距是由于向后传递和跨Pod梯度归约之间缺乏重叠造成的。我们期望在未来的工作中解决这个问题。

4.1 训练效率

先前报告的大多数有关加速器效率的数字都使用了我们称之为硬件 FLOPs 利用率（HFU）的指标。这通常反映了在给定设备上观察到的 FLOPs 数与其理论峰值 FLOPs 的比率的估计。然而，硬件 FLOPs 利用率存在几个问题。首先，执行的硬件 FLOPs 数量取决于系统和实现，编译器中的设计选择可能导致不同数量的操作。重生成是一种广泛使用的技术，用于在内存使用和计算之间进行权衡。为了有效地计算大多数神经网络架构的反向传播，必须在内存中存储批处理的许多中间激活。如果无法全部容纳，可以重新计算一些前向传播操作（使一些激活重新生成而不是存储）。这创建了一个权衡，使用额外的硬件 FLOPs 可以节省内存，但培训系统的最终目标是实现每秒标记的高吞吐量（因此是快速培训），而不是尽可能多地使用硬件 FLOPs。其次，测量观察到的硬件 FLOPs 取决于用于计数或跟踪它们的方法。已根据分析会计（Narayanan等人，2021b）以及使用硬件性能计数器（Xu等人，2021）报告了观察到的硬件 FLOPs。

鉴于这些问题，我们认识到 HFU 不是 LLM 训练效率的一致且有意义的度量标准。我们提出了一种新的效率度量标准，即模型 FLOPs 利用率（MFU），该度量标准与实现无关，并允许对系统效率进行更清洁的比较。MFU 是观察吞吐量（每秒标记数）与在峰值 FLOPs 下操作的理论最大吞吐量之间的比率。至关重要的是，“理论最大”吞吐量仅考虑了计算前向+后向传播所需的操作，而不考虑重生成。因此，MFU 允许在不同系统上进行训练运行的公平比较，因为分子仅仅是观察到的每秒标记数，而分母仅取决于模型架构和给定系统的已发布最大 FLOPs。我们在附录 B 中详细阐述了计算 MFU 的数学公式。

我们在表3中展示了PaLM 540B模型的模型FLOPs利用率，并将其与先前的大型模型进行了对比。在不同模型参数数量、架构和模型质量的背景下，MFU对比模型和系统非常有用。根据OpenAI通过Patterson等人（2021）报告的数据，GPT-3的MFU数字为21.3%，基于每个GPU的24.6个非注意力模型TFLOP/s，而Gopher的MFU数字为32.5%，基于每秒0.0152步的训练速度。Megatron–Turing NLG 530B的MFU数字是29.7%，没有自注意力，或者是30.2%，有自注意力，基于65.43K个tokens/秒的训练吞吐量（Smith等人，2022）。 相比之下，PaLM 540B在批量大小为2048时实现了平均训练吞吐量为238.3K个tokens/秒。PaLM 540B的训练使用了重新计算技术，因为使用重新计算可以实现更高的训练吞吐量。PaLM 540B的MFU是45.7%，没有自注意力，或者是46.2%，有自注意力。我们通过分析计算出的硬件FLOPs利用率，包括重新计算的FLOPs，为57.8%。PaLM通过其并行策略和其他几个因素，包括XLA TPU编译器优化和“并行层”的使用（见第2节），实现了高加速器利用率。我们认为PaLM在LLM训练效率方面代表了重要的进步。

5 训练设置

模型训练遵循了大型Transformer语言模型的标准设置。具体细节如下：

• 权重初始化 - 内核权重（即除了嵌入和层归一化尺度之外的所有内容）采用“输入方差缩放”进行初始化，即W ∼ N（0, 1/sqrt(n_in）)，其中n_in是内核的输入维度。输入嵌入被初始化为E ∼ N（0, 1），因为嵌入层归一化不适用于嵌入。由于输入和输出嵌入层是共享的，因此我们通过1/sqrt(n)对预softmax输出对数进行缩放，其中n是嵌入的大小。

• 优化器 - 模型使用Adafactor优化器（Shazeer＆Stern，2018）进行训练，没有分解。这实际上相当于Adam（Kingma＆Ba，2014）与“参数缩放”的等效，它通过参数矩阵的均方根来缩放学习率。由于权重初始化与1/sqrt(n)成正比，这种效果与Brown等人（2020）中手动缩放Adam学习率的效果类似。然而，参数缩放的好处在于，操作在不同尺度上运行的参数矩阵（嵌入和层归一化尺度）的学习率不会以相同的速率缩小。

• 优化超参数 - 我们使用Adafactor学习率为10-2进行前10000步的训练，然后以1/sqrt(k)的速率进行衰减，其中k是步数。我们使用β₁ = 0.9的动量。第二阶矩插值值计算为β₂ = 1.0 - k-0.8，其中k是步数。我们发现，在训练大型语言模型时，这比标准β₂ = 0.99更稳定，因为稀有嵌入标记可能在较短窗口内估计第二时刻时有很大的误差。我们在训练过程中使用全局范数梯度裁剪（Pascanu等人（2012）），对于所有模型，值为1.0。我们在训练期间使用动态权重衰减lr^2.0，其中lr是当前学习率。

• 损失函数 - 模型使用标准的语言建模损失函数进行训练，即所有标记的平均对数概率，不使用标签平滑。我们另外使用辅助损失z损失= 10^-4·log²Z，以鼓励softmax归一化器log（Z）接近0，我们发现这样增加了训练的稳定性。

• 序列长度 - 所有模型的序列长度都为2048。输入示例被连接在一起，然后被分割成确切的2048个标记的序列，以便没有填充标记，但是示例可能被分割在中间。输入示例通过特殊的[eod]标记与彼此区分开来。

• 批量大小 - 对于所有模型，我们在训练期间增加批量大小。对于最大的模型，我们在50k步时使用批量大小512（100万个标记），然后将其加倍到115k步时的1024（200万个标记），最后再次加倍到255k步时的2048（400万个标记）直到训练完成。较小的模型遵循类似的计划。使用这种批量大小计划的原因有两个：（1）在训练初期，较小的批量大小更有效率（即看到的标记的损失更好），而在训练后期，较大的批量大小由于更好的梯度估计（Smith等人，2018；McCandlish等人，2018）而受益，（2）较大的批量大小导致更大的矩阵乘法维度，从而增加了TPU的效率。

• 位确定性 - 模型可以从任何检查点完全位确定地再现。换句话说，如果模型在单次运行中训练到17,000步，然后我们从检查点15,000重新开始，那么训练框架保证在检查点15,000到17,000之间的两次运行中产生相同的结果。这通过两种方式实现：（1）由JAX + XLA + T5X提供的位确定性建模框架，以及（2）确定性数据集管道，其中打乱的数据以随机访问格式写入，因此给定训练批次的内容仅取决于步数（Roberts等人，2022）。 •

退出 - 模型在训练时没有使用退出，尽管在大多数情况下，微调使用0.1的退出。

5.1 训练不稳定性

对于最大的模型，在训练过程中我们大约观察到了20次损失的尖峰，尽管梯度裁剪已经启用。这些尖峰发生在高度不规则的间隔内，有时发生在训练后期，并且在训练较小的模型时并未观察到。由于训练最大模型的成本，我们无法确定一个原则性的策略来缓解这些尖峰。

相反，我们发现了一个简单有效地缓解问题的策略：我们从尖峰开始前大约100步的检查点重新开始训练，并跳过大约200-500个数据批次，其中包括在尖峰之前和之后看到的批次。通过这种缓解，损失在相同的时间点不再出现尖峰。我们不认为这些尖峰本质上是由“坏数据”引起的，因为我们进行了几项消融实验，其中我们取了围绕尖峰的数据批次，然后从不同的、更早的检查点开始对这些相同的数据批次进行训练。在这些情况下，我们并没有观察到尖峰。这意味着尖峰只会在特定的数据批次与特定的模型参数状态的组合下发生。在未来，我们计划研究更加原则性的策略来缓解非常大型语言模型中损失尖峰的问题。

6 评估

6.1 英语自然语言处理任务

为了与之前的大型语言模型进行比较，我们在与 Du 等人 (2021) 和 Brown 等人 (2020) 相同的 29 个英语基准测试上评估 PaLM 模型。这些基准测试包括：

• 开放领域闭卷问答任务： TriviaQA (Joshi 等人，2017)，自然问题 (Kwiatkowski 等人，2019)，Web 问题 (Berant 等人，2013) • 填空和补全任务：LAMBADA (Paperno 等人，2016)，HellaSwag (Zellers 等人，2019)，StoryCloze (Mostafazadeh 等人，2016)

• Winograd 风格任务：Winograd (Levesque 等人，2012)，WinoGrande (Sakaguchi 等人，2020) • 常识推理：PIQA (Bisk 等人，2019)，ARC (Clark 等人，2018)，OpenBookQA (Mihaylov 等人，2018) • 上下文阅读理解：DROP (Dua 等人，2019)，CoQA (Reddy 等人，2018)，QuAC (Choi 等人，2018)，SQuADv2 (Rajpurkar 等人，2018)，RACE (Lai 等人，2017) • SuperGLUE (Wang 等人，2019a)

• 自然语言推理 (NLI)：对抗性 NLI (Nie 等人，2019)

表 4 包括了 PaLM 540B 的结果以及其他大型语言模型的先前最优结果。在这个表格中，我们只考虑预训练语言模型的单一检查点结果。任何使用微调或多任务适应 (Wei 等人 (2022a)，Sanh 等人 (2021)) 的模型都不包含在表格中。

在 1-shot 设置中，PaLM 540B 在 29 项任务中有 24 项超越了先前的最优结果，在 few-shot 设置中有 28 项任务超越了先前的最优结果。有趣的是，PaLM 540B 在一些阅读理解和 NLI 任务的 few-shot 设置中超过先前的最优结果超过 10 分。虽然模型大小在实现这些结果中起到了重要作用，但是 PaLM 540B 在所有基准测试中均超过了相似规模的模型 (Megatron-Turing NLG 530B)。这表明，预训练数据集、训练策略以及训练过程中观察到的令牌数量也在实现这些结果中起着重要作用。

表5列出了自然语言理解（NLU）和自然语言生成（NLG）任务的平均得分。PaLM 540B 在两个类别的平均得分上都提高了超过5分。正如表中所示，每个类别的平均得分也随着PaLM模型的规模而提高。有趣的是，PaLM 62B 在两个类别中均超过了GPT-3 175B。8B和62B模型的所有任务结果可以在附录H.1中找到。

6.1.1 大规模多任务语言理解

我们还评估了PaLM模型在大规模多任务语言理解（MMLU）（Hendrycks等人，2021年）基准测试上的表现。这个多项选择问答基准测试涵盖了57个不同主题的任务，包括数学、历史和法律等。如表6所示，PaLM 540B将MMLU基准测试的平均得分提高了约2个点。PaLM 540B在除了其他任务类别之外的所有类别上都优于Chinchilla模型。

表7报告了在超级粘合任务的微调中使用任务比例混合的验证结果。在超级粘合任务上，我们与T5-11B（Raffel等人，2020年）和ST-MoE-32B（Zoph等人，2022年）等最先进的模型进行比较，并表明PaLM获得了竞争性的接近最先进水平的性能。值得注意的是，超级粘合的两个表现最佳模型都是使用跨度破坏目标训练的编码器-解码器模型。已经证明，当训练成本相等时，这种架构通常会在分类任务微调上优于仅解码器的自回归模型（Raffel等人，2020年）。这些结果表明，规模可以帮助弥合差距。表8还表明，少样本微调结果与微调结果之间仍然存在显着差距。最后，表9报告了超级粘合排行榜测试集上的结果。我们表明，PaLM在与最先进技术的竞争中表现出色，同时在排行榜上优于最佳的仅解码器自回归语言模型。

BIG-bench是一个合作基准，旨在为大型语言模型提供具有挑战性的任务（BIG-bench合作，2021年）。它包括150多个任务，涵盖了各种语言建模任务，包括逻辑推理、翻译、问答、数学等等。在本节中，我们介绍和分析了PaLM模型系列在BIG-bench上的少样本评估结果。BIG-bench包括文本任务和程序任务。在这次评估中，只考虑了文本任务。BIG-bench数据发布包括人类表现度量标准，其中在众包平台上的若干工作者（通常为10人）被要求解决每个任务。这些工作者被给予与少样本模型看到的类似的任务描述和示例，并且被允许使用任何外部工具来提供答案（搜索引擎、计算器等）。模型和人类的表现都使用相同的指标针对黄金标签进行度量，这些标签由每个任务的原始创建者提供。通过使用具有最高每个示例指标分数的人类生成的答案计算“最佳”人类表现，而通过使用所有人类生成的答案的平均指标分数来计算“平均”人类表现。

图3左显示了对PaLM模型系列在BIG-bench上进行评估的结果，与先前发表的结果进行了比较（Rae等人，2021年）。请注意，由于BIG-bench数据集最近发布，并且先前的模型仅在部分任务和样本上进行了基准测试，因此与其他更成熟的基准相比，比较更为有限。在58个常见任务上评估了三个模型，因此此图仅呈现了这58个任务的结果。我们看到，PaLM显著优于GPT-3、Gopher和Chinchilla，而5-shot PaLM 540B在相同任务上的平均得分高于人类被要求解决相同任务的平均得分。在58个常见任务中，PaLM 540B 5-shot优于先前的SOTA的44个任务，具体任务结果如图4所示。此外，PaLM模型的性能随规模的增加似乎遵循对数线性行为，表明进一步扩展可能会导致性能提升。图3右显示了PaLM在BIG-bench文本任务集合（150个任务）上的结果，其性能特征类似。

接下来，我们重点介绍了几个来自BIG-bench的任务，其中PaLM表现出特别有趣的性能特征。这些任务的个别评估结果显示在图5中。我们将简要描述以下每个任务：

• goal step wikihow - 目标是推理事件之间的目标步骤关系。例如：

输入：为了“清洁银器”，应该首先执行哪个步骤？ (a) 将银器擦干 (b) 用手洗银器

答案：(b) 用手洗银器

• logical args - 目标是从段落中预测正确的逻辑推理。例如：

输入：学生告诉代课老师他们正在学三角函数。代课老师告诉他们，与其教他们关于三角形的无用知识，他宁愿教他们如何处理概率。他在暗示什么？(a) 他认为数学不需要有用才有趣。(b) 他认为理解概率比学三角函数更有用。(c) 他认为概率理论是一个无用的学科。

答案：(b) 他认为理解概率比学三角函数更有用。

• english proverbs - 目标是猜测哪个谚语最能描述文本段落。例如：

输入：Vanessa在当地无家可归救助中心周末帮助了许多年。最近，当她失去工作时，中心立即准备给她提供一份新工作。以下哪个谚语最适合这种情况？(a) 诅咒，就像鸡一样回家了。(b) 哪里有烟就有火。(c) 种瓜得瓜，种豆得豆。

答案：(c) 种瓜得瓜，种豆得豆。

• logical sequence - 目标是将一组“事物”（月份、行动、数字、字母等）按照其逻辑顺序排列。例如：

输入：以下哪个列表按时间顺序正确排序？(a) 喝水、感到口渴、封住水瓶、打开水瓶 (b) 感到口渴、打开水瓶、喝水、封住水瓶 (c) 封住水瓶、打开水瓶、喝水、感到口渴

答案：(b) 感到口渴、打开水瓶、喝水、封住水瓶

• navigate - 目标是按照一组简单的导航指示前进，并弄清楚最终会到达何处。例如：

输入：如果按照这些说明，您会回到起点吗？始终朝前。向左走6步。向前走7步。向左走8步。向左走7步。向前走6步。向前走1步。向前走4步。

答案：不会

• mathematical induction - 目标是执行逻辑推理数学归纳规则，即使它们与现实世界的数学相矛盾。例如：

输入：已知将2添加到任何奇整数会产生另一个奇整数。2是一个奇整数。因此，6是一个奇整数。这是一个正确的归纳论证吗（尽管一些假设可能是错误的）？

答案：是的

从图5中可以看出，goal step wikihow 和 logical args 的表现遵循对数线性的缩放曲线，PaLM 540B 模型的准确性接近最佳人类表现。而英语谚语和逻辑顺序的表现也非常强大，但它们遵循的是不连续改进曲线——从 62B → 540B 的改进远远大于从 8b → 62B 的改进。这些任务在这里特别引人关注，因为这样的缩放曲线意味着模型的某些能力只有在达到一定规模后才会显现出来。例如，英语谚语需要非常高水平的抽象推理能力来理解复杂的比喻，因此，从 PaLM 62B 的 25% 提高到 PaLM 540B 的 87% 是一个非常令人兴奋的结果。

为了进一步量化 BIG-bench 中不连续的改进，我们首先用一个例子来解释我们对不连续性的定义。逻辑顺序任务的标准化准确率分别为 PaLM 8b、62b 和 540b 的 13%、25% 和 87%。因此，使用 8b → 62b 的对数线性预测，预计 540b 的准确率大约为 37%（= 25% +（25% - 13%））。实际准确率为 87%，所以不连续性为 +50%（= 87% - 37%）。在所有 150 个任务中，25% 的任务的不连续性大于 +10%，15% 的任务的不连续性大于 +20%。这表明，从规模上来看，不连续的改进是在具有挑战性的少样本语言任务中的一种常见现象。

然而，并非所有任务都受益于规模的扩大。对于 navigate 和 mathematical induction 任务，PaLM 540B 仅略微优于 PaLM 62B，而且两者距离最佳人类表现还有很大差距。这表明任务在示例级难度上存在很高的方差。例如，在数学归纳任务中，一些示例具有正确的假设（例如，“2 是一个偶数。”），而一些示例具有不正确的假设（例如，“2 是一个奇数。”）。通过分析这个任务的几个示例，似乎所有的 PaLM 模型（以及被要求解决任务的人类）在假设不正确的情况下都很难，尽管指令说明假设的正确性并不重要。

在图 6 中，我们展示了比较 PaLM 540B 与人类评估的平均表现分数时任务改进的分布情况。我们可以看到，尽管 PaLM 540B 在总体上优于平均人类表现，但在 35% 的单个任务中，平均人类表现仍然高于 PaLM 540B（请参见附录中的表 43 中的示例）。这表明在 BIG-bench 上仍有很大的改进空间。

我们考虑了几个任务，其中 PaLM 540B 的表现超过了平均人类表现（请参见附录中表 42 的详细结果）。我们观察到其中几个任务展示了 PaLM 在许多语言中表现良好的能力，例如，persian idioms 和 swedish to german proverbs，评估这些任务的人类池可能对所有语言都不太熟悉。其中一个任务，periodic elements，依赖于大型语言模型的记忆能力。大多数其他任务，如 common morpheme、sufficient information 和 logical args，强调了 PaLM 540B 令人印象深刻的自然语言处理能力。为了进一步说明这一点，我们考虑了因果关系任务，该任务要求模型确定两个呈现的事件中的哪一个导致了另一个事件。该任务有两个版本：

• 因果关系（一句话无提示） - 在一句话无提示的子任务中，事件以两种不同的顺序合并成一句话，并使用模型对每个句子的对数似然进行评分。不提供提示。例如：

输入 A：I washed the car because my car got dirty.

输入 B：My car got dirty because I washed the car.

更高似然的句子：I washed the car because my car got dirty.

• 因果关系（两句话） - 在两句话的子任务中，模型展示两个事件，需要选择哪个句子对应于导致另一个事件的事件。例如：

输入：对于每个示例，给出两个事件。哪个事件导致了另一个事件？(a) 我的车变脏了。(b) 我洗了车。

正确预测：(a) 我的车变脏了。

我们发现所有的 PaLM 模型在这个任务的一句话无提示版本上表现良好，其中 8B 模型的准确率达到了80%以上，但较小的 PaLM 模型在这个任务的两句话版本上表现不佳，8B 模型的得分接近随机猜测。相比之下，540B 模型能够解决两句话版本，并且准确率达到了90%以上，展示了规模可以释放的通用语言建模能力。

最后，图 7 展示了 BIG-bench Lite 的详细评估结果，这是 BIG-bench 的一个精选子集，包含 24 个任务，作为轻量级评估目标。尽管其中一些 BIG-bench Lite 任务已经解决或接近解决，但与人类评估的最佳表现分数相比，其他任务距离解决仍有很大差距。

我们采取了几项措施来确立这些结果的有效性，特别是排除模型通过记忆 BIG-bench 数据来达到这些结果的可能性。首先，BIG-bench 任务文件包含一个唯一的标记字符串；我们确保该字符串不会出现在 PaLM 的训练数据中。其次，BIG-bench 数据集在训练数据收集时并未在互联网上公开，而绝大多数 BIG-bench 任务都是由任务作者专门为 BIG-bench 构建的全新基准。最后，我们在几个模型表现强劲的任务上进行了抽查，手动验证了解码过程中未泄露任何来自金标签的信息。

The post 论文翻译：PaLM: 使用Pathways扩大语言模型(1) first appeared on 南极滑稽的博客.