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

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

왜 중요한가

핵심 기여

액션 프리픽스 예측을 통한 병렬 다중-호라이즌 잠재 예측

Fast-LeWM은 후보 액션 시퀀스의 각 접두사를 요약하는 프리픽스 토큰을 생성하고 현재 관찰 잠재와 결합하여 모든 미래 호라이즌 잠재를 한 번의 순전파로 예측한다. 이 설계는 각 호라이즌을 독립적 질의로 다루어 예측이 순차적으로 의존하지 않게 만든다. 따라서 반복 롤아웃에서 발생하던 중간 예측 오차의 누적을 줄이고 후보 평가 속도를 크게 향상시켰다.

프리픽스 수준의 촘촘한 감독 손실과 상태 토큰 조건화

학습 시 각 프리픽스 토큰에 대응하는 실제 미래 잠재를 타깃으로 하는 평균 MSE 손실을 적용하여 접두사별 누적 행동 효과를 직접 학습시켰다. 현재 관찰 잠재를 두 층 MLP로 토큰화하여 프리픽스 인코더에 선행 토큰으로 제공함으로써 동일한 개방형 액션이 초기 상태에 따라 다른 결과를 낳을 때 이를 구분할 수 있게 했다. 이 조합은 중간 단계 표현의 의미성을 강화하여 최종 예측 성능을 높였다.

계획 효율성 개선과 오픈루프 오차 성장 완화

동적 모듈 평가 시간을 31.4초에서 8.0초로 줄이며 전체 CEM 해결 시간을 54.4초에서 28.3초로 감소시켰다. 평균 성공률은 LeWM의 85.8%에서 Fast-LeWM으로 90.5%까지 향상되었고, 일관성 항을 더하면 92.0%에 도달했다. 다양한 환경에서 초기 예측 오차와 손실의 기울기 모두 감소하여 장기 예측의 견고성이 개선되었다.

선택적 자기일관성 페널티를 통한 후보 안정성 보완

계획 시 길이-H 접두사로 직접 예측한 터미널 잠재와 중간 접두사를 통한 재예측 터미널의 차이를 β 가중치로 페널티에 추가하는 일관성 항을 도입했다. β 값을 조절하면 CEM이 서로 다른 접두사 분해에서 예측이 안정적인 후보를 선호하도록 유도할 수 있다. 논문 결과에서 이 항을 추가하면 평균 성공률이 추가로 상승했다.

핵심 아이디어 이해하기

기존 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 설정은 전체 모델 성능에 비해 저하가 발생함이 확인되었다.

실무 활용

Fast-LeWM은 샘플링 기반 잠재 계획을 사용하는 시스템에서 후보 평가 비용을 크게 줄이고 장기 예측의 견고성을 높여 실시간 또는 근실시간 계획이 요구되는 제어 애플리케이션에 적용 가능하다. 메서드는 기존 LeWM 인코더와 호환되며 동적 모듈 호출 횟수를 줄이므로 동일한 하드웨어에서 더 많은 후보를 평가하거나 더 빠른 주기로 재계획할 수 있다. 공개된 코드로 재현성과 통합이 비교적 용이하다.

실시간 목표-조건화 비전 기반 로봇 제어에서 CEM 예산을 유지하면서 플래닝 주기를 단축하는 용도
데이터 효율적 오프라인 궤적 집합으로 학습한 잠재 모델을 활용하는 시뮬레이터 기반 정책 탐색
대규모 후보 샘플링을 요구하는 모델 기반 MPC에서 계획 안정성 향상을 통한 실패율 저감

코드 공개 여부: 공개

코드 저장소 보기

키워드

JEPA(Joint-Embedding Predictive Architecture)latent world model(잠재 세계 모델)action-prefix prediction(액션 프리픽스 예측)CEM(Cross-Entropy Method)SIGReg(SIGReg 정규화)

용어 해설

Joint-Embedding Predictive Architecture(JEPA): — JEPA는 고해상도 픽셀 재구성 없이 관찰을 잠재 임베딩으로 매핑하고 미래 임베딩을 예측하는 방식이다. 입력 이미지에서 인코더가 잠재 z를 생성하고 모델은 이 잠재 공간에서 다음 상태를 예측하여 픽셀 복원이 아닌 제어 관련 표현을 학습한다. Reconstruction-free 설계로 불필요한 시각적 세부 정보 보존을 억제하여 제어·계획에 적합한 잠재를 형성한다.
액션 프리픽스 예측(Action-Prefix Prediction): — 액션 프리픽스 예측은 후보 액션 시퀀스의 부분 접두사마다 그 접두사를 실행했을 때 도달하는 미래 잠재를 직접 예측하는 방식이다. 각 접두사는 누적된 액션 효과를 요약하는 토큰으로 인코딩되고 현재 관찰 잠재와 함께 병렬로 미래 잠재를 생성하여 시퀀스 전체를 순차적으로 롤아웃하지 않는다. 이 접근은 반복적 한 단계 전이 예측에서 발생하는 오차 누적과 계산 비용을 줄여 장기 예측을 더 견고하게 만든다.
SIGReg 정규화(SIGReg): — SIGReg는 reconstruction-free 잠재 학습에서 표현 붕괴(collapse)를 방지하기 위한 가우시안 기반 정규화 기법이다. 배치 내 잠재의 통계적 특성을 제어하여 표현이 의미 있는 분산을 유지하도록 강제하고 예측 손실과 결합해 안정적인 임베딩 학습을 유도한다. 본 논문에서는 접두사 수준 예측 손실과 병합하여 잠재의 표현력을 보전한다.
Cross-Entropy Method(CEM): — CEM은 확률적 최적화 기법으로 샘플링 기반으로 후보 액션을 생성하고 하위 성능 샘플에 맞춰 분포를 갱신하는 반복 절차이다. 본 논문에서는 목표 잠재와의 거리로 후보를 평가하며 CEM이 많은 후보를 평가하므로 각 후보 평가 비용이 전체 계획 시간에 큰 영향을 미친다. Fast-LeWM은 후보 평가 당 동적 모듈 호출 횟수를 줄여 CEM 기반 계획 비용을 감소시켰다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

순수 C로 구현한 Gemma 4 E2B 추론 엔진 gemma4.c

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

순수 C로 구현한 Gemma 4 E2B 추론 엔진 gemma4.c