MambaAct, Mamba로 플로 매칭 VLA 추론 가속

배경

비전-언어-행동(VLA) 모델은 로봇이 자연어 지시와 시각적 관찰을 연결해 조작 작업을 수행하도록 하는 핵심 기술로 주목받고 있습니다. 그러나 실제 제어에서는 작업을 이해하는 능력만큼 빠른 행동 생성도 중요합니다. 플로 매칭 정책은 노이즈가 있는 행동을 여러 단계에 걸쳐 디노이징하며, 이 과정에서 행동 전문가를 반복 호출합니다. 결과적으로 행동 전문가의 계산량이 실시간 제어의 주요 병목이 됩니다.

MambaAct는 이 반복 호출 경로를 직접 줄이는 방법입니다. 전체 VLA를 다시 설계하는 대신 π0.5의 사전학습된 시각·언어 프리픽스 인코더를 유지하고, Transformer 기반 행동 전문가만 소형 Mamba 기반 전문가로 교체합니다.

핵심 설계

• 행동 시퀀스 처리에 Mamba 적용: Mamba를 행동 믹서로 사용해 노이즈가 있는 연속 행동 시퀀스의 시간적·구조적 의존성을 모델링합니다. 이를 통해 행동 생성 부분을 더 작은 구성으로 만들 수 있습니다.
• 시각·언어 조건 압축: 모든 디노이징 단계에서 전체 시각·언어 토큰과 토큰 단위 교차 어텐션을 반복하면 계산 비용이 커집니다. Context Compressor는 프리픽스 토큰을 압축된 조건 표현으로 바꿉니다.
• Mamba 블록 조절: 압축된 조건 신호는 각 Mamba 행동 블록의 게이팅과 변조에 사용됩니다. 전체 토큰 시퀀스를 매번 처리하지 않아도 지시와 장면 정보가 행동 생성에 반영됩니다.
• 기존 인코더 재사용: 시각·언어 표현을 새로 학습하는 대신, 반복적으로 실행되는 행동 전문가를 집중적으로 경량화합니다.

공개된 결과

프로젝트에 따르면 MambaAct의 행동 전문가는 기존 전문가보다 약 10배 작습니다. LIBERO에서는 플로 매칭 디노이징 지연이 64.8% 줄었고, 실제 로봇 배포에서는 58.8% 감소했습니다. 성능도 유지되거나 개선되었습니다. 서로 다른 세 개의 미세조정 시드에서 LIBERO 평균 성공률은 1.08%포인트 상승했으며, LIBERO-Plus와 실제 환경 벤치마크에서도 개선 추세가 유지되었습니다.

이 결과의 의미는 한 번의 호출 비용을 줄이는 효과가 디노이징 전체 과정에서 누적된다는 데 있습니다. 또한 행동 생성 단계에서 매번 전체 시각·언어 토큰에 어텐션을 수행하지 않고도 조건 정보를 전달할 수 있음을 보여줍니다.

의미와 남은 과제

MambaAct는 시각·언어 인코더와 행동 생성기의 역할을 분리합니다. 인코더는 장면과 지시를 풍부하게 표현하고, 행동 전문가는 압축된 문맥을 받아 연속 제어 신호를 효율적으로 생성합니다. 이는 연산 자원이나 제어 주기에 제약이 있는 로봇에서 플로 매칭 VLA를 활용하는 데 유용한 방향입니다.

다만 제공된 자료만으로는 하드웨어 구성, 전력 소비, 디노이징 단계 수에 따른 변화, 다른 경량 행동 전문가와의 체계적인 비교까지 확인하기 어렵습니다. 다양한 로봇 형태와 작업 분포로의 일반화도 추가 검증이 필요합니다. 코드는 공개되어 있으며, 모델 가중치의 제공 일정은 프로젝트 공지에 따릅니다.

출처: Hugging Face Daily Papers