MOSAIC: Multi-Objective Search for Allocating Inference Cache Across LLM Decoder Layers
MOSAIC 是一个多目标神经架构搜索框架，通过优化长文本大语言模型推理中的逐层 KV 缓存预算，在多种模型和驱逐方法中一致地优于均匀分配方案，实现了更高的任务得分并显著降低了内存使用量。
原始论文根据 CC0 1.0（http://creativecommons.org/publicdomain/zero/1.0/）发布到公有领域。 这是对下方论文的AI生成解释。它不是由作者撰写的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
大语言模型是当今许多最先进人工智能工具背后的引擎，它们能够阅读长篇文档或进行长时间的对话。然而，这些模型在处理长文本序列时面临着一个物理限制：它们必须在计算机内存中存储不断增长的临时信息，以追踪已经阅读过的内容。这个被称为键值缓存（key-value cache）的存储区域会随着文本长度的增加而线性扩张。对于试图理解整本书籍长度的文档的模型来说，这种内存需求可能会变得如此庞大，以至于超过标准计算机硬件的容量，迫使系统要么停止运行，要么丢弃信息。
为了解决这个问题，研究人员开发了关于保留哪些信息的方法，使其具有选择性。这些技术就像一个过滤器，决定哪些文本部分足够重要而值得保留在内存中，哪些可以被安全地遗忘。直到现在，这些过滤器对模型内部处理的每个部分都应用相同的严格规则。当时的假设是，模型大脑的每一层都需要相同数量的内存来完成工作。这项新研究挑战了这一假设，认为不同的层实际上需要不同数量的空间，而给予它们相等的份额是一种资源的浪费。
一个研究小组推出了一种名为 MOSAIC 的新方法，它可以自动计算出为大型语言模型的每一层分配的最佳内存量。MOSAIC 不再将模型视为一个每个部分都获得相等内存份额的统一整体，而是将各层视为具有独特需求的个体。该系统使用一种复杂的搜索过程来测试成千上何种不同的内存分配方式。它在寻找一种特定的模式，即在保持总内存限制的同时，让某些层获得更多空间，而另一些层获得的更少。其目标是找到一种排列方式，使模型在回答问题或总结文本时最为准确，且不会耗尽内存。
研究人员在两个流行的语言模型上测试了这种方法，并将其与给每一层分配相同内存的标准方法进行了对比。他们在各种困难的任务上对模型进行了评估，包括复杂的阅读理解、代码生成和长文档总结。结果显示，在绝大多数情况下，这种新方法始终优于旧方法，在 154 个测试设置中有 146 个胜过了均匀分配法。在某些情况下，改进是显著的，模型在长上下文理解的标准测试中提升了近 50 分。或许更重要的是，新方法使模型能够以比以前少多达四倍的内存实现同样的高性能。这意味着模型可以处理更长的文本，或者在更小、更便宜的计算机上运行。
最令人惊讶的发现之一是，最佳的内存分布并非随机的。搜索揭示了一个一致的模式：模型中间的一小部分层始终比开头或结尾的层需要更多的内存。开头的层和结尾的层可以用较少的空间正常运作，而中间的层则是信息需要被保留的关键瓶颈。通过将内存从不需要的层转移到需要的层，系统变得更加高效。这一发现表明，这些模型的内部结构并不是均匀的；网络的不同部分扮演着不同的角色，并且需要不同的资源才能发挥最佳性能。
研究人员还发现，这种最优模式可以被发现一次后，便可应用于许多不同的内存限制。他们不需要为用户可能想要的每一种内存大小都进行一次新的、昂贵的搜索。相反，他们找到了一个用于内存分布的“获胜形状”，并只需对其进行缩放即可。这意味着，一次性的搜索就能为未来的任何内存约束提供解决方案，使该方法在实际应用中具有可行性。该系统通过在模型用于回答问题之前进行离线学习理想的平衡，因此在实际使用过程中不会增加额外的延迟。
与最近解决此问题的其他尝试相比，这种新方法探索了更为广阔的可能性空间。虽然其他方法一次只测试几种变化，但该系统评估了总可能配置中的极小一部分——在扫描比以往方法大数十亿倍的搜索空间时，仅评估了几百个候选方案以找到最佳解。尽管考察了如此多的选项，它所需的计算时间却更少。这项研究证实，将每一层视为相同对象的旧观念是低效的。通过认识到某些部分比其他部分更关键并据此分配资源，我们可以使这些强大的工具变得更快、更便宜，并能理解比以往认为的要长得多的上下文。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。