【极速前进】20240524-20240526：Meta多模态模型Chameleon、参数高效微调MoRA、跨层注意力降低KV Cache、Octopus v4、SimPO

一、Chameleon：混合模态早期融合基础模型

论文地址：https://arxiv.org/pdf/2405.09818

​ Chameleon是一个基于token的早期融合多模态模型，能够理解和生成任意序列的文本或者图像。

1. 预训练

1.1 Tokenization

​ 图像Tokenization。从大小为8192的codebook中奖512*512的图像编码为1024个离散tokens。使用授权图片来训练tokenizer。由于生成人脸的重要性，对人脸图像进行了上采样。该tokenizer的核心弱点是重构具有大量文本的图像，因此遇到需要大量OCR相关的任务会导致模型效果较差。

​ 文本Tokenizer。使用训练数据的子集来训练新的BPE tokenizer，词表大小为65536，其中包含8192个图像codebook tokens，使用sentencepiece库。

1.2 预训练数据

​ 将预训练分为两个阶段，第一个阶段占80%，第二阶段20%。

​ 第一阶段。第一阶段包含大量无监督数据的混合数据集，包括：(1) 大约2.9T tokens的纯文本数据集；(2) 使用大约1.4 B的图文对，其能够产生1.5T的图文tokens；(3) 从互联网收集大约400B tokens的图文交错数据。

​ 第二阶段。降低第一阶段数据的权重为50%并混合高质量数据集，同时维护图像文本token的比例相近。此外，包含了指令调优的部分数据集。

1.3 稳定性

​ 当Chameleon超过8B参数和1T tokens时会遇到稳定性问题，而且不稳定性出现在训练后期。为了解决稳定性问题，采用措施：

​ 架构。模型结构大致类似于LLaMa-2。由于训练中后期输出范数的缓慢增长，LLaMa架构会出现复杂的分歧。最终定位到是由softmax的平移不变性导致的。由于模型权重跨模态共享，每个模态会通过轻微增加范数来与其他模态竞争。虽然训练开始时没有问题，但是随着训练进行超过bf16有效表示范围，就会出现分歧。在单一模态中，该问题称为logit漂移问题。

​ 上图5(a)是训练过程中最后Transformer层输出的范数，大约在训练的20%处出现分歧。范数出现不受控增长与loss出现分歧高度相关。在Transformer中Softmax出现在两处：核心注意力机制和logit之后。这里使用QK-Norm来控制softmax输入的范数。此外，除了QK-norm以外，也有必要在注意力和前馈层前引入dropout。但是，该方案对于Chameleon-34B并不足够，其需要额外的范数重排。具体来说，使用Liu et al.的normalization策略。设$h$表示输入$x$经过自注意力机制后的隐藏向量，则
Chameleon-34B:   h = x + attention_norm ( attention ( x ) ) output = h + ffn_norm ( feed_forward(h) ) Llama2:   h = x + attention ( attention_norm ( x ) ) output = h + feed_forward ( ffn_norm ( h ) )

$$\begin{align} \textbf{Chameleon-34B: }&h=x+\text{attention\_norm}(\text{attention}(x)) \\ &\text{output}=h+\text{ffn\_norm}(\text{feed\_forward(h)}) \\ \textbf{Llama2: }&h=x+\text{attention}(\text{attention\_norm}(x)) \\ &\text{output}=h+\text{feed\_forward}(\text{ffn\_norm}(h)) \\ \end{align}$$ \\ Chameleon-34B: Llama2: ​h=x+attention_norm(attention(x))output=h+ffn_norm(feed_forward(h))h=x+attention(attention_norm(x))output=h+feed_forward(ffn_norm(h))​​
​ 此外，发现这种类型的normalization与dropout不能很好的结合，因此训练Chameleon-34B不使用dropout。

​ 优化。应用QK-Norm有助于注意力内的softmax，但对于最后输出层的logit shift问题没有帮助。因此，应用z-loss正则化。具体来说，对softmax的配分函数进行正则化，即对于$\sigma (x{)}_i=\frac{e^{x_i}}{Z}$中的$Z={\sum }_i{e}^{x_i}$加$10^{-5}{\log }_{2}Z$。Chameleon-7B使用dropout和z-loss实现稳定性，而Chameleon-34B仅需要z-loss。

2. 对齐

​ 数据。将SFT数据分为以下类别：文本、代码、视觉对话、图像生成、文本/图像交替生成、安全。

​ 微调策略。SFT阶段平衡模态数据比例至关重要。

二、MoRA：用于参数高效微调的高秩更新

论文地址：https://arxiv.org/pdf/2405.12130

​ LoRA的低秩更新机制可能会限制LLM学习和记忆新知识的能力。本文提出MoRA，能够实现相同可训练参数情况下的高秩更新。

1. 分析低秩更新的影响

​ 给定预训练权重矩阵$W_0\in {\mathbb{R}}^{d\times k}$，LoRA利用两个低秩矩阵来计算权重更新$\Delta W$：
$$h=W_{0}x+\Delta Wx=W_{0}x+BAx$$
其中$A\in {\mathbb{R}}^{r\times k}$和$B\in {\mathbb{R}}^{d\times r}$表示低秩矩阵。为了确保训练开始时$\Delta W=0$，LoRA使用高斯分布初始化$A$并初始化B为0。由于$\Delta W$被低秩分解为$BA$，则$\text{rank}(\Delta W)\le r$。LoRA更新在文本分类或者指令微调上表现出与全秩微调相当的效果，但是在复杂推理或者持续预训上则效果会变差。

​ 基于这些观察，提出假设：低秩更新能够更好利用LLM的原始知识和能力，但是对于需要增强LLM知识和能力的任务来说就很挣扎。经验实验验证，LoRA确实难以学习新知识，即使可以通过增大秩来缓解。

2. MoRA

​ 给定预训练权重$W_0\in {\mathbb{R}}^{d\times k}$，LoRA使用两个低秩矩阵$A$和$B$来实现低秩更新，总可训练参数为$(d+k)r$。在相同可训练参数下，方阵$M\in {\mathbb{R}}^{\hat{r}\times \hat{r}}$的$\hat{r}=\lfloor \sqrt{(d+k)r}\rfloor$能够实现最高的秩，因为$r\ll \min (d,k)$。

​ 为了实现该目标，应该缩减$M$的输入微调并增加输出维度。正式来说，
$$h=W_{0}x+f_{\text{decomp}}(M{f}_{\text{comp}}(x))$$
其中$f_{\text{comp}}:{\mathbb{R}}^k\to {\mathbb{R}}^{\hat{r}}$表示将$x$输入维度从$k$降低至$\hat{r}$的函数，$f_{\text{decomp}}:{\mathbb{R}}^{\hat{r}}\to {\mathbb{R}}^d$则是将输出维度从$\hat{r}$增加到$d$的函数。此外，这两个函数应该是非参数化的操作且能在线性时间内执行。对于将$M$转换为$\Delta W$，对应函数$f_{\frac{}{\text{comp}}}:{\mathbb{R}}^{\hat{r}\times \hat{r}}\to {\mathbb{R}}^{\hat{r}\times k}$和$f_{\frac{}{\text{decomp}}}:{\mathbb{R}}^{\hat{r}\times k}\to {\mathbb{R}}^{d\times k}$。对于任意的$x$，均有
$$f_{\text{decomp}}(M{f}_{\text{comp}}(x))=\Delta Wx,\forall x\in {\mathbb{R}}^k$$
其中$\Delta W=f_{\frac{}{\text{decomp}}}(f_{\frac{}{\text{comp}}}(M))$。若上式成立，$M$可以无损扩展为$\Delta W$。这就能像LoRA那样合并回LLM。

​ 对于$f_{\text{comp}}$和$f_{\text{decomp}}$，探索了几种实现函数。一种直接的方法是截取维度，随后在填充相应的维度。即
f comp ( x ) = x 1 : r ^ f decomp ( x ) = [ x 0 ]

$$\begin{align} f_{\text{comp}}(x)=x_{1:\hat{r}} \\ f_{\text{decomp}}(x)=\begin{bmatrix}x \\ \textbf{0}\end{bmatrix} \\ \end{align}$$ \\ fcomp​(x)=x1:r^​fdecomp​(x)=[x0​]​​
对应的$\Delta W$为
$$\Delta W=\left[\begin{array}{ccc}M& & \text{0}\\ \text{0}& & \text{0}\end{array}\right]$$
但这个方法会造成信息丢失。为了改善效果，共享$M$的行和列来实现更高效的压缩和解压缩。正式来讲，其可以表示为
$$\begin{array}{r}{f}_{\text{comp}}(x)=[\sum _{j\in g_i}{x}_j{]}_{i=1}^r\\ f_{\text{decomp}}(x)=[x_{{\tilde{g}}_i^{\prime }}{]}_{i=1}^d\end{array}$$
其中，$g$和$g^{\prime }$表示M中共享行和列的预定义组。$j\in g_i$表示第$j$维度属于$g$中第$i$个组。${\tilde{g}}_i^{\prime }$表示$g_i^{\prime }$的反转，表示第i个维度关联$g^{\prime }$中的第${\tilde{g}}_i^{\prime }$个组。对应的$\Delta W$为：
$$\Delta W_{i,j}=M_{{\tilde{g}}_i^{\prime },{\tilde{g}}_j}$$
共享行和列对于较大的秩更有效，因为$\Delta W$中仅有少量的行或列共享公共行或列。例如，对于$\Delta W\in {\mathbb{R}}^{4096\times 4096}$且$r=128$，则有$\hat{r}=1024$且$M\in {\mathbb{R}}^{1024\times 1024}$。在这种情况下，仅有4行或列共享相同的行或者列。

为了增强在小秩上的表现，不再直接压缩$x$，而是对其进行reshape。这种情况下，$f_{\text{comp}}(x):{\mathbb{R}}^k\to {\mathbb{R}}^{n\times \hat{r}}$且$f_{\text{decomp}}:{\mathbb{R}}^{n\times \hat{r}}\to {\mathbb{R}}^d$。对应的$f_{\text{comp}}$、$f_{\text{decomp}}$和$\Delta W$如下：
f comp ( x ) = [ x 1 : r ^ , x r ^ : 2 r ^ , … , x ( n − 1 ) r ^ : n r ^ ] f decomp ( x ) = concat ( x ) Δ W = [ M 0 … 0 0 M … 0 ⋮ ⋮ ⋱ ⋮ 0 0 … M ]

$$\begin{align} f_{\text{comp}}(x)&=[x_{1:\hat{r}},x_{\hat{r}:2\hat{r}},\dots,x_{(n-1)\hat{r}:n\hat{r}}] \\ f_{\text{decomp}}(x)&=\text{concat}(x) \\ \Delta W&=\begin{bmatrix} M && \textbf{0} &&\dots&&\textbf{0} \\ \textbf{0} && M &&\dots&&\textbf{0} \\ \vdots && \vdots && \ddots && \vdots \\ \textbf{0} && \textbf{0} && \dots && M \\ \end{bmatrix} \\ \end{align}$$ \\ fcomp​(x)fdecomp​(x)ΔW​=[x1:r^​,xr^:2r^​,…,x(n−1)r^:nr^​]=concat(x)= ​M0⋮0​​0M⋮0​​……⋱…​​00⋮M​ ​​​
其中$\text{concat}(x)$是指将$x$的行拼接至向量中。

三、使用跨层注意力降低KV Cache

论文地址：https://arxiv.org/pdf/2405.12981

​ KV cache在自回归LLM加速解码中扮演重要角色。但是，对于长序列和大batch size情况下，存储KV cache需要大量的显存。本文提出了一种相邻层之间共享key和value的新注意力机制，称为跨层注意力(CLA)。

1. MQA和GQA

​ 在原始Transformer中使用MHA，其中每个query都会关注由不同key/value头所生成的keys和values。在MHA中，每个key/value头的KV激活都必须单独存储在KV cache中，每个token的总存储开销为$2\cdot n_{query}\cdot d_{head}$，其中$n_{query}$是query头的数量，$d_{head}$是每个头的嵌入维度。

​ 为了减少KV cache的显存占用，提出了MQA(Multi-Query Attention)和GQA(Grouped-Query Attention)。GQA对单层Transformer的query头进行分组，每组query头共享单个key/value头。KV cache的大小只和key/value头有关，而和query头无关，GQA将KV cache的显存开销降低为$2\cdot n_{group}\cdot d_{head}$，其中$n_{group}$表示分组数量且$n_{group}<n_{query}$。MQA可以看作是GQA在$n_{group}=1$情况下的特例。

​ MQA和GQA能够显著降低KV cache的大小和解压延时，同时只会导致精度的小幅降低。

2. 跨层共享KV激活

​ 受MQA和GQA启发，其在单个层内跨不同query头共享key/value头，这里提出跨层共享key/value头。称这种注意力机制为CLA(Cross-Layer Attention)，如上图1所示。CLA仅计算模型中部分层的key/value投影，没有计算key/value投影的注意力块则重用前一层的KV激活。

​ CLA与MQA/GQA/MHA正交，可以与任意方法合并。CLA2表示相邻层之间共享KV投影，CLA3则表示每3层一组共享KV投影。

3. 系统设计的影响

​ CLA的主要作用是减少KV Cache的显存占用，对模型训练和推理消耗的其他资源影响很小。这里从系统工程的视角总结了一些影响：

• KV Cache显存占用：显著减少显存占用；
• 训练显存占用：CLA减少了训练期间KV激活张量的显存占用，虽然这部分激活要比模型的hidden states和MLP激活小。
• 模型并行：CLA完全兼容标准的张量并行，流水线并行的话则需要共享KV cache的层在同一个流水线阶段。
• 参数和FLOPs：CLA能够轻微降低模型的参数量和前后向传播的FLOPs。
• 解码延时：能够支持更大的batch size和更长的KV cache，能够改善推理延时。
• 核心注意力延时：不同于MQA和GQA，CLA对于每个解码步内注意力机制所消耗内存带宽没有直接影响，因为每个注意力层都会独立的读取内存。因此，CLA对于核心注意力的延时没有直接影响。

4. 一些实验结论

• 与GAQ或者MQA相比，CLA能够更好的实现准确率/显存的权衡；
• 实验中，每2层共享的效果最好；
• CLA和MQA相结合能够更好的减低KV cache；
• CLA模型需要比非CLA模型更高的学习率；

四、Octopus v4：语言模型图

论文地址：https://arxiv.org/pdf/2404.19296

1. Octopus v2中用于分类的语言模型

​ 在Octopus v2中，引入了一种称为functional token的分类方法，其能够高效处理任务
$$P(f,params|q)$$
其中$f$表示某个候选集合中的选项，$params$对query $q$中的信息进行重新格式化。在Octopus v2中将其运用在function calling中，$f$是某个特定的函数，$params$是从$q$中抽取的函数参数。functional token也可以运用在类似的场景，即从特定候选池中选择最优项，然后对query进行格式化后传递给后续节点。

2. 语言模型作为图中节点

考虑有向异构图定义为：
$$G=(N,E)$$
其中$N$表示图中节点，$E$表示连接节点的边。节点分为两种类型：(1) master节点$N^m$：将查询导向合适的worker节点；(2) worker节点$N^w$：传递执行任务所需要的信息。worker节点从master节点接收信息并执行任务。节点信息传递的过程如上图2所示。处理用户queries $q$并生成应答$y$可以建模为：
$$P(y|q)=P(y|q;G)$$
由于单步任务仅涉及单个worker节点，过程可以定义为：
$$P(y|q;G)=P(N^w,q_h|q;N^m)P(y|q_h;N^w)$$
其中$P(N^w,q_h|q;N^m)$使用Octopus v2模型来选择$N^m$的最优邻居worker节点以及重新格式化的query $q_h$。$P(y|q_h;N^w)$是由合适的worker节点计算的似然概率。

​ 对于多步任务，例如多智能体workflows，整个过程涉及到多个节点的序列交互，即
$$P(y|q;G)=\prod _{i=0}^{k-1}P(N_i^w,q_{h_i}|q;N_i^m)P(u|q_{h_i};N_i^w)$$
其将单步任务扩展至多步，每个步骤可能由不同节点处理，然后由各种的master节点进行协调。每个步骤处理query的不同部分并为最终的输出做出贡献。

​ 在一个基于各种语言模型的预定义图中，每个worker模型也可以是一个Octopus模型。若要执行并行函数调用，master节点可以将query导向相同节点多次来执行并行调用。相比于像GPT-4这样的模型，这种设计方式的优势是回答用户问题时，仅需要激活两个小模型，而不是一个巨大的模型。

3. 使用图作为多步操作的任务规划

​ 在多步任务规划中，合并多个步骤必不可少。通常，所有可用功能都在上下文中列出并输入到模型中，然后基于用户的queries生成计划。但是对于少于10B的模型来说，处理冗长的函数描述很难。而且这种方法没有考虑不同函数描述之间的关联性。为了应对这些挑战，构建一个映射各种节点相关性的图是可行方案。这种方式有两个层次的抽象：(1) 对于每个语言模型，可以应用functional token来使得其成为单个AI agent；或者单语言模型可以是一个用于问答、写作的普通语言模型；(2) 可以创建另一个Octopus v3模型，以便从不同的节点中进行选择。

4. Functional token和数据收集

​ 将每个模型概念化为不同的函数，利用functional tokens来激活特定的模型。该方法简化了语言模型的函数设计，仅需要一个输入参数和输出结果。此外，对于特定的模型，可以在函数的doc string中详细说明所需的prompt模板。

5. 语言模型图的系统设计

​ 在复杂图的系统架构中，每个节点表示语言模型，利用多个Octopus模型进行协调。将整个系统划分为几个可管理组件：

• Worker节点部署：每个worker节点$N^w$都对应一个独立的语言模型。利用无服务架构来部署这些节点，且限制模型参数在10B内。
• Master节点部署：master节点采用小于10B的模型，以便在边缘设备部署。
• 通信：worker节点和master节点可能分布在不同的设备上。因此，网络通信对节点间数据传输至关重要。虽然master节点位于智能设备上，但是worker节点能够托管在云上或者备用设备上。

五、SimPO：一种Reference-Free的简单偏好优化

论文地址：https://arxiv.org/pdf/2405.14734

1. DPO

​ DPO使用最优策略的封闭解来重参数化奖励函数$r$：
$$\begin{gathered} r(x,y)=\beta \log \frac{{\pi }_{\theta }(y|x)}{{\pi }_{\text{ref}}(y|x)}+\beta \log Z(x) \\ \text{\hspace{1em}(1)} \end{gathered}$$
其中${\pi }_{\theta }$是策略模型，${\pi }_{\text{ref}}$是reference策略，$Z(x)$是配分函数。将该奖励函数合并至Bradley-Terry目标函数$p(y_w\succ y_l|x)=\sigma (r(x,y_w)-r(x,y_l))$，从而DPO利用策略模型而不是奖励模型来表达偏好数据的概率，目标函数为：
$$\begin{gathered} {\mathcal{L}}_{\text{DPO}}({\pi }_{\theta };{\pi }_{\text{ref}})=-{\mathbb{E}}_{(x,y_w,y_l)\sim \mathcal{D}}\left[\log \sigma \right(\beta \log \frac{{\pi }_{\theta }(y_w|x)}{{\pi }_{\text{ref}}(y_w|x)}-\beta \log \frac{{\pi }_{\theta }(y_l|x)}{{\pi }_{\text{ref}}(y_l|x)}\left)\right] \\ \text{\hspace{1em}(2)} \end{gathered}$$
其中$(x,y_w,y_l)$是由prompt、正样本和负样本构造的偏好样本对。

2. 一个简单的Reference-Free奖励

​ 奖励和DPO生成的差异。使用等式(1)作为潜在的奖励存在以下确定：(1) 训练时需要reference模型${\pi }_{\text{ref}}$，增加了额外的显存和计算开销；(2) 在训练时被优化的奖励和推理时的生成存在差异。具体来说，在生成时，策略模型${\pi }_{\theta }$被用于生成近似最大平均对数似然的序列，即
$$\begin{gathered} p_{\theta }(y|x)=\frac{1}{|y|}\log {\pi }_{\theta }(y|x)=\frac{1}{|y|}\sum _{i=1}^{|y|}\log {\pi }_{\theta }(y_i|x,y_{<i}) \\ \text{\hspace{1em}(3)} \end{gathered}$$
在解码时直接最大化该度量比较困难，各种解码策略可以用于近似，例如：贪心解码、beam search和top-k采样等。在DPO中，任意满足奖励排序$r(x,y_w)>r(x,y_l)$的三元组$(x,y_w,y_l)$并不意味着满足似然排序$p_{\theta }(y_w|x)>p_{\theta }(y_l|x)$。事实上，DPO训练时仅有50%的三元组满足这个条件。

​ 长度归一化奖励公式。直觉上，使用公式(3)中的$p_{\theta }$替换DPO中的奖励公式，这样就能够与指导生成的似然度量保持一致。这会带来一个长度归一化奖励：
$$\begin{gathered} r_{\text{SimPO}}(x,y)=\frac{\beta }{|y|}\log {\pi }_{\theta }(y|x)=\frac{\beta }{|y|}\sum _{i=1}^{|y|}\log {\pi }_{\theta }(y_i|x,y_{<i}) \\ \text{\hspace{1em}(4)} \end{gathered}$$
其中$\beta$是控制奖励差异的缩放因子。实验发现长度归一化奖励至关重要，若删除长度归一化导致模型生成更长但质量较差的序列。

3. SimPO目标函数

​ 目标奖励间隔。在Bradley-Terry目标中引入了目标奖励间隔项$\gamma >0$，确保正样本的奖励$r(x,y_w)$能够比负样本的奖励$r(x,y_l)$至少大$\gamma$：
$$\begin{gathered} p(y_w\succ y_l|x)=\sigma (r(x,y_w)-r(x,y_l)-\gamma ) \\ \text{\hspace{1em}(5)} \end{gathered}$$
两个类别的间隔会影响分类器的泛化能力。在实践中发现，随着目标间隔的增大能够观察到生成质量的改善，但是当间隔太大时则会下降。

​ 目标函数。最终的SimPO的目标函数为：
$$\begin{gathered} {\mathcal{L}}_{\text{SimPO}}({\pi }_{\theta })=-{\mathbb{E}}_{(x,y_w,y_l)\sim \mathcal{D}}\left[\log \sigma \right(\frac{\beta }{|y_w|}\log {\pi }_{\theta }(y_w|x)-\frac{\beta }{|y_l|}\log {\pi }_{\theta }(y_l|x)-\gamma \left)\right] \\ \text{\hspace{1em}(6)} \end{gathered}$$

4. 实验

​ 如上表4所示，SimPO始终显著优于现有的偏好优化方法。此外，消融实验显示SimPO的两个关键设计都至关重要。

​