Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Look Every Frame All at Once: Video-Ma$^2$mba for Efficient Long-form Video Understanding with Multi-Axis Gradient Checkpointing

Hosu Lee, Junho Kim, Hyunjun Kim, Yong Man Ro

TL;DR

This work tackles the memory bottlenecks of long-form video understanding in transformer-based LMMs by introducing Video-Ma$^2$mba, which replaces attention with State Space Models in the Mamba-2 framework to achieve linear time and space complexity with respect to sequence length. A novel bi-axis gradient checkpointing strategy, MA-GC, further reduces memory by storing activations along both layer and sequence axes, enabling training on sequences up to $O(S)$ memory and up to $0.8$M context tokens on a single GPU. The model is trained in three stages—cross-modal alignment, long video knowledge learning, and supervised fine-tuning—using dense long-video data (SceneWalk) and diverse video QA datasets, and evaluated on long-video benchmarks (Video-MME, LongVideoBench) and general video benchmarks (ActivityNetQA, VideoChatGPT, MVBench). Results show Video-Ma$^2$mba achieves competitive or superior performance compared to larger models while using significantly less memory, demonstrating the practicality and effectiveness of linear-time long-video modeling for real-world deployment.

Abstract

With the growing scale and complexity of video data, efficiently processing long video sequences poses significant challenges due to the quadratic increase in memory and computational demands associated with existing transformer-based Large Multi-modal Models (LMMs). To address these issues, we introduce Video-Ma$^2$mba, a novel architecture that incorporates State Space Models (SSMs) within the Mamba-2 framework, replacing the attention mechanisms. This allows the LMMs to scale linearly in terms of time and memory requirements, making it feasible to handle long-duration video content. Furthermore, we enhance the memory efficiency introducing the Multi-Axis Gradient Checkpointing (MA-GC) method, which strategically manages memory by retaining only essential activations across multiple computational axes. Our approach significantly reduces the memory footprint compared to standard gradient checkpointing. Empirical analyses show that Video-Ma$^2$mba can process extensive video sequences-equivalent to millions of tokens or over two hours of continuous sequences at 1 FPS-on a single GPU. By maintaining a detailed capture of temporal dynamics, our model improves the accuracy and relevance of responses in long video understanding tasks, demonstrating substantial advantages over existing frameworks.

Look Every Frame All at Once: Video-Ma$^2$mba for Efficient Long-form Video Understanding with Multi-Axis Gradient Checkpointing

TL;DR

This work tackles the memory bottlenecks of long-form video understanding in transformer-based LMMs by introducing Video-Mamba, which replaces attention with State Space Models in the Mamba-2 framework to achieve linear time and space complexity with respect to sequence length. A novel bi-axis gradient checkpointing strategy, MA-GC, further reduces memory by storing activations along both layer and sequence axes, enabling training on sequences up to memory and up to M context tokens on a single GPU. The model is trained in three stages—cross-modal alignment, long video knowledge learning, and supervised fine-tuning—using dense long-video data (SceneWalk) and diverse video QA datasets, and evaluated on long-video benchmarks (Video-MME, LongVideoBench) and general video benchmarks (ActivityNetQA, VideoChatGPT, MVBench). Results show Video-Mamba achieves competitive or superior performance compared to larger models while using significantly less memory, demonstrating the practicality and effectiveness of linear-time long-video modeling for real-world deployment.

Abstract

With the growing scale and complexity of video data, efficiently processing long video sequences poses significant challenges due to the quadratic increase in memory and computational demands associated with existing transformer-based Large Multi-modal Models (LMMs). To address these issues, we introduce Video-Mamba, a novel architecture that incorporates State Space Models (SSMs) within the Mamba-2 framework, replacing the attention mechanisms. This allows the LMMs to scale linearly in terms of time and memory requirements, making it feasible to handle long-duration video content. Furthermore, we enhance the memory efficiency introducing the Multi-Axis Gradient Checkpointing (MA-GC) method, which strategically manages memory by retaining only essential activations across multiple computational axes. Our approach significantly reduces the memory footprint compared to standard gradient checkpointing. Empirical analyses show that Video-Mamba can process extensive video sequences-equivalent to millions of tokens or over two hours of continuous sequences at 1 FPS-on a single GPU. By maintaining a detailed capture of temporal dynamics, our model improves the accuracy and relevance of responses in long video understanding tasks, demonstrating substantial advantages over existing frameworks.

Paper Structure

This paper contains 42 sections, 2 theorems, 13 equations, 6 figures, 10 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Context Extension Methods
  4. Long Video Understanding with LMMs
  5. Gradient Checkpointing Techniques
  6. Video-Ma$^2$mba
  7. Overview.
  8. Preliminary: Mamba-2 and Simplification
  9. Multi-Axis Gradient Checkpointing
  10. Analysis of Upper Bound of Memory Reduction
  11. Model Architecture
  12. Training Stages
  13. Stage 1: Cross-modal Alignment.
  14. Stage 1.5: Long Video Knowledge Learning.
  15. Stage 2: Supervised Fine-Tuning.
  16. ...and 27 more sections

Key Result

Theorem 1

If $f$ is continuous on a closed interval $[a, b]$, then $f$ achieves both maximum and minimum values on $[a, b]$.

Figures (6)

  • Figure 1: Memory usage comparison across sequence lengths for Mamba-2-2.7B with different checkpointing methods, demonstrating the memory-saving capability of Multi-Axis Gradient Checkpointing (MA-GC).
  • Figure 2: Overview of MA-GC grid structure. Checkpoints are stored every $l$ layers and $s$ steps. The blue, red, and green arrows indicate forward propagation, activation restoration, and gradient propagation, respectively. This grid design optimizes memory by selectively restoring activations as needed. The below table shows comparison of checkpointing usage, maximum sequence length on 80GB VRAM, and peak activation memory in BFloat16 at sequence length 16384.
  • Figure 3: The overall summarization for the training stages of Video-Ma$^2$mba.
  • Figure 4: Qualitative examples on Video-MME fu2024video with Video-Ma$^2$mba-3.1B.
  • Figure 5: Qualitative examples from the Generative Subset of VideoChatGPT maaz2023video with Video-Ma$^2$mba-3.1B.
  • ...and 1 more figures

Theorems & Definitions (2)

  • Theorem 1: Extreme Value Theorem
  • Theorem 2: Fermat's Theorem