Technical Summary: OmniMem

Problem Statement

Audio-visual Large Language Models (av-LLMs) show significant promise for long-form video understanding but face a fundamental bottleneck during streaming inference: the linear growth of video tokens and Key-Value (KV) caches. As video length increases (e.g., hour-long videos), the memory requirements quickly exceed GPU limits. While chunked streaming processing has emerged as a practical solution, existing compression methods suffer from two primary limitations:

1. Uniform Token Treatment: Current methods often apply a unified token selection metric (e.g., cosine similarity or generic proxy query attention) to both audio and visual tokens. This fails to account for the severe token imbalance between modalities (visual tokens often outnumber audio tokens by 30–40x), leading to the systematic under-weighing or discarding of critical audio information.
2. Suboptimal Selection Criteria: Existing streaming methods typically rely on "generic proxy queries" (external to the model's computation graph) to estimate token importance. These scores may not accurately reflect the actual perturbation caused by evicting a token, potentially leading to the removal of informative states that do not score highly on a generic proxy but are crucial for the model's specific output.

Methodology

The authors propose OmniMem, a memory-efficient, perturbation-aware, and modality-aware streaming framework designed to address these limitations. The workflow processes video chunk-by-chunk, maintaining separate KV caches for audio and visual tokens.

1. Perturbation-Aware KV Cache Selection

Instead of using external proxy queries, OmniMem selects KV pairs to retain based on a criterion that minimizes the distortion to the attention output after eviction. The selection score ($\psi_k$) for a token $k$ combines two factors:

• Importance (Attention Score): The accumulated attention mass ($a_k$) directed toward the token from subsequent queries. High values indicate the token is frequently used by the next layer.
• Redundancy (Cosine Similarity): The cosine similarity ($s_k$) between the token's hidden state and its adjacent positions. High similarity indicates redundancy.

The selection score is calculated as:
$$\psi_k = a_k^\lambda (1 - s_k)$$
where $\lambda$ is a normalization factor. This multiplicative combination ensures that the system retains tokens that are both frequently used (high importance) and non-redundant (low similarity to neighbors). To compute similarity without excessive overhead, the method utilizes hidden states (or value vectors) rather than the full KV cache.

2. Audio-Visual Budget Allocation (AVBA)

To address modality imbalance, OmniMem introduces AVBA, a mechanism that allocates separate memory budgets for audio and visual KV caches rather than applying a uniform global budget.

• Compressibility Metrics: The system calculates a compressibility score ($C$) for each modality using normalized attention entropy and average cosine similarity.
• Dynamic Ratio: A prior visual-to-audio token ratio ($r$) is established (set to 5 in the paper). The actual budget weights ($w_v, w_a$) are then adjusted based on the relative compressibility scores of the audio and visual streams. This ensures that less compressible modalities (often audio, such as human speech) receive adequate memory allocation even if they constitute a small fraction of the total token count.
• Layer-wise Allocation: The total budget is distributed non-uniformly across Transformer layers based on layer-specific compressibility, with earlier layers generally receiving larger budgets.

3. Budget-Aware Fine-Tuning

While OmniMem functions as a training-free method, the authors explore supervised fine-tuning (SFT) with the allocated budget scheme. During this process, the model is trained to consolidate useful information into the retained memory tokens. To manage computational costs for long sequences, back-propagation is truncated to the last $L-l$ layers while ensuring gradient flow through all layers of the final chunk.

Key Contributions

The paper summarizes its contributions as follows:

• OmniMem Framework: A perturbation-aware, modality-aware KV cache selection mechanism that combines attention-based importance with cosine-similarity-based redundancy to minimize output distortion after eviction.
• AVBA Mechanism: The first audio-visual budget allocation mechanism for av-LLMs, which addresses modality representation discrepancies and prevents audio tokens from being marginalized due to token count imbalance.
• Performance Gains: Demonstrated consistent accuracy improvements over generic proxy query and cosine-similarity baselines, with further gains achieved through budget-aware fine-tuning.

Experimental Results

Experiments were conducted on VideoMME Long, LVBench, and LVOmniBench using video-SALMONN 2+ (4B and 8B parameters) and Qwen2.5-Omni (7B parameters).

• Training-Free Performance: Under identical memory budgets (e.g., 8K KV pairs per layer), OmniMem consistently outperformed strong baselines (InfiniPot-V, StreamMem, HERMES) by 2–4% in absolute accuracy.
  • On video-SALMONN 2+ (8B), OmniMem achieved improvements of 4.5% on VideoMME Long, 3.0% on LVBench, and 2.2% on LV-OmniBench compared to the best baseline (HERMES).
  • Improvements were particularly notable on the "Omni-source" partition of StreamingBench, where audio-visual balance is critical.
• Fine-Tuning Gains: Applying budget-aware fine-tuning provided an additional 1–2% absolute accuracy improvement. The most significant gains were observed on LVBench (up to 4.1% for the 4B model), suggesting that the model learns to utilize the fixed memory budget more effectively for longer videos.
• Efficiency: Runtime memory overhead was minimal (<1GB on H800 GPUs), primarily due to the retention of hidden states for similarity computation. First-token latency remained comparable to baselines (~0.2s).

Significance and Claims

The paper positions OmniMem as a necessary evolution for streaming av-LLMs, moving beyond heuristic token scoring to a principled approach that explicitly optimizes for the impact of token eviction on model outputs.

• Modality Awareness: The authors claim that treating audio and visual tokens uniformly is a fundamental flaw in current streaming methods. By introducing AVBA, they demonstrate that preserving the integrity of the less frequent but information-dense audio modality is critical for long-form understanding.
• Perturbation Awareness: The shift from generic proxy queries to perturbation-aware selection (minimizing output distortion) is presented as a more robust criterion for determining token importance.
• Practical Deployment: The method is designed for realistic deployment constraints, operating without the need for user prompts to determine budgets and functioning effectively within strict memory limits for hour-long videos.

The authors conclude that while OmniMem introduces a slight memory overhead for similarity computation, the trade-off yields consistent and significant accuracy gains, establishing a new standard for memory-constrained long-form audio-visual inference.