HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling

HCLSM: 객체 중심 세계 모델링을 위한 계층적 인과 잠재 상태 기계

1. 문제 정의 (Problem)

기존의 비디오 기반 세계 모델 (World Models) 은 다음과 같은 근본적인 한계를 가지고 있습니다:

• 평탄한 잠재 표현 (Flat Latent Representations): 모든 객체, 시간 척도, 인과 관계를 단일 벡터에 얽매어 (entangle) 표현합니다. 이는 개별 객체의 상태를 분리하거나 물리적 인과 관계를 추론하는 데 실패합니다.
• 시간적 역학의 단일화: 물체의 연속적인 운동 (밀리초 단위), 이산적인 사건 (충돌 등), 그리고 장기적인 목표 (분 단위) 를 구별하지 않고 단일 스케일로 처리합니다.
• 인과성 부재: "만약 그리퍼가 더 세게 밀었다면?"과 같은 반사실적 (counterfactual) 추론이 불가능합니다.
• 기존 모델의 한계: Slot Attention 과 같은 객체 중심 모델은 시계열 역학을 다루지 못하거나, DreamerV3 와 같은 세계 모델은 객체 분해 없이 평탄한 상태를 사용합니다.

2. 방법론 (Methodology)

HCLSM 은 객체 분해, 계층적 시간 역학, 인과 구조 학습을 통합한 5 층 아키텍처를 제안합니다.

A. 5 계층 아키텍처

1. 지각 (Perception): ViT(Vision Transformer) 인코더를 사용하여 비디오 프레임을 패치 임베딩으로 변환합니다.
2. 객체 분해 (Object Decomposition):
   • Slot Attention: 학습된 가우시안 분포에서 초기화된 슬롯들이 패치 토큰을 위해 경쟁하여 객체를 발견합니다.
   • 공간 브로드캐스트 디코더 (SBD): 각 슬롯이 자신의 이미지 영역을 재구성하도록 강제합니다. 재구성 목표는 픽셀이 아닌, EMA 타겟 인코더의 고정된 ViT 특징 (DINOSAUR 방식) 입니다.
   • 생성/소멸 (Birth/Death): 잔여 관심도가 임계값을 넘으면 잠자는 슬롯이 활성화됩니다.
   • 관계 그래프 (GNN): 슬롯 간의 상호작용을 인코딩하는 그래프 신경망을 사용합니다.
3. 계층적 역학 (Hierarchical Dynamics): 3 단계 엔진을 통해 다양한 시간 척도를 모델링합니다.
   • Level 0 (연속 물리): 선택적 상태 공간 모델 (Selective SSM) 을 사용하여 프레임 간 연속적인 물리 운동을 처리합니다.
   • Level 1 (이산 사건): 이벤트 감지기가 상태 변화를 포착하면, 희소 Transformer(Sparse Transformer) 가 해당 시점만 처리하여 계산 효율성을 높입니다.
   • Level 2 (추상 목표): 압축된 Transformer 가 이벤트 시퀀스를 요약하여 장기적 목표 추론을 수행합니다.
4. 인과 구조 (Causal Reasoning): GNN 엣지 가중치와 NOTEARS(DAG 제약) 정규화를 통해 객체 간 인과 관계를 학습합니다.
5. 지속적 기억 (Continual Memory): Hopfield 네트워크와 EWC(Exponential Weighted Consolidation) 를 활용합니다.

B. 2 단계 학습 프로토콜 (Two-Stage Training)
모델이 분산된 슬롯 코드로 쉽게 예측하는 것을 방지하고, 객체별 특수화를 유도하기 위해 도입된 핵심 기법입니다.

• Stage 1 (재구성 전용): 학습 초기 40% 동안 SBD 재구성 손실만 활성화합니다. 이는 슬롯들이 공간적 영역을 독점하도록 강제하여 객체 분해를 유도합니다.
• Stage 2 (역학 예측): 나머지 60% 에서 JEPA 스타일의 미래 상태 예측 손실을 추가합니다. 이미 분해된 객체 상태에 대해 역학을 학습합니다.

3. 주요 기여 (Key Contributions)

1. 통합 아키텍처: 객체 슬롯, 3 단계 시간 계층 (SSM + 희소 Transformer + 목표 Transformer), GNN 기반 인과 상호작용을 단일 미분 가능 모델로 통합했습니다.
2. 공간 브로드캐스트 디코더 (SBD): 고정된 ViT 특징을 재구성 목표로 사용하여, 실제 로봇 비디오에서 비지도 객체 분해를 가능하게 했습니다.
3. 2 단계 학습 프로토콜: 공간 재구성을 통한 슬롯 특수화 (Stage 1) 와 그 이후의 역학 예측 (Stage 2) 을 순차적으로 수행하여 구조적 학습을 보장합니다.
4. 고성능 Triton 커널: 선택적 SSM 스캔을 위한 커스텀 Triton 커널을 개발하여 순차적 PyTorch 구현 대비 38 배의 속도 향상을 달성했습니다.
5. 오픈 소스 및 실증: Open X-Embodiment 데이터셋의 PushT 로봇 조작 태스크에서 68M 파라미터 모델을 훈련하고, 전체 코드베이스 (8,478 줄, 51 모듈) 를 공개했습니다.

4. 실험 결과 (Results)

• 데이터셋: Open X-Embodiment 의 PushT 태스크 (로봇이 T 자 모양 블록을 목표 지점으로 밀어 넣는 작업).
• 성능 지표:
  • 예측 오차 (MSE): 0.008 (2 단계 학습 적용 시). SBD 없이 분산 표현만 사용한 경우보다 예측 오차는 높았으나 (0.002), 이는 개별 슬롯이 특정 공간 영역을 학습했기 때문입니다.
  • SBD 재구성 손실: 0.0075 로, 슬롯들이 명확한 공간적 영역을 재구성함을 보여줍니다.
  • 다양성 (Diversity): SBD 적용 시 슬롯 간 차이가 더 명확해졌습니다.
• 공간 분해: 학습 후 슬롯들이 T 자 블록, 그리퍼, 테이블 등 서로 다른 공간 영역을 할당받는 것을 시각적으로 확인했습니다.
• 이벤트 감지: 16 프레임 시퀀스당 2~3 개의 주요 상태 변화 (예: 그리퍼와 블록의 접촉) 를 정확히 감지했습니다.
• 속도 향상: NVIDIA T4/H100 환경에서 SSM 스캔 속도가 38 배 이상 향상되었습니다.

5. 의의 및 한계 (Significance & Limitations)

의의:

• HCLSM 은 세계 모델이 실제 세계의 구조 (객체, 계층적 시간, 인과성) 를 반영해야 함을 증명했습니다.
• 단순한 비디오 예측을 넘어, 계획 (Planning) 및 반사실적 추론에 필요한 구조화된 지식을 제공하는 세계 모델의 새로운 패러다임을 제시합니다.
• 실제 로봇 조작 데이터에서 객체 중심 표현이 학습 가능함을 입증했습니다.

한계 및 향후 과제:

• 슬롯 수: 현재 32 개의 슬롯이 모두 활성화되어 있으며, 3 개의 객체가 여러 슬롯에 분할되는 현상이 발생합니다. 슬롯 수를 동적으로 조절하는 메커니즘이 필요합니다.
• 인과성: 명시적인 인과 adjacency matrix 학습은 sparsity 정규화 하에서 실패했습니다. GNN 엣지 가중치가 실제 인과 관계를 반영하는지 검증이 필요합니다.
• 확장성: 262M 이상 파라미터의 Large 모델에서 bf16 정밀도 시 NaN 문제가 발생하여 확장성이 제한적입니다.
• 데이터: 단일 태스크 (PushT) 에서만 훈련되었으며, 다양한 다중 객체 환경으로의 확장이 필요합니다.

결론적으로, HCLSM 은 구조화된 세계 모델링을 위한 강력한 기초를 제공하며, 향후 사전 학습된 비전 백본과 수치적 안정성 해결을 통해 완전한 자율 에이전트로 발전할 가능성을 보여줍니다.