0.13
0.10
LWM-Text-1M (Ours)
0.67
0.84
0.69
Figure 5: Multiple needles retrieval task with LWM-1M. NN is the number of facts in the context, and RR is the number of given facts model is asked to retrieve.
3.3.4 Evaluation on LOFT
Table 3: Evaluations on some benchmarks in the LOFT dataset.
Setting: 512K Context
LWM (512K)
GPT-4o (128K)
Claude 3 Opus (200K)
Quora
0.38\bf 0.38
0.230.23
0.370.37
NQ
0.37\bf 0.37
0.220.22
0.370.37
HotPotQA
0.72\bf 0.72
0.210.21
0.320.32
We further evaluate our model on a coverage of the LOFT (Lee et al., 2024) dataset collection, we provides a more natural set of benchmarks that examine capabilities for long-context models in the context of document retrieval, and RAG. The benchmark includes tasks such as duplication detection (Quora 11
1
https://quoradata.quora.com/First-Quora-Dataset-Release-Question-Pairs), document retrieval (HotpotQA (Yang et al., 2018)), and retrieval-based question-answering (NQ). Each dataset contains a corpus of 1000s of documents, and the model is asked to retrieve a set of document ids pertaining to its specific task (Quora, HotpotQA). For RAG (NQ dataset), the model is asked to answer the question using the given context. Table 3 shows evaluations results on 512K context length against various language model baselines.
Takeaway: Long context capability enables LWM to outperform state-of-the-art text models at multiple benchmarks. This demonstrates the effectiveness of our methods for enabling long context.
4 Stage II: Extending to Long-Context Vision-Language
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 4.
4.1 Architectural Modifications For Vision
We use the pretrained VQGAN (Esser et al., 2021) from aMUSEd (Patil et al., 2024) that tokenizes 256×256256\times 256 input images to 16×1616\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 mechanisms 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 (see Figure 3).
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 8 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 (Schuhmann et al., 2022) and COYO-700M (Byeon et al., 2022). 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 modal