Rápido LeWorldModel

Resumo

Arquiteturas Preditivas de Incorporação Conjunta (JEPAs), incluindo o recente LeWorldModel (LeWM), tornaram-se uma base promissora para modelos de mundo visual livres de reconstrução. No entanto, para o planejamento visual, o LeWM avalia sequências de ações candidatas aplicando repetidamente um modelo de transição latente local de um passo. Esse desdobramento autorregressivo torna o planejamento computacionalmente caro e expõe a trajetória prevista a erros latentes acumulados à medida que o horizonte cresce. Propomos o Fast LeWorldModel (Fast-LeWM), um modelo de mundo latente rápido que substitui o desdobramento local repetido pela predição de prefixos de ação. Dado o estado latente atual e uma sequência de ações candidata, o Fast-LeWM codifica seus prefixos e prevê os estados latentes futuros alcançados após executar esses prefixos em paralelo. Ao tornar os prefixos de ação a unidade básica de predição, o Fast-LeWM modela diretamente os efeitos das ações acumulados em diferentes graus ao longo de múltiplos horizontes. Essa supervisão em nível de prefixo força o modelo a aprender como os estados evoluem continuamente sob diferentes prefixos de ação, em vez de apenas ajustar transições de estado de um passo. Durante o planejamento, o preditor pode usar o último token do prefixo da sequência de ações codificada para avaliar o estado latente futuro correspondente, sem precisar percorrer explicitamente cada estado imaginado intermediário. Em múltiplas tarefas, o Fast-LeWM melhora a taxa média de sucesso em relação ao LeWM, ao mesmo tempo que reduz substancialmente o tempo de planejamento, alcançando uma perda latente em malha aberta menor, cujo crescimento se torna significativamente mais lento à medida que o horizonte de desdobramento aumenta.

English

Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the last prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.