LinMU: Multimodal Understanding Made Linear

TL;DR

This work tackles the quadratic $O(N^2)$ complexity of self-attention in Vision-Language Models by introducing LinMU, which replaces all attention layers with the linear-cost M-MATE blocks that combine a global Flex-MA branch and a local Local-Swin branch to achieve $O(N)$ per layer. It couples this architectural shift with a three-stage distillation framework that reuses pretrained attention weights and progressively trains the dual branches, aided by LoRA fine-tuning. Empirically, LinMU matches or closely approaches teacher performance on diverse image- and video-based benchmarks while delivering substantial efficiency gains (e.g., up to $2.7\times$ faster TTFT and $9\times$ higher token throughput for minute-length videos). This demonstrates that high-quality multimodal reasoning can be achieved with linear-complexity networks, enabling long-context, high-resolution perception on more modest hardware and expanding the practical deployment of VLMs.

Abstract

Modern Vision-Language Models (VLMs) achieve impressive performance but are limited by the quadratic complexity of self-attention, which prevents their deployment on edge devices and makes their understanding of high-resolution images and long-context videos prohibitively expensive. To address this challenge, we introduce LinMU (Linear-complexity Multimodal Understanding), a VLM design that achieves linear complexity without using any quadratic-complexity modules while maintaining the performance of global-attention-based VLMs. LinMU replaces every self-attention layer in the VLM with the M-MATE block: a dual-branch module that combines a bidirectional state-space model for global context (Flex-MA branch) with localized Swin-style window attention (Local-Swin branch) for adjacent correlations. To transform a pre-trained VLM into the LinMU architecture, we propose a three-stage distillation framework that (i) initializes both branches with self-attention weights and trains the Flex-MA branch alone, (ii) unfreezes the Local-Swin branch and fine-tunes it jointly with the Flex-MA branch, and (iii) unfreezes the remaining blocks and fine-tunes them using LoRA adapters, while regressing on hidden states and token-level logits of the frozen VLM teacher. On MMMU, TextVQA, LongVideoBench, Video-MME, and other benchmarks, LinMU matches the performance of teacher models, yet reduces Time-To-First-Token (TTFT) by up to 2.7$\times$ and improves token throughput by up to 9.0$\times$ on minute-length videos. Ablations confirm the importance of each distillation stage and the necessity of the two branches of the M-MATE block. The proposed framework demonstrates that state-of-the-art multimodal reasoning can be achieved without quadratic attention, thus opening up avenues for long-context VLMs that can deal with high-resolution images and long videos.

Figures (7)

• Figure 1: By replacing quadratic-complexity attention layers with our proposed M-MATE layers, LinMU achieves linear complexity without sacrificing model performance.
• Figure 2: The overview of LinMU, in which the Vision Encoder is untouched but all the attention layers in the Language Model (LM) Decoder are replaced with our proposed M-MATE blocks. Residual/gate connections are not shown in this figure for brevity. The M-MATE block contains two parallel branches: (1) an Flex-MA branch built around a bidirectional Mamba2 augmented with reverse direction token mask and RMS token rearrangement, and (2) a Local-Swin branch that implements a local 3D Swin Attention with a fixed window size. The Flex-MA branch provides efficient long-range sequence mixing with linear complexity, while the Local-Swin branch focuses on adjacent spatial/temporal correlations to preserve local consistency. Both branch outputs are summed to produce the M-MATE block output that replaces the original self-attention output. Other parts of the LM Decoder (e.g., multi-layer perceptron, layer norms, etc.) remain unchanged.
• Figure 3: The receptive field comparison between native Mamba2 (suitable for text-only tasks), bidirectional Mamba2 (suitable for vision-only tasks), and our proposed masked flexible bidirectional Mamba2 (suitable for multimodal tasks).
• Figure 4: Window division, shifting, and re-combination in the local 3D Swin attention. After window shifting, tokens around the image/video edge are combined to form new attention windows, ensuring the total number of windows remains constant.
• Figure 5: Our proposed three-stage distillation pipeline to replace attention layers in pre-trained VLMs with linear-complexity M-MATE blocks.