LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破

文章圖片

LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破



本文共同第一作者蔣體通 , 清華大學直博五年級學生 , 研究方向是VLA、自動駕駛和人機交互等 。 共同第一作者蔣雪楓 , 中國科學院計算技術研究所直博五年級學生 , 研究方向聚焦弱監督學習 , 多模態大模型應用和生成式自動駕駛等 。 本文通訊作者朗咸朋 , 理想汽車智能駕駛副總裁 。



項目主頁:https://liauto-research.github.io/LightVLA/ 論文鏈接:https://arxiv.org/abs/2509.12594
核心創新

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 篩選的多樣性 , 我們給重要性評分注入采樣噪聲:



最后 , 篩選后的Token集可以通過以下公式得到:





為了在訓練前期鼓勵模型探索 Token 篩選的多樣性 , 而在訓練后期使 Token 篩選的策略收斂 , 我們對采樣噪聲的方差進行線性衰減 , 使噪聲方差從 1 逐漸衰減至 0 。

實驗結果





LightVLA 在指標上顯著超越現有基線:在 LIBERO 基準上的所有任務中 , LightVLA 均取得了最佳表現 , 平均成功率達到 97.4% , 全面超越包括 OpenVLA-OFT 在內的各類強基線模型 。 這表明 LightVLA 在兼顧效率的同時仍能保持領先的任務執行能力 。 Token 稀疏性揭示冗余視覺信息:與消耗 512 個視覺 token 的 OpenVLA-OFT 相比 , LightVLA 僅保留平均 78 個視覺 token , 卻仍實現更優性能 。 這一結果表明 , 大量視覺 token 并未貢獻有效信息 , LightVLA 成功捕捉關鍵語義 token , 證明了視覺模態的高度稀疏性 。 唯一兼顧性能與效率的加速方案:在與其他加速方法的對比中 , LightVLA 不僅將 FLOPs 與延遲分別減少 59.1% 與 38.2% , 同時還提升平均成功率 2.6% 。 值得注意的是 , LightVLA 是現有所有加速方法中唯一一個在加速的同時還能提升性能的方案 , 驗證了消除視覺冗余能夠同時優化性能與效率 。
剪枝過程可視化



為了說明 LightVLA 剪枝過程的可解釋性 , 我們隨機選擇了任務 “把咖啡壺放在爐子上” , 并展示任務執行過程中的 token 選擇可視化結果 , 每幀圖片的左右兩列分別代表第三人稱相機輸入和腕部相機輸入 , 第二行點亮的區域代表被選擇的視覺 token 。 關鍵幀代表操作任務的重要階段(物體交互 , 任務完成等) , 可以看出被保留的視覺 token 更多地關注咖啡壺、爐子、和機械臂本體等任務相關物體 , 并沒有關注無用的背景信息 。 進一步驗證了 LightVLA 在自適應 token 選擇上的能力 。

消融實驗

噪聲衰減的有效性:

引入噪聲提升模型探索能力:不引入噪聲的變體保留最少的視覺 token , 實現了次優的性能 , 說明噪聲的引入對訓練過程中模型主動探索任務相關的視覺 token 至關重要 , 否則會導致對語義密集場景視覺 token 的 “無感” 。

噪聲衰減讓模型變得更 “聰明”:固定噪聲的引入使得模型保留最多的視覺 token , 但模型對關鍵 token 的篩選能力不足 , 噪聲衰減讓模型學會對視覺 token 的有效取舍從而提升到最優性能 。

Token 選擇有效性:

保留無用 token 導致性能下降: 當在 LightVLA 已保留的 k 個 token 之外再補充 k 個隨機 token 時 , 整體性能反而下降 , 說明 LightVLA 已經捕捉到所有關鍵信息 , 額外的隨機 token 只會引入噪聲與干擾 。

丟棄有用 token 導致性能下降:當從 LightVLA 已篩選的 k 個 token 中隨機丟棄 10% 時 , 性能同樣下降 。 充分驗證 LightVLA 學會了選擇對任務成功率更相關的視覺 token , 并沒有保留無用信息 。

【LightVLA可微分token剪枝,首次實現VLA模型性能和效率的雙重突破】結論

我們研究了視覺 - 語言 - 動作(VLA)模型中固有的視覺冗余問題 , 并提出了一種無參數的可微分視覺 token 剪枝框架 LightVLA 。 通過基于無參數查詢的 token 剪枝過程 , 該方法能夠自適應地選擇最具信息量的視覺 token 。 在 LIBERO 基準上 , LightVLA 在顯著降低計算開銷的同時取得了當前最優的性能 。 我們還提出了另一種變體 LightVLA* , 相較于 LightVLA , 其引入了可學習查詢作為額外的可訓練參數 , 同樣在性能上優于同類方法 。 本工作為解決 VLA 模型中的視覺冗余挑戰提供了新的范式 , 在實現更低計算開銷與時延的前提下取得了更優性能 , 為未來 VLA 模型輕量化與部署提供了新穎的解決方案 。

    推薦閱讀