Source-linked AI summary

Past- and Future-Informed KV Cache Policy with Salience Estimation in Autoregressive Video Diffusion

Hanmo Chen, Chenghao Xu, Xu Yang, Xuan Chen, Cheng Deng

arXiv:2601.21896v3cs.CV

TL;DR

Heuristic KV-cache policies for long-horizon autoregressive video generation overlook token importance, risking lost context and redundant cache. PaFu-KV uses a salience head distilled from a bidirectional teacher to retain informative tokens, and experiments report smaller caches, faster inference, and preserved video quality.

  • Problem

    Heuristic KV-cache policies ignore differences in token importance during long-horizon video generation, risking critical spatiotemporal information loss and redundant cache.

  • Method

    PaFu-KV uses a lightweight Salience Estimation Head distilled from a bidirectional teacher to predict token salience and guide KV-cache retention under causal generation.

  • Results

    PaFu-KV significantly reduces KV Cache size and accelerates inference while preserving high-quality video generation.

  • Takeaways & Limitations

    Selective retention of salient tokens supports more efficient long-horizon autoregressive video generation while maintaining generation fidelity.

  • Takeaways & Limitations

    Memory reduction is relatively limited because FlashAttention already reduces the dominant attention-memory footprint.

Abstract

from arXiv · show

Video generation is pivotal to digital media creation, and recent advances in autoregressive video generation have markedly enhanced the efficiency of real-time video synthesis. However, existing approaches generally rely on heuristic KV Cache policies, which ignore differences in token importance in long-term video generation. This leads to the loss of critical spatiotemporal information and the accumulation of redundant, invalid cache, thereby degrading video generation quality and efficiency. To address this limitation, we first observe that token contributions to video generation are highly time-heterogeneous and accordingly propose a novel Past- and Future-Informed KV Cache Policy (PaFu-KV). Specifically, PaFu-KV introduces a lightweight Salience Estimation Head distilled from a bidirectional teacher to estimate salience scores, allowing the KV cache to retain informative tokens while discarding less relevant ones. This policy yields a better quality-efficiency trade-off by shrinking KV cache capacity and reducing memory footprint at inference time. Extensive experiments on benchmarks demonstrate that our method preserves high-fidelity video generation quality while enables accelerated inference, thereby enabling more efficient long-horizon video generation. Our code will be released upon paper acceptance.

1. Introduction

Autoregressive video diffusion improves real-time generation but heuristic KV-cache policies ignore token importance, risking lost spatiotemporal context and redundant cache. PaFu-KV estimates token salience using past and future information to retain useful tokens while reducing cache footprint and preserving quality.

  • Autoregressive DiT methods use KV caching to support low-latency, real-time sequential video generation despite bidirectional models’ memory and latency costs.
  • Heuristic FIFO-style policies ignore token importance, causing critical spatiotemporal information loss, error accumulation, and degraded quality and efficiency.
  • PaFu-KV estimates token salience from past, present, and future interactions to identify globally useful tokens for long-video generation.
  • A lightweight Salience Estimation Head distilled from a bidirectional teacher predicts token importance under causal-generation constraints.
  • During inference, the policy retains informative tokens and evicts marginal ones, preserving fidelity while reducing redundant-cache accumulation.
  • PaFu-KV reduces KV Cache size and accelerates inference while preserving high-quality video generation.

2. Related Work

Prior work on autoregressive video diffusion seeks efficient long-video generation through teacher forcing, diffusion forcing, distillation, and inference strategies. Related KV-cache research in language models similarly reduces memory and computation through attention-based token selection and hierarchical compression.

  • Autoregressive Video Diffusion: Autoregressive video diffusion research targets efficient long-video generation through teacher forcing, diffusion forcing, and distilled few-step generators.
  • Autoregressive Video Diffusion: CausVid distills an autoregressive video generator from a pretrained bidirectional DiT teacher, while prior approaches face training-inference discrepancy and error accumulation.
  • KV Cache in Large Language Models: Large language model methods reduce KV-cache memory and computation using attention-driven token selection, temporal locality, and hierarchical compression.

3. Method

PaFu-KV estimates token salience from past, present, and future interactions, then maintains a compact KV Cache by retaining the most salient entries during training and inference.

  • 3.2. Spatialtemporal-balanced Salience Estimation: PaFu-KV estimates token salience from spatially and temporally balanced attention interactions, capturing influence across past, present, and future contexts.The method partitions attention into lower-triangular, diagonal, and upper-triangular components, then aggregates their values.
  • 3.3. Salience Estimation Head: The method uses a lightweight Salience Estimation Head to infer future-relevant token importance despite causal inference lacking direct future information.The head concatenates contextualized Q, K, and V features and predicts a scalar score for ranking and cache retention.
  • 3.3. Salience Estimation Head: Layerwise analysis at the final denoising step measures overlap between each intermediate layer’s top-k tokens and the final layer’s selections.The analysis motivates placing salience estimation in the final Transformer layer.
  • 3.3. Salience Estimation Head: During training, PaFu-KV distills salience scores from a real score network and trains the generator-side head with SmoothL1 supervision.Index-based cache management supports gradient checkpointing by recording selected token indices instead of explicitly reorganizing the cache.
  • 3.4. Inference: During inference, the method generates video chunk by chunk and evicts lower-scoring entries whenever the KV Cache reaches capacity.Current and historical salience scores are combined, and the top-k tokens update the cache for subsequent chunks.

4. Experiments

Experiments show that PaFu-KV improves short- and long-video generation efficiency while preserving or improving key quality measures, and ablations identify effective cache and salience-design choices.

  • Quantitative Results: PaFu-KV achieves higher short-video scores than Self Forcing and LongLive, higher semantic score than Rolling Forcing, and higher throughput through reduced KV Cache size.
  • Quantitative Results: On long-video benchmarks, PaFu-KV achieves the best subject and background consistency with lower quality drift and only marginal motion-smoothness degradation.
  • Qualitative Results: Qualitative comparisons show substantially improved temporal consistency and comparable robustness to error accumulation relative to existing long-video methods.
  • Ablation Study: Simpler salience pooling degrades consistency and quality drift, while increased SEH architectural complexity also fails to improve long-video generation.
  • Ablation Study: Cache sizes of 9360 and 6240 tokens match the full model’s performance, whereas reducing below 4680 tokens noticeably degrades generation quality.
  • Ablation Study: Removing frame sink causes only marginal degradation, indicating it is not a primary contributor to PaFu-KV’s performance gains.

5. Conclusion

PaFu-KV is a past- and future-informed KV Cache policy that uses principled salience estimation to support stable long-horizon video generation. Its salience score separates past, present, and future attention contributions to reduce self-attention’s diagonal bias.

  • PaFu-KV uses a principled salience estimation framework to guide KV Cache policy design.
  • The method identifies token salience as a key factor governing long-horizon generation stability.
  • Its salience score disentangles past, present, and future attention contributions, mitigating self-attention’s strong diagonal bias.

A. Additional Implementation Details

The implementation uses a 4680-token KV Cache and a lightweight two-layer S-MLP Salience Estimation Head. Training reduces memory through gradient checkpointing and uses 5-second video segments per iteration.

  • The full PaFu-KV model uses a 4680-token KV Cache, equal to three times the latent tokens per video frame, during training and inference.
  • Training enables PyTorch gradient checkpointing to reduce GPU memory by recomputing intermediate activations during backpropagation.
  • Each training iteration processes a 5-second video segment.
  • The Salience Estimation Head uses a lightweight S-MLP with two linear layers and an intervening SiLU activation.

B. Streaming Long Tuning and KV Re-caching in LongLive.

PaFu-KV combines long-horizon training with one-time KV re-caching to support consistent interactive video generation and prompt switching.

  • Streaming Long Tuning: Streaming Long Tuning simulates long autoregressive diffusion rollouts instead of training on short clips and testing on long videos.It is designed to address error accumulation during long-horizon generation.
  • KV Re-caching: KV re-caching refreshes prompt-conditioned cache contents when prompts switch, avoiding delayed or incorrect adherence to new prompts.Prompt information is repeatedly injected through cross-attention and propagated through self-attention into the cache.
  • Qualitative evaluation: Qualitative comparisons evaluate PaFu-KV against representative open-source autoregressive models on 30-second videos.The figures mark inconsistent regions with red circles.
  • KV Re-caching: One-time KV re-caching adds minimal overhead while supporting rapid adaptation, smooth temporal transitions, and consistent visual quality after prompt changes.Subsequent frames use the refreshed cache without further intervention.

C. Algorithm

The paper presents complete training and inference workflows for PaFu-KV while noting that teacher-side salience computation can be memory-intensive.

  • Algorithm: The paper provides complete algorithmic workflows for PaFu-KV training and inference in Algorithms 1 and 2.These algorithms are included to support reproducibility and precise procedural description.
  • Algorithm: Teacher-model ground-truth salience computation explicitly calculates attention weights across all tokens.The teacher uses full bidirectional self-attention.
  • Algorithm: Full bidirectional attention during teacher training can impose substantial GPU memory overhead.

D. Additional Qualitative Results.

Additional qualitative comparisons report strong long-horizon consistency and visualize the effects of the ablation study.

  • Additional qualitative results: Qualitative results indicate that PaFu-KV achieves strong long-horizon consistency across additional video comparisons.Prompts appear below the corresponding figures, and noticeable inconsistencies are marked with red circles.
  • Additional qualitative results: The section provides additional qualitative comparisons in Figures 5, 6, and 7, with video-format results available online.
  • Ablation study: Figure 8 visualizes the ablation study, including the averaging strategy discussed in the section.The supplied passage states that averaging produces substantial error, without reporting its complete outcome.

E. Discussion and Future Work.

The discussion identifies limited memory savings under FlashAttention and outlines algorithmic procedures for salience-based KV selection and future deployment improvements.

  • E. Discussion and Future Work: FlashAttention makes further KV Cache reduction yield only marginal additional memory savings because attention memory is already reduced from O(N^2) to approximately O(N).The method trades recomputation for memory efficiency by discarding intermediate attention products.
  • E. Discussion and Future Work: The discussion concludes that reducing GPU memory footprint most effectively requires reducing model parameters rather than only shrinking the KV Cache.The authors identify this as a direction for future AR video diffusion research.
  • Algorithm: Training computes generator salience scores and selected-token indices alongside denoising, then optimizes distribution-matching and salience losses.The supplied algorithm fragments show the training inputs, salience computation, and loss update.
  • Algorithm: PaFu-KV inference initializes a cache, generates frames autoregressively, and retains top-k salient tokens when the cache reaches its capacity.The workflow uses a salience MLP and selects cache entries by salience scores.
  • Algorithm: Block-wise salience computation partitions queries into chunks and aggregates attention-derived scores across upper, diagonal, and lower regions.The procedure fuses the resulting scores into the output salience representation.
Loading 2601.21896v3…