Fast-LeWorldModel 액션-프리픽스 예측 기반 빠른 잠재 세계 모델

샘플링 기반 계획은 후보 액션 시퀀스마다 반복적인 잠재 롤아웃을 수행하므로 계산 비용과 누적 예측 오류가 빠르게 증가한다. Fast-LeWM은 액션 접두사 단위로 병렬 예측을 수행하여 후보 평가당 반복적 연산을 제거하고 오픈루프 오차 성장을 현저히 완화했다. 이 접근은 샘플링 예산이 제한된 환경에서 실시간에 가까운 계획과 더 안정적인 장기 예측을 가능하게 한다.

기존 LeWM 계열의 잠재 세계 모델은 한 단계 전이만 학습하는 로컬 전이 함수를 사용하고 후보 시퀀스 평가는 이 함수를 반복적으로 재귀 적용하여 전체 상상 궤적을 생성한다. 이 방식은 각 단계의 예측을 다음 단계 입력으로 사용하므로 초기 예측 오차가 이후 단계로 누적되고, 후보 평가 시 동일한 인코딩과 전이 모듈을 여러 번 호출해야 하므로 계산 비용이 크게 증가한다. 이러한 특성은 샘플링 기반 계획에서 성능 저하와 시간 소모의 주요 원인이었다. Fast-LeWM은 후보 시퀀스의 부분 접두사별 누적 영향을 요약하는 프리픽스 토큰을 도입하여 문제를 재정의한다. 인코더는 causal Transformer로 액션 시퀀스를 처리하고 각 위치 출력은 해당 길이의 접두사만을 반영하도록 마스킹된다. 예측기는 관찰에서 얻은 앵커 잠재 z_t와 각 프리픽스 토큰을 입력으로 받아 모든 호라이즌의 미래 잠재를 병렬로 생성하므로 중간 예측에 의존하지 않는 독립적 쿼리가 가능해진다. 이 구조는 학습과 추론에서 서로 다른 이점을 제공한다. 학습 단계에서는 각 프리픽스에 대응하는 실제 미래 잠재를 타깃으로 하는 촘촘한 다중-호라이즌 손실이 각 접두사 토큰이 누적 행동 효과를 정확히 요약하도록 강제한다. 추론 단계에서는 한 번의 프리픽스 인코딩과 한 번의 병렬 예측으로 모든 후보 호라이즌을 평가하므로 반복적 동적 호출이 제거되고 오픈루프 오차 성장률이 크게 완화되었다.

전체 파이프라인은 현재 관찰을 시각 인코더로 잠재 z_t로 변환하고 이 잠재를 상태 토큰으로 변환한 뒤 액션 토큰 시퀀스 앞에 붙인다. 이 상태-액션 토큰 시퀀스는 causal Transformer 기반의 Action-Prefix Encoder로 처리되어 k 길이 접두사에 대응하는 프리픽스 토큰 p_{t,k}들을 생성한다. Parallel Latent Predictor는 앵커 잠재 z_t와 모든 프리픽스 토큰을 입력으로 받아 각 호라이즌의 예측 잠재 hat{z}_{t+k}를 병렬로 출력하며 손실은 각 호라이즌의 MSE 평균과 SIGReg 정규화의 합으로 구성된다.

주요 벤치마크에서 Fast-LeWM은 LeWM 대비 평균 성공률을 85.8%에서 90.5%로 향상시켰고 자기일관성 항을 추가하면 평균 성공률이 92.0%로 올라갔다. 환경별로는 Two-Room에서 87%에서 98%로, Reacher나 PushT 등에서도 동등하거나 더 나은 성능을 보였다. 이 결과는 접두사 기반 학습이 목표-조건화 계획 성능을 개선함을 의미한다. 계산 효율성 측면에서 Fast-LeWM은 동적 모듈 평가 시간을 31.4초에서 8.0초로 줄여 약 3.9배 빠르게 동작했고 전체 CEM 시간은 54.4초에서 28.3초로 약 48% 감소했다. 모든 시간 측정은 NVIDIA 4090에서 동일한 CEM 예산으로 수행되었다. 이는 후보 평가 당 반복적 전이 호출을 제거한 것이 실질적인 계획 시간 절감으로 이어졌음을 의미한다. 오픈루프 잠재 예측 실험에서 Fast-LeWM은 초기 예측 오차가 더 작고 손실 곡선의 기울기가 작아 오차 성장률이 느려졌다. 물리 상태 프로빙에서는 선형 프로브 성능은 LeWM과 유사했고 MLP 프로브에서는 Fast-LeWM이 더 낮은 MSE와 더 높은 상관을 보여 잠재가 더 많은 물리 정보를 보전함이 확인되었다. Ablation에서 Long-Action LeWM은 단순히 액션 블록을 늘리는 방식으로는 성능 향상을 달성하지 못했고 접두사 밀집 감독이 핵심 요소임이 검증되었다.

전체 아키텍처는 LeWM의 시각 인코더와 동일한 재구성 없는 JEPA 스타일 표현을 사용하고 동적 모듈만 액션-프리픽스 기반으로 대체했다. 관찰 o_t는 f_\theta로 인코딩되어 z_t를 형성하고 이 z_t는 두 층 MLP로 상태 토큰으로 변환되어 액션 토큰 앞에 삽입된다. Action-Prefix Encoder는 causal Transformer로 구현되어 각 위치 출력은 해당 접두사만 열람 가능하도록 causal mask를 적용받는다. 수학적 흐름은 다음과 같다. 길이 k의 접두사 a_{t:t+k-1}가 주어지면 인코더 E_\psi^{(k)}는 접두사 토큰 p_{t,k}을 생성한다. 예측기 G_\phi는 입력으로 앵커 잠재 z_t와 접두사 토큰 p_{t,k}을 받아 연산을 통해 hat{z}{t+k}를 출력한다. 구체적으로 학습 손실은 각 k에 대해 |\hat{z}{t+k}-z_{t+k}|_2^2를 평균한 값이며 배치 수준에서 SIGReg 정규화를 더하여 표현 붕괴를 억제한다. 구현 세부는 다음과 같다. 액션-프리픽스 Transformer는 3개 레이어와 6개 어텐션 헤드를 사용하고 토큰 차원은 192로 설정되었다. 예측기는 6개 레이어의 action-modulated residual MLP로 설계되었고 잠재 차원 192, 히든 폭 2048, fusion 폭 768, AdaLN-zero 모듈과 드롭아웃 0.1을 사용했다. 전체 모델 파라미터 수는 약 17.9M로 공개된 LeWM 체크포인트와 유사하다. 계획 통합은 기존 LeWM과 동일한 CEM 기반 목표-조건화 프로토콜을 따르며 후보 평가 함수는 터미널 잠재와 목표 잠재의 제곱거리로 계산된다. 선택적 자기일관성 항은 길이-H 접두사로 직접 추정한 터미널과 중간 접두사를 통해 재추정한 터미널의 차이를 β로 가중하여 후보 비용에 더한다. β는 0이면 기존 목표 전용 비용이며 1로 설정한 실험에서 추가 성능 향상을 관찰했다.

논문에 명시된 제한 사항으로 먼저 액션 프리픽스만으로는 결과가 결정되지 않기 때문에 현재 관찰을 상태 토큰으로 조건화해야 정확한 접두사 효과를 예측할 수 있다. 학습 시 예측 호라이즌은 로딩된 궤적에 따라 적응적으로 결정되고 논문 실험에서는 [1,5] 범위로 클램프되어 있어 더 긴 연속 제어 단계에 대한 성능과 일반화 가능성은 추가 검증이 필요하다. 또한 촘촘한 프리픽스 감독이 핵심 구성 요소이므로 최종 상태만 감독하는 terminal-only 설정은 전체 모델 성능에 비해 저하가 발생함이 확인되었다.