세계 최초 S4 기반 월드모델, AI 에이전트의 '장기기억' 한계를 넘다

“Facing Off World Model Backbones: RNNs, Transformers, and S4”(Deng, Park, and Ahn, 2023) 논문은 모델 기반 강화학습에서 잠재 월드모델의 시퀀스 모델 백본으로 어떤 구조가 가장 효과적인지 탐구합니다. 연구진은 S4WM을 제안하여, 구조화된 상태공간 시퀀스(S4) 레이어를 고차원 시각적 잠재 상상에 최초로 적용했습니다.

관련 연구

• 잠재 월드모델 / 모델 기반 강화학습. 가장 직접적인 기반이 되는 연구는 PlaNet, Dreamer, DreamerV2, DreamerV3 같은 잠재 동역학 에이전트들입니다. 이들은 컴팩트한 확률적 상태 표현을 학습하고, 정책 최적화를 위해 잠재 미래 궤적을 예측합니다. 논문은 순환 RSSM 스타일 모델, 특히 Dreamer를 핵심 기준선(baseline) 계열로 설정합니다.

• 순환 신경망(RNN). RNN, LSTM, GRU 기반 동역학 모델은 온라인 순환 상태 업데이트를 제공하며 시각적 월드모델링의 중심이었습니다. 주요 한계는 시간 축을 따라 순차적으로 훈련해야 한다는 점으로, 이는 시계열 병렬 처리를 어렵게 만들고 긴 문맥 모델링에 불리합니다.

• 트랜스포머 월드모델. TransDreamer를 포함한 트랜스포머 기반 동역학 모델은 순환 구조를 자기 주의(self-attention)로 대체하여 병렬 훈련과 긴 히스토리에 대한 직접 접근을 가능하게 합니다. 그러나 시퀀스 길이에 대한 표준 2차 시간 복잡도(quadratic attention cost) 때문에, 긴 상상 롤아웃과 장기 시각적 모델링에 계산 비용이 크게 증가합니다.

• 구조화 상태공간 모델(S4). S4는 장거리 의존성을 위해 설계된 구조화 상태공간 시퀀스 모델입니다. 순환 상태공간 해석과 병렬 시퀀스 훈련을 가능하게 하는 컨볼루션 형태를 결합합니다. 논문은 S4가 중간 지점으로 유망하다고 주장합니다. 즉, 선형 시간의 순환 추론과 병렬화 가능한 훈련을 동시에 제공한다는 점입니다.

• S4WM의 차별성. 기존 연구가 S4를 일반적인 시퀀스 벤치마크에 적용한 것과 달리, S4WM은 S4를 잠재 시각 월드모델에 통합합니다. 이미지 인코딩/디코딩과 확률적 잠재 동역학을 결합하여, RNN 및 트랜스포머 백본과 동일한 월드모델링 프레임워크 안에서 직접 비교 평가합니다. 이 연구의 핵심 기여는 고차원 이미지 시퀀스를 잠재 상상을 통해 생성하는 최초의 S4 기반 월드모델이라는 점입니다.

문헌 검토를 위한 간결한 포지셔닝 요약:

기존 월드모델은 주로 순환 잠재 동역학이나 트랜스포머를 사용해 왔습니다. RNN 기반 모델은 추론에 효율적이지만 순차적으로 훈련하는 반면, 트랜스포머 월드모델은 훈련을 병렬화하지만 시계열 비용이 2차적으로 증가합니다. S4WM은 구조화 상태공간 레이어를 대안 백본으로 탐구하여, 병렬 훈련을 유지하면서도 효율적인 장기 동역학 모델링을 지원하는 것을 목표로 합니다.