LightVLA 是一個旨在提升 VLA 推理效率且同時提升性能的視覺 token 剪枝框架。當前 VLA 模型在具身智能領域仍面臨推理代價大而無法大規模部署的問題,然而大多數免訓練剪枝框架依賴于中間注意力輸出,并且會面臨性能與效率的權衡問題。為應對這些挑戰,LightVLA 引入了兩大核心創新:
- 無參數可微分token剪枝框架:創新的應用無參數查詢初始化和 Gumbel softmax 技術實現訓練時可微分 VLA 模型能夠根據多樣的文本任務輸入自適應地選擇對任務完成最有貢獻的關鍵視覺 token,驗證了性能和效率可以做到協同優化。
- 基于可學習查詢的 token 剪枝框架:相比于無參數的 LightVLA,LightVLA * 初始化一系列的可學習查詢(Learnable Query),可分別作用于視覺編碼器或 LLM 淺層解碼器層,借助額外參數引導 VLA 模型學習關鍵視覺 Token 選取,同樣實現了較好的性能提升。
我們研究了 VLA 模型中廣泛存在的視覺 token 冗余,設計了一種在微調中實現可微分視覺 token 剪枝的加速框架,創新實現 Gumbel-softmax 引導的無參數 token 選擇過程,強化 LightVLA 對關鍵視覺 token 的選擇能力,為 VLA 模型的推理加速提供新的范式。
在 LIBERO 上的實驗證明,LightVLA 不僅取得了當前最佳性能(SOTA),超越了 Pi-0 、Openvla-OFT 等經典 VLA 模型,并且實現了高效的推理加速。并且通過可學習的參數初始化 query 選擇方法 LightVLA * 驗證微調實現推理加速的可行性。消融實驗充分驗證了 LightVLA 自適應 token 選擇的有效性,并證實效率和性能并非是此消彼長的零和博弈,為構建更高效的具身智能大模型提供了新的思路。
研究動機與核心挑戰 (Motivation)
讓 VLA 學會更聰明地選擇關鍵視覺 token
當前視覺 - 語言 - 動作(VLA)模型在機器人任務中展現了令人矚目的認知推理和操作等能力,但龐大的計算開銷和較高的推理延遲限制了其廣泛部署,如家用機器人。
我們認為計算瓶頸的核心大部分來源于視覺 token 的固有冗余,而傳統基于 Token 剪枝的加速方法往往面臨 “效率 vs 性能” 的權衡困境,現有工作為了提升效率而剪枝視覺 token,不可避免地造成模型性能下降。我們認為對于 VLA 的視覺輸入,冗余的視覺 token 不僅會帶來額外的計算開銷,而且稀釋了模型對于關鍵區域的注意力,同時造成性能和效率的下降。
因此,我們認為效率和性能并非天生矛盾,但需要引入更聰明的剪枝方法,而非以固定比例或固定個數限制保留 token 的數量,讓模型學會主動、任務自適應地 “關注” 信息最豐富的視覺區域,并忽略無關的冗余信息。基于此,我們實現了兩種 token 選擇機制:
- LightVLA: 它不再依賴任何啟發式的超參數預設剪枝比例,基于無參數查詢的方式,進一步引入 Gumbel-softmax 實現在微調過程實現 token 的可微分選擇,使模型自發學會保留對提升任務性能更重要的 “關鍵 token”,從而實現性能和效率的雙重提升。
- LightVLA*: 為了驗證在微調過程中剪枝的可行性以及剪枝位置的影響,我們進一步探索了基于可學習查詢的剪枝框架,在引入可訓練參數后,LightVLA * 仍然可實現性能和效率的較好提升。

上圖展示了 LightVLA 在 LIBERO 中和主流 VLA 模型、剪枝類方法在視覺 token 數量和任務性能上的對比。從圖中可以看出,LightVLA 實現了在保留 token 最少情況下最好的性能,不僅說明了視覺 token 的大量冗余,同時說明通過微調可以實現性能和效率優化的最優解。
方法詳解
總體架構示意圖

可微分的 Token 剪枝
我們提出了可微分的 Token 剪枝算法,以實現自適應剪枝。我們使用一系列 Query 來實現 Token 篩選,具體來說,LightVLA 構建了與 Visual Token 數量相同的 Query,并由每個 Query 獨立選擇一個最重要的 Visual Token。沒有被 Query 選中的 Token 被剪除,而所有被 Query 選中的 Visual Token 組成剪枝后的 Token 集。可微分的 Token 剪枝算法具體流程如下:
- Query 生成
LightVLA 使用一組 Query Token 來識別有用和無用 Token。一個 Visual Token 是否有用,由攜帶的視覺信息和 VLA 輸入的文字指令共同決定。因此,LightVLA 取 Visual Token 對 Language Token 的 Cross Attention,在視覺表征中融合任務信息,作為 Query。

- Token評分
計算Query Token與Visual Token之間的內積,作為每一個Query Token對每一個Visual Token的重要性評分。

- Token 篩選
每個 Query 獨立地選擇重要性評分最高的 Visual Token,所有被選中的 Visual Token 保留下來,而沒有被選中的 Visual Token 被剪除。

注意到 argmax 是不可導運算,在這里,我們使用 Gumbel-softmax 技巧將 argmax 變為可導運算,以實現訓練時梯度的反向傳播。首先,為了提高訓練過程中 Token 篩選的多樣性,我們給重要性評分注入采樣噪聲:
