Apple Intelligence揭秘：端侧AI大模型的落地之道

发布日期：2024-08-04 18:46:39 浏览次数： 5280

作者：大模型沉思录

微信搜一搜，关注“大模型沉思录”

2024年6月苹果全球开发者大会上，正式推出了Apple Intelligence（苹果智能），这是一个深度集成到iOS 18、iPadOS 18和macOS Sequoia中的个人智能系统。

Apple Intelligence由多个高效的生成模型组成，这些模型专门针对用户的日常工作任务进行了微调，能够即时适应用户当前的行为活动，可以在文本创作、通知总结、图像生成，以及跨应用交互等任务中提供更为优秀的交互体验。

近期，苹果也发布了自己大模型的技术报告，公布了大量细节，为业界提供了一条清晰的端侧大模型落地应用方案，受到大量关注。

在这份报告中，详细说明其中的两个模型—AFM-on-device（AFM代表Apple Fundation Model，苹果基础模型），一个大约30亿参数的语言模型，以及AFM-server，一个更大的基于服务器的语言模型—是如何构建和适应的，以高效、准确、负责任地执行专业任务。这两个基础模型是苹果创建的大模型家族的一部分，旨在支持用户和开发者，包括一个编程大模型（基于AFM语言模型），用于在Xcode中辅助代码编写；一个扩散模型，帮助用户在各类App中实现视觉创作。

图1. 苹果基础大模型训练流程图

下面我们将从模型架构、预训练、后训练、系统以及评测等方面为大家一一揭晓。

01 模型架构

AFM基础模型采用了经典的自回归模型架构，并选择了如下参数设置：

使用共享的输入/输出嵌入矩阵来减少参数的内存使用
预归一化结合RMSNorm以提高训练稳定性
查询/键归一化以改善训练稳定性
分组查询注意力（GQA），具有8个键值头，以减少键值缓存的内存占用
SwiGLU激活以提高效率
RoPE位置嵌入，将基础频率设置为500k以支持长上下文

下表是AFM-on-device模型采用的参数设置。

表1. AFM-on-device模型的参数设置

02 预训练

预训练采用的训练策略同主流模型一样，主要包括数据、训练策略以及训练Infra三部分：

2.1 数据 (Data)

数据来源：预训练数据集由多种来源组成，包括从出版商那里授权的数据、公开可用或开源的数据集，以及通过Apple的网络爬虫Applebot抓取的公共信息。
数据多样性：数据集包含多种类型的数据，例如网页、编程代码、数学问答等。
数据质量：重视数据质量而非数量，进行了严格的数据清洗，包括去除不安全内容、亵渎语言、个人信息以及对常见评估基准的去重。
数据集细节：

网页：使用Applebot爬取，并通过一系列步骤进行质量过滤和纯文本提取。
授权数据集：从出版商处授权的高质量数据，用于预训练的持续和上下文延长阶段。
代码数据：从GitHub上的开源仓库获取，涵盖多种编程语言。
数学数据：整合了来自网络的高质量数学问答数据和数学内容。
公共数据集：选用了高质量且允许用于训练语言模型的公共数据集。

2.2 训练策略 (Training Strategy)

预训练过程：分为三个阶段：核心预训练、持续预训练和上下文延长预训练。
核心预训练：AFM-server从头开始训练，而AFM-on-device则通过从较大的模型中进行知识蒸馏和结构修剪来初始化。
持续预训练：在更长的序列长度上进行，增加了数学和代码数据的权重。
上下文延长：在更长的序列长度上进行预训练，并包含合成的长上下文问答数据。
训练技术：

使用了知识蒸馏和结构修剪来提高模型性能和训练效率。
采用了分步训练，其中包括线性预热和余弦衰减学习率计划。
利用了μParam简化版本进行正则化。

2.3 训练基础设施 (Training Infrastructure)

硬件平台：AFM模型在v4和v5p Cloud TPU集群上进行预训练，使用基于JAX的深度学习库AXLearn。
并行策略：训练结合了张量并行、完全分片的数据并行和序列并行，以实现大规模模型参数和序列长度的高效训练。
效率：AFM-server在8192个TPUv4芯片上训练，实现了大约52%的Model FLOPs utilization (MFU)，AFM-on-device则使用了2048个TPUv5p芯片。

03 后训练

苹果的后训练策略同样是和主流大模型一样，也包括了两个阶段：监督微调（SFT）和人类反馈强化学习（RLHF）。值得注意的是，AFM在强化学习阶段，提出了两个自己的工作：：（1）一种带有教师委员会的拒绝采样微调算法（iTeC）；（2）一种使用镜像下降策略优化（Mirror Descent Policy Optimization）和留一法优势估计器（Leave-One-Out Advantage Estimator）MDLOO的基于人类反馈的强化学习算法，显著提高了模型的效果。

3.1 SFT

监督微调阶段，AFM并未有自己特别的工作：

数据上，采用了人工标注和合成数据两种方式，数据种类除了通用对话数据外，还包括数学、工具使用以及代码等数据，并再次强调了数据质量的重要性（高于数据数量）；
训练策略上，包括数据的选择、配比以及模型训练的参数（AFM-server模型采用了5e−6的常量学习率，而AFM-device模型则采用了 2e−5，Dropout参数为0.1）

3.2 RLHF

人类反馈强化学习这个阶段，主要是通过人类偏好数据来对模型进行优化，包括奖励模型（Reward Modeling）和强化学习（Reinforcement Learning）两个阶段。

奖励模型
采用收集的人类偏好数据训练奖励模型，训练数据的标签包括两部分：

同一个问题的偏好回答以及偏好等级，分为接受回答比拒绝回答显著更好、更好、好一点、可以忽略不计的好四种
每个回答的单边评价，包括指令跟随能力、真实性、无害性等

训练阶段，包括以下两个创新点：

设计了一个软标签损失函数,考虑了人类偏好级别
将单边评分作为正则化项纳入奖励建模中

一种带有教师委员会的拒绝采样微调算法-iTeC

苹果提出了一种创新的迭代式RLHF框架,以充分利用多轮RLHF优化模型的能力，主要内容包括: 提出了一种新颖的迭代RLHF框架,有效结合了各种偏好优化算法,包括rejection sampling(RS)、直接偏好优化(DPO)及其变种IPO,以及在线强化学习(RL)。这样能够将RLHF的优势应用到不同大小的AFM模型中，明显提升了模型的对齐效果。

‍迭代式委员会：在每一批次的人类偏好数据收集中，使用一系列最新最优秀的模型（来自SFT、RS、DPO/IPO、RL以及上一轮迭代的最佳模型）组成"模型委员会"。从该委员会采样响应并收集配对的人类偏好数据。根据新获得的人类偏好数据，更新奖励模型，并使用各种偏好优化算法训练一组新模型。然后用新的模型委员会进行下一轮迭代RLHF数据收集。‍
委员会蒸馏：使用最新的奖励模型对模型委员会进行rejection sampling(蒸馏)。不同于全局级别的重排序，在提示级别对每个提示从委员会中采样多个响应，使用奖励模型选择最佳响应。这结合了不同偏好优化算法训练模型的优势。
扩大蒸馏规模：为了将RLHF改进应用到所有大小的AFM模型，扩大了来自模型委员会的蒸馏规模。小模型在扩大蒸馏提示数量时可获得巨大改进，最终AFM-On-Device模型在100多万高质量响应上进行训练。

在线RLHF算法-MDLOO

苹果在这部分工作中提出了一种名为MDLOO的在线强化学习算法，它在模型训练过程中解码模型回答，并利用强化学习算法来最大化奖励。

与RLHF中常用的PPO等方法类似，同样采用了基于信任域的策略迭代方法。此外，主要做了两方面的改进：

使用留一法（Leave-One-Out, LOO）估计器来估计提示-响应对的优势
使用Mirror Descent Policy Optimization (MDPO)来优化策略，与常用的基于剪切的PPO方法不同

因此，该算法被命名为带有留一法估计的镜像下降（Mirror Descent with Leave-One-Out estimation, MDLOO），实验证明了该方法要比PPO更有效。

04 系统

介绍了这里，苹果的模型训练整体上与业务主流方案保持一致，主要在强化学习阶段提出了两个自己的两个方法。那么，苹果的大模型最终是要部署到设备端的，那么它是如何处理各种任务以及云端模型是如何配合的呢，下面我们一起来看下。

从图中可以看出，苹果的整体设计还是比较务实的：

多任务：苹果为常用的任务每个任务均通过Lora训练了一个Adapter，这样可以在基础模型的基础上可以更灵活的适配每个任务，进而提升用户在不同任务上的体验。苹果自身也能根据用户的反馈，灵活轻量级的来迭代各个Adapter
端云结合以及不同任务的路由：从图中可以看到，苹果提供了语义索引以及意图工具箱两种方法来路由自己的模型能力。虽然没有详细介绍，可以猜测到语义索引指的是根据用户的问题或者请求通过端侧模型将请求路由到能够给出最优回答的模型上（具体实现可以参考[2]里的实现）；而意图工具箱则指的是对于那些已经明确意图的请求，例如在邮箱里点击文本总结，直接调用文本总结的模型。

此外，为了保证端侧模型的推理性能，苹果采用了量化方法对模型进行了低精度量化：