predominantly limited to processing short sequences, whether in the form of language, images, or video clips (Brown et al., 2020, Touvron et al., 2023a, Touvron et al., 2023b, OpenAI, 2023, Brooks et al., 2024, Team et al., 2023). As a result, these models fall short when tasked with understanding complex, long-form language and visual contexts.
However, training models to process sequences that exceed millions of tokens is a significant challenge due to the high memory and computational costs, as well as the lack of long-context data. In this work, we address these challenges by leveraging Blockwise RingAttention (Liu et al., 2024, Liu and Abbeel, 2023), a technique that scales context size without approximations or overheads, enabling efficient training on long sequences. We curate an extensive dataset of long-form videos and books from public sources, covering a wide variety of activities and narrative structures. To address the scarcity of long-form conversational datasets, we developed a model-based question-answering technique, where a short-context model generates training data from books, significantly enhancing the model’s chat capabilities over long sequences. To mitigate computational costs, we gradually extended context size from an initial 4K tokens to 1M tokens, achieving a cost-effective and scalable approach for long-context modeling.
Following this, we further train our long-context language model to incorporate visual modalities, such as image and video. Contrary to existing popular vision-language models (Liu et al., 2023a, OpenAI, 2023, Chen et al., 2023a), we opt to additionally optimize next-token prediction losses for image and video (generation) with a VQGAN (Esser et al., 2021) encoder. We encountered various challenges training on mixed modalities (video, image, text). To balance their unique characteristics - sequential information, visual detail, and linguistic content - we implement an efficient masked sequence packing strategy, as well as introduce careful loss balancing to retain short context accuracy. This approach handles varying sequence lengths more effectively than standard methods. We also optimized the ratio of image, video, and text inputs in each batch, proposing an empirically effective balance for cross-modality learning.
Since our model aims to model both textual and visual projections of the world through a large context window, drawing inspiration from prior work on world models (Brooks et al., 2024, Ha and Schmidhuber, 2018), we name our work as Large World Model (LWM).
Our contributions are threefold: (a) we train one of the largest context size transformers to date on long text documents and videos and achieved competitive results on long video understanding and long context fact retrieval. (b) We discover a range of challenges associated with training on long sequences and propose solutions for them: masked sequence packing to effectively train with different sequence lengths and synthetic model-generating question-answering for effective attention. (c) We provide an open-source and optimized implementation for training with millions of tokens in context, as well as a family of Llama-based 1M context models capable of processing long documents (LWM-Text, LWM-Text-Chat) and videos (LWM, LWM-Chat) of 1M tokens.
2 Method overview
We train a large autoregressive transformer model with a large context window of up to one million tokens, building upon Llama2 7B (Touvron et al., 2023b). To achieve this goal, we implement a two-stage training strategy. In Stage I (Section 3), we extend the context to 1M tokens using book-length texts. This is followed by Stage II (Section 4), where we conduct joint training on diverse long multimodal sequences, incorporating text-image data, text-video data, and book-length texts. Our model architecture is the standard autoregressive transformer design, as illustrated in Figure 3. For a comprehensive overview of our training stages and the datasets emp