オフライン隠れ状態蒸留による強剪定VLAモデルの復元

本論文は、Vision-Language-Action (VLA) モデルの言語バックボーンを構造的プルーニングで圧縮した際に失われる性能を、オフラインの Hidden-State Distillation で回復する手法を提案する。OpenVLA-OFT の 63% 削減で LIBERO-Long 成功率が 93.2% から 0.8% に低下する問題に対し、オンライン rollouts や強化学習を必要とせず、教師の 1 回の forward pass で作ったキャッシュを用いて生徒を訓練する。約 8 GPU-hours で教師との差 3.5 ポイント以内まで回復し、6-DoF マニピュレータで 72% 削減時に 77.5% 成功、教師より 2.23 倍高速、メモリ 62% 削減を達成する。

先行研究では、教師あり微調整 (supervised fine-tuning) と強化学習 (reinforcement learning) を組み合わせて回復する手法があり、オンライン rollouts と数百 GPU-hours を要した。本手法は完全にオフラインで、教師の 1 回のパスで構築したキャッシュのみを用い、約 8 GPU-hours で回復する。また、9 つの削減率をスイープし、45% 削減までは回復目的が有意に効かないこと、63% から 87% で +2.1〜+4.5 ポイント、CogACT では 63% 以降で +9.4〜+22.1 ポイントの改善を示す。81% 削減の CogACT では、回復予算を 3 倍にすると蒸留生徒の教師との差が平均 3.9 ポイントに縮まる一方、教師あり回復は 20 ポイント以上下回る。

幅プルーニング (width pruning) はブロックを狭めるが residual stream を元のサイズに保つため、教師と生徒の hidden states が同じ形状となり、projector なしで直接マッチングできる。教師の 1 回の forward pass でキャッシュを構築し、それに対して訓練する。これによりオンライン rollouts 不要で、オフラインで効率的に蒸留を行う。深さプルーニング (depth pruning) と比較して、同じ圧縮率では幅プルーニングの方が成功率が高く、深さプルーニングはレイテンシが低いという特性がある。

LIBERO-Long ベンチマークで OpenVLA-OFT の 63% 削減時の成功率を評価し、教師の 93.2% から 0.8% に落ちたものを、蒸留により教師との差 3.5 ポイント以内まで回復させた。9 つの削減率をスイープし、45% まで有意差なし、63%〜87% で +2.1〜+4.5 ポイント、CogACT で 63% 以降 +9.4〜+22.1 ポイントを確認。81% 削減の CogACT で回復予算 3 倍時に平均 3.9 ポイント差まで縮小。6-DoF マニピュレータで 72% 削減時に 77.5% 成功 (教師あり回復は 59.5%)、教師より 2.23 倍高速、メモリ 62% 削減を実証。

幅プルーニングと深さプルーニングのトレードオフが議論されている。同じ圧縮率では幅プルーニングの方が成功率が高く、深さプルーニングはレイテンシが低い。また、回復目的が有効になる削減率の閾値 (45% まで有意差なし) が示され、それ以上の圧縮で蒸留の効果が顕著になる。CogACT では 63% 以降で大きな改善が見られ、回復予算を増やすことで教師との差をさらに縮められるが、教師あり回復では限界があることが示唆されている。

要旨で参照/比較されている研究として、OpenVLA-OFT、CogACT、LIBERO-Long ベンチマーク、および教師あり微調整と強化学習を組み合わせた回復手法が挙げられる。関連手法として、構造的プルーニング (structured pruning)、幅プルーニング (width pruning)、深さプルーニング (depth pruning)、Hidden-State Distillation がある。同分野の定番として、Vision-Language-Action (VLA) モデル、Vision-Language Model (VLM) の圧縮手法も次に読むべき候補である。