Llmcad: Fast and scalable on-device large language model inference

题目：Llmcad: Fast and scalable on-device large language model inference 发表于2023.09

链接：https://arxiv.org/pdf/2309.04255

声称是第一篇speculative decoding+边缘设备的论文（不一定是绝对的第一篇），不开源，没有github地址

动机

生成性任务，在移动应用领域占据重要的地位。由于对隐私问题的敏感性，越来越多的需求是直接在移动设备上执行这些任务。目前，这些生成性任务的执行严重依赖于LLM。然而，这些设备内存容量十分有限。
在本研究中，引入了LLMCad，这是一个专为高效生成性NLP任务设计的设备端推理引擎。

注意，它的目标模型也是在移动设备上的，只不过不放在内存中，只在验证时期从存储中加载到内存；小LLM是常驻内存的。

核心思想

LLMCad的核心思想围绕模型协作：一个紧凑的LLM驻留在内存中，负责生成最简单的令牌，而一个高精度的LLM则介入以验证这些令牌并纠正识别出的错误。LLMCad融合了三种新技术：

• LLMCad不是以顺序方式生成候选令牌，而是使用较小的LLM构建一个token树。随后，较大的LLM可以同时验证所有这些路径。

• 它采用自适应回退策略，每当较小的LLM生成错误令牌时，迅速启动验证过程。

• 为了确保token生成的连续流动，LLMCad在验证过程中通过实施计算-IO管道来推测性地生成令牌。

论文基于以下观察：虽然较小的LLM不足以生成令人满意的端到端句子，但它们通常能够正确生成大多数简单的标记（例如，限定词、代词和标点符号）。目标LLM验证N个令牌可以在目标模型的一次推断中完成，因此比顺序生成N个令牌要快得多。

 策略

论文指出遇到的挑战以及解决策略：

• 小LLM可能产生的一些次优token，实际上它认为的次优token是目标LLM的最优输出。这就需要在验证的过程中，也考虑到那些次优token，即使用tree attention

• 需要确定何时启动验证过程，设备上的验证是耗时的，例如Jetson TX2上需要7.1秒。过早验证会导致检测不到错误，过晚验证会生成无用token，浪费了移动设备计算资源。先前的工作通常依赖于单个标记或标记序列长度，这可能无法准确地定位最佳验证时机。本文中token tree中比较累积不确定性、基于历史数据评估是否需要进行验证。

• 目标LLM的验证会阻塞小LLM的推理。所以在验证过程中，继续让小LLM执行推理。同时为了保证小模型推理不影响验证，继续推理只在不影响大LLM内存占用的情况下运行。

下图是整个流程图：

设计

论文的一些设计细节：

• 在小LLM生成token树的过程中，任何置信度高于阈值的token都会被生成一个分支(例如0.3)

• 树累计置信度(tree cumulative-confidence)，列的公式和字符看起来挺吓唬人，实际上就是同一个序列的每个token的置信度累乘

• 关于如何并行地验证树，论文也是给了较长的介绍，但实际上应该是用了其他论文中出现的tree attention机制。详情可以看第一个提出tree attention机制的specInfer论文，在我的博客里也有这篇文章的解读。