Source-linked AI summary
World Model on Million-Length Video And Language With Blockwise RingAttention
Hao Liu, Wilson Yan, Matei Zaharia, Pieter Abbeel
TL;DR
Long-context models must overcome short-sequence limits, high training costs, and scarce long-context data. The paper develops 1M-token language and video-language models through curated long-form data, progressive context extension, and multimodal training, achieving competitive retrieval and long-video understanding results while releasing open-source 7B models.
Problem
Existing sequence models are limited in processing long language and visual contexts, while million-token training is constrained by computational costs and scarce long-context data.
Method
The paper combines long-form data curation, progressive context extension, Blockwise RingAttention, and joint training over text, image, and video sequences.
Results
The resulting 7B models support over 1M tokens and achieve competitive long-context fact retrieval and long-video understanding results.
Takeaways & Limitations
The work provides an open-source family of long-context models and an implementation for scalable training on million-token text and video sequences.
Takeaways & Limitations
Because the models are much smaller than current 100B+ language models, extrapolation to larger scales should be done cautiously.
Abstract
from arXiv · showhide
Enabling long-context understanding remains a key challenge in scaling existing sequence models -- a crucial component in developing generally intelligent models that can process and operate over long temporal horizons that potentially consist of millions of tokens. In this paper, we aim to address these challenges by providing a comprehensive exploration of the full development process for producing 1M context language models and video-language models, setting new benchmarks in language retrieval and new capabilities in long video understanding. We detail our long context data curation process, progressive context extension from 4K to 1M tokens, and present an efficient open-source implementation for scalable training on long sequences. Additionally, we open-source a family of 7B parameter models capable of processing long text documents and videos exceeding 1M tokens.
1 INTRODUCTION
The paper targets long-context understanding by addressing the costs and data scarcity of million-token training, then develops multimodal models and training methods for long text and video.
- Long-context understanding is difficult because existing models primarily process short language, image, and video sequences.These limits hinder understanding of complex, long-form language and visual contexts.
- Million-token training faces high memory and computational costs alongside a shortage of long-context data.The work addresses these challenges with Blockwise RingAttention and curated long-form data.
- The training pipeline progressively extends context from 4K to 1M tokens and uses model-generated question-answering data for long-sequence conversations.The data includes long-form videos and books from public sources.
- The multimodal model jointly handles text, images, and video while optimizing next-token prediction for visual generation and using masked packing and loss balancing.These choices address differing sequence lengths and modality characteristics while retaining short-context accuracy.
- The paper contributes 1M-context models, an open-source scalable training implementation, and competitive long-context retrieval and video-understanding results.The released family includes Llama-based 7B models for long documents and videos.
2 METHOD OVERVIEW
LWM is a Llama2 7B-based autoregressive transformer trained on multimodal token sequences, supporting text, images, video, and text-video inputs and outputs.
- The model uses a two-stage strategy: first extending context to 1M tokens with book-length texts, then jointly training on long multimodal sequences.Stage II combines text-image data, text-video data, and book-length texts.
- LWM is an autoregressive transformer that combines visual and textual tokens to predict the next token.Video frames are tokenized with VQGAN into 256 tokens each, while text uses BPE tokenization.
- The architecture supports text, image, video, and text-video input-output modalities using special delimiters to distinguish modalities and sequence endings.The delimiters include <vision>, </vision>, <eof>, and <eov>.
3 STAGE I: LEARNING LONG-CONTEXT LANGUAGE MODELS
Stage I develops long-context language models through progressive context extension, efficient attention, curated book data, and model-generated question-answering. The resulting models retain short-context performance and support retrieval over contexts up to 1M tokens.
- 3 STAGE I: LEARNING LONG-CONTEXT LANGUAGE MODELS: Blockwise RingAttention addresses the memory and computational costs of million-token training by scaling context without approximations or overheads.The approach is described as theoretically allowing context bounded by available devices and is fused with FlashAttention for efficiency.
- 3.1 PROGRESSIVE TRAINING TOWARDS LONG CONTEXT: The training strategy progressively extends context from 32K to 1M tokens across five stages, initializing each stage from the previous one.The stages are 32K, 128K, 256K, 512K, and 1M tokens, using filtered Books3 data.
- 3.2 MODEL-GENERATED QUESTION-ANSWERING FOR EFFECTIVE CONTEXT: Model-generated question-answering data supports long-context chat by generating question-answer pairs from 1,000-token book chunks and placing them after concatenated long sequences.Chat fine-tuning mixes the custom question-answering data with UltraChat at approximately a 7:3 ratio, while keeping the datasets separate during packing.
- 3.3.1 SHORT CONTEXT TASKS: 32K-to-1M context expansion does not significantly degrade performance on shorter-context language tasks compared with Llama-2 7B using a 4K context window.The comparison covers multiple language tasks and reports equal or better performance for larger-context models.
- 3.3.3 RETRIEVAL TASK: MULTIPLE INFORMATION: The 1M-context model achieves nearly perfect single-needle retrieval across its context and remains competitive with GPT-4 on multi-needle retrieval.Retrieval accuracy degrades as the number of requested needles increases, indicating remaining room to improve 1M-context utilization.
- 3.3.4 EVALUATION ON LOFT: Long-context capability enables LWM to outperform state-of-the-art text models on multiple LOFT benchmarks.LOFT evaluates document retrieval, duplication detection, and retrieval-based question answering across corpora containing thousands of documents.
4 STAGE II: EXTENDING TO LONG-CONTEXT VISION-LANGUAGE
Stage II extends the 1M-context language model to jointly process text, images, and videos through progressive multimodal training. The resulting LWM models support long-video understanding and unified image/video understanding and generation, while retaining important short-context capabilities.
- Multimodal representation: The multimodal pipeline tokenizes 256 × 256 images into 16 × 16 discrete VQGAN tokens and concatenates per-frame codes for videos.Special markers distinguish text and vision generation and indicate modality transitions.
- Training data: The training data covers text-image and text-video pretraining, mixed image/video sequences, and chat tasks spanning generation and understanding.The first two stages also mix 16% pure text data to preserve language capabilities during vision training.
- Evaluation: LWM processes thousands of video frames in 1M-token contexts, enabling fine-grained retrieval from long videos and strong performance on long-video understanding.On a one-hour YouTube compilation containing more than 500 clips, LWM correctly answers questions that baseline methods generally find difficult.
- Evaluation: On standard image and short-video benchmarks, LWM performs comparably to baselines but remains below state-of-the-art models.The paper attributes the gap partly to discrete-token information loss and the absence of large-scale CLIP-based pretraining.
- Capabilities: A unified any-to-any architecture supports image and video captioning, question answering, and text-conditioned image and video generation.Figure 7 provides examples of both static image and dynamic video generation.
5 RELATED WORKS
The paper relates its approach to three lines of work: context-window extension, instruction tuning, and vision-language integration. It differs by combining progressive long-context training with multimodal long-sequence modeling.
- Context extension: Prior context-extension methods commonly extrapolate pretrained positional encodings and finetune models on longer-context data.The paper instead describes incrementally increasing the RoPE θ parameter alongside training expansion.
- Instruction tuning: The work extends instruction-tuning ideas from conversational language tasks to long-sequence understanding in video and language.It uses comprehensive datasets including books and long videos to develop these capabilities.
- Vision-language integration: Unlike many vision-language systems that use continuous visual embeddings, the paper uses discrete visual tokens and learns text-image alignment from scratch.The paper identifies CLIP-based continuous representations as a contrasting approach.
6 CONCLUSION
The paper concludes that its 7B LWM family demonstrates million-token text and video processing through progressive training and scalable infrastructure. It also identifies tokenization and model scale as important boundaries on generalization.
- Limitations: The study’s tokenization remains limited because it uses a vanilla image tokenizer and frame-by-frame video tokenization.Future work could model temporal redundancy and add continuous embeddings.
- Limitations: The findings may not directly transfer to much larger models because the study uses models far smaller than current 100B-plus systems.The paper recommends caution when extrapolating to larger scales because different scaling behaviors could emerge.
- Scalable training: The implementation uses TPUv4-1024, FSDP, and BlockwiseRingAttention while maintaining good MFU at very large context sizes.Figure 8 reports MFU for language and vision-language training stages.
- Scalable inference: Inference for million-length sequences requires at least a v4-128 TPU configuration with tensor and sequence parallelism.The inference implementation uses RingAttention for decoding and pure single precision.
- Training lessons: Correct masking and loss re-weighting are crucial for image understanding, because naive packing substantially degrades accuracy.The authors hypothesize that naive packing down-weights shorter text answers.
B.2 MIXING SYNTHETIC AND CHAT DATA
The mixing of chat and fact-retrieval data creates a trade-off between conversational ability and retrieval accuracy. Increasing chat data improves MT-Bench performance while reducing the emphasis on fact retrieval.
- Data mixture trade-off: Increasing the proportion of chat data improves MT-Bench scores as the proportion of fact-retrieval data decreases.MT-Bench measures conversational performance in this comparison.
- Evaluation: The experiment evaluates how chat and fact-retrieval mixtures jointly affect MT-Bench score and Needle Retrieval accuracy.The relationship is summarized in Table 11 across different task mixtures.
C MORE SINGLE-NEEDLE RETRIEVAL RESULTS
The LWM-Text-Chat models maintain near-perfect single-needle retrieval across positions and context lengths up to 1M tokens.
- Near-perfect retrieval accuracy is maintained across positions and context lengths for the 1M-token LWM-Text-Chat-1M model.
- Near-perfect retrieval performance is maintained across positions and context lengths for the 256K LWM-Text-Chat model.
- Near-perfect retrieval performance is maintained across positions and context lengths for the 512K LWM-Text-Chat model.
D DETAILS OF THE MULTI-NEEDLE EVALUATIONS
The multi-needle evaluation extends single-number retrieval by inserting multiple city-associated needles and asking for a random subset, with correctness checked by string matching.
- The evaluation extends Needle in a Haystack by inserting multiple city-associated magic numbers into a long context.
- N denotes the number of inserted needles, while R denotes the random subset of magic numbers requested for retrieval.
- Retrieved numbers are checked against the requested city-associated values using string matching.
F MORE VIDEO UNDERSTANDING EXAMPLES
The examples show LWM answering questions about long and varied videos, including details that other systems miss, while also supporting multimodal generation.
- LWM correctly identifies three lemons in a person’s car during a 1-hour video, while GPT-4V, Gemini Pro Vision, and Video-LLAVA fail.
- LWM correctly identifies a cat standing on a piano in a 1-hour video, unlike the compared systems’ incorrect or incomplete answers.
- LWM answers questions across videos about activities, briefings, and humorous events, including a man changing a car tire.
- Qualitative evaluation samples YouTube videos across egocentric footage, how-to videos, interviews, and animations at 1 FPS.
- LWM also generates image and video sequences from text alongside image and video captioning and question-answering.
J TRAINING HYPERPARAMETERS
The listed tables provide training-stage and training-detail information for LWM-Text and LWM models, including chat variants.
- Table 12 is titled “LWM-Text Training Stages.”
- Table 13 is titled “LWM-Text-Chat Training Details.”
- Table 14 is titled “LWM / LWM-Chat Training Stages.”