PaperlayerView exact revision on arXiv

Source-linked AI summary

OmniMem: Scalable and Adaptive Memory Retrieval for Long Video Generation

Lin Zhao, Yushu Wu, Yifan Gong, Yanzhi Wang, Pu Zhao

arXiv:2605.30519v1cs.CV

TL;DR

Long-video autoregressive generation must manage a growing historical KV cache, while truncation or compression can limit access to query-relevant details. OmniMem performs explicit sparse retrieval and addresses selection bias and memory-side Union Explosion. On VBench-Long, it improves Dynamic Degree by 52.3% over the strongest baseline while preserving visual consistency with 1.7% additional memory overhead.

  • Problem

    Long-video autoregressive generation grows the KV cache, while truncation and compression can limit access to query-relevant historical details, motivating explicit full-range retrieval.

    Source: §1 · p. 2

  • Method

    OmniMem performs explicit sparse historical-KV retrieval, combining Adaptive Window Exclusion, Query-Shared KV Selection, and Per-Head Scattered KV Access.

    Source: §1 · p. 2

  • Results

    52.3%: OmniMem improves Dynamic Degree over the strongest baseline on VBench-Long while preserving visual consistency with 1.7% additional memory overhead.

    Source: §1 · p. 3

  • Takeaways & Limitations

    OmniMem provides explicit long-range retrieval for chunk-based generation while improving long-video quality, temporal consistency, and motion dynamics with practical latency and memory footprint.

    Source: §6 · p. 10

    Evaluation uses one open-sourced DiT backbone, Wan2.1-T2V-1.3B; validation on larger backbones and other pipelines remains future work, and hardware-dependent gains may vary.

    Source: Limitation and Broader Impact · p. 13

Abstract

from arXiv · show

Autoregressive (AR) video generation extends videos by producing latent chunks sequentially, but scaling to long videos requires repeated access to a growing historical KV cache. Existing methods reduce this cost by truncating the KV cache or compressing it into implicit memory, but both lose explicit access to query-relevant historical details. We propose OmniMem, an explicit full-range memory retrieval framework that performs sparse KV retrieval over the historical cache. To make this practical for chunk-based AR video generation, OmniMem addresses two issues: (i) local bias in sparse KV selection and (ii) Union Explosion in memory access. Adaptive Window Exclusion removes local-window blocks from the selection candidates when sufficient long-range history is available, preserving the sparse budget for informative long-range retrieval. Query-Shared KV Selection reduces cross-query diversity, while Per-Head Scattered KV Access avoids expanding head-specific selections into a large selected KV buffer. This allows each attention head to retrieve non-contiguous KV blocks according to its own selection pattern. Experiments on long-video generation show that OmniMem improves Dynamic Degree by 52.3% and preserves strong consistency over strong baselines, while maintaining comparable memory usage.

1 Introduction

Long-video autoregressive generation strains a growing KV cache, while truncation and compression can sacrifice long-range consistency or fine detail. OmniMem retrieves explicit historical context while addressing local selection bias and memory-side Union Explosion.

1 Introduction

  • As generated sequences lengthen, KV-cache growth increases memory use and per-step attention computation; discarding distant states can cause identity drift.Fixed anchor banks can stabilize identity but may produce repetitive motion, while compressed history may lose fine-grained details needed for long-horizon consistency.Source: §1 · p. 2
  • Explicit full-range sparse retrieval is motivated by local bias in block selection and Union Explosion from divergent query and head selections.Thousands of chunk queries can each select few blocks yet create a large union and current-step KV footprint.Source: §1 · p. 2
  • OmniMem combines Adaptive Window Exclusion, Query-Shared KV Selection, and Per-Head Scattered KV Access for scalable long-range retrieval.The components filter near-window blocks, reduce cross-query selection diversity, and avoid materializing a large merged cache across heads.Source: §1 · p. 2
  • 52.3%: OmniMem improves Dynamic Degree over the strongest baseline on VBench-Long while preserving visual consistency with 1.7% additional memory overhead.Source: §1 · p. 3

2 Related Work

Prior long-video work addresses consistency through training, denoising, cache compression, and sparse memory strategies. OmniMem explores explicit retrieval over the historical KV cache as a distinct approach.

Long Video Generation

  • Long-horizon consistency methods include rollout-aware training, rolling-window denoising with attention sinks, and long-context flow matching with sliding-window distribution matching.Source: §2 · p. 3

KV Cache Compression and Memory Retrieval in Long Video Generation

  • KV-cache management approaches include quantization, learned or SSM-based history compression, cache sparsification, on-demand activation, and head-wise or salience-based policies.Source: §2 · p. 3
  • OmniMem differs by dynamically managing the historical cache through explicit retrieval during chunk-based long-video generation.Source: §2 · p. 3

3 OmniMem: Explicit Memory Retrieval and Access

OmniMem makes long-range memory retrieval explicit in chunk-based AR video generation, where sparse selection faces local bias and Union Explosion. It combines adaptive candidate exclusion with shared query selections and head-specific scattered access.

3.1 Framework

  • The framework combines compression, selection, and sliding-window attention over pooled history, retrieved full-resolution blocks, and recent cache, respectively.Selection uses compression-attention scores to identify top-K historical blocks; coefficients fuse the three branch outputs.Source: §3.1 · p. 4

3 OmniMem: Explicit Memory Retrieval and Access

  • OmniMem addresses local bias and Union Explosion with Adaptive Window Exclusion, Query-Shared KV Selection, and Per-Head Scattered KV Access.Together, the methods target which historical blocks are retrieved and how selected blocks are organized for access.Source: §3 · p. 3, §3.1 · p. 4, §3.3 · p. 6

3.2 Addressing Local Bias with Adaptive Window Exclusion

  • Nearby blocks dominate top-K selection, and the sliding-window span can exceed the sparse budget, leaving long-range history unselected.Without exclusion, K must exceed ⌈NwS/B⌉ to guarantee retrieval of even one block outside the local window.Source: §3.2 · pp. 4–5
  • Adaptive Window Exclusion removes sliding-window-covered blocks from selection candidates only when at least K blocks remain outside the window.When outside-window history is insufficient, selection instead uses the full historical cache.Source: §3.2 · p. 5

3.3 Mitigating Union Explosion with Query-Shared KV Selection and Per-Head KV Access

  • Across a chunk, query and head preferences can diverge, making the selected-block union grow rapidly despite each query's small top-K selection.Different queries can produce a union as large as S·K, while different heads may favor distinct historical regions.Source: Figure 3 · p. 4, §3.3 · pp. 5–6
  • Query-Shared KV Selection assigns one top-K selection to adjacent query groups, reducing the worst-case query-dimension union from S·K to R·K.Selections remain independent across groups, preserving inter-group retrieval diversity.Source: §3.3 · p. 6
  • Per-Head Scattered KV Access reads each head's selected non-contiguous blocks from native storage, avoiding a padded cross-head selected-KV buffer.This keeps the representation proportional to the sum of per-head selections rather than the cross-head union.Source: §3.3 · p. 6

4 Experiments

In multi-prompt long-video generation, OmniMem reports stronger visual quality and consistency, with CLIP scores remaining stable over time. Ablations and scalability results support its use of complementary memory branches and sparse long-range access.

4 Experiments

  • The implementation uses Wan2.1-T2V-1.3B, while the design ablations use a 20-second video setting.The model is fine-tuned with ODE initialization on VidProM prompts and then trained with Self-Forcing and long-video tuning.Source: §4 · p. 7, §4.2 · p. 8

4.1 Main Results

  • CLIP scores drop only 2.44 from 0–10s to 50–60s, while OmniMem achieves superior visual-quality and consistency performance in multi-prompt 60-second generation.The evaluation uses CLIP scores per 10-second segment to assess prompt adherence.Source: §4.1 · p. 8, Table 2 · p. 8

4.2 Ablation Study

  • Combining SWA, CMP, and SLC achieves the best performance, with coarse global memory and fine-grained retrieval complementing the local window.Using SWA alone gives the weakest results; adding either CMP or SLC improves results.Source: §4.2 · p. 8, Table 3 · p. 8
  • Excluding more recent chunks improves performance over no exclusion, and removing the full 3-chunk SWA window yields the best overall result.The 3-chunk exclusion is the default setting.Source: §4.2 · p. 8, Table 4 · p. 8
  • Increasing query-group size from 1 to 15 leaves metrics nearly unchanged, while sizes 30 or 60 degrade them.Grouping adjacent queries reduces cross-query selection diversity and redundant selection computation; a single shared selection cannot cover diverse content in larger groups.Source: §4.2 · p. 8, Table 5 · p. 9

4.3 Scalability of Long-Range Memory Access

  • At 717 frames, OmniMem uses 3.3× less memory and runs 2.7× faster than full-history attention, matching Naive Sparse latency.Figure 5 reports nearly constant memory usage for OmniMem, unlike Naive Sparse, whose selected-block union grows close to full-history attention.Source: Figure 5 · p. 9, §4.3 · p. 9

4.2 Ablation Study

  • Sharing KV selection across all 12 heads significantly degrades all metrics, and even groups of 3 heads trail per-head selection.The results indicate that different attention heads attend to different historical regions.Source: §4.2 · pp. 8–9, Table 6 · p. 9

5 Discussion and Further Analysis

Further analysis finds that dynamic retrieval can access useful historical context beyond fixed sink regions, and that full-resolution selection can use historical memory more effectively. The paper also reports head-level specialization and potential zero-shot generalization to longer videos.

Dynamic Retrieval over Fixed Sink Tokens

  • Adding fixed sink anchors to OmniMem brings no additional gain in this setting, where dynamic retrieval can select sink-like regions when useful.Many heads use historical context beyond the sink region.Source: §5 · p. 9

Selection Attention Uses Historical Tokens More Effectively

  • SWA+SLC achieves better performance with fewer historical tokens than SWA+Sink under the same recent-window context.SWA+SLC selects full-resolution historical blocks per query and head rather than reserving fixed anchors.Source: §5 · p. 9

Further Analysis

  • Different heads exhibit distinct gating preferences across the three memory branches, and explicit retrieval shows potential for zero-shot generalization to longer videos.Source: §5 · p. 9

6 Conclusion

OmniMem retrieves selected full-resolution historical KV blocks for chunk-based AR video generation and addresses local-biased selection and Union Explosion with three mechanisms. Experiments on 60-second generation report gains in quality, temporal consistency, and motion dynamics over SOTA baselines while maintaining practical latency and memory footprint.

6 Conclusion

  • OmniMem retrieves selected full-resolution historical KV blocks and addresses local bias and Union Explosion with Adaptive Window Exclusion, Query-Shared KV Selection, and Per-Head Scattered KV Access.Source: §6 · pp. 9–10
  • On 60-second generation, OmniMem improves long-video quality, temporal consistency, and motion dynamics over SOTA baselines while maintaining practical inference latency and memory footprint.Source: §6 · p. 10

Limitation and Broader Impact

OmniMem’s evaluation is limited to one DiT backbone, with validation on larger backbones and other pipelines left for future work. The preliminaries describe chunk-wise generation and the growing historical KV cache that makes long-video generation difficult to scale.

Limitation and Broader Impact

  • Evaluation uses only Wan2.1-T2V-1.3B; larger backbones and other video-generation pipelines remain unvalidated, and kernel efficiency may vary across hardware.The authors also note that the Triton selection-attention kernel is tuned for current GPUs.Source: Limitation and Broader Impact · p. 13
  • OmniMem may reduce compute costs for creative tools, simulation, and world-model research, while higher-quality long-video generation also carries disinformation and impersonation risks.The paper suggests content filtering, watermarking, and dataset controls for practical deployments.Source: Limitation and Broader Impact · p. 13

A Preliminaries

  • Chunk-based AR generation produces a sequence of latent chunks, with each chunk’s tokens processed jointly at each denoising step.Unlike token-wise autoregressive generation in LLMs, one chunk introduces thousands of query tokens simultaneously.Source: A Preliminaries · p. 13
  • The historical KV cache grows linearly with generated chunks, increasing attention computation and memory use.At each step, attention covers the current chunk and all previously generated chunks; prior keys and values are cached and reused.Source: A Preliminaries · p. 13

B Implementation Details

The implementation combines a three-stage training recipe with separate compressed and full-resolution KV caches. Selection accesses per-head, non-contiguous blocks through pointers, and attention visualizations show varied head preferences, including non-local history.

B Implementation Details

  • Training follows ODE initialization, Self-Forcing, and long-video tuning on Wan2.1-T2V-1.3B.The stages use VidProM prompts, then the CausVid and LongLive recipes; training runs on 4 nodes with 8× NVIDIA H100 GPUs.Source: B Implementation Details · p. 13
  • The memory module pools KV states, selects top-4 historical blocks per head per query group, and combines compression, selection, and sliding-window branches with learned sigmoid gates.Its pooling block is 15 × 2, the query-group size is 15, and the sliding window covers the 3 most recent chunks.Source: B Implementation Details · pp. 13–14

C Shared Chunk-level KV cache and Per-Head Selection Forward

D The Visualization of Attention Map

E Short Video Generation.

On short-video VBench, OmniMem matches Self-Forcing’s total score within 0.02 points and achieves the highest Quality score among the methods. Ablations report better metrics with finer pooling at a matched token budget and diminishing returns as Top-K increases.

E Short Video Generation.

  • 84.29 Total VBench score for OmniMem matches Self-Forcing (84.31) within 0.02 points, while OmniMem achieves the highest Quality score across the methods.The comparison uses the standard VBench prompt suite for 5-second video generation.Source: E Short Video Generation · p. 16, Table E1 · p. 17

F.1 Impact of Pooling Size

  • With the same 240-token budget, 15×2 pooling outperforms 15×4 on each metric; frame-level pooling performs worst despite accessing 6.5× more tokens.The tested configurations were restricted to those with kernel utilization above 90%.Source: §F.1 · p. 17, Table F1 · p. 17

F.2 Impact of Top-K Selection

  • Top-K metrics improve monotonically from K = 1 to 12, with the best performance at K = 12 and diminishing returns beyond K = 4.Increasing K from 1 to 4 yields a +0.55 gain in CLIP.Source: §F.2 · p. 17, Table F2 · p. 17

F.3 The Efficiency–Memory Trade-off of LRU Size

LRU cache performance saturates at a moderate size: size 4 reaches a 99.9% hit rate, while increasing size from 7 to 10 saves little latency relative to its added GPU memory.

F.3 The Efficiency–Memory Trade-off of LRU Size

  • From GLRU=7 to GLRU=10, latency improves by only 0.56 s (16.91 →16.35 s), while memory rises from 21.4 to 24.5 GB.GLRU = 4 already reaches a 99.9% hit rate; GLRU = 1 has a 0% hit rate and 2.2× higher latency than the optimum.Source: §F.3 · pp. 17–18, Figure F1 · p. 18

G Analysis and Discussion

The discussion connects OmniMem’s long-video behavior to head-specialized use of memory branches and consistency without long-video finetuning, and explains its retrieval and implementation choices for chunk-based video generation.

G Analysis and Discussion

  • Different attention heads favor different memory branches, and Fig. G1 shows fusion gates routing heads toward their preferred historical information.Compression provides a global summary, selection retrieves fine-grained long-range details, and the sliding window captures recent context.Source: G Analysis and Discussion · p. 18
  • OmniMem maintains long-video consistency without long-video finetuning, including a 20-second example from a model trained only on 5-second clips.As videos grow, each head ranks a larger pool of chunks while preserving its role, such as favoring recent frames or tracking anchor tokens.Source: G Analysis and Discussion · p. 18, Figure G2 · p. 19

H Benchmark for Triton Utilization Different Pooling Size

  • 93.8% utilization is achieved by the selected best tile-fit pooling group for the Wan-2.1-1.3B 480P configuration.Smaller blocks offer finer selection but low utilization, while larger blocks reduce sequence length more aggressively but may lose temporal or spatial granularity.Source: H Benchmark for Triton Utilization… · p. 18, Table H1 · p. 19

I Explicit Retrieval in Chunk-Based Video Generation

  • Chunk-based video retrieval must manage selected-block unions across thousands of query tokens and dedicated MHA heads; OmniMem combines window exclusion, proxy query sharing, and per-head scattered access.The scattered access avoids expanding head-specific selections into a large head-aligned KV buffer.Source: I Explicit Retrieval in Chunk-Based… · p. 19

J Permutation Equivariance of Reorder with RoPE

  • A one-time reorder makes 2D spatial blocks contiguous for mean-pooling, while jointly reordering tokens and RoPE positions preserves attention up to that same reorder.The reorder–compress–attend–restore pipeline is equivalent to applying 2D-block compression before attention on the unreordered sequence.Source: J Permutation Equivariance of Reorder… · p. 20

K More Qualitative Results

The 60-second qualitative comparisons show OmniMem maintaining subject consistency and natural scene evolution where baselines exhibit drift or repetition. In multi-prompt generation, it follows prompt transitions while preserving consistency.

Long Video Generation

  • OmniMem maintains subject consistency and natural scene evolution in 60-second generation, while SWA shows subject drift and LongLive [18] produces repetitive frames.SWA discards distant context; LongLive [18] preserves identity via attention sinks but exhibits repetition.Source: K More Qualitative Results · p. 20, Figure K1 · p. 21

Multi-Prompt Long Video Generation

Related papers

Loading 2605.30519v1…