World Model on Million-Length Video And Language With Blockwise RingAttention

## Abstract

Abstract Current language models fall short in understanding aspects of the world not easily described in words, and struggle with complex, long-form tasks. Video sequences offer valuable temporal information absent in language and static images, making them attractive for joint modeling with language. Such models could develop a understanding of both human textual knowledge and the physical world, enabling broader AI capabilities for assisting humans. However, learning from millions of tokens of video and language sequences poses challenges due to memory constraints, computational complexity, and limited datasets. To address these challenges, we curate a large dataset of diverse videos and books, utilize the Blockwise RingAttention technique to scalably train on long sequences, and gradually increase context size from 4K to 1M tokens. This paper makes the following contributions: (a) Largest context size neural network: We train one of the largest context size transformers on long video and language sequences, setting new benchmarks in difficult retrieval tasks and long video understanding. (b) Solutions for overcoming vision-language training challenges, including using masked sequence packing for mixing different sequence lengths, loss weighting to balance language and vision, and model-generated QA dataset for long sequence chat. (c) A highly-optimized implementation with RingAttention, Blockwise Transformers, masked sequence packing, and other key features for training on millions-length multimodal sequences. (d) Fully open-sourced a family of 7B parameter models capable of processing long text documents (LWM-Text, LWM-Text-Chat) and videos (LWM, LWM-Chat) of over 1M tokens.
This work paves the way for training on massive datasets of long video and language to develop understanding of both human knowledge and the multimodal world, and broader capabilities.

## 1 Introduction

Figure: Figure 1: LWM can answer questions over a 1 hour YouTube video. Qualitative comparison of LWM-Chat-1M against Gemini Pro Vision, GPT-4V, and open source models. Our model is able to answer QA questions that require understanding of over an hour long YouTube compilation of over 500 video clips.
Refer to caption: x1.png

Figure: Figure 2: LWM can retrieval facts across 1M context with high accuracy. Needle retrieval comparisons against Gemini Pro and GPT-4 for each respective max context length – 32K and 128K. Our model performs competitively while being able to extend to 8x longer context length. Note that in order to show fine-grained results, the x-axis is log-scale from 0-128K, and linear-scale from 128K-1M.
Refer to caption: extracted/5749045/figure/needle_comparison_baselines_log_linear_mix_rot.png

Current approaches on modeling the world are mostly restricted to short sequences of language sequences or short sequences of images and clips  [, , , , ].
resulting in models which lack understanding about parts of the world that are hard to represent in texts or short clips, and are unable to process complex long-form language and visual tasks. Temporal structure in video sequences provides helpful information that is missing from language or far less obvious in static images and short clips. Long language sequences encode information that short sequences cannot, crucial for various applications such as long document retrieval or coding. Long videos provide a rich context that short clips cannot grasp, showing how scenes connect, the development of events, and the cause and effect of actions within the temporal dimension of the video. This exposure to diverse long language and video scenarios also broadens the AI systems to generalize across various real-world situations. By jointly modeling both long videos and books, the model can develop an understanding of both the multimodal world and long sequences of texts and videos, leading to more advanced AI systems with a multimodal understanding, capable of assisting humans in a broader range of tasks.

To learn from video and language sequences, we need to train a model that is capable of processing more than millions of tokens per sequence and train it on a very large dataset. However, modeling millions of tokens is extremely difficult due to high memory cost, computational complexity, and lack of suitable datasets. Luckily, we have Blockwise RingAttention [, ], a technique for scaling up context size arbitrarily without approximations or overheads, allowing for scalable training on long sequences. We curated a large dataset of videos and languages from public book and video datasets, consisting of videos of diverse activities and long-form books. Considering the high compute cost, we grow context size from a smaller 4K context to a larger 1M context size gradually to reduce this cost, and this approach performs well in extending context effectively. Furthermore, we identify challenges associated with training on video and language: we discovered that training on a mixture of video, image, and text is crucial for optimal performance, due to images represent higher visual quality, videos offer sequential information, and text retains language understanding. To achieve this, we implemented an efficient masked sequence packing to effectively train with different sequence lengths, rather than standard sequence packing mechanism. Moreover, determining the right balance between image, video, and text training is crucial for cross modality understanding, and we suggest a ratio that proved effective. Furthermore, to address the lack of long-form chat datasets, we developed a model-generated question-answering (QA) approach by using a short-context model to generate a QA dataset from books. We found this to be crucial for chat abilities over long sequences.

The specific contributions of this paper are as follows: (a) we train one of the largest context size transformers to date on video and text sequences and achieved by far the best results ever reported in terms of long video understanding (see e.g., Figure 1) and long context fact retrieval (see e.g., Figure 2). (b) We discover a branch of challenges associated with training on video and text sequences, and propose solutions for them: loss weighting to balance language and vision, masked sequence packing to effectively train with different sequence lengths, and model-generated QA dataset for long sequence chat. (c) A highly-optimized, open-source implementation with RingAttention with Blockwise Transformer, masked sequence packing and other key features for millions-length multimodal training. (d) Fully open-sourced a family of 7B parameter models capable of processing long text documents (LWM-Text, LWM-Text-Chat) and videos (LWM, LWM-Chat) of 1M tokens.
Our work paves the way for training on massive datasets of long video and language, and is useful for future development of AI systems with an understanding of both human knowledge and the multimodal world, and broader capabilities.

Figure: Figure 3: This figure illustrates the multimodal training of a Large World Model. Stage 1, LLM Context Extension, focuses on expanding context size using the Books3 dataset, with context size growing from 32K to 1M. Stage 2, Vision-Language Training, focuses on training on visual and video contents of varying lengths. The pie chart details the allocation of 495B tokens across images, short and long videos, and 33B tokens of text data. The lower panel shows interactive capabilities in understanding and responding to queries about complex multimodal world.
Refer to caption: extracted/5749045/figure/data.png

## 2 Overview

We train a large autoregressive transformer model with a very large context window of up to one million tokens, building upon Llama2 7B []. To achieve this goal, we leverage several strategies: extending the context to 1M using books (Section 3), followed by joint training on long multimodal sequences, including text-image, text-video data, and books (Section 4).

Our training stages and datasets are shown in Figure 3 and the model architecture is shown in Figure 4.

## 3 Stage I: Learning Long-Context Language Models

This stage aims at first developing LWM-Text and LWM-Text-Chat, a set of long-context language models learned by training on progressively increasing sequence length data with RingAttention and Blockwise Transformer, and modifying positional encoding parameters to account for longer sequence lengths (see Section 3.1). The training steps for growing context size are shown in Section 3.2. In Section 3.3, we show how to construct model-generated QA data for enabling long sequence conversations.

### 3.1 Extending Context

Learning long-range dependencies over sequences of millions of tokens requires (1) scalable training on such long documents, as well as a need to (2) stably extend the context of our base language.

Scalable Training on Long Documents. Training on long documents becomes prohibitively expensive due to memory constraints imposed by the quadratic complexity of computing the attention weights. In order to address these computational constraints, we use the Blockwise RingAttention [, ] implementation that leverages block-wise transformer with sequence parallelism to theoretically extend to an infinite context, bounded only by the number of devices available. We further fuse Blockwise RingAttention with FlashAttention [, ] using Pallas [] to optimize performance compared with using XLA compiler. In general, given a large enough tokens per device, the communication cost during Blockwise Transformer and RingAttention fully overlap with computation, and does not add any extra overhead.

Progressive Training on Increasing Context Length. Although our implementation allows us to train on long documents of millions of tokens, it still remains costly since the quadratic computational complexity of attention remains, where gradient step time scales roughly linearly with context size (given a fixed number of tokens per batch). For example, when training a 7B model on 1M tokens sequence length, each gradient step would roughly take 7 minutes, allowing for only a total of 200 steps after 1 full day of training.

Therefore, we adopt a training approach inspired from [], where our model is trained on progressively longer sequence lengths, starting from 32K tokens and ending at 1M tokens in increasing powers of two. Intuitively, this allows the model to save compute by first learning shorter-range dependencies before moving onto longer sequences. By doing this, we are able to train on orders of magnitude more tokens compared to directly training on the maximum target sequence length. The progressive training of growing context sizes is shown in Figure 3.

Positional Extrapolation for Long Contexts. For extending positional embeddings on longer contexts, we adopt a simple, scaled up version of the approach explored in [], where the $\theta$ for RoPE [] is scaled up with context length. We generally found this approach to be a stable method to extend positional embeddings with context lengths due to its relatively simple nature of only needing to tune a single hyperparameter. We scale up the $\theta$ for RoPE along with context window sizes – the values are shown in Table 1.

Figure: Figure 4: LWM is a autoregressive transformer on sequences of millions-length tokens. Each frame in the video is tokenized with VQGAN into 256 tokens. These tokens are concatenated with text tokens and fed into transformers to predict the next token autoregressively. The input and output tokens’ order reflect the varied training data formats, including image-text, text-image, video, text-video, and purely text formats. The model is essentially trained in an any-to-any manner using multiple modalities. To differentiate between image and text tokens, and for decoding, we surround video and image tokens with the special delimiters <vision> and </vision>. We also include <eof> and <eov> vision tokens to mark the end of intermediate and final frames in images and videos. For simplicity, these delimiters are not shown.
Refer to caption: extracted/5749045/figure/model.png

### 3.2 Training Steps

We initialize from LLaMA-2 7B [] and progressively increase the effective context length of the model in 5 stages - 32K, 128K, 256K, 512K, and 1M. For each stage, we train on different filtered versions of the Books3 dataset from The Pile []. Table 1 details each information about each training stage, such as number of tokens, total time, and the Books3 dataset filtering constraints. Each successive run is initialized from the run of the prior sequence length.

### 3.3 Chat Fine-tuning for Long-Context Learning

Constructing QA data for Long Context Reasoning.  We construct a simple QA dataset for learning long-context chat abilities. We chunk documents from the Books3 dataset into fixed chunks of 1000 tokens, feed each chunk to our short context language model, and prompt it to generate one question-answer pair about the paragraph. Then, given a context length such as 32K, we construct a single 32K token example by concatenating adjacent chunks together, as well as appending relevant QA pairs towards the end of the sequence in chat form.

Training Details.  For chat fine-tuning, we train each model on a mix of UltraChat [] and our custom QA dataset, with roughly a 7:3 ratio. We found it crucial to pre-pack the UltraChat data to the training sequence length, and keep them separate from examples with our QA data, as UltraChat data generally has a much higher proportion of loss tokens (densely packed, short chat sequences), whereas our QA data has a much lower percentage of loss tokens per sequence (< 1%) since there is no loss on the long documents that are in the given context. Table 2 shows further training details for each run. Note that progressive training is not performed very each of the chat models, and instead are initialized from their respective pretrained models at the same context length.

Figure: Figure 5: Needle retrieval task. Our LWM-Text-Chat-1M have near perfect accuracy across different positions in 1M context window.
Refer to caption: extracted/5749045/figure/single_needle/single_needle_1M.png

Figure: Figure 6: Multiple needles retrieval task with LWM-1M. $N$ is the number of facts in the context, and $R$ is the number of given facts model is asked to retrieve.
Refer to caption: extracted/5749045/figure/multi_needle.png

**Table 1: LWM-Text Training Stages**
|  | 32K | 128K | 256K | 512K | 1M |
| --- | --- | --- | --- | --- | --- |
| Parameters | 7B | 7B | 7B | 7B | 7B |
| Sequence Length | $2^{15}$ | $2^{17}$ | $2^{18}$ | $2^{19}$ | $2^{20}$ |
| RoPE $\theta$ | 1M | 10M | 10M | 25M | 50M |
| Tokens per Batch | 4M | 4M | 4M | 4M | 4M |
| Total Tokens | 4.8B | 12B | 12B | 3B | 1.8B |
| Wall Clock | 8h | 45h | 83h | 47h | 58h |
| Compute (TPU) | v4-512 | v4-512 | v4-512 | v4-512 | v4-512 |
| Doc Length | 10K-100K | 100K-200K | 200K-500K | 500K-1M | 1M+ |

### 3.4 Language Evaluation Results

#### 3.4.1 Single Needle Retrieval

We evaluate on the popular Needle In A Haystack task [] - more specifically an easier to evaluate version []
that finds and retrieves random numbers assigned to randomized cities from the context. Figure 5 shows nearly perfect retrieval accuracy over the entire context of our 1M context model. In addition, Figure 2 shows that we can scale to far larger contexts compared to the current best available LLMs. Appendix A shows more single needle retrieval results for our other shorter context length models.

**Table 2: LWM-Text-Chat Training Details**
|  | 128K | 256K | 512K | 1M |
| --- | --- | --- | --- | --- |
| Parameters | 7B | 7B | 7B | 7B |
| Sequence Length | $2^{17}$ | $2^{18}$ | $2^{19}$ | $2^{20}$ |
| RoPE $\theta$ | 10M | 10M | 25M | 50M |
| Tokens per Batch | 4M | 4M | 4M | 4M |
| Total Tokens | 1.2B | 1.2B | 1.2B | 1.2B |
| Wall Clock | 6h | 10h | 20h | 40h |
| Compute (TPU) | v4-512 | v4-512 | v4-512 | v4-512 |

#### 3.4.2 Multi-Needle Retrieval

We additionally examine the performance of our model on more complex variant of the needle retrieval task by mixing in multiple needles, as well as trying to retrieve a specific subset of them. Figure 6 shows multi-needle retrieval results under different settings. Our model generalizes well when retrieving a single needle from multiple needles in context, with slight degradation when asked to retrieve more than one needle. Table 3 shows multi-needle comparisons between our model, Gemini Pro, and GPT-4, where our model is able to perform competitively or better than GPT-4 at retrieving one needle, or slightly lower performance when retrieving more than one needle. Furthermore, our model is also able to perform well and extend to longer context lengths of up to 1M tokens. However, we note that we see degradation in accuracy while increasing the difficulty of the needle retrieval task, suggesting that there is still more room to improve on the 1M context utilization of our model. We believe that our released model will provide a foundation for future work on developing longer context models, as well as encourage more challenging benchmarks that contain difficult long-range tasks that require higher levels of synthesis, rather than pure fact retrieval.

**Table 3: Multi-Needle Retrieval Accuracy Baseline Comparison**
| Context Length | Model | $N=2,R=2$ | $N=4,R=1$ | $N=4,R=2$ |
| --- | --- | --- | --- | --- |
| 32K | Gemini Pro | 0.34 | 0.44 | 0.6 |
| GPT-4 | 0.97 | 0.95 | 0.9 |  |
| LWM-Text-1M Ours) | 0.84 | 0.97 | 0.84 |  |
| 128K | Gemini Pro | - | - | - |
| GPT-4 | 0.92 | 0.8 | 0.82 |  |
| LWM-Text-1M (Ours) | 0.83 | 0.98 | 0.83 |  |
| 1M | Gemini Pro | - | - | - |
| GPT-4 | - | - | - |  |
| LWM-Text-1M (Ours) | 0.67 | 0.84 | 0.69 |  |

#### 3.4.3 Short Context Language Evaluation

Table 4 presents a comparative analysis between the Llama2-7B model with a 4K context and its context-expanded counterparts, ranging from 32K to 1M. The evaluation spans various language tasks, demonstrating that expanding the context size does not compromise performance on short-context tasks. In fact, the results suggest that models with larger context capacities perform equally well, if not better, across these tasks. This evidence indicates the absence of negative effects from context expansion, highlighting the models’ capability to adapt to different task requirements without losing efficiency in shorter contexts.

**Table 4: Evaluation of language tasks: Comparison between Llama2-7B (4K context) and context-expanded versions of LWM-Text: 32K to 1M. Results indicate that expanding context does not negatively impact performance on short-context tasks.**
|  |  | LWM-Text |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Task / Metric | Llama-2 7B | 32k | 128k | 256k | 512k | 1M |
| arc_challenge/acc | 0.4 | 0.43 | 0.45 | 0.44 | 0.44 | 0.43 |
| arc_challenge/acc_norm | 0.43 | 0.47 | 0.47 | 0.46 | 0.46 | 0.46 |
| hellaswag/acc | 0.57 | 0.57 | 0.57 | 0.57 | 0.56 | 0.57 |
| hellaswag/acc_norm | 0.77 | 0.76 | 0.76 | 0.76 | 0.75 | 0.75 |
| mmlu | 0.39 | 0.4 | 0.41 | 0.41 | 0.36 | 0.35 |
| openbookqa/acc | 0.32 | 0.33 | 0.31 | 0.32 | 0.33 | 0.3 |
| openbookqa/acc_norm | 0.44 | 0.44 | 0.44 | 0.43 | 0.41 | 0.41 |

#### 3.4.4 Chat Evaluation

We additionally evaluate the our model on MT-Bench [] to test its conversation ability. Table 6 shows the MT-Bench scores of for each of our models. Table 6 illustrates the relationship between the mix of chat and fact retrieval tasks and the performance on MT-Bench score and Needle Retrieval accuracy. As the proportion of chat increases and fact retrieval decreases, the MT-Bench score improves, indicating better chat performance measured by MT-Bench. Conversely, Needle Retrieval accuracy decreases, suggesting a trade-off where increasing chat interaction capabilities may reduce the system’s precision in retrieving specific information or ’needles’ from input context. Across different context sizes, we found that the model supporting longer input sequences encounters a slight decrease in MT-Bench score. We hypothesize that this is because we chose to train with fewer examples on longer sequence training and can be improved by simply training on more data. In addition, this trade-off may be resolved by acquiring higher quality long-context chat data that is closer to the chat distribution of the UltraChat dataset.

**Table 5: Results on MT-Bench across different context sizes. Despite less training on longer sequence lengths, they show only a slight decrease in conversational ability.**
| Model | MT-Bench |
| --- | --- |
| LWM-Text-Chat-128k | 4.62 |
| LWM-Text-Chat-256k | 5 |
| LWM-Text-Chat-512k | 4.83 |
| LWM-Text-Chat-1M | 4.19 |

## 4 Stage II: Learning Long-Context Vision-Language Models

Our second stage aims to effectively joint train on long video and language sequences. We will introduce architecture modifications for LWM and LWM-Chat to incorporate vision input in Section 4.1. Training on varying sequence lengths is discussed in Section 4.2. The evaluation results are shown in Section 4.3.
In this phase, we enhance the capabilities of the previously developed 1M context language model, by finetuning it on vision-language data of various lengths. The datasets used and the steps involved in the training process are illustrated in Figure 3.

### 4.1 Architectural Modifications For Vision

The model is illustrated in Figure 4. We use the pretrained VQGAN [] from aMUSEd [] that tokenizes $256\times 256$ input images to $16\times 16$ discrete tokens. Videos are tokenized by applying the VQGAN per-frame, and concatenating the codes together. In order to distinguish between modalities when generating, as well as knowing when to switch, we introduce mechanims to mark the end of text generation / beginning of vision generation, and vice-versa.
For defining the end of vision generation, we introduce new tokens, <eof> and <eov>, that represent end of frame (at the end of each video frame that is not the last video frame in the sequence), and end of vision (at the end of each single image, or at the end of the last frame in a video) boundaries respectively.
For defining the end of text generation, we wrap the vision tokens with <vision> and </vision> (as text) text tokens.
The model is trained with interleaved concatenations of vision and text tokens, and predicted autoregressively.

**Table 7: LWM and LWM-Chat Training Stages**
|  | 1K | 8K | Chat-32K | Chat-128K | Chat-1M |
| --- | --- | --- | --- | --- | --- |
| Parameters | 7B | 7B | 7B | 7B | 7B |
| Sequence Length | $2^{10}$ | $2^{13}$ | $2^{15}$ | $2^{17}$ | $2^{20}$ |
| RoPE $\theta$ | 50M | 50M | 50M | 50M | 50M |
| Tokens per Batch | 8M | 8M | 8M | 8M | 8M |
| Total Tokens | 363B | 107B | 10B | 3.5B | 0.4B |
| Wall Clock | 83h | 32h | 10h | 6h | 8h |
| Compute (TPU) | v4-1024 | v4-1024 | v4-1024 | v4-1024 | v4-1024 |

### 4.2 Training Steps

We initialize from our LWM-Text-1M text model, and perform a similar process of progressive training on a large amount of combined text-image and text-video data, with the exception that we do not additionally scale RoPE $\theta$, as it already supports up to 1M context.
Table 7 shows details for each training stage, where the model is initialized from the prior shorter sequence length stage. For each stage, we train on the following data:

- •
LWM-1K: We train on large set of text-image dataset comprising of a mix of LAION-2B-en [] and COYO-700M []. The datasets were filtered to only include images with at least 256 resolution – in total roughly 1B text-image pairs. During training, we concatenate the text-image pairs and randomly swap the order of the modalities to model both text-image generation, unconditional image generation, and image captioning. We pack text-image pairs to sequences of 1K tokens.
- •
LWM-8K: We train on a text-video dataset mix of WebVid10M [] and 3M InternVid10M [] examples. Similar to prior works [, , ], we jointly train on both images and video with a 50-50 ratio of each modality. We pack images to sequences of 8K tokens, and 30 frame videos at 4FPS. Similar to image training, we randomly swap the order of modalities for each text-video pair.
- •
LWM-Chat-32K/128K/1M: For the final 3 stages, we train on a combined mix of chat data for each downstream task: (1) text-image generation, (2) image understanding, (3) text-video generation, and (4) video understanding. We construct a simple version of text-image and text-video chat data by sampling random subsets of the pretraining data augmented with chat format. For image understanding, we using the image chat instruct data from ShareGPT4V []. Lastly, for the video understanding chat data, we use a combined mix of Valley-Instruct-73K [] and Video-ChatGPT-100K instruct data []. For all short context data (image generation, image understanding, video generation), we pack sequences to the training context length. During packing, we found it crucial to mask out the attention so that each text-vision pair only attends to itself, as well as re-weighting losses to make computation identical to training in a non-packed + padding training regime. For video understanding data, we uniformly sample a max number of frames to fit the training context length of the model if the video is too long. During training, We allocate 25% of each batch to each of the 4 downstream tasks.

For the first two stages of training (LWM-1K and LWM-8K), we additionally mix 16% of the batch to be pure text data from OpenLLaMA [], as we found it beneficial to preserve language capabilities while training on vision data.

Figure: ((a)) The assistant answers questions about a 3-minute video of animals.
Refer to caption: x2.png

### 4.3 Vision-Language Evaluation Results

#### 4.3.1 Long Video Understanding

Although vision-language model [, , ] can ingest long videos, this is commonly done by performing large temporal subsampling of video frames due to limited context length. For example, Video-LLaVA [] is restricted to uniformly sampling 8 frames from a video, no matter how long the original video may be. As such, models may lose more fine-grained temporal information that is important for accurately answering any questions about the video. In contrast, our model is trained on long sequences of 1M tokens, and as a result, can simultaneously attend thousands of frames of videos to retrieve fine-grained information over short time intervals. Figure 1 shows an example of our model correctly answering questions about a long, 1-hour YouTube compilation consisting of more than 500 individual clips. Our baseline methods, on the other hand, generally have difficulty answering the questions due to a limited number of frames. More results are shown in Figure 7 and Appendix B.

Although we demonstrate that our model can perform QA over complex, long-form videos, we note that there is still room to improve for better context utilization across all 1M tokens, as generated answers from our model may not always accurate, and the model still struggles with more complex questions that require higher-level understanding of the video. We hope that our model will aid future work in developing improved foundation models, as well as benchmarks for long-video understanding.

#### 4.3.2 Image Understanding and Short Video Understanding

Tables 8 and 9 show results on common benchmarks for image understanding and short video understanding. Figure 17 shows qualitative examples for image understanding. Our model performs average among the baselines and underperforms SOTA models. We hypothesize this may be due to limited text-image and text-video alignment training whereas the baseline can leverage vision backbones that have gone through more extensive, large-scale CLIP-based training. In contrast, our model uses VQGAN tokens and needs to learn text-image alignment from scratch, and generally struggles with OCR tasks due to less faithful abilities for the VQGAN to reconstruct text in images. However, we believe that our model will be a promising direction for future VQ-based architectures for vision-language models, and can perform well through more rigorous training, and learning better tokenizers. Appendix B shows more qualitative image understanding and Appendix B shows more qualitative video understanding examples.

Table: Table 8: Image Understanding Benchmarks

Table: Table 9: Video Understanding Benchmarks

#### 4.3.3 Image and Video Generation

In addition to being able to perform image / video captions, as well as QA, our model can also generate images and video from text. Figure 8 shows some examples of such capabilities. We use classifier-free guidance [] on the logits for autoregressive sampling in a manner similar to prior work [, ]. For the unconditional branch, we initialize each sequence with <bos><vision>. Appendix E and D show more image and video generation examples.

Figure: Figure 8: LWM can generate images and videos given text input. Examples of image and video generations. More examples are shown in Appendix E and Appendix D.
Refer to caption: x4.png

#### 4.3.4 Masked Sequence Packing Ablation.

As mentioned in Section 4.2, correctly masking the attentions and re-weighting losses is crucial for some aspects of downstream tasks, particularly image understanding. Table 10 shows a comparison of our model with and without packing corrections. Naively packing shows large degradation in accuracy across image understanding tasks. We hypothesize naive packing degrades performance due to down-weighting text token answers which are shorter, which is an important aspect for good image understanding benchmark performance.

**Table 10: Ablation study comparing standard and masked sequence packing mechanisms across three tasks. Masked sequence packing is crucial for performance.**
|  | VQAv2 | SQA | POPE |
| --- | --- | --- | --- |
| Naive Packing | 48.3 | 34.8 | 62.5 |
| LWM (Ours) | 55.8 | 47.7 | 75.2 |

## 5 Further Details

MFU.  We trained our models using TPUv4-1024, which is approximately equivalent to 450 A100s, with a batch size of 8M using FSDP [], Blockwise Transformer, and RingAttention for large contexts. Figure 9 shows the model FLOPS utilization (MFU) for each training stage. Blue color bars show language training and orange color bars show vision-language training. Our training achieves good MFUs even for very large context sizes.

Figure: Figure 9: High MFU training across sequence lengths. Model flops utilization (MFU) of each training stage for LWM-Text (top), and LWM / LWM-Chat (bottom)
Refer to caption: extracted/5749045/figure/mfu_lwm_text.png

Training Loss Curves.  Figures 10 and 11 show the training loss curves for each stage of training the language and vision-language models respectively.

Figure: Figure 10: Train loss curve for each training stage for LWM-Text models.
Refer to caption: extracted/5749045/figure/text_lwm_loss_curve.png

Figure: Figure 11: Train loss curve for each training stage for LWM and LWM-Chat models. Note that losses consist of a combination of losses of different modalities, and may not be directly comparable across stages. The sharp peak in the middle of 1K training is due to newly incporating EOF and EOV tokens into the vision codebook.
Refer to caption: extracted/5749045/figure/lwm_loss_curve.png

Training Hyperparameters. See Appendix F

Scaling Inference We additionally scale our inference code to support million-length sequences by implementing RingAttention for decoding. Inference for such long sequences requires a minimum of v4-128 with a TPU mesh sharding of 32 tensor parallelism, and 4 sequence parallelism (ring dimension). We perform inference in pure single precision, where additional improvements can be made through techniques in scalability such as quantization.

## 6 Related Works

Our work is related to efforts that extend language models’ context windows to allow more tokens [, , , inter alia], often by using novel extrapolation methods to expand pretrained positional encodings followed by finetuning the model on longer context data. Our model adopts a simple approach by gradually increasing $\theta$ in RoPE positional encodings along with the training context window sizes, and we found this method to be effective. There has been research into architectures that do not model pairwise interaction, such as sparse attention and sliding window [, ]. Parallelization along the sequence dimension is studied in [], but is not optimized for blockwise transformers and is incompatible with memory-efficient attention, which are crucial
to large context training. [] propose hybrid usage of tensor parallelism and sequence parallelism to better reduce memory cost on a single device.
Our work utilizes RingAttention and Blockwise Transformer [, ] to model exact pairwise interactions in very long sequences for optimal performance. Concurrent to our work, Gemini 1.5 [] achieves millions of tokens context size.
Further training performance improvement is also possible with the load balancing of skipping causal masked computation [, ].

Our work is also related to works on instruction tuning [, , , inter alia].
These studies focus on finetuning models using conversational data to enhance their capability across language tasks. Our approach seeks to advance models’ understanding of complex, long sequences of videos and languages. To achieve this, we extend the models’ context size by training on books and long videos, and finetuning on model-generated QA data to learn chat ability over long sequences.

Our work is also related to efforts combining vision and language [, , , , , , inter alia]. These efforts often use CLIP [] or BLIP [] to encode visual information into embeddings for inputting into language models. They have the potential advantages of leveraging CLIP’s cross-modal understanding for encoding textual information in images. However, they can only be trained to predict text given visual input but not vice versa, limiting their ability to learn from diverse formats of visual and language information. In cases where they do predict visual tokens [], a stronger diffusion decoder is generally required due to the relatively lossy nature of CLIP embeddings. Our work, on the other hand, is autoregressive "tokens in, tokens out", allowing us to flexibly model diverse forms of image-text, text-image, text-video, video-text, and purely video, image, or text formats.

## 7 Conclusion

In this paper, we address the challenge of learning models to understand the world better by combining language and video. We utilize RingAttention and Blockwise Transformers to scalably train on a massive dataset of long videos and books and gradually increase sequence length, from 32K to 1M tokens, to keep compute manageable. We develop masked sequence packing and loss weighting to effectively train on a diverse dataset of videos, images, and books. Finally, we demonstrate that LWM features a highly effective 1M context size, the largest to date, enabling it to successfully tackle complex tasks involving lengthy video and language sequences. We open source our optimized implementation of Blockwise RingAttention, masked sequence packing and other key features for training on millions-length sequences, as well as a 7B parameter model capable of processing over 1M multimodal tokens.
We hope this work paves the way for advancing AI models with a reliable reasoning and a grounded understanding of the world and broader capabilities.

##### Limitations and Future Work.

Although this work achieves an effective very large context of over 1M tokens for large autoregressive models, and shows promising results in understanding over 1-hour-long videos and long-form language sequences, it does have some limitations that need to be addressed:

- •
Better Video Tokenization. This work uses image tokenizer for videos – improving tokenization to be more compact could not only enhance video quality but also enable the processing of significantly more complex and longer videos, allowing more capable world model.
- •
More Modalities. Our work paves the road for learning from additional modality sources such as audio and other long sequences, allowing a deeper and more holistic understanding about the world.
- •
Better and More Video Data. Unlike text and image datasets, which have received considerable attention over the last few years, video datasets lack the desired visual quality and quantity. Future research can address this by sourcing YouTube videos.

## Acknowledgments

This project is supported in part by Office of Naval Research grant N00014-21-1-2769 and Multidisciplinary University Research Initiative (MURI) award by the Army Research Office (ARO) grant No. W911NF-23-1-0277. We thank Google TPU Research Cloud for granting us access to TPUs, and thank Google Cloud for granting us research credits for storage.

## Appendix A More Single-Needle Retrieval Results

Figure: Figure 12: Single needle retrieval accuracy for LWM-Text-Chat-256K
Refer to caption: extracted/5749045/figure/single_needle/single_needle_256k.png

Figure: Figure 13: Single needle retrieval accuracy for LWM-Text-Chat-512K
Refer to caption: extracted/5749045/figure/single_needle/single_needle_512k.png

## Appendix B More Video Understanding Examples

Figure: Figure 14: LWM demonstrates video understanding over 1 hour video.
Refer to caption: x5.png

Figure: Figure 15: LWM demonstrates video understanding over 1 hour video.
Refer to caption: x6.png

Figure: ((a)) The assistant answers questions about a 4-minute video of a man talks to the camera and shows a tree of apples.
Refer to caption: x7.png

## Appendix C More Image Understanding Examples

Figure: Figure 17: Question answering based on image input using LWM.
Refer to caption: x10.png

## Appendix D More Video Generation Examples

Figure: Figure 18: Video generation using LWM.
Refer to caption: x12.png

## Appendix E More Image Generation Examples

Figure: Figure 19: Image generation using LWM
Refer to caption: x14.png

## Appendix F Training Hyperparameters

**Table 11: LWM-Text Training Stages**
|  | 32K | 128K | 256K | 512K | 1M |
| --- | --- | --- | --- | --- | --- |
| Parameters | 7B | 7B | 7B | 7B | 7B |
| Initialize From | LLaMA-2 7B | Text-32K | Text-128K | Text-256K | Text-512K |
| Precision | float32 | float32 | float32 | float32 | float32 |
| Sequence Length | $2^{15}$ | $2^{17}$ | $2^{18}$ | $2^{19}$ | $2^{20}$ |
| RoPE $\theta$ | 1M | 10M | 10M | 25M | 50M |
| Tokens per Batch | 4M | 4M | 4M | 4M | 4M |
| Total Tokens | 4.8B | 12B | 12B | 3B | 1.8B |
| Total Steps | 1200 | 3000 | 3000 | 720 | 450 |
| LR Schedule | Constant | Constant | Constant | Constant | Constant |
| LR Warmup Steps | 100 | 200 | 200 | 50 | 25 |
| LR | $4\times 10^{-5}$ | $4\times 10^{-5}$ | $4\times 10^{-5}$ | $4\times 10^{-5}$ | $4\times 10^{-5}$ |
| Compute (TPU) | v4-512 | v4-512 | v4-512 | v4-512 | v4-512 |
| Mesh Sharding | 1,-1,4,1 | 1,-1,8,1 | 1,-1,16,1 | 1,-1,16,2 | 1,-1,16,4 |

**Table 12: LWM-Text-Chat Training Details**
|  | 128K | 256K | 512K | 1M |
| --- | --- | --- | --- | --- |
| Parameters | 7B | 7B | 7B | 7B |
| Initialize From | Text-128K | Text-256K | Text-512K | Text-1M |
| Precision | float32 | float32 | float32 | float32 |
| Sequence Length | $2^{17}$ | $2^{18}$ | $2^{19}$ | $2^{20}$ |
| RoPE $\theta$ | 10M | 10M | 25M | 50M |
| Tokens per Batch | 4M | 4M | 4M | 4M |
| Total Tokens | 1.2B | 1.2B | 1.2B | 1.2B |
| Total Steps | 300 | 300 | 300 | 300 |
| LR Schedule | Constant | Constant | Constant | Constant |
| LR Warmup Steps | 25 | 25 | 25 | 25 |
| LR | $4\times 10^{-5}$ | $4\times 10^{-5}$ | $4\times 10^{-5}$ | $4\times 10^{-5}$ |
| Compute (TPU) | v4-512 | v4-512 | v4-512 | v4-512 |
| Mesh Sharding | 1,-1,4,1 | 1,-1,8,1 | 1,-1,16,1 | 1,-1,16,2 |

**Table 13: LWM / LWM-Chat Training Stages**
|  | 1K | 8K | 32K | 128K | 1M |
| --- | --- | --- | --- | --- | --- |
| Parameters | 7B | 7B | 7B | 7B | 7B |
| Initialize From | Text-1M | 1K | 8K | 32K | 128K |
| Precision | float32 | float32 | float32 | float32 | float32 |
| Sequence Length | $2^{10}$ | $2^{13}$ | $2^{15}$ | $2^{17}$ | $2^{20}$ |
| RoPE $\theta$ | 50M | 50M | 50M | 50M | 50M |
| Tokens per Batch | 8M | 8M | 8M | 8M | 8M |
| Total Tokens | 363B | 107B | 10B | 3.5B | 0.4B |
| Total Steps | 45000 | 14000 | 1200 | 450 | 50 |
| LR Schedule | Cosine | Cosine | Cosine | Cosine | Cosine |
| LR Warmup Steps | 1000 | 500 | 100 | 50 | 5 |
| Max LR | $6\times 10^{-4}$ | $6\times 10^{-4}$ | $8\times 10^{-5}$ | $8\times 10^{-5}$ | $8\times 10^{-5}$ |
| Min LR | $6\times 10^{-5}$ | $6\times 10^{-5}$ | $8\times 10^{-5}$ | $8\times 10^{-5}$ | $8\times 10^{-5}$ |
| Compute (TPU) | v4-1024 | v4-1024 | v4-1024 | v4-1024 | v4-1024 |
| Mesh Sharding | 1,-1,1,1 | 1,-1,1,1 | 1.-1.4,1 | 1.-1.8,1 | 1,-1,16,4 |

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: 1. When training the models of LWM-1K and LWM-8K, there was an 16% mix of the pure text data added from OpenLLaMA in the batch, how was it determined and how sensitive is the model's performance to this parameter?
2. How did you achieve loss balancing across different tasks and contexts? Any specific strategy?
soundness: 3
strengths: 1. The paper addresses an important technical challenge in building foundational AI models, which is the long sequence understanding.
2. This work achieves multiple state-of-the-art results across different training stages
3. The open-sourced implementation and pre-trained models benefit and accelerate progress in the research community.
summary: This paper proposes a novel model architecture and training process that achieves significant advancements in long-context modeling for developing large language and vision-language models capable of processing sequences up to 1 million tokens in length. The authors propose a two-stage training process - in the first stage, they progressively train a base language model on increasingly longer text sequences, while in the second stage, the models incorporate vision capabilities through joint training on image and video data. The consequent Large World Model(LWM) family demonstrates impressive performance on a few challenging tasks such as long context retrieval and long video understanding.
weaknesses: The model's suboptimal performance on image and short video understanding tasks and the choice of visual encoding constrain the general applicability of this work.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: The core issue is the unclear relationship between the concept of a large world model and the multi-modal LLM with long context window presented in the paper. Additional suggestions include:
1. Provide quantitative evaluations for image and video generation, as mentioned in Section 4.3.3.
2. Explain the larger fluctuations in the purple curve compared to other curves in Fig. 9.
To improve the paper, the authors should:
● Clearly define and discuss the concept of a "large world model" and its relationship to the proposed multi-modal LLM.
● Conduct and present ablation studies to demonstrate the impact of individual components.
● Provide more in-depth analysis of the interaction between text and video modalities during training.
● Include comprehensive comparisons with other recent multimodal LLM works.
soundness: 2
strengths: 1. The paper provides extensive training details on scaling token sequences to 1M tokens, pushing the boundaries of context length for multimodal models.
2. Empirical results demonstrate strong performance on challenging tasks like long video understanding and retrieval across 1M token contexts.
3. The paper is well-structured and clearly written, with detailed explanations of the model architecture, training process, and evaluation results.
summary: This paper introduces the Large World Model (LWM), a LLaMA-based model capable of processing sequences of up to 1 million tokens, combining language and video. The paper details the training datasets, progressive training stages, and evaluation results. It incorporates techniques such as RingAttention, Blockwise Transformers, and others for efficient training on million-length multimodal sequences.
weaknesses: 1. The title "Large World Model" is misleading. The paper primarily explores training multi-modal (text and video) large language models with long context windows, without adequately defining or discussing the concept of a "large world model" or its relationship to MLLMs.
2. The paper fails to address fundamental questions about the necessity of long-context multi-modal LLMs for large world models and whether such models are sufficient for this purpose.
3. The paper largely describes training datasets and stages, lacking in-depth ablation studies and analysis. It relies heavily on previously proposed components like blockwise RingAttention, making it more akin to an experimental record than a research paper.
4. There is insufficient analysis of the potential synergies or conflicts between text and video in joint training. Do text and video in joint training mutually benefit each other, or do they have a negative influence on one another?
5. The paper lacks comprehensive comparisons with recent multimodal LLM works.
6. Despite claiming contributions in areas like masked sequence packing and loss weighting, the paper lacks ablation studies on those components

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 4
questions: - For the retrieval task in section 3.3.3, what is the exact evaluation benchmark used for the Multi Needle in a Haystack?
- What are the trainable parameters for each step in the stage 2?
soundness: 4
strengths: - The empirically feasible training process that integrates understanding and generation capabilities into one extensive model using VQ is quite useful, and the evaluations encompassing both linguistic and multimodal tasks are comprehensive.
- The settings and hyperparameters for training and evaluation are presented in a meticulous manner.
summary: This paper introduces a training process that gradually extends LLMs from a limited context length of 4K tokens to an extremely long context of millions of tokens. The two-stage training process — where the first stage focuses on long-context expansion and the second stage addresses image and video understanding and generation — enables the models to support multimodal understanding and generation. The empirically feasible training process, which integrates understanding and generation capabilities into one extensive model using VQ, is quite useful. The evaluations, which involves both linguistic and multimodal tasks, are comprehensive.
weaknesses: - One major contribution of this paper is extending the context length of existing LLMs to extremely long through progressive training, which is a valuable practical implementation. However, aside from the fundamental support of RingAttention, the technical improvements in this study are marginal.
- Building models that support both multimodal understanding and generation is interesting and meaningful. However, the simple incorporation of an off-the-shelf VQ model into the framework does not yield promising results. I believe that more in-depth analysis and improvements are needed in this direction.
- There is a lack of quantitative evaluation for multimodal lengthy content understanding. For example, can we evaluate understanding of image details by increasing resolutions with lengthy vision tokens? Additionally, there are long video understanding benchmarks that could be used to assess the proposed LWM-1M.
- Some crucial details in the main content of the paper, such as the evaluation data and the trainable parameters in each stage, are missing.
- Compared to existing and up-to-date VLMs (e.g., LLaVA-NeXT, VILA, Qwen-VL, CogVLM), the performance of LWM on image understanding and video understanding benchmarks is relatively low.

---

confidence: 5
contribution: 4
initial_rating: 6
presentation: 3
questions: 1. The three benchmarks used to report video performance is not long. It is better to show accuracy on long-video benchmarks such as VideoMME[1] and MLVU[2].
2. In Table 3, how about comparing with existing open-sourced LLM supporting context length of 128K?
3. The paper is the first to scale context length up to 1M. However, there is no baseline comparison in 1M context length. How about using linear scaling technique during inference stage to construct a naive baseline? E.g., scale 8x from a 128k model.

[1] Video-MME: The First-Ever Comprehensive Evaluation Benchmark of Multi-modal LLMs in Video Analysis. arXiv preprint arXiv:2405.21075.
[2] MLVU: A Comprehensive Benchmark for Multi-Task Long Video Understanding. arXiv preprint arXiv:2406.04264.
soundness: 4
strengths: + This paper is the first to scale sequence models up to 1 million tokens, pushing the boundaries of long-context processing in both language and multimodal (text-video) tasks.
+ The model use VQGAN tokenization to enable image/video understanding as well as generation.
+ In multi-needle retrieval tasks with a 128K context, the model achieves results comparable to or better than GPT-4, demonstrating its effectiveness in long-context retrieval tasks.
summary: The paper scales the LLM context length up to 1M tokens by utilizing Blockwise RingAttention. They progressively increase the context length from 32K to 1M using book-length text data and video-language data. It uses discrete token representation of VQGAN to facilitate image/video generation tasks. It demonstrates competitive performance in long-context retrieval tasks. It also reports comprehensive experimental results in both image/video understanding and generation tasks.
weaknesses: + Lack of hour-long video benchmark evaluation, please refer to questions.
+ No comparison with existing open-source LLMs in context length of 128K, for example Llama 3.1
+ The usage of VQGAN decrease the image understanding ability as Table 4 shows.
+ There is no quantitative result on image generation tasks.

Meta-review: ### Summary:
This paper introduces LWM (Large World Model), a novel approach for processing million-length sequences in both language and video modalities. The key contribution is demonstrating scalable training of models that can handle extremely long context windows (up to 1M tokens) while maintaining competitive performance, particularly in video understanding tasks.

### Strengths:
1. Novel technical solution for scaling context length
> "The paper addresses an important technical challenge in building foundational AI models, which is the long sequence understanding." - gZJr

2. Comprehensive empirical validation
> "Empirical results demonstrate strong performance on challenging tasks like long video understanding and retrieval across 1M token contexts." - gZJr

3. Open-source contribution with practical impact
> "The open-sourced implementation and pre-trained models benefit and accelerate progress in the research community." - gZJr

### Weaknesses:
1. Limited image/video understanding capabilities
> "The model's suboptimal performance on image and short video understanding tasks and the choice of visual encoding constrain the general applicability of this work." - gZJr

2. Insufficient comparative analysis 
> "The paper lacks comprehensive comparisons with recent multimodal LLM works." - cWVz

3. Questions about fundamental necessity
> "The paper fails to address fundamental questions about the necessity of long-context multi-modal LLMs for large world models and whether such models are sufficient for this purpose." - QidU

### Justification:

Despite the noted weaknesses, I recommend acceptance because:
1. The paper makes significant technical contributions in scaling context length - a fundamental challenge in AI systems - with demonstrated improvements over existing approaches
2. The comprehensive empirical validation and open-source implementation provide substantial value to the research community
3. The authors have thoroughly addressed most reviewer concerns during the discussion period with additional experiments and clarifications
4. While image/video understanding performance could be improved, this limitation is acknowledged and does not diminish the core contribution of enabling million-token context processing

The reviewers are have carefully evaluated the technical details. The authors have provided thorough responses with additional experimental validation.

Decision: Accept (Poster)