Transformers are Sample-Efficient World Learners
트랜스포머는 샘플 효율적인 월드 학습기이다
Vincent Micheli, Eloi Alonso, François Fleuret (2023)
관측을 이산 토큰으로 변환한 뒤 Transformer 세계 모델이 토큰 공간에서 다음 관측, 보상, 종료를 자기회귀적으로 예측하여 상상 속에서 정책을 학습하는 모델 기반 강화학습 알고리즘이다.
배경
기존 세계 모델(Dreamer 시리즈)은 연속 잠재 공간이나 범주형 잠재 공간에서 RNN 기반 시퀀스 모델로 역학을 학습했다. 한편 NLP에서 Transformer의 자기회귀 생성이 놀라운 성공을 거두었고, 이미지 분야에서도 VQ-VAE/VQGAN이 이미지를 이산 토큰으로 변환하는 효과적인 방법을 제공했다. Transformer의 강력한 시퀀스 모델링 능력을 세계 모델에 직접 활용하려는 시도가 자연스럽게 등장했다.
핵심 아이디어
IRIS(Imagination with auto-Regression over an Inner Speech)는 두 단계 아키텍처로 구성된다. 첫째, VQ-VAE(이산 오토인코더)가 이미지 관측을 이산 토큰 시퀀스로 인코딩한다. 둘째, GPT 스타일의 자기회귀 Transformer가 (관측 토큰, 행동, 보상, 종료) 시퀀스를 다음 토큰 예측으로 학습한다. 이 세계 모델은 현재 상태와 행동이 주어지면 다음 관측 토큰, 보상, 종료 여부를 자기회귀적으로 생성하여 완전한 상상 궤적을 만들어낸다. 이 상상 속에서 Actor-Critic을 학습하며, 실제 환경과의 상호작용 없이도 정책을 개선한다.
방법론
16x16 이미지 패치를 VQ-VAE로 인코딩하여 각 관측을 16개의 이산 토큰으로 변환한다. Transformer 세계 모델은 과거 맥락(관측 토큰 + 행동)에서 다음 관측 토큰, 보상, 종료를 순차적으로 예측한다. 환경에서 수집한 실제 경험과 세계 모델이 생성한 상상 경험을 모두 활용하여 Actor-Critic을 학습한다.
주요 결과
Atari 100K 벤치마크(환경 상호작용 100K 스텝 제한)에서 26개 게임 중 10개에서 인간 성능을 초과했으며, 평균 인간 정규화 점수에서 DreamerV2, SPR 등 기존 방법을 상회했다. 특히 시각적으로 복잡한 게임(Breakout, Asterix 등)에서 뛰어난 성능을 보여, 이산 토큰 기반 세계 모델의 효과를 입증했다.
임팩트
IRIS는 Transformer 기반 자기회귀 세계 모델이 RL에서 효과적임을 증명하여, 이후 DIAMOND, Genie 등 토큰 기반 세계 모델 연구의 흐름을 촉발했다. NLP의 대규모 언어 모델링 패러다임을 세계 모델링에 직접 적용할 수 있다는 가능성을 열었으며, 비디오 생성 모델과 세계 모델의 융합 연구에도 영향을 미쳤다.