기술 요약: World Machine: 시계열을 위한 생성적 세계 모델링을 향해

1. 문제 제기

현재의 생성형 AI 모델, 특히 트랜스포머 기반 모델들은 환경을 예측적으로 이해하고 제어 가능한 시뮬레이션이 가능한 진정한 "세계 모델"을 만드는 데 심각한 도전에 직면해 있습니다. 기존 트랜스포머는 컨텍스트 길이에 비례하는 2 차 계산 및 메모리 비용으로 인해 긴 시간적 시퀀스에 대한 확장성이 제한됩니다. 또한, 기존 접근 방식들은 종종 수동으로 설계된 규칙에 의존하거나, 해당 관측치를 생성하는 내부 잠재 상태를 명시적으로 모델링하지 않고 감각 스트림 간의 상관관계에만 초점을 맞추는 경향이 있습니다. 다음과 같은 아키텍처가 필요합니다:

• 다양한 양의 관측 데이터와 컨텍스트에 적응할 수 있어야 함.
• 감각 상관관계뿐만 아니라 내부 세계 상태 (잠재 표현) 를 명시적으로 모델링할 수 있어야 함.
• 전체 기록에 의존하지 않고 단축된 컨텍스트 (예: 단일 이전 상태) 에서 예측을 가능하게 함으로써 추론 비용을 줄일 수 있어야 함.
• 이질적인 데이터와 여러 감각 채널 (잠재적으로 멀티모달) 을 처리할 수 있어야 함.

2. 방법론

2.1. World Machine 아키텍처

World Machine 은 **잠재 세계 상태 ($ws$)**를 기반으로 작동하는 디코더 전용 트랜스포머에 기반한 생성형 멀티모달 시계열 모델링 아키텍처입니다.

• 잠재 세계 상태 ($ws$): 이는 표준 트랜스포머의 임베딩과 유사하지만 고정된 시뮬레이션 주기로 시간에 따라 진화하는, 시뮬레이션된 세계의 현재 구성을 인코딩하는 실수 벡터입니다.
• 감각 채널: 모델은 특정 채널을 통해 입력 및 출력 감각 데이터를 처리합니다. 입력 채널은 사전 인코더를 통과하는 반면, 모델은 $ws$를 디코딩하여 출력 감각 데이터를 생성합니다.
• 핵심: 시뮬레이션 엔진은 트랜스포머 블록으로 구성됩니다. 이는 감각 데이터를 통해 스케일링 및 시프트를 수행하여 잠재 상태를 조건화하는 감각 블록 (특히 adaLn-Zero 유형) 을 활용합니다. 특정 단계에서 감각 데이터가 누락된 경우, 해당 블록은 이러한 연산을 건너뜁니다.
• 위치 인코딩: 모델은 정현파 인코딩 대신 ALiBi(선형 편향을 가진 어텐션) 를 사용합니다. 이는 무제한 시퀀스 크기를 허용하며 위치 정보를 어텐션 행렬에 직접 주입하여 잠재 상태 임베딩을 방해하지 않습니다.
• 상태 활성화: 모델 안정성을 보장하기 위해 상태 출력에 tanh 활성화 함수가 적용됩니다.

2.2. 훈련 프로토콜: 상태 발견

주요 과제는 핵심이 실제 값 (ground truth) 으로 제공되지 않는 잠재 상태 ($ws$) 를 기반으로 작동한다는 점입니다. 이를 해결하기 위해 저자들은 상태 발견 (State Discovery) 프로세스를 제안합니다:

• 초기화: 잠재 상태는 균일 분포로 초기화됩니다.
• 반복적 업데이트: 훈련 중 핵심에서 예측된 상태 ($\hat{ws}$) 는 "왼쪽으로 이동 (shift to left)" 연산자를 사용하여 시간적으로 이동된 후 다음 에포크의 입력으로 피드백됩니다. 모든 시퀀스의 첫 번째 상태는 널 벡터로 재설정됩니다.
• 최적화: 모델은 출력 차원에 대한 가중 손실 합을 사용하여 다음 시간 단계의 감각 데이터를 예측하도록 훈련됩니다 (다음 토큰 예측과 유사). 이는 $ws$가 감각 데이터를 재구성하는 데 필요한 정보를 인코딩하도록 강제합니다.

2.3. 향상된 훈련 프로토콜

표현의 질과 견고성을 개선하기 위해 훈련 프로토콜에는 몇 가지 구체적인 단계가 포함됩니다:

• 감각 마스킹: 훈련 중 감각 데이터를 무작위로 숨겨 데이터가 누락되었을 때 모델이 잠재 상태에 의존하도록 강제합니다.
• 시퀀스 분할: 시퀀스를 무작위 세그먼트로 나눕니다. 모델은 선행 데이터 없이 시퀀스 중간 상태에 대처해야 합니다. "Fast Forward" 메커니즘은 한 세그먼트의 마지막 상태를 다음 세그먼트의 첫 번째 상태로 이동시켜 정보 기울기를 전파합니다.
• 노이즈 추가: $ws$ 및/또는 감각 차원에 가우시안 노이즈를 추가합니다.
• 단기 시간 회상: 현재 단계에 대한 과거 또는 미래 데이터를 나타내는 합성 채널을 생성하여 시간적 컨텍스트를 개선합니다.
• 로컬 모드: 세그먼트에 대해 어텐션 메커니즘을 무작위로 비활성화하여 모델이 현재 $ws$와 즉각적인 감각 데이터에만 의존하도록 강제합니다 (RNN 은닉 상태 처리와 유사).

2.4. 평가 프레임워크

저자들은 무작위 초기 상태와 임펄스/정사각파 입력을 가진 감쇠 물리 시스템을 나타내는 합성 1 차원 시계열인 Toy1D 데이터셋을 소개합니다. 평가 작업은 다음과 같습니다:

• ** shallow 예측 (Prediction Shallow):** 감각 데이터 없이 단순히 마지막 인코딩된 상태만을 사용하여 시퀀스의 미래 50% 를 추론합니다. 이는 컨텍스트 단축 능력을 테스트합니다.
• 상태 사용 (Use State): 감각 입력 없이 이전에 인코딩된 상태에서 출력을 생성합니다.
• 마스킹 감각 (Mask Sensory): 무작위 비율의 감각 데이터가 마스킹된 시퀀스를 생성합니다.
• 영향 테스트: 성능과 훈련 시간에 대한 특정 하이퍼파라미터 및 프로토콜 단계의 영향을 정량화하기 위한 통계적 방법 (윌콕슨 부호 순위 검정 사용).

3. 주요 기여

1. 새로운 아키텍처: 경험을 고차 마르코프 과정으로 취급하며 감각 상관관계가 아닌 내부 잠재 상태 ($ws$) 를 명시적으로 모델링하는 통합 엔드투엔드 아키텍처인 World Machine 의 도입.
2. 상태 발견 프로토콜: 실제 상태 레이블 없이 잠재 상태를 인코딩하고 진화시키는 방법을 학습할 수 있게 하는 훈련 메커니즘으로, 널 벡터에서 자기회귀 추론을 가능하게 함.
3. 컨텍스트 단축 능력: 단축된 컨텍스트 (특히 "Prediction Shallow" 작업) 로 추론을 수행할 수 있는 능력 입증으로, 기존 트랜스포머와 관련된 2 차 비용 절감을 위한 잠재적 경로 제공.
4. 종합적 영향 분석: 성능과 비용에 대한 구체적인 기여도를 결정하기 위해 영향 테스트 프레임워크를 사용하여 훈련 프로토콜 구성 요소 (마스킹, 시퀀스 분할, 로컬 모드 등) 를 체계적으로 평가.
5. 오픈소스 베이스라인: 최근 대규모 세계 모델의 불투명성을 해결하기 위해 아키텍처, 훈련 프로토콜 및 합성 데이터셋을 공개.

4. 결과

Toy1D 데이터셋에서 다양한 모델 구성과 훈련 프로토콜을 사용하여 실험이 수행되었습니다.

• 훈련 가능성: 손실 곡선의 감소로 입증된 바와 같이, 상태 발견 기술이 성공적으로 모델 훈련을 가능하게 함.
• 작업 성능:
  • 모델은 "일반" 자기회귀 추론에서 가장 좋은 성능을 보임.
  • Base 모델(상태 발견만 사용) 은 감각 데이터가 없는 작업을 수행하지 못함.
  • 완전 프로토콜(마스킹, 시퀀스 분할, 노이즈 등 포함) 은 특히 Prediction Shallow 작업에서 상당한 개선을 보였으나, 이는 여전히 가장 어려운 작업으로 남음.
  • 정성적 결과는 감각 데이터가 있을 때 특히 일관된 예측을 보였으나, 고주파 시퀀스는 여전히 어려웠음.
• 프로토콜 영향:
  • 안정성: tanh 상태 활성화 사용이 결정적이었음. 이를 제거하고 (MSE 정규화로 대체) 약 20% 의 발산율을 보인 반면, 이를 유지하면 발산율이 1% 미만으로 감소함.
  • 성능: "Prediction Shallow" 작업이 가장 큰 변이를 보임. 1,000 에포크 후 0.3435의 Prediction Shallow MSE 를 달성한 최상의 경험적 구성 (시퀀스 분할, 특정 블록 구성, 노이즈 추가, 로컬 모드 결합) 을 보임.
  • 비용/편익: 로컬 모드 변수는 비싼 어텐션 메커니즘을 우회하여 성능과 훈련 시간 모두에 긍정적인 영향을 보임.
  • 상관관계: "Use State" 및 "Prediction Shallow" 작업은 다른 작업과 낮은 상관관계를 보여, 세계 모델링 능력을 평가하는 고유하고 필수적인 평가자임을 시사함.

5. 중요성과 주장

이 논문은 World Machine 을 예측적 이해와 제어 가능한 시뮬레이션 사이의 균형을 맞춘 생성적 세계 모델링을 향한 기초적인 단계로 위치시킵니다.

• 효율성: 단일 잠재 상태 (Prediction Shallow) 에서 미래 상태를 추론할 수 있는 능력은 현재 트랜스포머의 2 차 확장성 문제를 완화할 가능성을 시사하며, 장기 계획에 대해 더 비용 효율적인 접근 방식을 제공함.
• 적응성: 이 아키텍처는 누락된 데이터와 다양한 감각 입력을 처리하도록 설계되어, 완전한 관측이 항상 가능하지 않은 복잡하고 역동적인 환경에 적합함.
• 투명성: 완전히 재현 가능한 오픈소스 베이스라인을 공개함으로써, 대규모 독점 시스템에서 종종 불투명했던 세계 모델의 메커니즘을 신비에서 해방시키려는 목표.
• 미래 가능성: 저자들은 이 접근 방식이 멀티모달 문제 (예: 시각 - 관성 오도메트리) 와 대규모 언어 모델 (LLM) 에 적용되어 추론 비용을 줄이고 품질/비용 트레이드오프에 대한 제어를 개선할 수 있다고 제안함. 또한 향후 실세계 멀티모달 데이터셋에서 테스트하고 데이터 생성 및 환경 비용과 관련된 윤리적 문제를 해결해야 함을 인정함.

이 논문은 합성 데이터에 대한 현재 결과가 유망하지만, 아키텍처가 훈련 가능하고 일관된 예측을 생성할 수 있음을 보여주며, 제안된 상태 발견 및 훈련 프로토콜의 실현 가능성을 검증한다고 결론지음.