미래 예측을 버리고 현재 관측 디노이징으로 생성 DiT를 로봇 제어에 적응시키는 NowWAM

무엇인가

이 논문은 사전학습된 생성형 Diffusion Transformer(DiT)를 로봇 제어에 어떻게 옮겨야 하는가라는 질문을 다룬다. 기존 생성형 로봇 정책들은 대개 미래 시각 예측을 그 인터페이스로 삼아 왔고, 최근에는 미래 예측을 추론 단계에서 빼고 공동학습(co-training)에만 남기는 흐름이 생겼다. 저자들은 여기서 더 근본적인 질문을 던진다. 미래가 더 이상 행동 추론에 쓰이지 않는다면, 미래를 예측한다는 사실 자체가 생성형 공동학습의 이득을 설명하는가. 그리고 사전학습된 생성형 DiT는 행동 학습에 실제로 무엇을 기여하는가.

어떻게 동작하나

제안 방법인 NowWAM은 별도의 미래 목표를 없애고, 제어에 쓰는 현재 시각 스트림 자체에 사전학습된 디노이징 과정을 그대로 적용한다. 기존 방식은 현재 관측 z_t와 별도의 목표 z_t^+를 함께 두고 [l | z_t | z_{t,σ}^+ | a] 형태의 두 스트림으로 학습했지만, NowWAM은 [l | z_{t,σ} | a] 단일 스트림만 쓴다. 즉 같은 시각 표현이 생성 목적과 행동 학습을 동시에 담당한다. 노이즈 레벨 σ는 u~U(0,1)에 대해 σ = su/(1+(s-1)u)로 샘플링하며, s=1이면 균등 샘플링, s<1이면 저노이즈 쪽으로 치우친다. 현재 잠재는 z_{t,σ} = (1-σ)z_t + σϵ (ϵ~N(0,I))로 만들고, 속도 목표는 v*_t = ϵ - z_t다.

무엇과 다른가

학습 손실은 L = λ_vis·L_denoise + λ_act·L_action이며, L_denoise = ||v̂_t - (ϵ - z_t)||²_2, L_action은 연속 행동 목표에 대한 마스크드 평균제곱오차다. 주 설정에서 λ_vis = 0.5, λ_act = 1.0이고 유효하지 않은 행동 차원과 패딩된 행동 토큰은 손실에서 제외한다. 구조는 사전학습 생성 DiT 백본과 별도 행동 DiT로 이루어진 DiT–MoT 형태다. 백본이 언어와 시각 스트림을 처리해 층별 시각-언어 컨텍스트를 행동 전문가에 전달하고, 행동 토큰과 파라미터는 분리된 채 마스크드 혼합 어텐션으로 상호작용한다. 중요한 점은 학습과 추론의 비대칭이다. 학습 중에는 행동 전문가가 z_{t,σ}에서 나온 컨텍스트를 읽어 디노이징 상태 전반에 걸친 표현에 결합되지만, 추론 시에는 σ=0, 즉 깨끗한 끝점 z_{t,0}=z_t만 단일 순전파로 평가한다. 보조 시각 스트림도, 시각 디노이징 롤아웃도 없다.

어떻게 쓰나

실험은 세 가지 설정에서 이뤄졌다. 표준 LIBERO에서는 평균 98.4% 성공률로 미래 목표 생성형 베이스라인과 동등했고, 다른 시뮬레이터와 과제군으로의 소수샷 전이를 보는 RoboCasa GR1 Tabletop에서는 과제당 100개 시연만으로 1,200 에피소드에서 64.9%를 기록해 같은 100샷 예산의 DIAL(58.3%), Fast-WAM(56.0%)을 앞섰다. 주 무대는 시각·환경 교란 7종을 다루는 LIBERO-Plus로, 10,030개 전체 에피소드 프로토콜에서 FLUX2-Klein 백본으로 87.7%를 달성해 미래 목표 공동학습 베이스라인의 81.6%를 6.1포인트 앞섰다. 개선은 카메라, 로봇, 배경 교란 같은 어려운 시각 변화에 집중됐고 언어 성능은 비슷했다. 순수 텍스트-투-이미지 DiT인 Z-Image 백본으로도 87.8%를 내 비디오 생성이나 이미지 편집 백본에 묶이지 않음을 보였다. 비교 대상에는 π0, π0.5, OpenVLA-OFT, X-VLA, ABot-M0, GR00T, LingBot-VA, Motus, Cosmos-Policy, Fast-WAM, ImageWAM 등이 포함된다.

전제와 한계

통제 실험이 논문의 핵심 근거다. 생성 초기화만으로도 강건성이 크게 올라, Edit 목적이 있을 때 58.14%에서 83.05%로, 시각 목적이 없어도 50.55%에서 75.81%로 향상됐다. 보조 시각 목표의 시간적 의미를 바꾼 실험에서는 미래 t+16이 82.89%, t+32가 83.70%, 과거 t-16이 83.80%로 사실상 비슷했다. 즉 순방향 시간 예측이 특별히 우월하지 않다. 반면 디노이징 궤적을 제한하면 성능이 떨어진다. 깨끗한 끝점에서만 학습하면 77.48%에 그쳤고, 균등 샘플링은 83.46%, 저노이즈로 치우친 샘플링은 84.97%였다. 노이즈 낀 현재를 쓰되 시각 감독을 제거하면 78.73%로 떨어졌다. 결론은 미래 예측이 아니라, 행동에 맞닿은 표현에 대해 네이티브 디노이징 궤적 전반의 생성 감독을 유지하는 것이 핵심이라는 것이다.

효율도 개선된다. 동일한 2×H200, 글로벌 배치 64 조건에서 두 번째 시각 스트림을 없애 학습 비주얼 토큰이 784에서 392로 절반이 되고, 스텝 시간은 2.85초에서 1.63초로 1.8배 빨라졌으며 최대 메모리는 78.9GiB에서 74.3GiB로 줄었다. 텍스트 임베딩과 VAE 잠재는 미리 계산해 캐시하므로 이 수치는 DiT-행동 공동학습만을 측정한 값이다. 실무 관점에서 이 논문이 주는 신호는 명확하다. 생성형 백본을 로봇 정책에 붙일 때 미래 프레임 예측 헤드를 별도로 유지할 필요가 없고, 현재 관측 스트림에 디노이징 목적을 걸고 노이즈 레벨을 궤적 전체에서 샘플링하면 된다. 다만 학습 시에는 노이즈 낀 표현을 봐야 하고 추론은 깨끗한 끝점 단일 순전파라는 비대칭을 코드에서 정확히 구현해야 하며, 백본이 비디오 생성 모델이어야 한다는 제약도 사라진다.

저자들이 밝힌 전제와 한계도 분명하다. 미래 모델링 자체는 명시적 동역학이나 장기 계획에는 여전히 유용할 수 있으며, 이 논문이 다룬 조작(manipulation) 설정에서만 미래 목표 없는 현재 중심 인터페이스가 충분하다는 주장이다. 또한 주 결과는 각 벤치마크의 전체 평가 프로토콜과 최종 EMA 체크포인트를 쓰지만, 통제 분석은 계산 비용 때문에 고정된 1,923 에피소드 LIBERO-Plus 부분집합에서 비EMA 체크포인트로 수행했다는 평가상 차이가 있다. 표준 LIBERO는 이미 포화 상태에 가까워 차별력이 낮고, 강건성 이득이 주로 시각 교란 범주에 집중된다는 점도 함께 봐야 한다.