기술 요약: 흐름 등변 세계 모델 (Flow Equivariant World Models)

1. 문제 정의

체화된 에이전트(Embodied agents)는 제한된 1인칭 시야만을 가진 동적인 환경에서 작동합니다. 이러한 **부분 관측성(partial observability)**은 기존 세계 모델에 큰 도전 과제를 안겨줍니다. 에이전트가 환경의 일부만을 관찰하는 동안, 관찰되지 않는 영역은 계속해서 변화하기 때문입니다. 대규모 잠재 디퓨전 트랜스포머(latent diffusion transformers)와 순환 모델(recurrent models)을 포함한 현재의 최첨단 방법론들은 이러한 설정에서 장기적인 일관성을 유지하는 데 종종 실패합니다.

세계의 기저에 깔린 대칭성을 존중하는 메모리 구조가 없다면, 이러한 모델들은 다음과 같은 문제를 겪습니다:

• 정보 손실(Information Loss): 에이전트가 이동함에 따라, 과거의 관련 정보가 컨텍스트 윈도우(예: 트랜스포머의 슬라이딩 윈도우)를 벗어나 폐기됩니다.
• 환각(Hallucination): 에이전트가 이전에 관찰했던 영역으로 다시 시선을 돌렸을 때, 모델은 해당 객체의 상태를 기억해내는 데 실패하고, 대신 새로운 객체를 환각하여 생성하거나 기존 객체를 사라지게 만듭니다.
• 구조적 사전 지식의 결여(Lack of Structural Priors): 기존 모델들은 관찰을 정적인 시퀀스나 시점 의존적 데이터로 취급하며, 자가 운동(self-motion)과 외부 객체 운동 모두가 매끄러운 시간 매개변수 대칭(flows)을 따른다는 사실을 간과합니다.

2. 방법론: 흐름 등변 세계 모델 (FloWM)

저자들은 FloWM을 제안하며, 이는 **흐름 등변성(flow equivariance)**을 세계 모델 메모리의 핵심에 통합하는 프레임워크입니다. 이들의 핵심 가설은 예측 표현(predictive representations)이 세계의 시간적 및 역학적 구조와 정렬될 때 더욱 강력해진다는 것입니다.

핵심 개념: 흐름 등변성 (Flow Equivariance)

이 프레임워크는 정적 군 등변성(static group equivariance)을 시간 매개변수 시퀀스 변환, 즉 "흐름(flows)"으로 확장합니다.

• 정의: 모델이 입력 시퀀스가 흐름(예: 시간이나 운동에 의한 이동)을 겪을 때, 출력 시퀀스 또한 동일한 흐름에 따라 변환된다면, 그 모델은 흐름 등변적이라고 합니다.
• 전역 등변성(Global Equivariance): 부분 관측 환경에서 모델은 전체 세계 상태(관찰된 영역과 관찰되지 않은 영역 모두)에 대한 잠재 메모리를 유지해야 합니다. 잠재 표현은 현재의 관찰뿐만 아니라 전역적인 세계 상태에 대해 등변적으로 변환되어야 합니다.
• 공동 이동 참조 프레임(Co-moving Reference Frame): 이를 달성하기 위해, 순환 메모리는 새로운 관찰을 통합하기 전에 동일한 흐름에 의해 전달(transported)됩니다. 이는 특징(features)들이 고정된 프레임에 남겨지는 것이 아니라, 움직이는 객체와 공간적으로 정렬된 상태를 유지하도록 보장합니다.

아키텍처 구성 요소

프레임워크는 일반화된 흐름 등변 순환 관계를 통해 구현됩니다:
$$h_{t+1}(\nu) = \psi_1(\nu) \cdot U_\theta [h_t(\nu); E_\theta [f_t; h_t] (\nu)]$$
여기서:

• $h_t(\nu)$는 서로 다른 벡터장 운동에 대응하는 "속도 채널(velocity channels)"($\nu$)로 구성된 은닉 상태를 나타냅니다.
• $E_\theta$는 관찰을 잠재 공간으로 매핑하는 인코더입니다.
• $U_\theta$는 업데이트 연산입니다.
• $\psi_1(\nu)$는 한 타임스텝 동안의 흐름 작용을 나타냅니다.

자가 운동 등변성(Self-Motion Equivariance):
모델은 에이전트의 자가 운동을 상대적 흐름으로 취급합니다. 에이전트의 행동($T_{a_t}$)에 따라 은닉 상태를 변환함으로써, 메모리는 효과적으로 에이전트의 공동 이동 참조 프레임 내에서 이동합니다. 이를 통해 모델은 에이전트의 행동을 기반으로 새로운 관찰이 올바른 상대적 위치에 통합되는 구조화된 환경 "지도"를 구축할 수 있습니다.

두 가지 구현체

1. 단순 순환 FloWM (2D): 2D 환경(예: MNIST World)을 위해 설계되었습니다. 은닉 상태와 공간적 입력에 대해 합성곱(convolutional) 연산을 사용합니다. 에이전트의 시야 주변으로 상태가 흐르는 동안, 은닉 상태에 읽고 쓰기 위해 고정된 윈도우를 사용합니다.
2. 트랜스포머 기반 FloWM (3D): 복잡한 3D 환경(예: Block World)을 위해 설계되었습니다. 비전 트랜스포머(ViT) 인코더와 디코더를 활용합니다. 잠재 메모리는 공간적으로 조직된 토큰 임베딩의 집합으로서 상위 수준의 추상적 지도로 작용합니다.
   • 참고 (등변성): 비용이 많이 드는 깊이 역투영(depth unprojection) 없이 3D-to-2D 투영에 대해 정확한 인코더 등변성을 달성하는 것은 어렵지만, 저자들은 순환 흐름 변환($T^{-1}_{a_t} \psi_1(\nu)$)이 훈련 중에 인코더가 등변성을 학습하도록 유도하여 명시적인 기하학적 제약 없이도 높은 충실도를 달성함을 보여줍니다.

3. 주요 기여

• 흐름 등변성의 공식화: 부분 관측 동적 세계 모델링을 위한 구조화되고 지속적인 메모리를 구축하기 위해 시간 매개변수 대칭(흐름)을 어떻게 사용할 수 있는지 공식화했습니다.
• 구조화된 메모리 아키텍처: 자가 운동 및 외부 객체 운동 모두에 따라 이동하고 변환되는 "속도 채널"을 가진 순환 메모리를 도입하여, 시야 밖 영역에 대한 정보가 시간이 지나도 정렬된 상태를 유지하도록 했습니다.
• 새로운 벤치마크: 부분 관측 세계 모델링을 위한 새로운 벤치마크를 도입했습니다:
  • 2D MNIST World: 무작위 속도로 움직이는 숫자들이 있는 부분 관측 환경.
  • 3D 동적 Block World: 에이전트가 이산적 행동을 통해 탐색하는 방 안에서 색깔 있는 블록들이 튀어 오르는 환경.
• 경험적 검증: FloWM이 장기 예측 및 길이 일반화(length generalization) 측면에서 최첨단 디퓨전(DFoT, DFoT-SSM) 및 순환(RSSM) 베이스라인을 크게 능가함을 입증했습니다.

4. 실험 결과

저자들은 DFoT, DFoT-SSM, 그리고 DreamerV3의 RSSM을 포함한 베이스라인들과 Flo-WM을 비교 평가했습니다.

성능 지표

• 장기 일관성(Long-Horizon Consistency): 2D MNIST World에서 FloWM은 훈련 호라이즌(20 타임스텝)을 넘어 150 타임스텝까지 일관된 예측을 유지했으나, 베이스라인들은 발산하거나 환각을 일으켰습니다. 3D Block World에서 FloWM은 210 프레임 동안 일관성을 유지한 반면, 베이스라인들은 객체를 놓치거나 새로운 객체를 환각했습니다.
• 오차 감소: FloWM은 모든 베이스라인에 비해 현저히 낮은 평균 제곱 오차(MSE)와 높은 PSIM/SSIM 점수를 달성했습니다. 예를 들어, 3D Block World에서 FloWM의 MSE는 약 0.0006(70 프레임)이었던 반면, DFoT는 0.0117였습니다.
• 학습 효율성: 자가 운동 등변성(SME)과 속도 채널(VC)을 결합한 모델은 절제된 버전들에 비해 수렴하는 데 필요한 훈련 단계가 수십 배 적었습니다.
• 다운스트림 플래닝(Downstream Planning): "빨간 블록 찾기" 플래닝 작업에서 FloWM은 목표물을 향해 성공적으로 이동했습니다. 베이스라인들은 종종 목표물의 위치를 환각하여 플래닝 성능이 저하되었습니다.
• 프로브 분석(Probe Analysis): 잠재 공간에서 블록 위치를 디코딩하도록 훈련된 프로브 네트워크는 DFoT의 경우 1% 미만의 정확도를 보인 반면, FloWM은 96%의 정확도를 보여 FloWM이 객체 궤적에 대한 구조적이고 등변적인 표현을 학습했음을 확인시켜 주었습니다.

절제 연구 (Ablation Studies)

속도 채널(VC) 또는 자가 운동 등변성(SME)을 제거하면 성능이 크게 저하되어 종종 베이스라인 디퓨전 모델과 비슷하거나 그보다 낮은 수준으로 떨어졌으며, 이는 두 구성 요소가 부분 관측성과 역학을 처리하는 데 매우 중요하다는 것을 강조합니다.

5. 의의 및 주장

본 논문은 예측 표현이 세계의 시간적 및 역학적 구조와 정렬될 때 근본적으로 더 강력하다고 주장합니다.

• 이론적 토대: 이 연구는 동적인 부분 관측 환경에서 구조화된 메모리의 필요성을 규정하는 흐름 등변성이라는 이론적 프레임워크를 제공합니다.
• 부분 관측성 해결: 컨텍스트 윈도우를 확장하거나 정적 메모리 뱅크에 의존하는 기존 방법과 달리, FloWM은 세계 상태와 함께 움직이는 메모리를 유지함으로써 관찰되지 않는 영역의 진화를 본질적으로 처리합니다.
• 일반화: 이 프레임워크는 훈련 중에 본 것보다 훨씬 긴 시퀀스로 일반화할 수 있게 하며, 이는 실제 체화된 에이전트에게 중요한 능력입니다.
• 제한된 범위: 저자들은 현재의 평가가 강체 기하학적 행동이 있는 통제된 환경에서 이루어졌음을 인정합니다. 이들은 이 작업을 오픈 월드 데이터셋, 연속적 행동 공간 및 더 복잡한 물리적 상호작용을 위한 향후 확장을 위한 기초로 설정합니다.

요약하자면, FloWM은 운동의 대칭성을 메모리 아키텍처에 명시적으로 인코딩함으로써, 부분 관측된 동적 환경에서 세계 모델이 안정적이고 정확하며 환각이 없는 예측을 달성할 수 있음을 보여줍니다.