KV-Efficient VLA: A Method to Speed up Vision Language Models with RNN-Gated Chunked KV Cache

一、论文概述

项目	内容
标题	KV-Efficient VLA: A Method to Speed up Vision Language Models with RNN-Gated Chunked KV Cache
作者	Wanshun Xu, Long Zhuang, Lianlei Shan
机构	University of Toronto, Tsinghua University
论文	arXiv:2509.21354
代码	KV-Efficient-VLA.github.io (项目主页，暂无公开代码仓库)
发布	2025-09-20 (v1), 2025-11-23 (v2)
许可	未明确说明

二、核心思想

问题定义

Vision-Language-Action (VLA) 模型为机器人感知和控制提供了统一框架，但在推理阶段面临严重的效率瓶颈。具体而言：

注意力计算开销大：标准 Transformer 注意力机制的时间和内存复杂度为 O(n)，其中 n 为上下文长度。当保留历史图像 token 作为上下文时，KV cache 会随序列长度线性增长。
实时控制要求高：灵巧操作和动态导航等实时机器人任务通常需要 >50-100 Hz 的控制频率，而现有 VLA 模型（如 OpenVLA 7B 约 6 Hz、HybridVLA 7B 约 6.1 Hz）远不能满足。
历史帧冗余严重：研究表明，直接包含未处理的历史帧会降低约 6% 的性能，说明需要更高效的时间上下文整合方法。

解决方案概述

KV-Efficient VLA 提出了一种模型无关的内存压缩方法，通过引入轻量级机制选择性保留高价值上下文。核心思想是：

分块策略（Chunked KV Strategy）：将 KV cache 分割为固定大小的块，通过 MLP 聚合每个块的表示。
LSTM 门控机制（LSTM Gating Mechanism）：使用 LSTM 门控模块评估每个聚合块的重要性，根据学习到的效用分数决定保留或丢弃。
保留近期细节：保留最近的 r 个 token 不压缩，确保细粒度的近期信息不丢失。

三、技术架构

整体框架图

输入序列: [历史 KV Cache (n tokens)] → [分块] → [MLP 聚合] → [LSTM 门控] → [选择性保留]
                                                    ↓
                                              [保留最近 W tokens 不压缩]
                                                    ↓
                                              [合并: 压缩块 M + 近期窗口 W]
                                                    ↓
                                              [注意力计算: n' = W + M << n]

框架包含两个主要步骤：

将 KV cache 分割为固定大小的段并聚合
使用 RNN 模块更新和选择保留/丢弃哪些块，实现高效的选择性压缩

核心公式

标准注意力机制：

$\text{Attention}(Q_t, K_{1:n}, V_{1:n}) = \text{softmax}\left(\frac{Q_t K_{1:n}^\top}{\sqrt{d_k}}\right) V_{1:n}$

其中在解码步骤 t 时，每个注意力层形成查询 $Q_t \in \mathbb{R}^{H \times d_k}$ 并与存储的键值对 $\{K_{1:t-3}, V_{1:t-3}\}$ 进行注意力计算。

KV Cache 分块聚合：

将 KV cache 分割为长度为 $C \in \mathbb{N}$ 的非重叠块。对于时间 t 的每个块 $C_t = \{(k_j, v_j)\}_{j=1}^{C}$ ，使用 MLP 计算代表性向量：

$\bar{K}_t = \text{MLP}_K(\{k_j\}_{j=1}^{C}), \quad \bar{V}_t = \text{MLP}_V(\{v_j\}_{j=1}^{C}), \quad \bar{K}_t, \bar{V}_t \in \mathbb{R}^{B \times 1 \times d}$

聚合后，每个块被压缩为单个键值对 $[\bar{K}_t, \bar{V}_t]$ ，形状为 $B \times H \times 1 \times d$ 。

LSTM 门控机制：

使用隐藏大小为 $d_g$ 的 LSTM 门控模块，维护隐藏状态 $h_t$ 并产生门控分数 $s_t \in [0, 1]$ ：

$h_t, s_t = \text{LSTM}([\bar{K}_t, \bar{V}_t], h_{t-1})$

缓存更新策略： 使用可学习阈值 $\tau$ 。如果 $s_t \geq \tau$ ，则保留压缩块；否则丢弃。最近的 $r < W$ 个 token 保持未压缩状态以保留细粒度的近期信息。

基准模型 FLOPs（每层每 token）：

$\text{FLOPs}_{\text{base\_attention}} = 4Bnd_{\text{model}}^2 + 2BHn^2 d_k$

$\text{FLOPs}_{\text{FFN}} = 3Bnd_{\text{model}}d_{ff}$

KV-Efficient 注意力 FLOPs：

$\text{FLOPs}_{\text{KV\_Eff}} = 4Bn'd_{\text{model}}^2 + 2BH{n'}^2 d_k + \frac{n}{C}BH(4d_{\text{sum}}d_{\text{model}} + 4d_g d_{\text{model}} + 4d_g^2)$

其中 $d_{\text{sum}}$ 为 MLP 聚合器的隐藏维度， $n' = W + M \ll n$ 为有效内存长度。

加速比：

$\text{Speed up} = \frac{\text{FLOPs}_{\text{base}} + \text{FLOPs}_{\text{SFF}}}{\text{FLOPs}_{\text{KV\_Eff}} + \text{FLOPs}_{\text{SFF}}}$

内存加速比： $\frac{n}{W+M}$ ，反映由于压缩和选择性保留带来的内存占用减少。

模型组件

组件	说明	关键参数
视觉编码器	DINO + SigLIP ViT-Large	ViT-SO400M-14-SigLIP 预训练权重
LLM 骨干	LLaMA-2 7B	$d_{\text{model}}=4096$ , $L=32$ , $H=32$ , $d_k=128$
注意力机制	Grouped Query Attention	8 个 KV heads
MLP 聚合器	2 层 MLP	$d_{\text{sum}}=128$
LSTM 门控	LSTM 门控模块	$d_g=128$
输出头	融合 GELU MLP	无额外对齐模块
微调方法	LoRA	通过 LLaMA Factory 框架

训练流程

数据集：

Open X-Embodiment 数据集：超过 500,000 个演示，涵盖 22 种机器人 embodiment 和 500+ 任务
每个样本包含：RGB 图像、语言指令、机器人状态、低层连续动作

仿真环境：

RLBench 环境（基于 CoppeliaSim 仿真器）：100+ 个任务，每个任务约 1,000 个专家演示，总计 100,000+ 轨迹
多模态观测：多视角 RGB 图像、深度图、分割掩码、本体感觉状态

实验设置：

基础 VLM：prism-dinosiglip-224px+7b
仅在 HybridVLA 模型上进行微调
使用混合精度训练，2 块 NVIDIA H800 GPU
图像输入：224×224 RGB
序列长度：平均约 n ≈ 20,000 tokens（包含视觉特征、任务指令、机器人状态、RGB 图像）

超参数：

最近窗口大小：W = 4096
块大小：C = 3,136
平均保留比率：r = 0.687
LSTM 隐藏大小： $d_g = 128$
MLP 聚合器隐藏维度： $d_{\text{sum}} = 128$

四、核心创新

创新点	说明	理论/实验依据
分块 KV 策略	将完整序列缓存分割为固定大小块，通过 MLP 聚合，大幅减少内存占用	理论分析：有效内存长度 n’ = W + M << n
LSTM 门控机制	使用 LSTM 评估每个聚合块的重要性，实现选择性保留	平均保留比率 r = 0.687，即仅保留约 68.7% 的历史上下文
模型无关设计	方法可无缝集成到 OpenVLA、CogACT、HybridVLA 等不同 VLA 架构	在三种基线模型上均验证有效
保留近期细节	最近的 r 个 token 保持未压缩，确保细粒度近期信息不丢失	平衡了压缩效率和信息保真度
因果性保证	步骤 t 的模型仅关注过去 token i < t，严格保持因果性	适用于自回归推理场景

五、实验结果

基准测试

理论分析结果（基于 LLaMA-2-7B）：

指标	基准模型	KV-Efficient	改进
注意力 FLOPs (每层每 token)	1,490,288,640,000	412,871,032,320	-72.3%
FFN FLOPs	1,351,756,800,000	不变	-
理论加速比	1.00×	1.61×	+61%
内存加速比	1.00×	2.44×	+144%

多任务设置下的实验结果（Table 1）：

模型	加速比	推理速度 (Hz)	总 FLOPs (T)
OpenVLA (7B)	1.00×	6.3	2.37
OpenVLA-KV-Efficient	1.22×	7.6	1.94
CogACT (7B)	1.00×	9.8	2.41
CogACT-KV-Efficient	1.33×	13.8	1.81
HybridVLA (7B)	1.00×	5.8	2.73
HybridVLA-KV-Efficient	1.47×	8.3	1.85

关键发现

FLOPs 节省：KV-Efficient 方法平均减少约 24.6% 的总 FLOPs
推理加速：平均推理速度提升约 1.34×
内存减少：KV cache 存储减少约 1.87×
参数不变：所有变体保持相同参数量，效率提升来自计算优化而非模型压缩
训练收敛：KV-Efficient 模块在前 100 次迭代内 loss 快速下降，随后稳定收敛

与现有方法对比

对比维度	OpenVLA	CogACT	HybridVLA	KV-Efficient 平均
原始推理速度	6.3 Hz	9.8 Hz	5.8 Hz	-
KV-Efficient 推理速度	7.6 Hz	13.8 Hz	8.3 Hz	9.9 Hz
加速比	1.22×	1.33×	1.47×	1.34×
原始 FLOPs (T)	2.37	2.41	2.73	2.50
KV-Efficient FLOPs (T)	1.94	1.81	1.85	1.87

六、总结

核心贡献

分块 KV 策略：提出将完整序列缓存分割为固定大小块并通过 MLP 聚合的方法，有效减少内存占用
LSTM 门控机制：设计基于 LSTM 的门控模块，根据学习到的效用分数选择性保留高价值上下文
模型无关架构：KV-Efficient 模块可无缝集成到多种 VLA 架构（OpenVLA、CogACT、HybridVLA），无需修改下游控制逻辑
理论与实验验证：通过理论成本模型分析和实证基准测试，验证了注意力 FLOPs 和 KV cache 内存使用的减少

技术影响

实时机器人控制：KV-Efficient VLA 使 VLA 模型更接近实时控制所需的 >50 Hz 频率要求
边缘部署：1.87× 的内存减少使得在资源受限的机器人平台上部署大型 VLA 模型成为可能
长时序任务：通过压缩历史上下文，模型能够处理更长的任务序列而不受内存限制
可扩展性：随着输入序列长度增加，效率优势更加显著

局限性

评估范围有限：由于计算资源限制，准确性和推理速度仅在单个基准任务上评估，可能不能完全反映模型的整体性能
缺少详细消融研究：论文提到计划进行 chunk size 和 retention gating 策略的详细消融研究，但目前尚未完成
真实世界验证缺失：尚未在真实世界的闭环机器人环境中进行部署和测试
精度与效率权衡：虽然论文展示了效率提升，但未详细讨论在不同任务和场景下精度的变化

七、参考资源

资源	链接
arXiv 论文	https://arxiv.org/abs/2509.21354
PDF 下载	https://arxiv.org/pdf/2509.21354
HTML 版本	https://arxiv.org/html/2509.21354v2
项目主页	https://KV-Efficient-VLA.github.io

文章目录

KV-Efficient VLA