Efficient Diffusion-Based Action Generation in Vision Language Action Model via State Space Model

Abstract

비전–언어–행동(VLA) 파운데이션 모델은 장기 로봇 조작을 위한 통합 인터페이스로 자리 잡아가고 있지만, action 생성은 여전히 어텐션 기반 확산 트랜스포머에 주로 의존한다. 본 연구는 행동 헤드에서 대부분의 트랜스포머 어텐션 블록을 선택적 상태공간 모델(SSM)로 대체하되, 하단에 소수의 어텐션 레이어를 유지하는 하이브리드 정책 아키텍처를 탐구한다. 고정된 비전–언어 모델이 시각 관측과 자연어 지시를 컨디셔닝 토큰으로 변환한다. 이후 행동 스트림은 두 가지 SSM 구성요소로 처리된다. Cross-SSM은 비전–언어 문맥을 SSM의 선택 메커니즘에 주입하여 교차-모달 융합을 수행하고, SSM 블록은 행동 시퀀스 내 시간적 의존성을 모델링한다. 최종 cross-attention과 self-attention 레이어만 보존하여, 비전–언어–행동 융합을 위한 하이브리드 SSM–attention 정책을 구성한다. 제안한 하이브리드 정책은 동일한 초기화, 학습 데이터, 최적화 설정 하에서 어텐션 전용 확산 트랜스포머 베이스라인과 비교한다. 시뮬레이션 조작 벤치마크에서 일관되게 더 높은 태스크 성공률을 달성했으며, 이는 성능 향상이 추가 사전학습이나 데이터가 아니라 아키텍처 설계에서 비롯됨을 보여준다. 또한 제안한 아키텍처는 행동 헤드의 파라미터 수와 활성 메모리를 줄인다. 소거 실험은 Cross-SSM의 핵심적 역할을 보여준다. Self-attention을 SSM으로만 대체하는 경우 이득이 제한적인 반면, SSM 내부에서의 명시적 교차-모달 융합은 일관된 성능 향상을 제공한다. 추가 분석에서는 모델 깊이와 성능 사이에 비단조적 관계가 관찰되었는데, 학습 손실이 안정적으로 수렴하더라도 과도하게 깊은 하이브리드 스택은 태스크 성공률이 저하된다. 이는 모델 용량 증가가 곧바로 정책 성능 향상으로 이어지지 않음을 시사한다. 이러한 결과는 교차-모달 융합 메커니즘과 깊이 스케일링 전략을 신중히 조율할 때, SSM 기반 VLA 정책이 자원 효율적인 로봇 파운데이션 모델을 위한 유망한 접근임을 보여준다.|Vision–language–action (VLA) foundation models are becoming a unified interface for long-horizon robot manipulation, yet their action generators remain predominantly built on attention-based diffusion transformers. This work investigates a hybrid policy architecture that replaces most transformer attention blocks in the action head with selective state-space models (SSMs), while retaining a small number of attention layers at the bottom. A frozen vision–language model processes visual observations and natural-language instructions into conditioning tokens. The action stream is then processed through two SSM components: Cross-SSM, which performs cross-modal fusion by injecting vision–language context into the SSM's selection mechanism, and SSM blocks that model temporal dependencies within the action sequence. Only the final cross-attention and self-attention layers are preserved, yielding a hybrid SSM–attention policy for vision–language–action fusion. The hybrid policy is compared against an attention-only diffusion transformer baseline under identical initialization, training data, and optimization settings. It consistently achieves higher task success rates on simulated manipulation benchmarks, demonstrating that the gains stem from architectural design rather than additional pretraining or data. The proposed architecture also reduces parameter count and activation memory in the action head. Ablation studies reveal that Cross-SSM plays a critical role. Replacing self-attention with SSMs alone provides limited benefit, whereas explicit cross-modal fusion within the SSM yields consistent improvements. Further analysis reveals a non-monotonic relationship between model depth and performance, where overly deep hybrid stacks exhibit degraded task success rates despite stable training loss convergence. This suggests that increased capacity does not directly translate to improved policy performance. These findings indicate that SSM-based VLA policies are a promising approach for resource-efficient robot foundation models, provided that cross-modal fusion mechanisms and depth-scaling strategies are carefully coordinated.