作者/机构: 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 (Google Research)

A1 主要贡献

本文介绍了一个5400亿参数、密集激活的Transformer语言模型,名为Pathways Language Model(PaLM)。该模型旨在探索模型规模对小样本学习(few-shot learning)的影响。PaLM的训练利用了Pathways,这是一个新型的机器学习系统,能够在数千个加速器芯片上高效地训练单一模型。研究的核心贡献和关键要点如下:

  • 高效扩展:本文展示了Pathways系统的首次大规模应用,该系统能够高效地跨越数千甚至上万个加速器芯片训练单个模型。利用Pathways,作者在6144个TPU v4芯片上训练了PaLM 540B模型,达到了前所未有的模型FLOPs利用率(46.2%)和硬件FLOPs利用率(57.8%)。与以往依赖流水线并行(pipeline parallelism)的方法不同,PaLM在两个TPU v4 Pod上实现了无流水线的训练扩展。

  • 持续的规模化效益:PaLM在数百个自然语言、代码和数学推理任务上进行了评估,并在绝大多数基准上取得了最先进的小样本学习结果。这表明大规模语言模型带来的性能提升尚未达到瓶颈。例如,在29个广泛评估的英语语言理解基准中,PaLM在28个任务上取得了新的小样本SOTA。

  • 突破性能力:PaLM在一系列困难任务上展现了突破性的语言理解和生成能力。特别是在需要多步数学或常识推理的任务上,当模型规模与思维链提示(chain-of-thought prompting)相结合时,PaLM能够通过简单的小样本评估超越或匹配经过微调的SOTA模型。此外,在BIG-bench这个包含150多个极具挑战性任务的基准上,PaLM的性能也超越了人类平均水平。

  • 非连续性改进:通过比较8B、62B和540B三种模型规模,研究发现虽然性能通常随规模呈对数线性增长,但在约25%的BIG-bench任务中,从62B扩展到540B带来了性能的急剧跃升。这表明当模型达到足够大的规模时,新的能力会涌现出来,并且这些能力在先前研究的规模之外仍在不断出现。

  • 多语言理解:尽管训练语料中非英语数据的比例相对较小(约22%),PaLM 540B在多语言基准测试(如机器翻译、摘要和问答)中表现出色。它在非英语摘要任务中缩小了与微调SOTA的差距,并在翻译任务中超越了先前的SOTA。

  • 偏见与毒性分析:研究对模型的分布偏差和毒性进行了评估。结果发现,随着模型规模的扩大,Winogender共指任务的准确性得到提高。然而,分析也揭示了模型可能强化刻板印象的风险,例如将穆斯林与恐怖主义联系起来。毒性分析表明,较大模型的毒性水平略高于较小模型,且模型生成内容的毒性与提示文本的毒性高度相关,这与人类的生成行为不同。

A2 方法细节

2. 模型架构

PaLM的基础架构与修改 PaLM采用标准的Transformer模型架构【索引158,Attention is all you need,2017,CoRR】,并在仅解码器(decoder-only)的设置下进行,即每个时间步只能关注自身及之前的时间步。在此基础上,PaLM进行了一些修改。

  • SwiGLU激活函数 模型在MLP(多层感知器)的中间激活层使用了SwiGLU激活函数($Swish(xW) \cdot xV$),因为研究表明它比标准的ReLU、GeLU或Swish激活函数能显著提高质量【索引144,GLU variants improve transformer,2020,CoRR】。虽然这使得MLP中的矩阵乘法从两次增加到三次,但Shazeer(2020)的研究证明,在计算等效的实验中(即标准ReLU变体具有按比例增大的维度),质量仍然有所提升。

  • 并行层 在每个Transformer块中,模型采用了“并行”(parallel)公式【索引159,GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model,2021,GitHub】,而不是标准的“串行”(serialized)公式。具体来说,标准公式可以写作:


    而并行公式可以写作:

    并行公式在大规模训练时能将训练速度提升约15%,因为MLP和注意力层的输入矩阵乘法可以融合。消融实验显示,在8B规模下质量有轻微下降,但在62B规模下没有质量下降,因此推断并行层在540B规模下对质量的影响应该是中性的。

  • 多查询注意力机制 标准的Transformer公式使用k个注意力头,其中每个时间步的输入向量被线性投影成形状为[k, h]的“查询”(query)、“键”(key)和“值”(value)张量,h是注意力头的大小。而在PaLM中,每个头的键/值投影是共享的,即“键”和“值”被投影到[1, h],但“查询”仍被投影到形状为[k, h]。研究发现,这对模型质量和训练速度没有影响【索引143,Fast transformer decoding: One write-head is all you need,2019,arXiv】,但在自回归解码时能显著节省成本。这是因为标准的多头注意力在自回归解码时,由于键/值张量不在样本间共享,且每次只解码一个token,因此在加速器硬件上的效率很低。

  • RoPE嵌入 模型使用了RoPE嵌入【索引153,Roformer: Enhanced transformer with rotary position embedding,2021,CoRR】,而不是绝对或相对位置嵌入,因为RoPE嵌入在长序列上表现更好。

  • 共享输入输出嵌入 模型共享了输入和输出的嵌入矩阵,这是过去工作中常用(但非普遍)的做法。

  • 无偏置项 所有的密集核(dense kernels)或层归一化(layer norms)中都没有使用偏置项。研究发现这能提高大型模型的训练稳定性。

  • 词汇表 模型使用了一个包含256k个token的SentencePiece词汇表【索引94,SentencePiece: A simple and language independent subword tokenizer and detokenizer for neural text processing,2018a,EMNLP】。选择这个规模是为了在不过度分词的情况下支持训练语料中的大量语言。词汇表是从训练数据中生成的,这被发现可以提高训练效率。该词汇表是完全无损和可逆的,这意味着空格被完全保留(对代码尤为重要),词汇表外的Unicode字符被拆分为UTF-8字节,每个字节都有一个对应的词汇表token。数字总是被拆分为单个数字token(例如,“123.5 → 1 2 3 . 5”)。

2.1 模型规模超参数

三种模型规模的配置 本文比较了三种不同规模的模型:540B、62B和8B参数。由于这些模型是标准的密集Transformer,每个token的FLOPs(浮点运算次数)约等于参数数量。这些模型使用了表1中的超参数进行构建。除了批量大小(batch size)外,这三种模型使用相同的数据和词汇表进行相同的训练。


表1:模型架构细节。表中列出了层数、d_model、注意力头数和注意力头大小。前馈网络大小d_ff总是4 × d_model,注意力头大小总是256。

4. 训练基础设施

PaLM的训练系统 我们的训练和评估代码库基于JAX【索引20,JAX: Composable transformations of Python+NumPy programs,2018,GitHub】和T5X【索引132,Scaling up models and data with t5x and seqio,2022,arXiv】,所有模型都在TPU v4 Pods【索引84,A domain-specific supercomputer for training deep neural networks,2020,Communications of the ACM】上进行训练。PaLM 540B模型通过模型和数据并行【索引167,GSPMD: general and scalable parallelization for ml computation graphs,2021,arXiv】的组合,在两个通过数据中心网络(DCN)连接的TPU v4 Pods上进行训练。我们在每个Pod中使用3072个TPU v4芯片,连接到768个主机。这个系统是迄今为止描述的最大的TPU配置,使我们能够高效地将训练扩展到6144个芯片,而无需使用任何流水线并行【索引79,GPipe: Efficient training of giant neural networks using pipeline parallelism,2019,NeurIPS】。

与先前大规模训练方法的比较 以往达到类似规模的模型训练报告主要采用两种方法。LaMDA【索引156,Lamda: Language models for dialog applications,2022,arXiv】和GLaM【索引42,GLaM: Efficient scaling of language models with mixture-of-experts,2021,arXiv】均在单个TPU系统上训练,没有利用流水线并行或DCN。Megatron-Turing NLG 530B【索引150,Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model,2022,arXiv】在2240个A100 GPU上使用了模型、数据和流水线并行的组合进行训练,而Gopher【索引128,Scaling language models: Methods, analysis & insights from training Gopher,2021,CoRR】则在四个通过DCN连接的TPU v3 Pods(每个Pod有1024个TPU v3芯片)上使用Pod间的流水线并行进行训练。

无流水线并行的优势与实现 流水线并行通常与DCN一起使用【索引150,Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model,2022,arXiv】,因为它对带宽要求较低,并能在模型和数据并行达到最大有效规模之外提供额外的并行化。流水线并行通常将训练批次分割为“微批次”(micro-batches),但这有几个重要缺点。首先,它会产生流水线“气泡”(bubble)的步骤时间开销,即在前向和后向传播的开始和结束时,许多设备完全空闲。其次,由于在每个小批次(mini-batch)内为每个微批次重新从内存加载权重,它需要更高的内存带宽。在某些情况下,它还增加了软件复杂性。我们通过以下策略,成功地将PaLM 540B的无流水线训练高效扩展到6144个芯片。

Pod内部的并行策略 每个TPU v4 Pod都包含模型参数的完整副本,其中每个权重张量通过12路模型并行和256路完全分片的数据并行(在【索引167,GSPMD: general and scalable parallelization for ml computation graphs,2021,arXiv】中称为“2D finalized”方法)被分区到3072个芯片上。在前向传播过程中,权重在数据并行轴上进行all-gather操作,并且每一层都会保存一个完全分片的激活张量。在后向传播过程中,其余的激活张量会被重新物化(rematerialized),因为与替代的重计算选择相比,这在较大批量下能获得更高的训练吞吐量。

利用Pathways系统进行跨Pod扩展 我们使用Pathways系统【索引9,Pathways: Asynchronous distributed dataflow for ML,2203,MLSys】将训练扩展到单个TPU v4 Pod之外。PaLM 540B利用Pathways的客户端-服务器架构,在Pod级别实现双向数据并行。具体来说,一个Python客户端将一半的训练批次分派给每个Pod,每个Pod使用标准的Pod内数据和模型并行来并行执行前向和后向计算以计算梯度。然后,各个Pod将梯度(在其一半批次上计算得出)与远程Pod进行交换,最后,每个Pod累积本地和远程梯度,并并行应用参数更新,以获得下一个时间步的比特级相同的参数。


图2:Pathways系统【索引9,Pathways: Asynchronous distributed dataflow for ML,2203,MLSys】利用Pod级别的双向数据并行将训练扩展到两个TPU v4 Pods。

Pathways系统的设计特点与跨Pod梯度传输 图2展示了Pathways系统如何执行双向的Pod级数据并行。一个Python客户端构建一个分片数据流程序(如图2左侧所示),该程序在远程服务器上启动JAX/XLA【索引166,XLA: Optimizing compiler for TensorFlow,2019,TensorFlow】工作,每个服务器包含一个TPU Pod。该程序包含用于Pod内前向+后向计算(包括Pod内梯度规约)的组件A,用于跨Pod梯度传输的传输子图,以及用于优化器更新(包括本地和远程梯度求和)的组件B。Pathways程序在每个Pod上执行组件A,然后将输出的梯度传输到另一个Pod,最后在每个Pod上执行组件B。Pathways系统的设计有几个特点,使其能够将程序执行扩展到数千个加速器芯片:首先,它通过在每个Pod调度器上的异步组调度(gang-scheduling)(如图2右侧所示)来掩盖从单个Python客户端到远程服务器分派JAX/XLA工作的延迟;其次,它通过分片数据流执行模型来分摊管理数据传输的成本(详情请参考【索引9,Pathways: Asynchronous distributed dataflow for ML,2203,MLSys】)。

应对跨Pod梯度传输的挑战 在6144个TPU v4芯片(横跨两个Pod,共1536个主机)的规模下,实现双向Pod级数据并行中跨Pod梯度传输的高训练吞吐量是一个有趣的挑战。需要注意的是,跨Pod梯度传输仅要求两个Pod上相应主机之间进行1:1的传输,因为每个核心只需要其模型分片参数的远程梯度。此外,两个Pod之间的主机通过谷歌数据中心网络连接【索引149,Jupiter rising: A decade of clos topologies and centralized control in google’s datacenter network,2015,ACM SIGCOMM computer communication review】。由于传输直到每个核心完成梯度计算后才开始(如图2所示),这导致了一个非常突发性的工作负载,所有主机同时通过数据中心网络链接传输它们的梯度。具体来说,每对主机在每个训练步骤中交换大约1.3GB的梯度,相当于所有主机的总突发流量达到81 Tbps。这种工作负载的突发特性带来了挑战,我们通过精心设计Pathways网络堆栈来解决,以实现最佳的DCN链接利用率。例如,为了减轻拥塞的影响,梯度传输的数据被分解成更小的块,并通过多个较小的流在不同的DCN链路上路由。通过这些优化,我们在训练期间实现了相对于单个Pod吞吐量约1.95倍的训练吞吐量(相当于完美弱扩展的97%,因为我们在两个Pod上的批量大小是单个Pod的两倍)。与理论上的2倍吞吐量相比,性能上的差距是由于后向传播和跨Pod梯度规约之间缺乏重叠。我们期望在未来的工作中解决这个问题。

4.1 训练效率

硬件FLOPs利用率(HFU)的问题 大多数先前报告的语言模型加速器效率使用的是我们称为硬件FLOPs利用率(HFU)的指标。这通常反映了给定设备上观察到的FLOPs与其理论峰值FLOPs的比率。然而,HFU存在几个问题。首先,执行的硬件FLOPs数量是系统依赖和实现依赖的,编译器的设计选择可能导致不同数量的操作。重物质化(Rematerialization)是一种广泛用于权衡内存使用和计算的技术。为了使用梯度下降高效地计算大多数神经网络架构的反向传播,批次的许多中间激活值必须存储在内存中。如果不能全部容纳,可以重新计算一些前向传播操作(从而使一些激活值被重物质化而不是存储)。这就产生了一种权衡,即使用额外的硬件FLOPs可以节省内存,但训练系统的最终目标是实现每秒高吞吐量的tokens(从而缩短训练时间),而不是尽可能多地使用硬件FLOPs。其次,测量观察到的硬件FLOPs依赖于用于计数或跟踪它们的方法。观察到的硬件FLOPs曾通过分析计算【索引113,Efficient large-scale language model training on gpu clusters using megatron-lm,2021b,SC'21】和使用硬件性能计数器【索引167,GSPMD: general and scalable parallelization for ml computation graphs,2021,arXiv】进行报告。

模型FLOPs利用率(MFU)的提出 鉴于这些问题,我们认识到HFU不是一个一致且有意义的LLM训练效率指标。我们提出了一种新的效率指标,该指标与实现无关,并允许更清晰地比较系统效率,称为模型FLOPs利用率(MFU)。这是观察到的吞吐量(tokens-per-second)与系统在峰值FLOPs下运行的理论最大吞吐量的比率。关键是,“理论最大”吞吐量只计算了计算前向+后向传播所需的操作,而不包括重物质化。因此,MFU允许在不同系统上对训练运行进行公平比较,因为分子只是观察到的每秒tokens数,而分母仅依赖于模型架构和给定系统公布的最大FLOPs。我们在附录B中详细阐述了计算MFU的数学公式。


表3:PaLM及先前大型模型的模型FLOPs利用率。PaLM由于在模型、编译器和并行策略上的多项优化,实现了显著较高的MFU。PaLM对应的硬件FLOPs利用率为57.8%。计算细节见附录B。

PaLM的MFU与其他模型的比较 我们在表3中展示了PaLM 540B模型的模型FLOPs利用率,并与先前的大型模型进行了对比。MFU有助于在不同模型参数数量、架构和模型质量的背景下比较模型和系统。根据Patterson等人(2021)通过OpenAI报告的每个GPU 24.6非注意力模型TFLOP/s,GPT-3的MFU为21.3%。而Gopher的MFU为32.5%,基于每秒0.0152步的训练速度。Megatron–Turing NLG 530B的MFU在不包含自注意力时为29.7%,包含时为30.2%,基于每秒65.43K tokens的训练吞吐量【索引150,Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model,2022,arXiv】。相比之下,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. 训练设置

权重初始化 模型的核心权重(即除嵌入和层归一化尺度外的所有权重)使用“fan-in方差缩放”进行初始化,即$W \sim N(0, 1/n_{in})$,其中$n_{in}$是核心的输入维度。输入嵌入初始化为$E \sim N(0, 1)$,因为层归一化不应用于嵌入。由于输入和输出嵌入层是共享的,我们通过$1/\sqrt{n}$来缩放softmax之前的输出logits,其中n是嵌入大小。

优化器 模型使用Adafactor优化器【索引145,Adafactor: Adaptive learning rates with sublinear memory cost,2018,ICML】进行训练,没有使用因式分解。这实际上等同于带有“参数缩放”的Adam【索引88,Adam: A method for stochastic optimization,2014,arXiv】,它通过参数矩阵的均方根来缩放学习率。因为权重初始化与$1/n$成比例,其效果类似于手动降低Adam学习率,如Brown等人(2020)的做法。然而,参数缩放的优点是,在不同尺度上操作的参数矩阵(嵌入和层归一化尺度)的学习率不会以相同的速率被缩减。

优化超参数 我们在前10,000步使用$10^{-2}$的Adafactor学习率,然后以$1/\sqrt{k}$的速率衰减,其中k是步数。我们使用$\beta_1 = 0.9$的动量进行训练。二阶矩插值计算为$\beta_2 = 1.0 - k^{-0.8}$,其中k是步数。我们发现在训练大型语言模型时,这比标准的$\beta_2 = 0.99$更稳定,因为罕见的嵌入token在较短的窗口内可能二阶矩估计不佳。我们对所有模型使用值为1.0的全局范数梯度裁剪【索引120,Understanding the exploding gradient problem,2012,CoRR】。在训练期间,我们使用动态权重衰减$lr^{2.0}$,其中lr是当前的学习率。

损失函数 模型使用标准的语言建模损失函数进行训练,即所有token的平均对数概率,不使用标签平滑。我们还使用了一个辅助损失$z_{loss} = 10^{-4} \cdot \log^2 Z$,以鼓励softmax归一化器$\log(Z)$接近0,我们发现这能增加训练的稳定性。

序列长度 所有模型都使用2048的序列长度。输入样本被连接在一起,然后被分割成长度恰好为2048的序列,这样就没有填充token,但样本可能会在中间被分割。输入样本之间用一个特殊的[eod] token来区分。

批量大小 对于所有模型,我们在训练期间增加批量大小。对于最大的模型,我们在前50k步使用批量大小512(1M tokens),然后在50k到115k步之间将其加倍至1024(2M tokens),最后再次加倍至2048(4M tokens)直到训练在255k步完成。较小的模型遵循类似的时间表。使用这种批量大小时间表的原因有二:(1)在训练早期,较小的批量大小样本效率更高(即,作为已见tokens函数的损失更好),而在训练后期,由于梯度估计更佳,较大的批量大小更有利【索引151,Don’t decay the learning rate, increase the batch size,2018,ICLR】;【索引105,An empirical model of large-batch training,2018,arXiv】(2)较大的批量大小导致更大的矩阵乘法维度,从而提高了TPU效率。

比特级确定性 模型从任何检查点开始都是完全比特级可复现的。换句话说,如果模型在一次运行中训练到了第17,000步,我们从第15,000步的检查点重新开始,那么训练框架保证在两次运行中从第15,000步到第17,000步产生相同的结果。这是通过两种方式实现的:(1)由JAX+XLA+T5X提供的比特级确定性建模框架,以及(2)一个确定性的数据集管道,其中打乱的数据以随机访问格式写出,因此给定训练批次的内容仅是步数的函数【索引132,Scaling up models and data with t5x and seqio,2022,arXiv】。

Dropout 模型在训练时没有使用dropout,尽管在大多数情况下,微调时会使用0.1的dropout。

5.1 训练不稳定性

损失峰值现象 对于最大的模型,尽管启用了梯度裁剪,我们在训练过程中观察到损失出现了大约20次峰值。这些峰值以非常不规则的间隔出现,有时在训练后期发生,并且在训练较小模型时没有观察到。由于训练最大模型的成本高昂,我们无法确定一个有原则的策略来缓解这些峰值。

缓解策略 相反,我们发现一个简单的策略可以有效缓解这个问题:我们在峰值开始前大约100步的检查点重新开始训练,并跳过大约200-500个数据批次,这些批次涵盖了峰值之前和期间看到的数据。通过这种缓解措施,损失在同一点没有再次出现峰值。我们不认为这些峰值是由“坏数据”本身引起的,因为我们进行了几个消融实验,我们取了峰值周围的数据批次,然后从一个不同的、更早的检查点开始在这些相同的数据批次上进行训练。在这些情况下,我们没有看到峰值。这意味着峰值仅在特定数据批次与特定模型参数状态的组合下才会发生。未来,我们计划研究更具原则性的缓解策略,以应对非常大的语言模型中的损失峰值。

A3 背景知识/关键Observation/设计原则

3. 训练数据集

高质量语料库 PaLM的预训练数据集包含一个高质量的7800亿token的语料库,代表了广泛的自然语言用例。该数据集是过滤后的网页、书籍、维基百科、新闻文章、源代码和社交媒体对话的混合体。该数据集基于用于训练LaMDA【索引156,Lamda: Language models for dialog applications,2022,arXiv】和GLaM【索引42,GLaM: Efficient scaling of language models with mixture-of-experts,2021,arXiv】的数据集。我们对所有三个模型都只在数据上训练一个epoch(对所有模型都进行了相同的洗牌),并选择混合比例以避免在任何子组件中重复数据。

源代码数据 除了自然语言数据,预训练数据集中还包含代码。预训练数据集中的源代码来自GitHub上的开源仓库。我们根据仓库中包含的许可证过滤了文件;排除了copyleft许可证。我们通过文件名扩展名过滤文件,以限制在24种常见的编程语言中,包括Java、HTML、Javascript、Python、PHP、C#、XML、C++和C,这导致了196GB的源代码。此外,我们根据文件之间的Levenshtein距离移除了重复项,因为众所周知,重复文件在源代码仓库中很常见【索引104,D´ej`aVu: a map of code duplicates on GitHub,2017,Proc. ACM Program. Lang.】;【索引4,The adverse effects of code duplication in machine learning models of code,2019,SPLASH Onward!】。

数据来源比例 表2列出了用于创建最终PaLM数据集混合的各种数据源的比例。我们在第8节检查数据污染,并展示了我们的训练数据集与评估数据之间的重叠分析。附录D包含一个数据表(datasheet)【索引55,Datasheets for datasets,2021,Commun. ACM】,其中包含更多信息,包括语言比例的细分。


表2:训练数据集中各来源数据的比例。多语言语料库包含来自100多种语言的文本,其分布见附录表29。

A4 实验环境

  • 数据集:PaLM预训练数据集,包含7800亿个token。这是一个高质量的混合语料库,来源包括经过筛选的网页(占50%)、社交媒体对话(占25%)、书籍(占8%)、GitHub代码(占8%)、维基百科(占4%)和新闻(占4%)。数据集覆盖超过100种语言,其中英语约占78%。
  • 模型架构:采用仅解码器(decoder-only)的Transformer架构,并进行了多项修改,包括SwiGLU激活函数、并行层、多查询注意力机制和RoPE位置嵌入。模型分为三个规模:80亿(8B)、620亿(62B)和5400亿(540B)参数。
  • 硬件配置:训练在Google的TPU v4 Pods上进行。540B模型使用了两个TPU v4 Pods,共计6144个TPU v4芯片,每个Pod包含3072个芯片,连接到768台主机。两个Pod通过数据中心网络(DCN)连接。
  • 软件配置:训练代码库基于JAX和T5X。使用了Pathways系统来协调跨多个TPU Pods的训练。优化器为Adafactor。

A4 实验结果

6.1 英语NLP任务

  • 实验内容: 在29个标准英语NLP基准上评估PaLM的小样本(few-shot)性能,涵盖问答、常识推理、阅读理解等多种任务。
  • 实验结果: PaLM 540B在1-shot设置下的24个任务和few-shot设置下的28个任务中超越了先前的SOTA(表4)。其在NLU和NLG任务上的平均得分也显著优于GLaM和GPT-3(表5)。
  • 分析结论: 模型规模的持续扩大带来了性能的显著提升。PaLM 540B的表现优于同等规模的Megatron-Turing NLG 530B,表明预训练数据、训练策略等因素也至关重要。

6.1.1 Massive Multitask Language Understanding (MMLU)

  • 实验内容: 在MMLU基准上进行评估,该基准包含57个涵盖数学、历史、法律等领域的多项选择问答任务。
  • 实验结果: PaLM 540B在MMLU上的平均分比之前的SOTA(Chinchilla)提高了约2个百分点(表6)。
  • 分析结论: PaLM在需要广泛世界知识和专业知识的任务上表现出色。

6.1.2 微调(Finetuning)

  • 实验内容: 在SuperGLUE基准上对PaLM模型进行微调。
  • 实验结果: 微调后的PaLM在SuperGLUE上取得了接近SOTA的性能,与专为分类任务设计的编码器-解码器模型(如T5、ST-MoE)相当(表7)。微调性能远高于小样本性能(表8)。在排行榜上,PaLM超越了最好的仅解码器模型(表9)。
  • 分析结论: 模型规模的扩大有助于弥合仅解码器架构在微调分类任务上的天然劣势。

6.2 BIG-bench

  • 实验内容: 在BIG-bench基准上评估PaLM的小样本性能。该基准包含150多个具有挑战性的新任务。
  • 实验结果: 5-shot的PaLM 540B在58个共同任务上的总分显著优于GPT-3、Gopher和Chinchilla,并超过了人类评估者的平均分(图3)。在150个文本任务中,PaLM 540B在约65%的任务上优于人类平均水平(图6)。
  • 分析结论: PaLM展现了强大的泛化能力。特别地,约25%的任务表现出“非连续性改进”,即性能从62B到540B规模时出现急剧跃升,表明某些能力只有在足够大的模型规模下才会涌现(图5)。

6.3 推理任务

  • 实验内容: 在需要多步算术或常识推理的任务上进行评估,并引入了“思维链提示”(chain-of-thought prompting)技术。
  • 实验结果: 通过结合模型规模和思维链提示,PaLM 540B在多个推理任务上仅用8-shot评估就达到或超过了经过微调的SOTA。例如,在GSM8K数学问题数据集上,PaLM 540B取得了58%的准确率,超过了先前SOTA的55%(图10,表10)。从62B扩展到540B修复了大量的推理错误(图9)。
  • 分析结论: 模型规模和思维链提示的结合是解锁大型模型复杂推理能力的关键。模型能够生成中间推理步骤,这不仅有助于解释,还能提高最终答案的准确性。

6.4 代码任务

  • 实验内容: 评估PaLM在文本到代码(HumanEval, MBPP)和代码到代码(TransCoder, DeepFix)任务上的性能。
  • 实验结果: 未经微调的PaLM 540B在HumanEval上的性能与专门在代码上训练的Codex 12B相当,尽管其训练的Python代码量少了50倍。经过代码数据微调后(PaLM-Coder),性能进一步提升,在多个基准上达到SOTA(表12)。在DeepFix代码修复任务上,编译成功率达到82.1%,创下新高(图12,表13)。
  • 分析结论: PaLM证明了单个大型模型可以在自然语言和代码任务上同时达到顶尖水平,这得益于跨领域知识的迁移以及大模型更高的样本效率。

6.5 翻译任务

  • 实验内容: 在多种语言对上评估PaLM的翻译能力,包括英语中心、非英语中心和极低资源语言对。
  • 实验结果: 在英语中心语言对上,PaLM的小样本性能全面超越GPT-3和FLAN,甚至在某些方向上超过了有监督的SOTA模型(表14)。在非英语和低资源语言对上,性能依然强大,例如在法德互译上能与有监督模型媲美(表15)。模型规模的扩大同样带来了翻译性能的非连续性跃升(图15)。
  • 分析结论: 尽管没有在并行的翻译数据上进行专门训练,PaLM仍然展现出强大的翻译能力,尤其是在翻译到英语时。

6.6 多语言自然语言生成(NLG)

  • 实验内容: 在多语言摘要(MLSum, XSum)和数据到文本生成(E2E NLG, WebNLG)任务上评估PaLM的小样本和微调性能。
  • 实验结果: 微调后的PaLM 540B在所有生成英语文本的任务上都达到或超过了SOTA,再次证明了规模可以弥补架构上的不足。但在生成非英语文本时,性能不如SOTA。小样本性能与微调性能之间仍有较大差距(表16)。
  • 分析结论: 这是首次对大型语言模型在这类条件生成任务上进行小样本评估。PaLM在处理非英语输入方面优于生成非英语输出。

6.7 多语言问答

  • 实验内容: 在TyDiQA-GoldP多语言问答基准上评估PaLM的小样本和微调性能。
  • 实验结果: 微调后的PaLM 540B取得了非常有竞争力的结果,性能优于mT5 XXL,但不及ByT5 XXL(表17)。
  • 分析结论: 即使非英语训练数据占比较小,PaLM通过规模扩展也能在多语言任务上达到接近SOTA的水平,在一定程度上克服了架构和数据上的劣势。

A5 结论

结论
本文通过训练5400亿参数的PaLM模型,进一步拓展了小样本语言理解和生成的边界。研究结果表明,通过规模扩展带来的性能提升尚未达到瓶颈。PaLM在广泛的英语NLP任务、挑战性的BIG-bench、源代码理解与生成、多语言NLP以及机器翻译任务上均取得了SOTA性能。

一个关键的发现是模型在推理任务上的突破性表现。通过结合模型规模和“思维链提示”(chain-of-thought prompting),PaLM在无需微调的情况下,在多个算术和常识推理任务上匹敌甚至超越了微调后的SOTA模型。这不仅证明了生成自然语言解释对用户的价值,更重要的是,它揭示了模型的生成能力本身可以极大地提升预测任务的质量。

性能提升不仅是线性的,在某些基准上还呈现出“非连续性”的跃升,这意味着某些复杂能力只有在模型达到足够大的规模时才会涌现,预示着未来更大规模的模型可能解锁更多未知的能力。

未来工作与开放性问题
尽管PaLM取得了成功,但关于理想的网络架构和训练方案仍有许多悬而未决的问题。PaLM是基于成熟的密集、仅解码器、全注意力Transformer架构,但Pathways系统的最终目标是探索更多样化的架构选择和训练方案。一个核心的开放问题是模型大小与训练数据量之间的权衡,即一个用更多数据训练的较小模型是否能匹敌用较少数据训练的更大模型。PaLM的成功为开发一个大规模、模块化、能够跨多种模态泛化的系统奠定了坚实的基础,这是向Pathways作为未来ML扩展范式的愿景迈出的第一步。

A6 附录

B. 计算资源使用和环境影响

MFU计算详解 模型FLOPs利用率(MFU)是衡量训练效率的指标,计算方式为观察到的吞吐量(tokens/秒)与系统在峰值FLOPs下运行的理论最大吞吐量的比率。对于一个N参数的仅解码器模型,每个token需要约6N的矩阵乘法FLOPs(前向2N,后向4N)。自注意力部分额外增加$6LH(2QT)$ FLOPs。若一组加速器的总理论峰值矩阵乘法吞吐量为P FLOPs/秒,则理论峰值吞吐量$R = P / (6N + 12LHQT)$。PaLM 540B的MFU达到46.2%,而分析计算出的硬件FLOPs利用率(HFU,包含重物质化开销)为57.8%。

计算用量 训练PaLM系列模型的计算用量如下表所示。PaLM 540B模型使用了重物质化技术,而8B和62B模型除自注意力外未使用。


表22:训练PaLM 8B和PaLM 540B至7800亿tokens以及PaLM 62B至7950亿tokens的计算用量。

环境影响 PaLM 540B的训练在Google的俄克拉荷马数据中心进行,该数据中心在训练期间89%的能源来自无碳能源。训练总共在6144个TPU v4芯片上运行1200小时,在3072个TPU v4芯片上运行336小时,总计排放了271.43吨二氧化碳当量(tCO2e)。相比之下,一次旧金山到纽约的单人往返航班排放约180 tCO2e,而GPT-3的训练排放估计为552 tCO2e。所有用于PaLM训练和实验的能源消耗都通过可再生能源进行了补偿。

C. 数据集分析

描述性统计 使用分类器对训练数据进行分析,发现1.75%的数据点包含地址,0.7%包含电子邮件,32%包含人名,1.5%包含电话号码。然而,这些不一定是个人身份信息(PII),例如10%的电话号码是免费热线,11%的人名来自书籍语料库。

主题建模 为了解数据集内容,我们运行了分层主题建模分类器。顶层主题类别包括“艺术与娱乐/在线媒体/在线图片库”、“参考/通用参考/指南、DIY和专家内容”、“游戏/电脑与视频游戏/游戏参考与评论”等。完整的层级结构如下图所示。


图25:数据集中检测到的分层主题。

关联性分析
- 性别代词: 数据集中,41%的数据点包含they/them代词,30%包含he/him代词,14%包含she/her代词。与不同代词共现频率最高的职业和TF-IDF词元如下表所示。


表23:与性别代词共现频率最高的职业。

表24:与性别代词共现频率最高的TF-IDF词元。

  • 宗教与种族/民族: 我们分析了宗教和种族/民族身份与其它词元的共现情况。下表展示了与选定的宗教及种族身份共现频率最高的TF-IDF词元。

    表25:与宗教(为匹配Brown等人(2020)的分析而选择的身份)共现频率最高的TF-IDF词元。

    表26:与种族和民族身份(为匹配Brown等人(2020)的分析而选择的身份)共现频率最高的TF-IDF词元。

F. 延长训练时间

延长训练的设置 如第13节所述,训练的tokens数量对小样本性能有显著影响。本节将PaLM 62B模型训练至1.325万亿tokens,并与PaLM 540B和Chinchilla进行比较。为了避免数据重复,我们使用了更新后的数据集,调整了子语料库的比例,并新增了arXiv和数学网页数据。

与Chinchilla的比较 延长训练后,PaLM 62B在BIG-bench上的性能差距与Chinchilla缩小,而在9个英语NLP任务上则超越了Chinchilla,其性能与插值缩放曲线相符。


图26:模型性能与总训练FLOPs的关系。实线表示基于0.78T tokens训练的PaLM 8B、62B和540B的插值性能曲线,虚线表示将PaLM 62B延长训练至1.325T tokens后的性能增益。

性能提升分析 延长训练使PaLM 62B在所有评估基准上都取得了显著提升,但仍未能完全弥补与PaLM 540B(训练FLOPs高5倍)的差距。


图27:经过额外训练至1.325T tokens,PaLM 62B在多种任务上取得了显著提升,但仍未弥补与PaLM 540B的差距。

表32:对应图27的任务特定结果,比较了PaLM 62B原始训练(约7950亿tokens)和延长训练(1.325万亿tokens)后与PaLM 540B在MMLU、BIG-bench、GSM8K、MBPP和HumanEval数据集上的表现。

更新数据集 vs. 重复数据集 消融实验证明,使用更新的数据集进行延长训练比使用含有重复子语料库的数据集效果更好。


表33:比较使用更新数据集与重复子语料库的数据集对PaLM 62B进行额外950亿tokens的训练。结果显示,PaLM 62B在英语NLP任务上使用更新数据集表现更佳。

G. 模型输出示例

推理示例 对于GSM8K数学应用题数据集,模型规模从62B扩展到540B显著减少了错误。以下是一些PaLM 540B修正了PaLM 62B错误的例子。


图28:PaLM 62B因语义理解错误而回答错误,但PaLM 540B修正了这些错误的示例问题。

图29:PaLM 62B因缺少一个步骤而回答错误,但PaLM 540B修正了这些错误的示例问题。

偏见与毒性示例


表34:用于各共现类别的提示。

表35:性别:针对“他/她非常...”等提示,PaLM 540B生成的第一完整句子中最常见的10个描述性词语。

表36:种族:针对“人们会形容{术语}的人为...”等提示,PaLM 540B生成的第一完整句子中最常见的10个描述性词语。

表37:宗教:针对“{术语}的信徒是...”等提示,PaLM 540B生成的第一完整句子中最常见的10个描述性词语。

表38:人类与PaLM 540B模型的续写对比。括号中的数字表示提示/续写的毒性概率。

记忆化示例


表39:PaLM 540B模型完全复现真实续写的几个训练样本。这些样本是精心挑选的,旨在展示被记忆的文本类型。

H. 补充结果

较小模型在英语NLP任务上的结果


表40:PaLM 8B和62B模型与PaLM 540B在29个NLP基准上的结果对比。对于62B模型,我们报告了在795B tokens和延长训练至1325B tokens后的1-shot结果。

BIG-bench补充结果


表41:PaLM 8B、62B和540B模型在24个BIG-bench Lite任务上的5-shot评估结果。报告的分数为未归一化分数,与图7中的数据一致。

表42:PaLM 540B性能超过人类平均性能的前十大BIG-bench任务的首选指标分数。

表43:人类平均性能超过PaLM 540B性能的前十大BIG-bench任务的首选指标分数。

多语言NLG补充结果


表44:GEM数据集的详细ROUGE-2结果,分为精确率、召回率和F1值,扩展了第6.6节报告的结果。

表45:GEM数据到文本和摘要数据集的详细ROUGE-L结果,分为精确率/召回率/F1值,扩展了第6.6节报告的结果。

表46:GEM数据到文本和摘要数据集的BLEURT-20结果,扩展了第6.6节报告的结果。