MambaAct、MambaでフローマッチングVLAの推論を高速化

導入

視覚言語アクション（VLA）モデルは、ロボットが指示と視覚情報を結び付け、物体操作を実行するための有力な構成になりつつあります。一方で、実時間制御においては、タスクを理解できることだけでは十分ではありません。フローマッチングを用いるポリシーでは、ノイズを含むアクションを段階的にデノイズする間、アクション専門家が何度も呼び出されます。そのため、専門家の計算量が全体のレイテンシーを押し上げます。

MambaActは、この繰り返し実行される部分に焦点を当てた設計です。π0.5の事前学習済み視覚言語プレフィックスエンコーダを残し、Transformerベースのアクション専門家だけを小型のMambaベース専門家へ置き換えます。

主な仕組み

• アクション混合にMambaを利用：Mambaによって、ノイズを含む連続アクション列の中にある依存関係をモデル化します。これにより、アクション生成向けのシーケンス処理をより小さな構成で実現します。
• 条件情報を圧縮：デノイジングの各反復で、視覚言語トークン全体とのクロスアテンションを行うのはコストが高くなります。Context Compressorは、プレフィックス出力を小さな条件表現へ変換します。
• 各ブロックを変調：圧縮された条件信号は、各Mambaブロックのゲートおよび変調に使われます。完全なトークン列を毎回処理しなくても、指示や観測に基づく文脈をアクション生成へ伝えられます。
• 既存の表現を活用：VLA全体を作り直すのではなく、繰り返し呼び出されるアクション経路を重点的に軽量化しています。

報告された結果

プロジェクトの報告によると、MambaActのアクション専門家は元の構成のおよそ10分の1です。LIBEROでは、フローマッチングのデノイジング遅延が64.8%減少し、実ロボット環境では58.8%減少しました。また、独立した3つの微調整シードでLIBEROの平均成功率が1.08ポイント向上し、LIBERO-Plusと実世界ベンチマークでも改善が維持されています。

アクション専門家はデノイジング中に繰り返し実行されるため、1回あたりの計算削減がエンドツーエンドの速度向上へ積み上がる点が重要です。同時に、視覚言語条件を毎回フルのトークン単位アテンションで処理しなくても、制御性能を保てる可能性を示しています。

意義と今後の確認点

MambaActは、視覚言語エンコーダとアクション生成器の役割を分離する設計として捉えられます。前者はシーンと指示の情報を抽出し、後者は圧縮された文脈を受けて連続アクションの構造を効率よく処理します。計算資源や制御周期に制約のあるロボットへの適用を考えるうえで、実用的な方向性です。

ただし、提示された資料だけでは、ハードウェア構成、消費電力、デノイジング回数ごとの挙動、他の軽量アクション専門家との体系的比較までは判断できません。異なるロボット形態やタスク分布への一般化も、今後の検証が必要です。コードは公開されており、モデル重みの提供時期はプロジェクトの告知に従います。

出典：Hugging Face Daily Papers