드래프트‑앤‑타깃 샘플링: 로봇 비디오 생성 정책을 위한 고속 확산 추론

초록

본 논문은 로봇 행동 예측에 사용되는 비디오 생성 정책의 추론 속도를 개선하기 위해, 하나의 확산 모델 내에서 대폭적인 스텝을 취하는 ‘드래프트 샘플링’과 작은 스텝으로 정밀 검증을 수행하는 ‘타깃 샘플링’이라는 두 개의 보완적 디노이징 경로를 결합한 Draft‑and‑Target Sampling (DTS) 을 제안한다. 토큰 청킹과 점진적 수용 전략을 도입해 불필요한 연산을 줄이고, 세 가지 벤치마크에서 최대 2.1배의 속도 향상과 성공률 감소 최소화를 달성하였다.

상세 분석

DTS는 기존의 speculative decoding(추측 디코딩) 아이디어를 비디오 생성용 확산 모델에 적용하면서, 별도의 경량 드래프트 모델을 훈련시키는 대신 동일 모델을 두 가지 샘플링 전략에 재활용한다는 점에서 혁신적이다. 구체적으로, 드래프트 샘플링은 DDIM 솔버의 큰 스텝 크기 (n_1) 을 사용해 초기 가우시안 노이즈부터 거친 디노이징 궤적을 빠르게 생성한다. 이때 생성된 토큰 집합 (D={x_T, x_{T-n_1}, …, x_0}) 은 전체 시퀀스를 근사하지만, 스텝이 크게 건너뛰어 오류가 누적될 가능성이 있다.

타깃 샘플링은 동일 모델에 동일 초기 노이즈를 입력하되, 작은 스텝 크기 (n_2 (< n_1)) 으로 각 드래프트 토큰을 세분화한다. 즉, (x_{T-k n_1}) 에 대해 (n_1/n_2) 개의 세부 스텝을 수행해 정밀 토큰 (\bar{x}) 시퀀스를 만든다. 이렇게 얻어진 타깃 궤적 (G_{all}) 은 드래프트 궤적보다 오류가 적고, 두 궤적을 1:1로 비교해 일치 여부를 판단한다.

핵심적인 두 가지 최적화가 추가된다. 첫째, 토큰 청킹(token chunking) 은 전체 디노이징 과정을 여러 청크로 나누어 순차적으로 처리함으로써 GPU 메모리와 연산을 효율적으로 활용한다. 청크 단위로 드래프트‑타깃 검증을 수행하면, 뒤쪽에서 반복적으로 거부되는 토큰에 대해 불필요한 연산을 방지할 수 있다. 둘째, 점진적 수용(progressive acceptance) 은 초기에는 완벽한 일치를 요구하지 않고, 허용 오차를 점차 감소시키는 방식으로 재샘플링 비용을 크게 낮춘다. 이는 비디오 생성 토큰이 고차원 연속값이기 때문에 엄격한 0‑거리 매칭이 현실적으로 불가능한 점을 고려한 설계이다.

실험에서는 iThor, Meta‑World, Libero 세 가지 로봇 시뮬레이션 벤치마크에 기존 최첨단 Diffusion‑Policy(예: DDPM‑based, DDIM‑based)와 비교하였다. DTS는 평균 2.1× 속도 향상을 보였으며, 성공률은 iThor에서 0.3% 미만, Meta‑World와 Libero에서는 1.2% 이하의 감소에 그쳤다. 이는 드래프트‑타깃 간 오류가 충분히 보정되었음을 의미한다.

한계점으로는 오류 누적에 대한 이론적 분석이 부족하고, 큰 스텝 크기 선택이 도메인에 따라 민감하다는 점이다. 또한, 현재 구현은 GPU 메모리 한계에 의해 청크 크기와 배치 규모를 조절해야 하며, 실시간 로봇 제어에 적용하려면 추가적인 하드웨어‑소프트웨어 최적화가 필요할 것으로 보인다.

전반적으로 DTS는 훈련 비용 없이 확산 기반 비디오 정책의 추론 효율을 크게 개선한 첫 사례이며, 토큰 청킹·점진적 수용이라는 두 가지 기법은 다른 고차원 연속 토큰 생성 작업에도 일반화 가능성이 높다.