Genie: Generative Interactive Environments 리뷰
https://arxiv.org/abs/2402.15391
https://sites.google.com/view/genie-2024
Intro
지난 2월, Google Deepmind가 공개한 생성적 상호작용 환경 (Generative Interactive Environments) 모델, Genie🧞 입니다. Text-to-Image 모델로 생성한 이미지, 손으로 그린 스케치, 실제 사진을 입력으로 받아 액션 제어가 가능한 게임 세계(Playable Worlds)를 생성하는 모델입니다.
지난 몇 년간 생성 AI의 엄청난 발전이 있었음에도, ChatGPT와 같은 언어 모델과 영상 생성 모델간의 사용자 상호작용 및 참여 수준의 격차는 여전히 존재합니다. 유저가 프롬프트를 입력하면 모델이 생성 결과를 출력하고, 그 결과를 바탕으로 유저가 또 다시 프롬프트를 입력하는 식의, 소위 ‘티키타카’라고 하는 상호작용이 생성 AI 사용에 있어 더욱 몰입감 있는 경험을 만들어내는 것도 사실입니다.
만일, 인터넷의 방대한 비디오 데이터가 주어진다면?
새로운 이미지 및 비디오 생성 뿐만 아니라, 완전히 인터랙티브한 경험(Entire Interactive Experience)을 만들어낼 수 있지 않을까요? 이러한 맥락에서, 구글 딥마인드는 생성 AI의 새로운 패러다임, Generative Interactive Environments를 제안합니다. 하나의 텍스트/이미지 프롬프트로부터 인터랙티브한 ‘환경’을 생성하는 것이죠.
🧞Genie
- 11B 개의 파라미터로 Foundation Model의 특성을 가집니다.
- 20만 시간이 넘는 인터넷 게임 영상 데이터셋으로 훈련되었습니다.
- (필터링을 거쳐 정제된 데이터셋은 약 3만 시간)
- 액션 레이블이 없는 Video-only 데이터로 훈련되었음에도, Latent Action Space(잠재 액션 공간)를 학습해 프레임 단위의 액션 제어가 가능합니다. (중요)
Genie의 구조
Genie는 크게 3 가지의 핵심 컴포넌트로 구성되어 있습니다.
- Latent Action Model(LAM): 프레임과 프레임 사이의 Latent Action을 학습합니다.
- Video Tokenizer: Raw 비디오 프레임을 이산 토큰(Discrete Token)으로 변환합니다.
- Dynamics Model: 이전 프레임들의 토큰과 Latent Action을 받아 다음 프레임을 예측합니다.
0. ST-transformer
핵심 컴포넌트를 살펴보기에 앞서, 먼저 ST-transformer에 대해 알아야합니다. Genie의 모든 컴포넌트가 이 아키텍처를 채택했기 때문입니다.
ST-transformer는 Spatial-Temporal Transformer의 약자로, 공간적(Spatial) 정보와 시간적(Temporal) 정보를 함께 처리할 수 있는 트랜스포머 기반 모델입니다.
ST-transformer는 Spatial attention 레이어, Temporal attention 레이어, Feed-forward 레이어로 구성된 L 개의 Spatiotemporal Block(시공간 블록) 으로 이루어져 있습니다.
- Spatial Layer → 각 timestep마다 1 X H X W 토큰 연산
- Temporal Layer → T timestep에 대해 T X 1 X 1 토큰 연산
이미지 처리에 많이 쓰이는 ViT(Vision Transformer)가 아닌, ST-transformer를 채택한 이유는 바로 높은 메모리 효율성입니다.
각 토큰이 시퀀스 내 다른 모든 토큰과 연결되는(Self-attention) 기존 Transformer 방식은 시간 레이어가 있는 비디오 데이터 특성상 연산 시 메모리 비용이 이차적으로(Quadratically) 증가해 효율성이 떨어진다는 문제가 있습니다.
반면 ST-transformer는 연산 복잡도의 지배적 요인인 Spatial Layer의 수가 프레임 수에 비례해 선형적으로 증가하기 때문에 인터랙션이 증가해도 일관성을 유지하는 효율적인 비디오 생성이 가능합니다.
1. Latent Action Model(LAM)
- 비지도 학습 방식: 잠재 액션 모델은 프레임 사이의 잠재 액션(Latent Action) $a_t$를 비지도 학습(Unsupervised Learning)합니다. 액션 제어가 가능한 비디오 생성을 위해서는 이전 프레임에서 행해지는 액션에 대한 다음 프레임 예측을 레이블을 통해 조건화(Conditioning)해야 하는데, 액션 레이블은 인터넷에서 구하기 힘들고 어노테이션 비용이 많이 들기 때문입니다.
- Encoder 학습
- 입력: 모든 이전 프레임 $x_{1:t}$ 와 다음 프레임 $x_{t+1}$
- 출력: 대응되는 연속적인 Latent Actions $ \tilde{a}_{1:t} $
- Decoder 학습
- 입력: 모든 이전 프레임 $x_{1:t}$ 와 Latent Actions $ \tilde{a}_{1:t} $
- 출력: 예측한 다음 프레임 $\hat{x}_{t+1}$
- ST-transformer 아키텍처 사용: Temporal 레이어의 Causal Mask를 통해 전체 비디오 $x_{1:T}$를 입력으로 받아 각 프레임 사이 모든 Latent Actions $\tilde{a}_{1:T-1}$를 생성합니다.
추론(Inference) 시에는 LAM 전체를 유저의 액션이 대체합니다.
2. Video Tokenizer
- VQ-VAE 방식: 토크나이저는 전 비디오 시퀀스에 대해 Standard VQ-VAE를 사용하여 훈련됩니다. T 프레임으로 구성된 비디오 $x_{1:T} = (x_1, x_2, … , x_T) \in \mathbb{R}^{T X HXWXC}$ 를 입력받아, 각 프레임에 대한 이산 표현 $z_{1:T} = (z_1, z_2, … , z_T) \in \mathbb{I}^{TXD}$을 생성합니다.
- ST-transformer 아키텍처 사용: 공간적인 압축에만 집중했던 기존 연구(Phenaki, MaskViT 등)와 달리, ST-transformer를 인코더와 디코더 모두에 사용해 인코딩 안의 Temporal Dynamics를 통합했고, 비디오 품질을 향상시켰습니다. 또한 ST-transformer의 특성에 따라, 각 이산 인코딩 $z_t$는 모든 이전 프레임 $x_{1:t}$의 정보를 포함합니다.
Phenaki도 Temporal-aware 토크나이저(C-ViViT)를 사용하지만 Genie에서 사용한 ST-transformer 기반 토크나이저 (ST-ViVIT)가 훨씬 컴퓨팅 효율적이라는 것을 Ablation 실험 결과로 증명했습니다.
3. Dynamics Model
- Decoder-only MaskGIT 구조: MaskGIT의 Decoder를 사용합니다. 매 timestep t 마다 토큰화된 비디오 $z_{1:t-1}$와 잠재 액션 $\tilde{a}_{1:t-1}$을 입력받아 다음 프레임 토큰 $\hat{z}_t$을 예측합니다.
- ST-transformer 아키텍처 사용: 아키텍처의 인과적(Causal) 구조로 인해 모든 (T-1)개의 프레임 $ z_{1:T-1} $의 토큰과 잠재 액션 $ \tilde{a}_{1:T-1} $ 을 입력으로 사용할 수 있으며, 모든 다음 프레임 $ \hat{z}_{2:T} $에 대한 예측을 생성합니다.
- Cross-entropy Loss: 예측된 토큰 $ \tilde{z}_{2:T} $ 와 Ground Truth 토큰 $z_{2:T}$ 사이의 크로스 엔트로피 손실 함수로 훈련됩니다.
추론 시에는, 0.5~1.0 사이로 균일하게 샘플링된 베르누이 분포 마스킹 비율(Bernoulli Distribution Masking Rate)에 따라 입력 토큰 $z_{2:T-1}$을 무작위로 마스킹합니다.
World Model을 훈련할 때 일반적인 관행은 시간 t에서의 액션을 해당 프레임과 연결시키는 것입니다. 그러나 연구진은 Latent Actions을 (Latent Action과 Dynamics Model 모두에 대한) 추가 임베딩(Additive Embeddings)으로 취급했고, 이것이 Controllability을 향상시킨다는 것을 발견했습니다.
Genie의 추론 과정
- 플레이어가 모델에 이미지 $x_1$을 넣어 프롬프팅 한다. (→ 첫번째 프레임으로 기능)
- 이미지가 비디오 인코더를 통해 $z_1$으로 토큰화 된다.
- 플레이어가 잠재 액션 $a_1$ (특정 범위의 정수)을 선택한다.
- Dynamics Model이 프레임 토큰 $z_1$과 대응되는 잠재 액션 $\tilde{a}_1$ 을 받아 다음 프레임 토큰 $z_2$를 예측한다.
- 액션이 모델에 계속 입력됨에 따라 Autoregressive 방식으로 나머지 시퀀스 $\hat{z}_{2:T}$를 생성한다.
- 동시에 토크나이저의 Decoder를 통해 토큰들이 비디오 프레임 $\hat{x}_{2:T}$로 Decode된다.
실험 결과
데이터셋 Genie는 공식적으로 사용가능한 2D 플랫폼 게임(Platformers)의 인터넷 비디오로부터 수집된 대규모 데이터셋으로 훈련되었습니다. 최종 데이터셋은 6.8M 16초 비디오 클립으로, 약 3만 시간에 해당합니다. 또한, Generality 입증을 위해 RT1에 사용된 로보틱스 데이터셋으로 학습된 모델을 따로 만들었습니다.
평가 지표 Video Fidelity, Controllability의 두 가지 요인으로 성능을 측정했습니다. 이 중 Controllability를 측정하기 위해 Peak Signal-to-Noise Ratio(PSNR) 기반 메트릭을 사용했는데, 1) Ground-truth로부터 추론된 잠재 액션으로 조건화된 $\hat{x}_t$ 과 2) Random Distribution으로부터 샘플링된 $\hat{x}\acute{}_t$ 가 얼마나 다른지를 측정했습니다.
$ \Delta{PSNR} $ 이 크다는 것은 랜덤 잠재 액션으로부터 생성된 비디오와 Ground-truth와의 차이가 더 크다는 것이고, 이는 잠재 액션에 대해 더 높은 수준의 Controllability를 의미합니다.
확장 결과 (Scailing Results)
연구진은 모델 크기와 배치 사이즈를 스케일링 했을 때 나타나는 영향에 대해 실험했습니다.
Scailing Model Size: 40M에서부터 2.7B 개의 파라미터를 가진 모델을 훈련시켰으며,모델 크기가 커질수록 최종 Training Loss가 지속적으로 감소하는 것을 확인할 수 있습니다. 이는 Genie의 Approach가 스케일링을 통해 이점을 얻을 수 있음을 시사합니다.
Scailing Batch Size: 2.3B 모델의 배치 사이즈를 128, 256, 448로 늘려가며 실험했으며, 배치 사이즈가 커질수록 최종 Training Loss가 감소하는 것을 볼 수 있습니다. 이는 배치 사이즈를 늘리면 모델 성능 측면에서 이득을 얻을 수 있음을 시사합니다.
Genie Model: 결과적으로, 10.1B 크기의 Dynamics Model을 배치 사이즈 512로 훈련했습니다. 여기에 Tokenizer 및 Action Model을 더하면 총 10.7B 파라미터, 942B 토큰이 되고, 이것이 최종적인 Genie Model입니다.
질적 결과 (Qualitative Results)
연구진은 두 종류의 모델을 훈련시켰습니다.
- Platformers 데이터셋으로 훈련한 11B 파라미터 모델
- Robotics 데이터셋으로 훈련한 더 작은 모델
Platformers-trained Model: Platformers 데이터셋으로 훈련한 모델에서 흥미로운 창발 능력(Emergent Capability)을 보입니다. 시차 에뮬레이션(Emulating Parallx) 능력으로, Scene 안의 Foreground에 있는 물체와 Background에 있는 물체의 움직임의 비율이 다른 것을 나타냅니다.
Robotics-trained Model: 로보틱스 데이터셋으로 훈련한 모델에서는 텍스트나 액션 레이블 없이도 분명하고 일관적인 액션을 잘 학습한 모습을 보입니다. 로봇 팔에 대한 컨트롤 뿐만 아니라, 다양한 오브젝트에 대한 Interaction 및 Deformation도 학습한 모습을 볼 수 있습니다. 언젠가는 Genie가 Generalist Agents를 훈련하는 Foundation World Model이 될 수 있을 것 같습니다.
마무리
생성 AI의 새로운 패러다임, 인터랙티브한 환경을 생성하는 모델 🧞Genie에 대해 살펴봤습니다. Video-only 데이터로만 학습되었음에도 프레임 수준의 액션 제어가 가능하다는 것은 아주 놀라운데요. 물론 이런 Genie에게도 앞으로 개선할 점이 많이 남아 있습니다.
- 타 Autoregressive Transformer 모델의 약점을 일부 이어받아 비현실적인 미래를 Hallucination할 가능성이 있습니다.
- 16 프레임의 메모리 제한으로 인해 긴 시간 동안 일관된 환경을 구현하는 데 어려움이 있습니다.
- 현재 1fps로 작동하고 있어, 원활한 인터랙션을 위해서는 향후 개선이 필요합니다.
반면, 앞으로의 AI 연구를 위한 잠재력도 가지고 있습니다.
- Genie의 범용성(Generality)으로 더 많은 인터넷 동영상으로 모델을 학습시켜, 더 다양하고 현실적인 상상 속 환경을 시뮬레이션 할 수 있습니다.
- 현재 강화학습의 한계 중 하나는 풍부하고 다양한 환경의 부족입니다. 이런 상황에서 Genie가 보다 범용적인 능력을 갖춘 에이전트 훈련의 새로운 길을 열 수 있습니다.
쓰다보니 조금 길고 어려운 글이 되었네요. 위 내용을 쉽게 요약한 카드뉴스를 여기에서 보실 수 있습니다. 제가 직접 포토샵으로 만들었습니다😅 그럼 여기서 마무리 하겠습니다. 읽어주셔서 감사합니다!
댓글남기기