Next-Latent Prediction Transformer
The official research codebase for a paper showing that predicting the next latent state, instead of the next token, trains transformers to learn more compact world models.
This repository is the official implementation accompanying the paper 'Next-Latent Prediction Transformers Learn Compact World Models,' including the Next-Latent Prediction (NextLat) algorithm and baseline algorithms such as next-token prediction (GPT-style), multi-token prediction, joint multi-token prediction and the Belief State Transformer, each implemented in its own model file for direct comparison through their compute_loss functions.
Training scripts support single-GPU and distributed runs, with YAML configuration files reproducing the paper's experiments and a sweep parameter for grid search across seeds and hyperparameters. Evaluation covers Manhattan taxi rides, TinyStories and FineWeb-Edu pretraining benchmarks. Requires PyTorch 2.6 or later, and 2.9 or later for the Muon optimizer.