谷歌：缓存增强优化冻结LLM的推理性能

标题：Deliberation in Latent Space via Differentiable Cache Augmentation
来源：arXiv, 2412.17747

摘要

通过生成和处理中间推理步骤，使大型语言模型（LLM）能够“思考更多”的技术在解决复杂问题方面显示出了希望。然而，标准方法在响应之前立即生成离散令牌序列，因此它们可能会产生巨大的延迟成本，并且难以优化。
在这项工作中，我们证明了冻结的LLM可以通过在模型的键值（kv）缓存上运行的离线协处理器来增强。该协处理器通过一组潜在嵌入来增强缓存，旨在提高后续解码的保真度。我们使用解码器在标准预训练数据上的语言建模损失来训练这个协处理器，同时保持解码器本身冻结。这种方法使模型能够以端到端可微的方式学习如何将额外的计算提取到其kv缓存中。由于解码器保持不变，协处理器可以离线异步运行，如果协处理器不可用或给定的缓存被认为不需要额外的计算，语言模型可以正常运行。
我们通过实验证明，当缓存被增强时，解码器对许多后续令牌的困惑度较低。此外，即使没有任何针对特定任务的训练，我们的实验也表明，缓存增强在一系列推理密集型任务中始终能降低困惑并提高性能。

️文章简介

研究问题：如何通过潜在嵌入增强冻结的大语言模型（LLM）的推理能力，特别是在需要复杂推理的任务中。
主要贡献：论文提出了一种可微的缓存增强方法，通过引入一个独立的协处理器模块来生成潜在嵌入，从而在不修改冻结LLM的情况下提升其推理和预测能力。

重点思路

模型架构：输入序列首先由冻结的LLM处理生成kv缓存，然后将kv缓存传递给协处理器模块，该模块输出潜在嵌入以增强kv缓存，最后将增强后的kv缓存输入LLM生成输出。
预训练策略：协处理器通过预测未来多个token（称为“ahead tokens”）进行训练，生成的潜在嵌入被插入到输入序列中，目标token随后被预测。训练过程中，协处理器和软token的权重被优化，而LLM的权重保持冻结。
实验设置：使用冻结的Gemma-2 2B模型，在包含2万亿token的预训练数据集上进行训练，训练过程中保持LLM的权重不变，仅训练协处理器模块。

分析总结

困惑度评估：模型在预测后续token时的困惑度显著降低，表明潜在嵌入增强了LLM的内部表示和生成能力。
公共基准测试：在GSM8K、MMLU、TriviaQA、NQ和MATH等推理基准测试中，模型性能显著提升，尤其是在需要复杂推理的任务中，潜在嵌入数量越多，性能提升越明显。
与现有方法比较：与Pause Token方法和零样本CoT方法相比，本文提出的方法在准确性和困惑度方面均表现更优，显示出动态生成潜在嵌入的有效性。
异步操作：由于协处理器的计算可以离线进行，且与LLM的解码过程并行，因此提高了计算效率。

个人观点

论文的核心在于训练了一个协处理器模块，为输入序列动态生成潜在嵌入，从而指导冻结LLM的推理。

附录