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

🔼 Figure 2 illustrates the Multi-Axis Gradient Checkpointing (MA-GC) method. It uses a grid structure to checkpoint activations at regular intervals along both the layer (every ’l’ layers) and sequence (every ’s’ steps) dimensions. The grid allows for selective restoration of activations only when needed during backpropagation, reducing memory usage. Arrows show the flow of forward propagation, activation restoration, and gradient propagation. The table compares MA-GC to other checkpointing methods, showing its advantages in terms of memory usage and maximum sequence length achievable with 80GB of VRAM. Peak activation memory at sequence length 16384 is also compared.

read the caption

Figure 2: Overview of MA-GC grid structure. Checkpoints are stored every l𝑙litalic_l layers and s𝑠sitalic_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.

🔼 This figure shows the three main stages of training for the Video-Ma2mba model. Stage 1 focuses on cross-modal alignment using image-text and video-text pairs to align visual and textual features. Stage 1.5 emphasizes long-video knowledge learning using the SceneWalk dataset to train the model on longer video sequences with detailed descriptions, enhancing its temporal understanding. Stage 2 involves supervised fine-tuning on a diverse video question-answering dataset to refine the model’s ability to respond accurately to various queries about video content. The figure highlights the architecture of the model at each stage, demonstrating its progression through different training phases. It also indicates the input modalities (image, video, and text) and the output (caption or detailed descriptions) at each stage.

read the caption

Figure 3: The overall summarization for the training stages of Video-Ma2mba.

🔼 This figure presents a table showing the experimental results of Video-Ma²mba on the Video Multimodal Evaluation (Video-MME) benchmark. It compares the performance of Video-Ma²mba with different model sizes (0.7B, 1.8B, and 3.1B parameters) against other state-of-the-art models on short, medium, and long video clips. The table likely displays metrics such as accuracy or F1-score, allowing for a comparison of Video-Ma²mba’s performance across various video lengths and against competing methods.

read the caption

(a) Experimental results on Video-MME

🔼 This figure presents a table showing the experimental results of Video-Ma²mba and other models on the LongVideoBench benchmark. It compares performance across different video lengths (8-15s, 15-60s, 180-600s, 900-3600s) and across different model sizes. The results are presented as scores on the LongVideoBench dataset, indicating the models’ ability to understand long-form video content.

read the caption

(b) Experimental results on LongVideoBench

🔼 Figure 4 showcases qualitative examples from the Video Multimodal Evaluation benchmark (Video-MME). It presents three example video questions, their options, the correct answer, and the prediction made by the Video-Ma2mba-3.1B model. This visualization demonstrates the model’s capacity for accurate and nuanced understanding of long-form video content within Video-MME’s diverse question types and contextual scenarios. Each example includes a series of video frames depicting the relevant portion of the video clip.

read the caption

Figure 4: Qualitative examples on Video-MME [13] with Video-Ma2mba-3.1B.

🔼 Table 2 presents a performance comparison of Video-Ma2mba against various baseline models on three distinct video question answering benchmarks: ActivityNetQA, VideoChatGPT, and MVBench. For each benchmark, the table shows the model size (in billions of parameters) and the accuracy score achieved by each model. This allows for a direct assessment of Video-Ma2mba’s performance relative to other models, highlighting its capabilities in handling various video understanding tasks.

read the caption

Table 2: Benchmark results for ActivityNetQA, VideoChatGPT, and MVBench, comparing Video-Ma2mba and baselines.

🔼 This table presents a comparison of memory usage (in gigabytes) for different gradient checkpointing methods applied to the Mamba-2-2.7B model across various sequence lengths. The sequence lengths are powers of 2 (from 2¹⁰ to 2²¹). The methods compared are: no checkpointing (‘GC off’), checkpointing per layer (‘GC on’), checkpointing layers in groups of the square root of the total number of layers (‘Sqrt GC’), and a multi-axis gradient checkpointing method optimized for sequence length (‘MA-GC’). The memory overhead shown for each method and sequence length represents the peak memory consumption during both the forward and backward passes, using BF16 precision, excluding the model weights and gradients themselves.

read the caption

Table 3: Memory overhead (GB) for GC methods in Mamba-2-2.7B across sequence lengths (S=2n𝑆superscript2𝑛S=2^{n}italic_S = 2 start_POSTSUPERSCRIPT italic_n end_POSTSUPERSCRIPT). “GC off” indicates no checkpointing; “GC on” applies checkpointing per layer; “Sqrt GC” groups layers by L𝐿\sqrt{L}square-root start_ARG italic_L end_ARG; and “MA-GC” optimizes based on sequence length. Each cell show peak memory during activation and backpropagation (BF16 precision), excluding model weights and gradients.

🔼 This ablation study investigates the impact of frame size and the inclusion of Stage 1.5 (Long Video Knowledge Learning) on the performance of Video-Ma2mba-3.1B using the Video Multi-modal Evaluation (Video-MME) benchmark. It compares different frame sampling rates (8 frames, 16 frames, and 32 frames) and the presence or absence of Stage 1.5 training. Results show performance across short, medium, and long video lengths, providing insights into the effect of temporal context on model accuracy.

read the caption

Table 4: Ablation study on frame size and Stage 1.5 effects in Video-MME using Video-Ma2mba-3.1B.

🔼 This table details the hyperparameters used during the three training stages of the Video-Ma²mba model. It includes specifications for the input modalities (video and image in Stage 1, video in Stages 1.5 and 2), frame rates, input resolution, the number of trainable parameters in different model sizes, learning rates for the language model (LLM) and vision components, optimizer used (AdamW), global batch sizes for each stage, training epochs, warmup ratio, weight decay, gradient clipping, precision, deepspeed stages, and the gradient checkpointing method used.

read the caption

Table 6: Hyperparameters for Training Stages.

🔼 This table presents model-specific constants used in calculating memory usage for Video-Ma²mba, a model for long video understanding. The constants are crucial for the memory optimization formula presented in the paper (Equation 12). Each constant represents the memory consumption of different parts of the model’s architecture under BFloat16 precision (except for SSM states, which use Float32 and thus count as two BFloat16 elements). The table breaks down these constants for three different sizes of the Mamba-2 model (370M, 1.3B, and 2.7B parameters), reflecting how memory usage scales with model size. It shows the memory requirements for layer-wise checkpoints (CL-ckpt), sequence-wise checkpoints (CS-ckpt), grid cells (Cgrid), and SSM states (Cstate).

read the caption

Table 7: Model-specific constants for memory estimation under BFloat16 precision. Constants reflect relative element counts, where SSM states in Float32 are equivalent to two BFloat16 elements.

🔼 This table presents a computational analysis comparing the throughput (tokens processed per second) and per-token processing time (milliseconds per token) across various gradient checkpointing methods. The experiments were conducted using the Mamba-2-2.7B model on an NVIDIA A100 GPU with 80GB of memory. Sequence lengths are denoted using the notation 2ⁿ, where n represents the exponent, indicating the number of tokens processed. The results help to demonstrate the trade-offs between memory efficiency and processing speed for different gradient checkpointing techniques.

read the caption

Table 8: Computational analysis of throughput and per-token processing time among gradient checkpointing methods. Results are measured using the Mamba-2-2.7b model on an A100 80GB GPU. The notation @bold-@@bold_@ 𝟐𝐧superscript2𝐧\mathbf{2^{n}}bold_2 start_POSTSUPERSCRIPT bold_n end_POSTSUPERSCRIPT specifies the sequence length (in tokens) used for measurement.