자가 지도 학습 기반의 강력한 월드 모델링 기법, StarWM
StarWM: Self-Supervised Trained Attention Routing for Robust World Models
arXivWorld Model의 한계를 극복하기 위해, StarWM은 크로스 어텐션 모듈을 활용하여 환경 역학 학습 시 재구성이 필요한 핵심 영역만 선별적으로 식별합니다.
- 이 기술은 어텐션 기반으로 재구성을 제한하여, 중요한 상태 정보는 보존하면서도 주변의 방해 요소(distractors)를 체계적으로 무시하는 것이 가능합니다.
- 동적 배경 환경 테스트에서 StarWM은 기존 재구성 기반 모델보다 강력한 추론 능력을 입증했으며, 최고 수준의 성능을 달성했습니다.
- StarWM은 AI가 단순히 데이터를 복원하는 것을 넘어, 환경의 역학적 관계를 이해하고 핵심 상태 정보를 유지할 수 있음을 보여주는 중요한 진전입니다.
월드 모델은 환경 역학을 충실히 포착하는 것과 불필요한 내용을 추상화하는 것 사이의 균형이 필요합니다. 기존 재구성 기반 모델들은 표현 능력을 픽셀 영역에 할당하여, 과제와 무관한 내용이 학습된 표상을 지배할 위험이 있습니다. StarWM은 이러한 문제를 해결하기 위해 자기 지도 역학(self-supervised dynamics)으로 훈련된 크로스 어텐션 모듈을 활용하여 재구성이 적용되어야 할 핵심 영역을 결정합니다.
StarWM은 이중 스트림 디코더를 통해 재구성을 오직 어텐션된 영역으로 제한하며, 스톱-그라디언트 장벽을 두어 두 가지 목표 간의 간섭을 방지합니다. 이러한 구조는 재구성이 어텐션된 영역의 시각적 콘텐츠만 감독하도록 하며, 잠재 공간에 예측과 무관한 정보가 오염되는 것을 막아줍니다.
실험 결과, 동적 비디오 배경이 있는 DeepMind Control 환경에서 기본(보상 없음) StarWM은 무작위 프레임 방해 요소가 존재하는 상황에서 재구성 기반 모델보다 가장 강력한 성능을 달성했습니다. 또한, 보상 증강 변형 모델은 순차 비디오에서 재구성-비제한 방식과 동등하거나 그 이상의 성능을 보여 모든 방해 요소 체제에서 최고 수준의 총 리턴을 기록했습니다.