AnyStep-WAM이 로봇 조작 추론의 디노이징 예산을 장면별로 줄인다

AnyStep-WAM: Budget-Aligned Distillation and Adaptive Inference for World Action Models

HF Daily2609.33748

Rui Wang, Xiangyu Wang, Donglin Yang2026-09-27조회 3

무엇인가

월드 액션 모델(WAM)은 관측과 지시로부터 미래 시각 정보와 행동 청크를 함께 생성하지만, 추론 시 대형 디퓨전 트랜스포머를 통해 노이즈를 반복 제거하기 때문에 새 관측에 대한 반응이 느리다. 문제는 조작 작업의 행동 청크가 오차에 대해 균일한 민감도를 갖지 않는다는 점이다. 접촉이 중요한 구간은 정밀도가 필요하지만, 그렇지 않은 구간은 적은 스텝으로도 충분하다. 기존 일관성 증류나 Flash-WAM 같은 방식은 1스텝 또는 고정 예산만 지원해 이런 가변적 정밀도 요구를 수용하지 못한다. 저자들은 하나의 모델이 여러 디노이징 예산에서 정확히 동작하게 만드는 것과, 현재 청크에 몇 스텝을 쓸지 온라인으로 결정하는 것, 이 두 가지를 분리된 난제로 규정한다.

어떻게 동작하나

첫 번째 난제에 대한 해법은 예산 정렬 적분 플로우맵 증류다. 동결된 교사 WAM을 세밀한 솔버로 돌려 얻은 궤적에서 두 시점 r<t의 상태 차이를 (t-r)로 나눈 값을 교사 평균 속도로 정의하고, 이를 학생의 회귀 목표로 삼는다. 두 끝점이 모두 동결 교사에서 나오므로 목표가 학생 파라미터에 의존하지 않고, 학생의 시간 미분을 추정할 필요가 없다. 학습 시에는 후보 추론 스케줄 집합에서 예산 K를 먼저 뽑고 그 안의 구간 인덱스 i를 뽑아 (t,r) 쌍을 만든다. 손실은 네 항으로 구성된다. 구간 감독 L_int, 순간 속도를 유지하는 표준 플로우매칭 L_FM, r=0으로 가는 끝점 감독 L_end(1스텝 프리뷰용 (1,0) 감독 포함), 그리고 학생이 스스로 만든 상태에서 이어가는 보정 연속 학습 L_rec이다. L_rec은 학생이 만든 첫 재활용 상태에서 교사를 한 번 재시작해 만든 예산별 참조 궤적을 목표로 쓴다. 구현은 어텐션·피드포워드·시간 조건화 투영에 LoRA를 넣고 사전학습 가중치는 모두 동결하며, 어댑터는 모든 예산에서 공유한다.

무엇과 다른가

두 번째 난제는 위험-이득 스케줄러로 푼다. 교사 궤적에서 스텝 간 속도의 방향 불일치와 상대적 크기 변화를 재어 난이도 대리값 ρ를 만들고, 각 예산 K에 대해 배포 시 쓰는 프리뷰 재활용 경로를 오프라인으로 돌려 실행 속도와 교사 참조 속도의 일치도 S_K를 계산한다. S_K는 앞쪽 디노이징 구간에 가중치 2^-i를 더 주는 가중합이다. 1스텝 프리뷰의 은닉 특징을 입력으로 받는 경량 예측기 g가 ρ와 S_K를 동시에 출력하도록 Smooth L1으로 학습한다. 추론 시에는 τ = τ_easy + (τ_hard - τ_easy)·ρ̂ 로 난이도에 따라 기준을 조절하고, 모든 활성 모달리티에서 Ŝ ≥ τ를 만족하는 가장 작은 예산을 고른다. 만족하는 예산이 없으면 예측 이득 합이 가장 큰 예산을 쓴다. K=1이면 프리뷰 결과를 그대로 반환하고, 아니면 원래 노이즈로 첫 중간 시점까지 재노이즈한 뒤 나머지 구간을 갱신한다. 프리뷰 계산이 선택된 궤적에 재사용되므로 추가 디노이징 패스나 온라인 교사 평가, 후보 롤아웃이 필요 없다.

어떻게 쓰나

실험은 RoboTwin 2.0의 양팔 조작 50개 과제(clean 및 무작위화 설정)와 실제 Unitree G1D 양팔 로봇의 6개 과제(Bowl Pouring, Clean Table, Put Block, Put Cup, Stack Blocks, Stack Bowls)에서 Motus, FastWAM, LingBotVA 세 백본으로 수행했다. 학습은 H100 4장에서 8,000 스텝, 전역 배치 32다. RoboTwin 평가는 A100 1장에서 과제당 100회, 실기는 RTX 5090에서 과제당 평균 20회다. 기준 Motus와 FastWAM은 10스텝 고정, LingBotVA는 비디오 DiT 25스텝·행동 DiT 50스텝을 쓴다. 제안 방법은 평균 성공률을 기준선 대비 0.24%p 이내로 유지하면서 평균 예산을 각각 3.98, 5.03, 3.68스텝으로 줄여 60.2%, 49.8%, 85.28%(LingBotVA 비디오/행동 92.64%)의 스텝 감소를 달성했다. A100 기준 호출당 지연은 Motus 1.932→0.859초, FastWAM 0.496→0.295초, LingBotVA 9.732→3.247초로 줄었고, 예산 선택 모듈 오버헤드는 예측당 2.881, 2.4996, 2.9772밀리초에 불과하다.

전제와 한계

1스텝 예산에서는 기준 1스텝 대비 평균 성공률이 Motus 7.07%p, FastWAM 12.08%p, LingBotVA 8.94%p 향상됐고, Flash-WAM 대비로도 3.19, 7.58, 1.30%p 앞선다. 학습 안정성 비교에서는 유한차분 목표가 표본별 분포가 넓고 그래디언트 스파이크가 잦아 전체 스텝의 45.9%가 클리핑 임계값 1.0을 넘긴 반면, 동결 교사 목표는 분포가 집중돼 훨씬 안정적이었다. 실기에서는 평균 성공률이 Motus 69.17%, FastWAM 71.67%, LingBotVA 70.00%로 오르고 평균 디노이징 스텝은 59.4%, 63.9%, 85.84% 감소했으며, 지연은 1.868→1.117초, 0.280→0.122초, 4.654→0.758초로 각각 1.67배, 2.30배, 6.14배 빨라졌다.

Motus 백본 절제 실험에서 L_int를 제거하면 1스텝 성공률이 1.40%p 떨어지고 2스텝 성공률이 86.21%에서 78.94%로, 적응적 성공률이 87.96%에서 78.09%로 하락하며 평균 예산이 3.98에서 5.27스텝으로 늘었다. L_end 제거 시 1스텝 성공률이 82.12%에서 75.53%로 내려가고 적응 예산이 4.74스텝으로 증가했으며, L_rec 제거 시 적응 성공률이 86.96%로 소폭 하락했다. 무작위화 샘플 없이 clean 데이터만으로 스케줄러를 학습해도 무작위 조건에서 87.56% 성공률과 4.05스텝을 기록해 완전판(87.80%, 3.94스텝)에 근접했다. 임계값을 5퍼센타일 낮추거나 높이면 87.60%·3.08스텝, 87.93%·5.85스텝으로 완전판과 0.36%p 이내 차이를 보여, 계산량은 크게 조절하면서 성공률 민감도는 제한적이었다.

개발자 관점에서 이 논문의 실용적 핵심은 재학습 비용이 낮다는 점이다. 사전학습 WAM을 동결한 채 LoRA 어댑터와 네 개의 회귀 손실만으로 여러 예산을 지원하는 단일 모델을 얻고, 스케줄러는 1스텝 프리뷰 특징만으로 동작해 추가 지연이 수 밀리초 수준이다. 로봇 정책 서빙에서 평균 스텝 수를 직접 줄이는 대신 구간별로 예산을 배분하는 방식을 검토한다면, 먼저 자신의 작업에서 접촉 민감 구간이 실제로 존재하는지, 그리고 프리뷰 재활용 경로가 배포 파이프라인과 맞는지 확인해야 한다. 다만 저자들이 명시한 전제를 유의해야 한다. 난이도 지표 ρ는 작업 실패 확률이 아니라 상대적 디노이징 난이도를 나타내고, 일치도 S_K는 최종 행동 오차나 과제 성공률을 직접 측정한 값이 아닌 대리 지표다. 또한 플로우맵으로 학습된 WAM에 대한 적응적 예산 선택을 다룬 선행 연구가 없어, 저자들이 직접 AnyStep+Gate라는 훈련 없는 백분위 게이팅 기준선을 만들어 비교했다는 점도 결과 해석의 범위를 제한한다.