Gist：探索语言模型的高效提示压缩之道

论文：https://arxiv.org/pdf/2304.08467v3 NeurIPS 2023

代码：https://github.com/jayelm/gisting

摘要（Abstract）

• 问题陈述：语言模型（LMs）的多任务能力主要通过提示（prompting）来实现，但提示占用了宝贵的输入上下文窗口空间，且重复编码相同的提示在计算上是低效的。

• 现有方法：微调（finetuning）和蒸馏（distillation）方法可以在不使用提示的情况下专门化语言模型，但需要针对每个任务重新训练模型。

• 提出的解决方案：本文提出了“gisting”方法，通过训练语言模型将提示压缩成更小的“gist”令牌集合，这些令牌可以被缓存和重用，以提高计算效率。

• 方法：通过简单地修改 Transformer 注意力掩码（attention masks）来鼓励提示压缩，可以在不增加额外成本的情况下训练 gist 模型。

• 结果：在解码器（LLaMA-7B）和编码器-解码器（FLAN-T5-XXL）语言模型上，gisting 实现了高达 26 倍的提示压缩率，从而减少了高达 40% 的浮点运算（FLOPs），提高了 4.2% 的墙钟时间（wall time）速度，并节省了存储空间，同时最小化了输出质量的损失。

引言（Introduction）

• 背景：每天处理数百万查询的 ChatGPT 如果不优化，将使用具有二次时间和内存复杂度的自注意力机制重复编码提示。

• 挑战：即使使用缓存策略，随着缓存提示数量的增加，仍然会产生内存和存储成本。

• gisting 方法：提出了一种新方法，通过压缩提示到“gist”令牌来减少提示长度，从而在大型规模下实现显著的计算、内存和存储节省。

Gisting 方法（Gisting）

• 指令微调：描述了指令微调的过程，即学习分布 pLM(y | t, x)，其中 t 是任务的自然语言提示，x 是任务的输入，y 是期望的输出。

• 压缩提示：提出了学习压缩版本的 t，即 G(t)，使得使用 pG(y | G(t), x) 的推理比 pLM(y | t, x) 更快。

• 上下文蒸馏视角：通过蒸馏的角度来看待 gisting，即通过修改注意力掩码来鼓励将提示压缩到 gist 令牌中。

通过掩码学习 Gisting（Learning Gisting by Masking）

• 方法：使用语言模型本身作为 gist 预测器 G，通过修改 Transformer 注意力掩码来强制执行提示压缩。

• 实现：在模型词汇表和嵌入矩阵中添加单个 gist 令牌 g，然后修改注意力掩码，使得输入令牌在 gist 令牌之后不能关注提示令牌之前的任何令牌。

实验（Experiments）

• 数据集：创建了一个名为 Alpaca+ 的数据集，结合了 Self-Instruct 和 Stanford Alpaca 指令微调数据集。

• 模型：在 LLaMA-7B 和 FLAN-T5-XXL 模型上进行了实验，使用不同数量的 gist 令牌 k ∈ {1, 2, 5, 10}。

• 评估：使用了 ROUGE-L、ChatGPT 和人类评估来评估模型性能。

结果（Results）

• 性能：在 Alpaca+ 分布内的未见提示上，gist 模型显示出与正面控制模型相当的性能。

• 人类评估：人类评估结果与 ChatGPT 评估一致，显示出 gist 压缩模型与正面控制模型相比具有竞争力。

讨论与局限性（Discussion and Limitations）

• 效率提升：gisting 可以显著提高计算效率，但可能在某些边缘情况下无法完美压缩提示。

• 未来工作：提出了将 gisting 集成到现有指令微调工作流程中的可能性，以及通过“gist 预训练”来提高压缩性能的想法。

这篇论文的核心贡献是提出了一种新的方法来压缩语言模型中的提示，以提高计算效率，同时保持输出质量。通过实验验证了该方法的有效性，并讨论了其在实际应用中的潜力和局限性。

本文首发于公众号：AI智语狂想，欢迎关注。