当前位置: 首页 » 资讯 » 科技头条 » 正文

Stability AI让AI"记忆"更省空间

IP属地 中国·北京 科技行者 时间:2026-05-29 00:23:42


这项由Stability AI研究团队开发的技术成果以预印本形式发布于2026年5月,论文编号为arXiv:2605.21226,感兴趣的读者可通过该编号检索完整技术文献。

每当你和AI聊天、看AI生成的视频或者听AI创作的音乐时,AI的大脑深处都在飞速运转一套叫做"注意力机制"的系统。这套系统有点像人类的短期记忆——它需要随时记住之前说过的每一句话、每一帧画面、每一个音符,然后在生成新内容时参考所有这些记忆。这些临时记忆在技术上叫做"KV缓存"(Key-Value Cache),可以简单理解为AI的"工作记忆本"。

麻烦在于,这个工作记忆本实在太占地方了。当AI需要处理超长文本(比如一百万字的小说)或者高质量视频时,这个记忆本会变得极其庞大,以至于撑爆AI服务器的内存,导致速度变慢、费用飙升、甚至直接崩溃。这就像你需要在一张小便条纸上记下整部《红楼梦》——根本写不下。

Stability AI的研究团队提出了一套他们称为"OCTOPUS"的技术方案,专门解决这个"记忆本太厚"的问题。这个名字是"通过八面体参数化和最优平方误差量化优化Transformer的KV缓存"的英文缩写,听起来相当复杂,但核心思路其实可以用一个非常直观的比喻来理解:**把记忆本里的内容压缩成速记符号,而且是经过数学精心设计的、损失最小的速记方式。**

一、AI的"工作记忆"为什么这么难压缩

要理解这项研究解决的问题,先要搞清楚AI在处理长内容时究竟在做什么。

每当AI模型生成一个新词、一帧新画面或一个新音符时,它都必须"回看"之前所有的内容,判断接下来应该生成什么。这个"回看"的过程用到的就是KV缓存——"K"代表"键"(Key),可以理解为每段记忆的"索引标签";"V"代表"值"(Value),是记忆的实际内容。每生成一步,AI都要把当前的查询信息(Query)与缓存中所有的K进行比对,找到最相关的内容,再结合对应的V生成输出。

问题是,随着内容越来越长,这个缓存会持续膨胀。一个处理百万字上下文的大语言模型,其KV缓存可能占据整个推理过程内存消耗的大半壁江山。而且每生成一步,系统都要把整个庞大的缓存从内存里读出来一次——这个读取操作本身就成了速度瓶颈。

过去研究者们想了很多办法。有人提出"丢掉不重要的记忆"(Token Eviction),就像做笔记时只记关键词;有人提出对每个数值进行粗略的四舍五入(标量量化);还有一类更精妙的方案,叫做"旋转预处理量化"(Rotation-preconditioned Quantization),OCTOPUS就属于这一类,并且在这条路上走得更远。

二、旋转魔法:让数据变得"均匀"再压缩

旋转预处理的思路来自一个巧妙的数学观察。原始的KV缓存数据往往分布不均匀——有些数值很大,有些很小,直接压缩会造成很大误差。但如果先对数据做一个随机旋转(类似于把所有信息在高维空间里均匀搅拌一下),旋转后的每个数值就会服从一种数学上已知的规律,从而可以用专门匹配这种规律的压缩器来处理,达到近乎最优的压缩效果。

这个"旋转操作"在实现上用的是一种叫做"符号翻转沃尔什-哈达玛变换"(sign-flipped Walsh-Hadamard Transform)的技巧,计算效率极高,每次操作的复杂度只与维度成对数关系,而不是平方关系。可以把它理解为一种超高效的"信息均匀混合器"。

TurboQuant和PolarQuant是OCTOPUS之前这条路上的两个代表性工作。TurboQuant证明了旋转后每个坐标独立服从一种叫做"对称Beta分布"的规律,并针对这个规律设计了最优的逐坐标压缩器;PolarQuant则把旋转后的方向信息用极坐标的方式逐层分解再压缩。这两种方法的共同点是:**每次只处理一个坐标(或一个角度)**。

OCTOPUS的核心突破是:**一次处理三个坐标组成的"三元组",把三维方向信息和长度信息分开,用不同的精度来压缩。**

三、八面体折叠:把三维方向"压平"成二维地图

OCTOPUS的名字里有"八面体"(Octahedral),这个词来自计算机图形学领域一个非常实用的技巧:如何用最少的数据精确表示空间中的一个方向向量。

考虑这样一个问题:空间中有无数个方向(就像地球上有无数个地理位置),如何用两个数字来表示其中任意一个方向,同时让误差尽可能小?

传统的方法是用经纬度(对应极坐标),但经纬度有个缺陷:靠近两极的地方被严重"拉伸",同样大小的面积在极点附近被表示成很窄的条带,而在赤道附近则很宽——信息密度严重不均匀。

八面体映射解决了这个问题。它把球面上的点先投影到一个八面体(由八个等边三角形组成的立体图形)的表面,再把这个八面体展开铺平,得到一个正方形。这个展开过程就像剥开一个橙子皮并压平——每块三角形皮的面积在展开前后保持不变,信息密度更均匀。

具体来说,对于空间中一个单位方向向量 (x, y, z),先计算 l = |x| + |y| + |z|,然后把向量除以 l 投影到八面体表面;如果 z 方向朝上(正半球),直接取 (px, py) 两个坐标;如果 z 方向朝下(负半球),就做一次"折叠"变换,把下半球的四个三角形翻折到正方形的四个角落。最终得到的两个数 (ξ, η) 都落在 [-1, 1] 的范围内,完整覆盖了所有可能的方向。

还原时也很简单:给定 (ξ, η),计算 r = 1 - |ξ| - |η|,如果 r 不小于零就直接构造方向向量并归一化;如果 r 小于零就先做反折叠再归一化。整个编解码过程只用加减乘除和取绝对值,没有任何三角函数,计算速度极快。

这种映射的优秀特性在于:每个八面体面内,从球面到正方形的变换是"等面积"的,也就是说球面上均匀分布的方向映射到正方形后依然近似均匀。这让后续针对这两个坐标的压缩器能够最大限度地利用每个比特。

四、三元组分解:长度和方向分开压,精度分配有讲究

OCTOPUS处理数据的完整流程像是一条精心设计的流水线。

第一步,对每个原始的K向量,先把它分解为整体长度 γ 和单位方向 u。整体长度用32位浮点数存储(相对于总数据量极小),后续所有的压缩预算都用于压缩这个单位方向向量。

第二步,对单位方向向量施加前面提到的符号翻转沃尔什-哈达玛旋转,得到旋转后的方向 u,它的每个坐标都服从数学上已知的分布规律。

第三步,把旋转后的方向向量切割成若干个"三元组"(每组三个相邻坐标),如果总维度不是3的倍数就在末尾补零。对于每个三元组 t?,再次分解为三元组长度 ρ? 和三维单位方向 n?。

这里有个重要的数学结论:根据论文中的Lemma 3.1,当原始向量维度 d 很大时,每个三元组的长度 ρ? 都集中在一个很小的数值 √(3/d) 附近——维度越高,三元组长度就越小且越集中。这意味着三元组长度的"变化范围"(方差)随维度增大而缩小,因此不需要太多比特就能精确描述它;相比之下,三维方向 n? 始终分布在整个球面上,变化范围始终是量级为1的,需要更多比特才能描述好。

第四步,对三维单位方向 n? 施加八面体映射,得到两个坐标 (ξ?, η?),再对这两个坐标和三元组长度 ρ? 分别用"劳埃德-马克斯量化器"(Lloyd-Max Quantizer,一种针对特定概率分布设计的最优压缩器)进行压缩。

劳埃德-马克斯量化器的工作原理类似于把一段连续的数轴切割成若干段,每段用一个"代表值"(码本中的一个质心)来替代。切割点的位置不是均匀分布的,而是根据数值出现的频率来优化的——出现频率高的区域切割得更细、代表值更密集,出现频率低的区域则相对稀疏。这样每个比特都用在了"刀刃上"。

五、非均匀比特分配:数学优化得出的"偏科"策略

OCTOPUS最精妙的创新之一是它的比特分配策略。

假设对每个三元组的总比特预算是固定的,应该如何在两个方向坐标 (ξ, η) 和一个长度 ρ 之间分配这些比特?最直觉的做法是三等分,每个分量用相同的比特数。但OCTOPUS的研究者通过严格的数学推导证明,这种均匀分配并不最优。

核心论据来自两个数量的量级差异:三元组长度的方差 σ?_ρ 随维度 d 的增大以 O(1/d) 的速度缩小,而方向误差贡献的方差 σ?_n 始终是 O(1) 的量级。把这两个数量代入拉格朗日优化框架(一种求解约束条件下最优化问题的数学工具),可以推导出最优的比特差值:

方向比特数 - 长度比特数 = log?(3·C?·σ?_n / (2d·C_ρ·σ?_ρ))

虽然式子里有很多符号,但把 σ?_ρ = O(1/d) 和 σ?_n = O(1) 代入后,d 互相抵消,结论是:**这个比特差值是一个常数,与维度 d 无关,也与总比特预算无关。**

研究者在维度 d=128 的合成数据上进行了系统性验证,在每个总预算 b(代表均匀分配时每个分量用 b 比特)周围,尝试了从 b-2 到 b+2 的所有偏移量 δ(即方向用 b+δ 比特、长度用 b-δ 比特)。结果非常清晰:在 b=2、3、4 三种情况下,均匀分配(δ=0)不是最优,最优点始终出现在 δ=+1,即**方向多分配一个比特、长度少分配一个比特**,记为 (b+1, b-1) 分配方案。

相比均匀分配,这个 (b+1, b-1) 方案把均方误差降低了31%到41%;而如果反过来让方向少分配一个比特(δ=-1),误差则会上升超过200%;δ=+2 会使误差上升44%到73%。数据的凸性非常清晰,最优点不偏不倚地落在 δ=+1 处。

六、联合取整:让三个数值协同工作而不是各自为政

量化中还有一个微妙的问题:即使有了最优的码本,如何从码本中选取最终输出的值?

最简单的方法是对每个数值独立地找最近的码本条目(标量取整)。但OCTOPUS的解码过程是非线性的——方向坐标经过八面体逆变换得到三维单位向量,再与长度相乘得到最终的三维向量。这个非线性组合意味着:独立地将每个坐标取整到最近码本值,并不能保证最终组合结果与真实值最近。

换句话说,在压缩"红"、"绿"、"蓝"三个颜色通道时,如果最终目标是让混合出来的颜色最准确,那么对每个通道单独取最近值并不是最好的策略——可能稍微偏离某个通道的最近值,但最终混合颜色反而更准。

OCTOPUS提出了一种"联合取整"方案。推导过程颇为精妙:把三元组的重建误差展开,可以证明对于任意固定的方向候选值,最优的长度量化值并不是最接近真实长度 ρ? 的码本条目,而是最接近**当前方向候选向量与真实三元组向量的点积** s? 的码本条目。这意味着方向选择和长度选择可以解耦:先找最优方向,再根据点积选最优长度。

寻找最优方向理论上需要遍历码本中所有 2^(2×b_dir) 个候选对,代价很高。但研究者发现,在所有测试的比特宽度下,标量取整得到的初始方向索引 (iξ, iη) 与真正最优的索引相差最多一个位置。因此,只需要检查以初始索引为中心的3×3邻域(9个候选对)就足够了,计算量极小。跨越10000个随机三元组的验证表明,这个3×3局部搜索的结果与全码本穷举搜索完全一致,没有任何区别。

七、实验结果:从文本到视频到音频的全面验证

OCTOPUS的研究者在四个不同场景下验证了这套方案的效果,对比的基准包括TurboQuant-MSE(逐坐标最优压缩)、TurboQuant-QJL(在TurboQuant基础上加1比特残差修正)和PolarQuant(用极坐标分解方向信息)。

**合成数据验证**:在维度128的各向同性高斯数据上,用1024个随机键向量和16个查询向量测试,每个实验重复64次取平均。结果显示,OCTOPUS在每个比特宽度下都达到最低的均方误差:在4比特时,OCTOPUS的均方误差(0.0071)比TurboQuant-MSE(0.0094)低约24%,比PolarQuant(0.0145)低约51%;在最极端的2比特压缩时,OCTOPUS的误差(0.0897)仅是PolarQuant(0.2197)的约41%,差距更为显著。余弦相似度和内积误差的趋势完全一致。在"针中找针"测试(在2048个随机干扰项中找到目标向量)中,2比特时OCTOPUS保留了0.92的softmax概率质量,而PolarQuant只有0.87,TurboQuant-MSE只有0.86——这意味着OCTOPUS在极端压缩下仍然能大致保持检索能力。

**大语言模型测试**:在Qwen2.5-7B-Instruct-1M(一个70亿参数的大型语言模型)上测试WikiText-2和C4数据集的语言建模困惑度(perplexity,数字越低越好,代表模型对语言的掌握更准确)。在4比特时,OCTOPUS的WikiText-2困惑度为10.306,相比基准(10.033)只上升了2.7%;TurboQuant-MSE上升了3.1%,PolarQuant上升了4.4%,TurboQuant-QJL上升了8.0%。

在2比特的极端压缩下,差距变得触目惊心:OCTOPUS的困惑度为13.517,相比基准上升34.7%;而TurboQuant-MSE上升63%,PolarQuant上升187%,TurboQuant-QJL则飙升772%,完全失控。简单来说,TurboQuant-QJL在2比特时的语言理解能力几乎归零,而OCTOPUS还保持着相当的可用性。

在"长文本找针"测试(在4000到128000个词的上下文中找到隐藏的随机8字符密码)中,4比特时所有方案都能完美找到;3比特时OCTOPUS保持完美(1.00),PolarQuant下降到0.86平均;2比特时只有OCTOPUS(0.81)和OCTOPUS-QJL(0.83)还保持可用水平,PolarQuant和TurboQuant-QJL完全崩溃(分别为0.04和0.01)。

**视频生成测试**:在两个基于Wan-1.3B架构的自回归视频生成模型上测试。CausVid以3帧为单位分块生成,Causal Forcing逐帧生成。在100个相同提示词和相同初始噪声下运行所有方案,测量生成视频与未压缩版本的感知差异(LPIPS,越低越好)、信噪比(PSNR,越高越好)等指标。

在4比特时,各方案差异不大(都在3%以内)。在2比特时,Causal Forcing上,TurboQuant-QJL的最坏情况LPIPS达到0.997(几乎等于随机噪声),均值为0.816;而OCTOPUS的均值仅为0.581,最坏情况0.821——虽然有质量损失,但至少还保持了图像的基本可辨识性。从论文中提供的视频帧对比图来看,2比特时TurboQuant-QJL和PolarQuant的输出变成了彩色像素噪点,完全无法辨识内容,而OCTOPUS的输出虽然有些模糊和色彩偏差,但场景内容依然清晰可辨。

**音频生成测试**:在AAR(一个基于下一尺度预测的自回归音频生成模型)上,对100个随机AudioSet音频片段进行测试,衡量生成音频与未压缩版本在对数谱距离(LSD,越低越好)和信噪比(SNR,越高越好)上的差异。在2比特时,TurboQuant-MSE、TurboQuant-QJL和PolarQuant的平均LSD都在12.6到13.2分贝之间,平均SNR全部为负值(即噪声比信号还大);而OCTOPUS的平均LSD仅为6.75分贝,平均SNR为+1.07分贝——仍然保持正向的信号质量。

八、可选的残差修正:让内积估计更无偏

除了基础版OCTOPUS,论文还提供了一个可选的增强版本OCTOPUS-QJL,在压缩结果的基础上附加了一个1比特的"残差修正"。

这个修正基于一种叫做Johnson-Lindenstrauss变换(JL变换)的技术,来自之前的QJL工作。压缩后的向量与真实向量之间存在一个残差 r,OCTOPUS-QJL对这个残差再做一次独立的随机旋转,只存储旋转结果的符号(每个坐标正还是负,占1比特)以及残差的长度(用16位浮点数存储)。在计算查询与键的内积时,利用这个1比特符号向量对内积估计进行修正,使修正后的估计值在数学期望意义上是无偏的。

代价是每个坐标额外增加1比特,相当于总压缩率略有下降(通常降低约0.5比特/坐标)。从实验结果来看,OCTOPUS-QJL在内积绝对误差上比OCTOPUS进一步降低,但在重建质量(均方误差、LPIPS等)上几乎没有区别——因为QJL修正只影响内积计算路径,不改变重建的键向量本身。研究者因此建议:OCTOPUS-QJL适合那些用压缩键直接计算注意力分数的部署场景,而如果还需要实际重建键向量,OCTOPUS基础版就够用了。

九、工程实现:让解码在"寄存器"里完成,不留痕迹

OCTOPUS不只是一个理论方案,还有配套的高效实现。研究者用Triton(一种专门为GPU编写高性能内核的编程语言)实现了压缩和解压缩的全部流程。

特别值得一提的是解码器的设计:它与闪存注意力(FlashAttention)的"分块计算"思路深度融合,在GPU的寄存器(最快速的片上存储)中直接完成解压缩、内积计算和在线Softmax的全部运算,整个解码过程中**从未把完整的未压缩键向量写入内存**。这意味着OCTOPUS没有为解压缩操作额外增加内存读写带宽——唯一需要从内存读取的是压缩后的紧凑比特流,而这正是压缩本来就要节省的东西。

编码器同样经过优化,使用基于Kronecker因子的沃尔什-哈达玛变换实现就地蝴蝶运算,每个键向量的编码时间在维度128时约为0.08微秒,对于每次只生成一个新词元的自回归大语言模型来说完全可以忽略不计。

从速度测试来看,在NVIDIA H200 GPU上,OCTOPUS的单步解码时间约为0.52到0.66毫秒(视具体配置而定),约为未压缩的PyTorch SDPA基准路径(0.06毫秒)的8到11倍。这个额外开销是内在的——任何在线解压缩方案都需要额外的计算来还原信息。但这与2.8到4.8倍的内存节省之间的权衡,在内存或带宽成为瓶颈的场景(超长上下文、大批量推理)下往往是值得的。

归根结底,OCTOPUS解决的是一个非常实际的工程与数学交叉问题:如何在有限的内存预算下,让AI的"工作记忆"既小又准。它的答案是:把记忆中的每个小片段分解为"朝哪个方向"和"有多长"两部分,用八面体折叠把三维方向压平成二维,再根据数学推导把有限的比特预算向"方向"那边多倾斜一个比特——因为长度信息天然更容易被少数比特描述,而方向信息更难被粗略量化。

这项技术对普通用户的影响是间接但真实的:它让AI服务在相同硬件上能处理更长的文本、生成更长的视频、理解更多的上下文,同时保持相近的质量水平。在极端压缩(2比特)的场景下,OCTOPUS是目前已知唯一一个不会彻底崩溃的旋转预处理量化方案,这对于内存极度受限的边缘部署场景尤为重要。

当然,这项技术也有其局限:解码计算量更大,在内存并不是瓶颈的场景下可能反而降低效率;该技术目前主要针对旋转预处理量化家族,与其他完全不同的压缩思路(如稀疏编码、Token驱逐)的比较还需要进一步研究。

对这项技术感兴趣的读者,可以通过arXiv编号2605.21226找到完整的论文和数学推导,研究团队也提供了项目主页(octopus-quant.github.io)供参考。

Q&A

Q1:KV缓存压缩和普通的视频压缩有什么区别?

A:KV缓存压缩针对的是AI推理过程中产生的中间计算结果,而不是最终的输出视频或文本。它压缩的是AI在生成内容时需要"记住"的键值信息,目标是让AI在内存有限的情况下处理更长的输入。普通视频压缩(如H.264)针对的是像素数据,两者在应用场景和技术原理上完全不同,不能互相替代。

Q2:OCTOPUS的(b+1, b-1)比特分配为什么是最优的,用更多比特给方向不会有更多收益吗?

A:这是因为方向误差和长度误差在总误差中的贡献受数学结构制约。当方向比长度多了超过一个比特后,方向的量化误差已经足够小,继续加比特的边际收益急剧下降;而此时长度只有一个比特(b-1),长度误差反而成为瓶颈。实验数据显示,(b+2, b-2)分配比均匀分配的误差反而更高,说明比特差值存在一个明确的最优点,而这个最优点恰好在差值为1处。

Q3:OCTOPUS在哪些实际AI产品场景中最有用?

A:最有价值的场景是那些内存带宽或容量成为推理瓶颈的地方。典型例子包括:超长上下文的大语言模型服务(处理几十万词的文档)、高分辨率长视频的自回归生成(每帧都需要参考之前所有帧的缓存)、以及在算力受限的设备上部署大模型。在这些场景中,OCTOPUS能在牺牲较少质量的前提下将内存占用压缩到原来的三分之一甚至四分之一,让原本不可能完成的任务变得可行。

免责声明:本网信息来自于互联网,目的在于传递更多信息,并不代表本网赞同其观点。其内容真实性、完整性不作任何保证或承诺。如若本网有任何内容侵犯您的权益,请及时联系我们,本站将会在24小时内处理完毕。