想象一下你正在读一本非常长的书,每当你翻一页时,你都需要记住之前读过的所有内容,才能理解下一句话。对于像这篇论文中的计算机人工智能(AI)来说,这种“记忆”被称为 KV Cache。
随着故事变得越来越长,这种记忆会变得巨大。这就像是为了读再一页,就得背着一整个图书馆去旅行。最终,背包变得太重(耗尽了计算机的所有内存),以至于你无法继续阅读了。
这篇论文提出了一个巧妙的两部分解决方案,让你的背包变得更轻、更好用。
第一部分:“分形地图”(存储技巧)
通常,当计算机试图节省空间时,它们会将数据压缩成一个巨大的、混乱的团块。为了以后找到特定的句子,它们必须解开整个团块,这非常缓慢。
作者建议了另一种方式:分形地图 (The Fractal Map)。
想象你有一张巨大的、神奇的城市地图。
- 规则: 每当你向记忆中添加一个新词时,你就在这张地图上迈出小小的一步。
- 神奇之处: 这张地图经过设计,如果你为“苹果”这个词迈出一步,你会落在某个特定的微型社区里。如果你接着为“派”迈出一步,你会落在“苹果”社区内部的一个特定位置。
- 结果: 你对整个故事的记忆不再是一串单词列表;它只是这张地图上的一个点。
- 如果你想知道最后一个词,你看那个点,看看它位于哪个微型社区。
- 如果你想知道最后两个词,你就看那个点,找出倒数第二个社区,以此类推。
为什么这很酷?
- 它是无损的: 你可以从那一个点完美地重建出原始的精确单词。
- 它很快: 你可以瞬间跳转到故事中的任何一点(随机访问),而无需先阅读整张地图。
- 它是可搜索的: 因为地图是基于几何构建的,如果你在寻找像“猫坐在”这样的短语,你只需寻找在特定模式下彼此靠近的点即可。你不需要阅读文本来寻找模式;那个点的形状本身就是模式。
第二部分:“智能收缩”(压缩技巧)
在将记忆转化为地图上的一个点之前,AI 需要先缩小数据。论文测试了如何压缩 AI 记忆中的“键”(Key)和“值”(Value)两部分。
把 AI 的记忆想象成两个人的对话:
- 键 (Keys): 这些像是“问题”或“标签”,决定了应该关注什么。
- 值 (Values): 这些像是“答案”或实际的内容。
论文发现了一个有趣的失衡现象:
- “键”很脆弱: 如果你搞砸了“问题”(过度压缩它们),AI 就会对该看什么感到困惑。这就像给某人一张模糊的地图;他们可能会看错街道。
- “值”很强韧: 如果你把“答案”弄得稍微有点模糊,AI 通常仍能理解大意。这就像听到一个略显模糊的声音;你仍然能明白意思。
解决方案: 作者创建了一个“混合背包”。他们非常小心地打包“问题”(键)(使用更多空间),而对“答案”(值)则打包得比较松散(使用更少空间)。这节省了大量的空间——比原始数据小了 36 倍——而仅仅让 AI 在预测下一个词时的准确度下降了一点点(大约降低了 11%)。
大局观
这篇论文结合了这两个想法:
- 使用“智能收缩”方法(区别对待问题和答案)来缩小数据。
- 将缩小的数据存储在“分形地图”上。
超能力:
由于数据存储在这个分形地图上,AI 可以做到一件了不起的事情:它可以在不“解压”文件的情况下搜索自己的过去。
如果 AI 需要找到 500 页前读到的某个特定句子,它不需要加载整本书。它只需查看地图,找到匹配的点,就能瞬间知道那个句子的位置。这就像拥有一个图书馆,你只需通过观察书架上灰尘的颜色就能找到特定的书,而无需把书从书架上取下来。
结论摘要
- 存储: 他们创造了一种存储 AI 记忆的方法,这种方法具有完美的准确性、极快的访问速度且易于添加内容。
- 压缩: 他们发现压缩“问题”(键)比压缩“答案”(值)要困难得多,并利用这一点节省了 36 倍的空间。
- 搜索: 存储方法本身就是一个搜索引擎,允许 AI 瞬间找到其过去记忆中的模式。
- 范围: 他们在特定的、小型 AI 模型(GPT-2)及其 1,000 词的上下文环境下进行了测试。他们尚未在大型模型或现实世界任务中测试此方法,但在标准笔记本电脑上,其数学逻辑和代码运行完美。
技术摘要:分形 KV 缓存档案 (Fractal KV-Cache Archives)
问题陈述
在长上下文环境下为 Transformer 语言模型提供服务,已从计算密集型挑战转向内存密集型挑战。键值(KV)缓存是自回归注意力机制所必需的,其大小随上下文长度线性增长,往往超过模型权重并限制了批处理大小。虽然现有的解决方案侧重于有损压缩(量化、剔除、卸载),但本文解决了一个下游且通常是隐性的问题:一旦 KV 状态通过向量量化被简化为符号流,该流应该如何存储?
作者认为,存储层不应是一个惰性的数据块,而应是一个能够支持不断增长的缓存之特定访问模式的主动组件:O(1) 摊还追加(随着新 token 到达)和 O(1) 随机访问(以检索过去的位置),同时可能提供诸如原地检索等附加功能。
研究方法
1. 存储原语:收缩迭代映射码 (Contractive Iterated-Map Codes)
本文建议使用一种无损收缩迭代映射码(具体为混沌游戏表示法 Chaos Game Representation 和通用序列映射 Universal Sequence Maps 的推广)作为量化 KV 索引的序列化层。
- 机制: 将一个字母表大小为 N 的符号序列映射到二维平面上的轨迹。固定一个正 N 边形,每个顶点对应一个符号。每个符号 ck 通过收缩映射 pk=V(ck)+r(pk−1−V(ck)) 将当前点 pk−1 映射为新点 pk,其中 r 是收缩率。
- 特性:
- 无损性: 只要收缩率 r 足够小以确保每个符号拥有不相交的“单元格”,该映射就是可逆的。
- 高效性: 编码为线性时间。解码允许通过从最近的存储锚点重建必要的跨度来随机访问任何位置。
- 自相似性: 可达到的点集形成一个分形吸引子(例如,N=3 时为谢尔宾斯基三角形)。
- 实现: 为了减轻反转过程中的浮点误差放大,编解码器存储每个固定长度跨度的一个点,并通过锚定后向搜索重建跨度。
2. 量化策略
作者在 GPT-2 (124M) 上进行了 1024 个 token 上下文的受控研究,以确定喂给该档案的最佳量化方案:
- 精确窗口: 为了保持性能,保留一个小的精确窗口(4 个“注意力汇聚”token + 32 个近期 token)。其余上下文将被归档。
- 量化变体:
- 池化 vs 每头 (Pooled vs. Per-Head): 比较由所有注意力头共享的单一码本与每个头独立的码本。
- 残差向量量化 (RVQ): 将向量存储为来自 s 个码本的索引之和。
- 键/值不对称性 (Key/Value Asymmetry): 研究量化键 (K) 与量化值 (V) 对性能影响的差异。
3. 原地检索
本文利用了迭代映射码的一个几何特性:共享相同长度为 s 的后缀的点在空间上是接近的(距离 ≈2rs)。这允许直接通过 2D 向量执行近似子串查询,而无需解压周围的文本。
关键结果
1. 编解码器性能
分形编解码器作为一个高性能的存储原语,独立于模型:
- 编码: 0.68 µs/字符(线性时间)。
- 随机访问: 对于 1M 字符的文档,单次查找耗时 311 µs。
- 追加: 175 µs/token(摊还)。
2. 量化发现 (GPT-2, 1024 tokens)
- 每头优于池化: 在相同的比特预算下,每头码本的表现显著优于池化码本。与池化 RVQ 相比,每头 RVQ(深度为 2)将困惑度增加控制在 +15.0%,而池化 RVQ 为 +29.2%。
- 键/值不对称性: 量化键对困惑度的破坏力大约是量化值的 4 倍。
- 仅量化值(键精确):+4.0% 困惑度。
- 仅量化键(值精确):+14.5% 困惑度。
- 这证实了键决定了注意力路由(高敏感度),而值是被平均化的(误差衰减)。
- 最优混合方案: 为键分配更多比特(4 级 RVQ)并为值分配较少比特(2 级)能实现最佳权衡:
- 压缩率: 相对于 fp16 缓存实现了 36 倍 的缩减。
- 困惑度代价: +11.2%(特别是在档案占据主导地位的“后半部分”上下文中为 +16.6%)。
- 存储: 1025 字节/token。
3. 检索能力
- 精确召回: 通过对存储点之间的距离进行阈值处理,可以实现对匹配长度为 ℓ 的后缀的所有位置的 1.00 召回率。
- 精确度: 由数值精度控制。双精度索引实现 0.89–1.00 的精确度;单精度则作为粗略过滤器。
- 效率: 在 2D 空间中对 100K 个位置进行暴力最近邻搜索耗时 ~0.9 ms/查询。
- 解码: 匹配的上下文是从匹配的向量向后解码,无需访问外部文本。
意义与主张
本文声称其主要贡献有三点,并将其与关注码本几何或向量存储引擎的同期工作(如 FibQuant, RetroInfer)区分开来:
- 统一的档案格式: 它将收缩迭代映射码定义为一个实用的 KV 档案,该档案同时是一个无损存储、随机访问解码器以及原地搜索索引。这种将存储与检索统一在单一表示中的做法被视为其创新贡献。
- 量化见解: 它量化了键/值不对称性作为一个特定的困惑度预算,证明了比特不对称混合方案(键 4 倍比特,值 2 倍比特)是长上下文压缩的最优运行点。
- 可搜索内存: 它展示了该档案支持直接在压缩表示上进行后缀匹配,从而实现检索增强解码,使模型能够在不实例化完整缓存的情况下定位相关过去的上下文。
作者承认的局限性
- 规模: 实验局限于 GPT-2 (124M) 和单个 1024-token 上下文。尚未在更大规模的模型上测试结果(尽管这些模型可能对量化更具容忍度)。
- 码本训练: 码本是针对特定语料库训练的,而非在数据间摊还。
- 检索范围: 检索严格限于序列结构匹配(后缀),而非基于嵌入的语义检索。
- 压缩率: 分形编解码器的原始压缩率与标准字节打包器(byte-packers)相当。其优势在于访问模式(随机访问/搜索)而非纯粹的大小缩减。
- 基准测试: 本文并未针对 FibQuant 等同期工作的特定码本设计进行基准测试,并承认将这些量化器与此存储格式结合是未来的工作方向。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。