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

Authors: Hao Liu, Wilson Yan, Matei Zaharia, Pieter Abbeel
Primary area: generative models
Keywords: long context, world model, attention

Abstract:
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.

Decision: accept (poster)
Reviewer scores: 8, 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1729118477529

thread_type:
official_review

<thread_object>
creator: reviewer_cwvz

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730560946487

thread_type:
official_review

<thread_object>
creator: reviewer_zobj

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596352060

thread_type:
official_review

<thread_object>
creator: reviewer_qidu

<content>
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.

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.

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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730712311636

thread_type:
official_review

<thread_object>
creator: reviewer_gzjr

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733258483027

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion Period Summary by Authors

comment:
We thank the reviewers for taking the time to review our paper and considering our rebuttals. Here, we provide a summary of the reviews, rebuttals, and responses over the discussion period.

&nbsp;

Reviewers gave positive evaluations on our paper, saying that it:
* Addresses important technical challenges in learning general purpose foundation models by being the **first to scale to 1M context** **(gZJr, QiDU, cWVz)**
* Provides a public codebase, as well as detailed experimental and training details that will **greatly benefit the open-source community** **(gZJr, QiDU, zoBJ)**
* Provides a **simple, unified model design** that supports both understanding and generation over language, image, and video **(zoBJ, cWVz)**
* Demonstrates **state-of-the-art** results in **long-context** tasks **(gZJr, QiDU, cWVz)**

&nbsp;

During the rebuttal period, we addressed reviewers’ concerns with further discussion and experimental results:

* **Long VideoMME benchmark (QiDU, zoBJ, cWVz)**: LWM-1M achieves scores of 63.7% and 60.8% on medium and long videos respectively, __outperforming other similar sized models by substantial margins__ (e.g., +12.3% over Long-LLaVA on medium-length videos). These results establish it as **the highest performance in the 7B model category**, approaching the capabilities of 70B+ models and competitive with proprietary solutions like GPT-4V and Gemini-1.5-Pro.
* **Multi-Needle evals on state-of-the-art models (cWVz)**: Although baseline models perform well at lower contexts, extrapolating to 1M context proves difficult, showing on average 0% - 32% accuracy. In contrast, our model **significantly outperforms these recent state-of-the-art methods at 1M context**, achieving 67% - 84% accuracy across different settings.

&nbsp; 

Best Regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
### 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.