RealtimeWAM이 월드 액션 모델 추론을 25배 빠르게 만든다

무엇인가

World Action Model(WAM)은 미래 시각 관측과 로봇 액션을 함께 모델링해 조작 정책을 학습하는 패러다임이다. VLA가 관측과 언어 지시를 곧바로 액션으로 매핑하는 것과 달리, WAM은 미래 영상 예측을 추가 학습 목표로 삼아 물리 동역학과 시간 구조를 정책에 반영한다. 구조는 비디오와 액션을 한 백본에서 함께 디노이징하는 shared-backbone 계열과, 비디오 전문가와 액션 전문가를 분리한 Mixture-of-Transformers(MoT) 계열로 나뉜다. Fast-WAM 같은 효율적 MoT WAM은 비디오 표현을 한 번만 계산해 액션 전문가가 재사용하도록 했지만, H100에서도 한 번 추론에 300ms 이상이 걸린다. 저자들은 병목을 두 가지로 지목한다. 하나는 액션 전문가의 다단계 디노이징(intra-expert iteration)으로 전체 지연의 약 90%를 차지하고, 다른 하나는 비디오 전문가가 끝난 뒤에야 액션 전문가가 시작하는 순차 실행(inter-expert waiting)이다.

어떻게 동작하나

첫 번째 병목을 없애기 위해 제안하는 것이 Teacher-Anchored Consistency Distillation(TACD)이다. 기존 consistency distillation(CD)은 같은 교사 궤적 위의 잡음 상태들을 하나의 깨끗한 종점으로 매핑하도록 학습해 소수 스텝 생성을 가능하게 하지만, 저자들은 국소 일관성 오차(local consistency error)만 줄일 뿐 교사의 깨끗한 목표 액션과의 전역 편차(global target error)는 잘 줄지 않는다는 local-global error gap을 지적한다. 총 오차를 e_total = e_local + e_global로 분해하면 CD는 e_local에만 직접 벌점을 주고 e_global은 명시적으로 다루지 않는다. TACD는 동결된 교사의 다단계 롤아웃 종점 a_0^T를 정확한 종점 a_0*의 대리값으로 삼아 구간 평균 속도 u_θT(a_t,t) = (a_t - a_0^T)/t를 계산하고, 학생의 속도 예측 v_θS를 이 값에 직접 정렬한다. 최종 목적함수는 L = L_CD + λ·L_TA이며 실험에서 λ=0.2를 쓴다. 저자들은 e_total = -t[v_θS - u_θT] + δ (δ는 교사의 수치 적분 오차)라는 관계를 통해 이 교사 앵커 항이 국소·전역 오차의 합을 직접 제약하고, 기대 제곱 종점 오차가 2L_TA + 2E||δ||²로 상한이 잡힌다는 것을 보인다.

무엇과 다른가

두 번째 병목은 Cross-Expert Wavefront Pipelining(CEWP)으로 푼다. 저자들은 각 DiT 블록을 Q/K/V 투영 P_i, 어텐션 Attn_i, 나머지 계산 R_i로 분해하고 의존 관계를 분석한다. 액션 블록 i의 어텐션은 비디오 KV 캐시 중에서도 같은 인덱스의 (K_i^v, V_i^v)만 필요로 하며, 액션 투영 P_i^a는 비디오 KV를 기다릴 필요 없이 이전 액션 은닉 상태만 있으면 실행할 수 있다. 기존 구현은 비디오 전문가의 마지막 블록이 끝난 뒤에야 액션 전문가를 시작하는 불필요한 대기를 강제한다. CEWP는 비디오와 액션 전문가를 두 개의 논블로킹 CUDA 스트림에 배치하고, 비디오 스트림이 P_i^v 직후에 이벤트 e_i를 기록하면 액션 스트림은 Attn_i^a 직전에서만 그 이벤트를 기다린다. 그 결과 액션 블록들이 비디오 스트림이 전진하는 만큼 KV 캐시를 점진적으로 소비하는, 모델 깊이 방향의 엇갈린 웨이브프런트가 형성된다.

어떻게 쓰나

실험은 Fast-WAM 기반 RealtimeWAM*과 Faster-WAM 기반 RealtimeWAM† 두 변형으로 LIBERO, LIBERO-Plus, RoboTwin 2.0에서 수행했다. 비디오 전문가는 동결하고 액션 전문가의 rank-128 LoRA만 학습한다. RoboTwin 2.0에서 RealtimeWAM*은 1스텝 생성으로 90.84% 성공률을 기록해 10스텝 Fast-WAM 대비 0.67%만 떨어졌고, RealtimeWAM†는 92.64%로 10스텝 Faster-WAM 대비 0.29% 감소에 그쳤다. 같은 조건에서 1스텝 Flash-WAM은 81.41%, DMD*는 88.22%였다. LIBERO에서는 RealtimeWAM*이 전체 97.0%로 10스텝 Fast-WAM과 동일한 성능을 냈고, Fast-WAM을 그냥 1스텝으로 줄인 95.0%보다 2.0%p 높았다. 1스텝 Flash-WAM 95.1%, DMD* 96.1%보다 우수하고 Light-WAM 97.2%와는 비슷한 수준이다. LIBERO-Plus 분포 이동 평가에서 RealtimeWAM†는 73.0%로 Faster-WAM의 73.6%보다 0.6%p 낮았고, 카메라·센서 노이즈·레이아웃 교란에서는 오히려 소폭 향상됐다.

전제와 한계

절제 실험에서 비디오 전문가를 함께 파인튜닝하면 성공률이 오히려 떨어지고 학습 시간과 최대 메모리 사용량이 늘어난다. 비디오 전문가를 동결한 상태에서 L_CD에 L_TA를 더하면 전체 성공률이 89.85%에서 90.84%로 올라간다. 교사 롤아웃 스텝 K는 5일 때 90.39%, 10일 때 90.84%, 20일 때 90.77%로 K=10 이후 이득이 없어 K=10을 채택한다. 지연 시간은 TACD만으로 Fast-WAM 299.7ms에서 65.8ms(4.56배), Faster-WAM 218.9ms에서 57.1ms(3.83배)로 줄고, CUDA Graph와 CEWP를 더하면 Fast-WAM 23.3→17.4ms(1.34배), Faster-WAM 26.2→22.4ms(1.17배)가 추가로 줄어든다. 모든 최적화를 켠 최종 추론 지연은 12.2ms와 16.1ms로 각각 24.55배, 13.56배 가속이다. 지연 측정에는 VAE 인코딩 시간이 포함되고 텍스트 인코딩은 에피소드당 한 번만 수행되므로 제외했다.

개발자 관점에서 이 논문의 실용적 가치는 기존 MoT WAM 체크포인트를 버리지 않고 후처리(post-training)만으로 실시간 제어 대역에 끌어올린다는 점이다. 30Hz 제어 주기(33ms)를 맞추려면 정책 추론이 그보다 짧아야 하는데 12.2ms는 그 조건을 만족한다. 적용 조건은 명확하다. 대상이 MoT 구조 WAM이어야 하고, 비디오 전문가를 동결한 채 액션 전문가에 LoRA를 붙이는 방식이므로 비디오 백본 자체를 바꾸려는 경우에는 그대로 쓸 수 없다. 또한 CEWP는 블록 단위 의존성을 CUDA 스트림과 이벤트로 직접 스케줄링하는 시스템 최적화라서, 사용하는 프레임워크의 스트림·이벤트 지원과 CUDA Graph 호환성을 먼저 확인해야 한다. 재현 시에는 λ, K, EMA decay(0.995), 학습 스텝(30,000) 같은 하이퍼파라미터와 지연 측정 범위(VAE 포함, 텍스트 제외)를 원문과 맞춰 비교하는 것이 좋다.

한계와 전제는 저자들이 밝힌 범위 안에서 다음과 같다. TACD의 교사 앵커는 정확한 종점 a_0*가 아니라 교사의 수치 적분 결과 a_0^T를 대리값으로 쓰기 때문에 종점 오차 상한에 교사의 적분 오차 δ 항이 남는다. δ는 교사 디노이징 스텝 K를 늘려 줄일 수 있지만 K=20에서는 성공률 이득이 사라지고 교사 연산 비용만 늘어나, 저자들은 K=10을 성능과 비용의 절충점으로 선택한다. 또한 이 방법은 MoT 기반 WAM을 후처리로 가속하는 데 한정되며, 실험은 LIBERO, LIBERO-Plus, RoboTwin 2.0과 부록의 실세계 양팔 조작 과제에 국한된다. 그 밖의 로봇 구성이나 과제 지평에 대한 일반화 수치는 이 본문 범위에서 제시되지 않았다.