Improving Token-Based World Models with Parallel Observation Prediction

• URL: http://arxiv.org/abs/2402.05643v5
• Date: Wed, 29 May 2024 07:16:28 GMT
• Title: Improving Token-Based World Models with Parallel Observation Prediction
• Authors: Lior Cohen, Kaixin Wang, Bingyi Kang, Shie Mannor,
• Abstract summary: token-based world models (TBWMs) were recently proposed as sample-efficient methods. During imagination, the sequential token-by-token generation of next observations results in a severe bottleneck. We devise a novel Parallel Observation Prediction (POP) mechanism to resolve this bottleneck. POP augments a Retentive Network (RetNet) with a novel forward mode tailored to our reinforcement learning setting.
• Score: 55.41770427527391
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Motivated by the success of Transformers when applied to sequences of discrete symbols, token-based world models (TBWMs) were recently proposed as sample-efficient methods. In TBWMs, the world model consumes agent experience as a language-like sequence of tokens, where each observation constitutes a sub-sequence. However, during imagination, the sequential token-by-token generation of next observations results in a severe bottleneck, leading to long training times, poor GPU utilization, and limited representations. To resolve this bottleneck, we devise a novel Parallel Observation Prediction (POP) mechanism. POP augments a Retentive Network (RetNet) with a novel forward mode tailored to our reinforcement learning setting. We incorporate POP in a novel TBWM agent named REM (Retentive Environment Model), showcasing a 15.4x faster imagination compared to prior TBWMs. REM attains superhuman performance on 12 out of 26 games of the Atari 100K benchmark, while training in less than 12 hours. Our code is available at \url{https://github.com/leor-c/REM}.

Related papers

• OmniSAT: Compact Action Token, Faster Auto Regression [70.70037017501357]
  We introduce an Omni Swift Action Tokenizer, which learns a compact, transferable action representation.<n>The resulting discrete tokenization shortens the training sequence by 6.8$times$, and lowers the target entropy.
  arXiv  Detail & Related papers  (2025-10-08T03:55:24Z)
• Test-Time Training Done Right [61.8429380523577]
  Test-Time Training (TTT) models context by adapting part of the model's weights (referred to as fast weights) during inference.<n>Existing TTT methods struggled to show effectiveness in handling long-context data.<n>We develop Large Chunk Test-Time Training (LaCT) which improves hardware utilization by orders of magnitude.
  arXiv  Detail & Related papers  (2025-05-29T17:50:34Z)
• Reviving Any-Subset Autoregressive Models with Principled Parallel Sampling and Speculative Decoding [55.2480439325792]
  In arbitrary-order language models, it is an open question how to sample tokens in parallel from the correct joint distribution. We find that a different class of models, any-subset autoregressive models (AS-ARMs), holds the solution. We show that AS-ARMs achieve state-of-the-art performance among sub-200M parameter models on infilling benchmark tasks, and nearly match the performance of models 50X larger on code generation.
  arXiv  Detail & Related papers  (2025-04-29T06:33:13Z)
• FIRP: Faster LLM inference via future intermediate representation prediction [54.897493351694195]
  FIRP generates multiple tokens instead of one at each decoding step. We conduct extensive experiments, showing a speedup ratio of 1.9x-3x in several models and datasets.
  arXiv  Detail & Related papers  (2024-10-27T15:53:49Z)
• Fluid: Scaling Autoregressive Text-to-image Generative Models with Continuous Tokens [53.99177152562075]
  Scaling up autoregressive models in vision has not proven as beneficial as in large language models. We focus on two critical factors: whether models use discrete or continuous tokens, and whether tokens are generated in a random or fixed order using BERT- or GPT-like transformer architectures. Our results show that while all models scale effectively in terms of validation loss, their evaluation performance -- measured by FID, GenEval score, and visual quality -- follows different trends.
  arXiv  Detail & Related papers  (2024-10-17T17:59:59Z)
• Diffusion Auto-regressive Transformer for Effective Self-supervised Time Series Forecasting [47.58016750718323]
  We propose a novel generative self-supervised method called TimeDART. TimeDART captures both the global sequence dependence and local detail features within time series data. Our code is publicly available at https://github.com/Melmaphother/TimeDART.
  arXiv  Detail & Related papers  (2024-10-08T06:08:33Z)
• TokenUnify: Scalable Autoregressive Visual Pre-training with Mixture Token Prediction [61.295716741720284]
  TokenUnify is a novel pretraining method that integrates random token prediction, next-token prediction, and next-all token prediction. Cooperated with TokenUnify, we have assembled a large-scale electron microscopy (EM) image dataset with ultra-high resolution. This dataset includes over 120 million annotated voxels, making it the largest neuron segmentation dataset to date.
  arXiv  Detail & Related papers  (2024-05-27T05:45:51Z)
• Parallel Decoding via Hidden Transfer for Lossless Large Language Model Acceleration [54.897493351694195]
  We propose a novel parallel decoding approach, namely textithidden transfer, which decodes multiple successive tokens simultaneously in a single forward pass. In terms of acceleration metrics, we outperform all the single-model acceleration techniques, including Medusa and Self-Speculative decoding.
  arXiv  Detail & Related papers  (2024-04-18T09:17:06Z)
• Fast Sequence Generation with Multi-Agent Reinforcement Learning [40.75211414663022]
  Non-autoregressive decoding has been proposed in machine translation to speed up the inference time by generating all words in parallel. We propose a simple and efficient model for Non-Autoregressive sequence Generation (NAG) with a novel training paradigm: Counterfactuals-critical Multi-Agent Learning (CMAL) On MSCOCO image captioning benchmark, our NAG method achieves a performance comparable to state-of-the-art autoregressive models, while brings 13.9x decoding speedup.
  arXiv  Detail & Related papers  (2021-01-24T12:16:45Z)
• LAVA NAT: A Non-Autoregressive Translation Model with Look-Around Decoding and Vocabulary Attention [54.18121922040521]
  Non-autoregressive translation (NAT) models generate multiple tokens in one forward pass. These NAT models often suffer from the multimodality problem, generating duplicated tokens or missing tokens. We propose two novel methods to address this issue, the Look-Around (LA) strategy and the Vocabulary Attention (VA) mechanism.
  arXiv  Detail & Related papers  (2020-02-08T04:11:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.