【论文笔记】LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models

🍎个人主页：小嗷犬的个人主页
🍊个人网站：小嗷犬的技术小站
🥭个人信条：为天地立心，为生民立命，为往圣继绝学，为万世开太平。

基本信息

标题: LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models
作者: Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, Jiaya Jia
发表: ICLR 2024
arXiv: https://arxiv.org/abs/2309.12307

基本信息

摘要

我们提出了LongLoRA，一种高效的微调方法，它通过有限的计算成本扩展了预训练大型语言模型（LLM）的上下文大小。

通常，使用长上下文大小训练LLM在计算上非常昂贵，需要大量的训练时间和GPU资源。例如，在 $8192$ 个上下文长度的训练中，自注意力层的计算成本是 $2048$ 个上下文长度的 $16$ 倍。

在本文中，我们从两个方面加速了LLM上下文扩展。

一方面，尽管在推理过程中需要密集的全局注意力，但通过稀疏局部注意力可以有效地进行模型微调。提出的移位稀疏注意力（S2-Attn）有效地实现了上下文扩展，与使用标准注意力微调具有相似的性能，同时实现了显著的计算节省。特别是，它可以在训练中仅用两行代码实现，而在推理中是可选的。

另一方面，我们重新审视了参数高效的上下文扩展微调机制。值得注意的是，我们发现LoRA在可训练嵌入和归一化的前提下，对于上下文扩展效果良好。LongLoRA将这种改进的LoRA与 $S^2$ -Attn相结合。

LongLoRA在Llama2模型（从7B/13B到70B）的各种任务上展示了强大的实证结果。LongLoRA将Llama2 7B的上下文从4k扩展到100k，或将Llama2 70B扩展到32k，在单个8×A100机器上完成。

LongLoRA在保持原始架构的同时扩展了模型的上下文，并且与大多数现有技术兼容，如Flash-Attention2。

此外，我们还使用LongLoRA和我们的长指令遵循LongAlpaca数据集进行了监督微调。

我们所有的代码、模型、数据集和演示代码都可在github.com/dvlab-research/LongLoRA上找到。

简介

LongLoRA缩小了传统LoRA和全量微调之间的精度差距，同时保持了比全量微调低1.8倍的内存成本。此外，LongLoRA通过 $S^2$ -Attn 将LoRA的训练速度提高了高达1.8倍。Llama2-7B使用Flash-Attention2和DeepSpeed的第二阶段进行微调，并在 proof-pile 测试集上评估了困惑度。

$S^2$ -Attn的示意图涉及三个步骤。首先，它将特征沿头部维度分为两个部分。其次，其中一个部分中的token向右移动了组大小的一半。第三，我们将token分成组，并将它们重塑为批量维度。在我们的模型中，注意力仅在每组中计算，而信息通过移动在组之间流动。移动可能会引入潜在的信息泄露，但通过在注意力掩码上进行微小修改可以轻松防止。

LongLoRA

背景

Transformer

大型语言模型（LLMs）通常是基于 Transformer 构建的。例如，以 Llama2 为例，一个 LLM 模型由一个嵌入输入层和若干解码器层组成。每个解码器层包含一个自注意力模块。它通过带有权重矩阵 ${W_q, W_k, W_v\}$ 的线性投影层将输入特征映射为一组查询、键和值 ${q, k, v\}$ 。给定 ${q, k, v\}$ ，它计算输出 $o$ ：

$\text{softmax}(qk^T)v$

输出随后通过一个权重矩阵 $W_o$ 的线性层进行投影，接着是多层感知机（MLP）层。在自注意力模块之前和之后，会应用层归一化。所有解码器层完成后，还会进行一次最终归一化。

对于较长的序列，自注意力在计算成本方面表现出困难，其计算复杂度与序列长度成平方关系。这大幅减慢了训练过程，并增加了 GPU 内存的使用成本。

Low-rank Adaptation

LoRA假设预训练模型中的权重更新在适配期间具有较低的内在秩（intrinsic rank）。对于一个预训练权重矩阵 $\in \mathbb{R}^{d \times k}$ ，它通过低秩分解 $\Delta W = W + BA$ 来更新，其中 $\in \mathbb{R}^{d \times r}$ ， $\in \mathbb{R}^{r \times k}$ 。秩 $\ll \min(d, k)$ 。在训练期间， $W$ 被冻结（没有梯度更新），而 $A$ 和 $B$ 是可训练的。这就是 LoRA 训练比完全微调更高效的原因。

在 Transformer 结构中，LoRA 仅适配注意力权重 ${W_q, W_k, W_v, W_o\}$ ，并冻结所有其他层，包括 MLP 和归一化层。这种方式简单且参数高效。然而，我们通过实验证明，仅在注意力权重中的低秩适配并不能很好地适用于长上下文扩展任务。

Shifted Sparse Attention

标准的自注意力计算成本为 $\mathcal{O}(n^2)$ ，这使得长序列上的LLM具有高内存成本和低速。为了在训练期间避免这一问题，我们提出了移位稀疏注意力（ $S^2$ -Attn），如图2所示。接下来，我们将进行一项初步研究，并逐步解释我们的设计。

Pilot Study

在表1中，我们建立了一个标准基线，该基线经过完整注意力和微调训练和测试，在各种上下文长度下表现出一致的良好质量。第一次试验是使用短注意力进行训练，仅模式1如图2所示。正如我们所知，在长上下文中，高昂的成本主要来自自注意力模块。因此，在这次试验中，由于输入很长，我们在自注意力中将其分为几个组。例如，模型在训练和测试阶段都以8192个token作为输入，但在每个组中进行自注意力操作，组大小为2048，组数为4。这种模式效率很高，但在非常长的上下文中仍然不起作用，如表1所示。随着上下文长度的增加，困惑度变大。其背后的原因是没有不同组之间的信息交换。

Effectiveness of S2-Attn under different context lengths

为了引入组之间的通信，我们包括了一个移位模式，如图2所示。我们在半注意力头中将组分区移位半个组大小。以总体8192个上下文长度为例，在模式1中，第一组从第1个到第2048个token进行自注意力。在模式2中，组分区移位1024。第一个注意力组从第1025个开始到第3072个token结束，而前1024个和最后1024个token属于同一组。我们在每个半自注意力头中分别使用模式1和模式2。这种方式不会增加额外的计算成本，但能够实现不同组之间的信息流。我们在表1中展示了它接近标准注意力基线的结果。

Consistency to Full Attention

现有的高效注意力设计也可以提高长上下文LLM的效率。然而，大多数这些设计并不适合长上下文微调。因为这些从头开始训练的Transformer与预训练中使用的标准全注意力存在差距。在表6中，我们展示了 $S^2$ -Attn 不仅能够实现高效的微调，还支持全注意力测试。尽管其他注意力机制也可以用于长上下文微调，但模型必须使用微调期间使用的注意力进行测试。移位防止了模型对特定注意力模式的过度拟合。

Easy Implementation

$S^2$ -Attn易于实现。它仅涉及两个步骤：

在半注意力头中移位token；
将特征从token维度转置到批次维度。

两行代码就足够了。我们在算法1中提供了一个PyTorch风格的代码示例。

Algorithm 1: Pseudocode of S2-Attn in PyTorch-like style.

Improved LoRA for Long Context

LoRA是一种高效且流行的将LLMs适应其他数据集的方法。与全微调相比，它节省了大量的可训练参数和内存成本。然而，将LLMs从短上下文长度适应到长上下文长度并不容易。我们观察到LoRA与全微调之间存在明显的差距。如表2所示，随着目标上下文长度的增加，LoRA与全微调之间的差距逐渐增大。并且，具有更大秩的LoRA无法缩小这个差距。

Finetuning normalization and embedding layers is crucial for low-rank long-context adaptation