TurboWAM이 필요한 월드모델 레이어만 골라 로봇 정책 추론을 빠르게 한다

무엇인가

월드 액션 모델(WAM)은 미래 시각 생성을 행동 예측과 하나의 프레임워크로 묶어, 로봇이 상호작용 중 환경이 어떻게 변할지 예측하게 한다. 문제는 추론 비용이다. 대규모 비디오 생성 모델을 반복 디노이징해야 하므로 지연이 커지고, 실시간 제어에는 쓰기 어렵다. 기존 연구들은 테스트 시점의 미래 프레임 상상(Fast-WAM)을 없애거나, 모든 월드모델 레이어의 키·밸류를 학습된 융합 가중치로 모아 단일 DiT 블록 액션 헤드에 넘기는 방식(Faster-WAM)으로 지연을 줄여 왔다. 이 논문은 여기서 두 가지를 지적한다. 첫째, 모든 레이어를 균일하게 처리하지만 레이어마다 행동 예측에 대한 관련성이 다르고 디노이징 타임스텝에 따라 그 중요도도 달라진다. 둘째, 행동 손실이 각 월드모델 레이어의 표현 학습에 어떻게 작용하는지 조절하는 장치가 없다.

어떻게 동작하나

TurboWAM의 뼈대는 학습 시 비디오 공동 학습을 유지하되 추론 시 미래 시각 예측을 제거하는 것이다. 월드모델은 현재 관측, 언어 지시, 고유수용 상태에 더해 현재의 노이즈 낀 행동 잠재변수를 입력받아, 디노이징 스텝 t마다 행동 조건화된 월드 표현을 만든다. 액션 헤드는 노이즈 낀 행동 잠재변수를 쿼리로, 이 표현을 키·밸류로 삼아 어텐션을 수행하고 속도를 예측하며, 샘플러가 이를 갱신한다. 즉 월드모델 백본을 매 스텝 다시 돌리는 대신 표현 집계만 갱신하는 구조다.

무엇과 다른가

첫 번째 구성 요소는 구조적 월드 표현 집계다. 저자들은 Figure 3에서 중간 레이어가 큰 융합 가중치를 받고, 레이어별 공간 특징 관계가 다르며, 각 레이어가 행동 예측에 미치는 영향이 디노이징 스텝에 따라 변한다는 관찰을 제시한다. 이를 근거로 레이어 예산 m을 두고 앞쪽 m개 레이어의 키·밸류만 집계에 쓴다(실험에서는 m=24). 융합 점수는 액션 청크의 노이즈 낀 행동 잠재변수를 평균한 벡터를 학습된 사영 행렬로 변환하고, 레이어별 라우팅 벡터와 내적한 뒤 차원의 제곱근으로 나눠 얻는다. 여기에 레이어별 정적 로짓과 학습 가능한 스칼라 게이트를 더해 소프트맥스로 가중치를 만든다. 추론 시에는 앞 m개 레이어만 컨트롤 결정마다 한 번 실행하고 키·밸류를 캐싱하며, 융합 가중치만 디노이징 스텝마다 갱신한다. 남은 n-m개 레이어는 건너뛴다.

어떻게 쓰나

두 번째는 구조적 행동 지도다. 융합 가중치가 행동 잠재변수에 조건화되어 있어, 행동 손실이 각 레이어로 흘러가는 직접 그래디언트는 정방향 융합 가중치에 비례한다. 즉 가중치가 작은 레이어는 행동 지도에서 소외된다. 저자들은 그래디언트를 레이어 공통 성분과 레이어별 잔차 성분으로 분해하고, 두 계수의 비율 ρ로 학습된 융합 분포와 균등 분포 사이를 보간한다. 정방향 계산은 그대로 두고 역전파에만 개입하므로 추론 비용은 늘지 않는다. 실험 설정은 μ_com=0.25, μ_res=0.05다.

전제와 한계

실험은 LIBERO(4개 스위트, 40개 태스크, 2,000회 평가), LIBERO-Plus(7종 교란에 대한 10,030개 제로샷 인스턴스), RoboTwin2.0(양팔 50개 태스크, 태스크당 100회 평가, 클린 2,500개·랜덤화 25,000개 시연으로 학습)에서 수행됐고, 실제 로봇 실험은 부록에 담겼다. 비교 대상은 Motus, LingBot-VA, Cosmos Policy, Efficient-WAM, Light-WAM, Faster-WAM, Fast-WAM이며 VLA 계열 X-VLA와 π0.5도 참조로 포함됐다. 모든 지연은 약 140GB 메모리의 동일 단일 데이터센터 GPU에서 측정됐다.

결과는 LIBERO 평균 성공률 99.0%에 지연 49.7ms로, 비교 방법 중 가장 높은 성공률과 가장 낮은 지연을 동시에 기록했다. RoboTwin2.0에서는 평균 90.5%, 47.1ms로 역시 최저 지연이다. LIBERO-Plus에서는 원 LIBERO로 학습한 모델을 추가 적응 없이 그대로 평가해 7개 교란 범주 전반에서 경쟁력 있는 성능을 보였다. 소거 실험에서 레이어 표현을 단순 평균하면 성공률이 크게 떨어지고, Faster-WAM의 DoT 집계로 바꾸면 성능이 하락하며, 정적 로짓을 제거하면 소폭 하락한다. 행동 지도도 그래디언트 조절 없음, 그래디언트 노름 정합 등방 스케일링, 그래디언트 절연과 비교해 제안 방식이 우세했다. 다만 본문에 인용된 범위에서는 소거 실험의 개별 수치가 제시되지 않았다.

개발자 관점에서 이 논문의 실용적 메시지는 월드모델 전체를 매 디노이징 스텝 돌리는 것이 필수가 아니라는 점이다. 앞단 레이어의 키·밸류를 한 번 계산해 캐싱해 두고, 액션 상태에 따라 융합 가중치만 다시 계산하는 패턴은 비디오 백본을 정책에 재활용하는 다른 파이프라인에도 그대로 옮길 수 있다. 도입 전에는 레이어 예산 m을 실제 백본 깊이 대비 얼마로 잡을지, 캐시가 유효한 관측 변화 범위가 어디까지인지, 그리고 학습 시 그래디언트 조절 강도(ρ)가 성능에 얼마나 민감한지를 확인해야 한다. 저자들도 조작 태스크 벤치마크 중심으로 평가했다는 점을 전제로 읽어야 한다.

저자가 밝힌 한계는 두 가지다. 복잡한 공간 추론과 정교한 물체 상호작용이 필요한 장면에서는 어려움을 겪을 수 있으며, 공간 파운데이션 모델의 기하 사전지식이나 비전-언어 모델의 의미 지식을 결합하는 것을 향후 방향으로 제시한다. 또한 낮은 추론 지연을 유지하면서 성공률을 더 끌어올리는 것이 남은 과제라고 명시한다.