로봇 확산 트랜스포머의 잠재 공간을 객체 슬롯으로 바꾼 SlotDiT

무엇인가

텍스트 조건 잠재 확산 모델은 영상 생성에서 강력한 성능을 보이며 로봇 응용의 백본으로 기대를 모은다. 그런데 기존 접근은 픽셀 수준 표현이나 VAE 기반 잠재를 쓰기 때문에 명시적인 의미 구조가 없고, 잠재 공간 설계가 최종 성능에 어떤 영향을 주는지는 거의 탐구되지 않았다. 로봇에서는 잠재 공간이 다운스트림 예측과 계획, 제어에 어떤 정보가 남는지를 직접 결정하므로 이 공백이 특히 치명적이다. 저자들은 장면을 객체 단위 잠재 벡터인 슬롯으로 분해하는 객체 중심 표현이 동역학 모델링과 계획에서 성과를 냈지만, 확산 기반 생성 모델링의 잠재 공간으로는 아직 쓰이지 않았다는 점을 문제로 삼는다.

어떻게 동작하나

SlotDiT의 첫 단계는 객체 중심 표현 모듈이다. DINOSAUR 프레임워크를 영상으로 확장해, DINOv2-ViT가 각 프레임을 특징 맵 h_t로 인코딩하고 Slot Attention이 이전 슬롯 S_{t-1}을 시각 특징에 반복적으로 정련한다. 슬롯은 cross-attention으로 특징에 주목하되 어텐션 가중치를 슬롯 축으로 정규화해 각 특징 위치를 두고 슬롯들이 경쟁하도록 만들고, 다시 특징 위치 축으로 정규화한 가중 평균 U를 GRU에 넣어 슬롯을 갱신한다. 디코딩은 2단계로, MLP 기반 broadcast decoder가 슬롯마다 객체 특징 맵과 알파 마스크를 만들고 슬롯 간 마스크 정규화 후 가중합으로 특징을 재구성한 다음, CNN 이미지 디코더가 프레임을 렌더링한다. 이 모듈은 이미지 재구성과 특징 재구성 손실 L_Slot = (1/τ) Σ_t (||X̂_t - X_t||² + ||ĥ_t - h_t||²)로 학습된다.

무엇과 다른가

두 번째 단계에서 SlotDiT는 이 슬롯 잠재 공간 위에서 확산을 수행한다. 참조 이미지 X_1과 언어 지시 C가 주어지면 장면을 N_S개 슬롯 S_1으로 파싱하고, T5 인코더로 지시를 텍스트 임베딩 C로 바꾼 뒤, 지시와 관측 장면 맥락에 조건화해 가우시안 노이즈에서 미래 슬롯 궤적을 반복적으로 디노이징한다. 모델은 고정 길이 시간 윈도 N개 슬롯 세트 위에서 동작하며, 최근 M개 맥락 슬롯과 텍스트 임베딩, 나머지 N-M개 미래 프레임용 노이즈를 입력받아 다음 슬롯 세트들을 예측하고, 마지막 M개 예측을 다음 단계의 맥락으로 재사용해 자기회귀적으로 전체 미래를 생성한다. 백본은 N_P개 블록으로 구성된 DiT로, 양방향 multi-head self-attention과 텍스트-슬롯 cross-attention, MLP를 LayerNorm과 잔차 연결로 감싸고 프레임별 노이즈 레벨은 adaLN으로 주입한다. 위치 인코딩은 RoPE를 쓰는데, 시간 축에만 적용해 슬롯 순열 등가성을 보존하는 변형과 시간 및 슬롯 두 축에 모두 적용해 등가성을 깨는 변형을 함께 연구한다.

어떻게 쓰나

학습은 2단계다. 객체 중심 모듈을 먼저 학습하고 슬롯 인코더를 동결한 채 SlotDiT를 학습한다. SlotDiT는 Diffusion Forcing Transformer(DFoT) 프레임워크를 따르며, 노이즈를 마스킹으로 보는 패러다임에 따라 각 프레임에 독립적인 확산 노이즈 레벨 k를 부여한다. 이 덕분에 모든 토큰이 학습 목표에 기여하고 추론 시 가변 길이 히스토리를 쓸 수 있다. 전방 확산 과정은 S_t^(k_t) = sqrt(ᾱ_{k_t}) S_t + sqrt(1-ᾱ_{k_t}) ε_t이고, v-prediction 파라미터화로 v_t = sqrt(ᾱ_{k_t}) ε_t - sqrt(1-ᾱ_{k_t}) S_t를 목표로 삼아 fused min-SNR 재가중 손실 L_SlotDiT = E[ (1/N) Σ_t w(k_t) ||v_θ(S_{1:N}^(k), k, C)_t - v_t||² ]를 최소화한다. classifier-free guidance를 위해 학습 중 텍스트 임베딩을 학습된 null 임베딩으로 무작위 대체하고, 추론은 DDIM 샘플러로 한다.

전제와 한계

실험은 CLIPort, LanguageTable-Synthetic, LanguageTable-Real, Bridgev2 네 로봇 데이터셋에서 텍스트 유도 영상 생성으로 수행됐고, 두 시뮬레이션 환경에서 로봇 제어도 평가했다. 비교는 두 갈래다. 같은 객체 중심 표현을 쓰지만 확산 대신 자기회귀 트랜스포머 예측기를 쓰는 TextOCVP와 OC-WM, 그리고 동일한 DiT 아키텍처와 학습 절차, 추론 방식을 공유하고 잠재 공간만 바꾼 SD-VAE, ImageVAE, VA-VAE, VideoVAE, DINOv2 기반 RAE 스타일 의미 정렬 잠재다. 결과적으로 SlotDiT는 네 데이터셋 모두에서 더 높은 작업 성공률을 냈다. CLIPort와 LanguageTable-Synthetic에서는 DiT+VideoVAE가 LPIPS와 FVD, JEDi 같은 시각 품질 지표에서 가장 강했지만 작업 성공률은 크게 낮았고, SlotDiT는 시각 품질에서 경쟁력 있으면서 성공률이 훨씬 높았다. 실제 환경 데이터인 LanguageTable-Real에서 SlotDiT가 최고 성공률을 기록했고, Bridgev2에서는 최고 VAE 베이스라인을 약 8퍼센트포인트 앞섰다. 저자들은 표준 영상 생성 지표가 로봇 환경의 작업 해결과 지시 따르기를 충분히 반영하지 못한다고 해석한다.

로봇 제어에서는 학습된 inverse dynamics model로 예측 잠재를 행동에 매핑해 평가했는데, LanguageTable-Synthetic의 in-distribution 설정에서 SlotDiT가 학습된 예측기 중 최고 성공률을 기록하며 정답 슬롯을 쓰는 오라클에 거의 근접했다. 8개 블록의 복잡한 장면이나 상대 위치 추론이 필요한 미학습 지시 템플릿 같은 분포 밖 설정에서도 일관되게 최고 성공률을 냈고, CLIPort의 개루프 실행에서도 마찬가지였다. 효율에서는 CLIPort 기준 9.33 FPS로 DiT+SD-VAE 대비 5.68배 빨랐는데, 주된 이유는 DiT가 처리하는 토큰 수가 프레임당 256개에서 슬롯 10개로 줄었기 때문이다. 절제 실험에서는 합성 데이터에서 슬롯 순열 등가 RoPE가, 실제 데이터에서 비등가 RoPE가 더 좋았고, T5-small을 T5-XL로 키워도 성능이 오르지 않았으며, Bridgev2에서 DiT 깊이를 N_P=14로 늘리면 개선됐다. 객체 중심 슬롯을 단일 고차원 잠재 벡터로 바꾼 Non-OC 변형은 두 데이터셋 모두에서 SlotDiT보다 나빴고, factorized attention이나 adaLN 시간 조건 제거, 평균 풀링 텍스트 임베딩 추가도 성공률을 떨어뜨렸다.

실무 관점에서 이 논문이 주는 신호는 두 가지다. 첫째, 로봇 조작 파이프라인에서 확산 기반 월드모델이나 영상 예측기를 쓴다면 잠재 공간 선택이 아키텍처 미세 조정보다 성공률에 더 큰 영향을 줄 수 있다. VAE 잠재를 그대로 쓰는 구성이라면 슬롯 기반 잠재로 바꿨을 때 토큰 수와 추론 속도, 그리고 작업 성공률이 어떻게 변하는지 비교해 볼 가치가 있다. 둘째, LPIPS나 FVD 같은 지표가 좋다고 해서 지시 수행이 잘 되는 것이 아니므로, 로봇 응용에서는 VLM 판정 작업 성공률 같은 과제 지향 지표를 함께 봐야 한다. 다만 SlotDiT는 슬롯 인코더를 먼저 학습해 동결하는 2단계 파이프라인이고, Slot Attention과 DINOv2 특징 추출, T5 인코더, DDIM 샘플러 등 여러 구성 요소를 함께 맞춰야 하므로 도입 비용은 작지 않다.

논문이 명시적으로 밝힌 한계는 제한적이다. 제공된 본문에는 별도의 한계 절이 없고, 저자들은 압축된 객체 중심 잠재가 고용량 VAE 잠재보다 담을 수 있는 시각적 세부 정보의 양을 제한할 수 있다고 인정한다. 또한 OC-WM 비교는 재실행이 아니라 원 논문의 수치를 동일 설정에 맞춰 인용했고, 작업 성공률 평가는 VLM 판정에 의존한다. 실험은 두 장의 NVIDIA A6000 48GB에서 T5-small, 224×224 해상도(CLIPort는 336×336), 슬롯 8~10개, 슬롯 차원 128 또는 256, DiT 블록 8개(Bridgev2는 14개), 은닉 차원 512, 시간 윈도 10, 확산 스텝 1000이라는 특정 구성에 한정되어 있으므로, 다른 규모나 도메인으로의 일반화는 별도 검증이 필요하다.