Calibrated token merging improves video generation cache efficiency
Offline Makes Online Better: Calibrated Cache Compression for Video Generation
Computer Vision and Pattern RecognitionArtificial IntelligenceMachine Learning
Summary
Videos generated by AI use a memory cache to keep track of important details from past frames. The authors found that treating all parts of the AI the same when compressing this cache wastes memory and hurts video quality. They measured how different parts react to compression and made a smart plan to compress more where it matters less. This approach, called CALM, helps AI models create better videos with less memory needed.
What this means in practice
- •For video generation engineers: Improve memory efficiency and video quality in autoregressive video generation models by using calibrated token merging per layer and head.
- •For machine learning platform developers: Build optimized cache compression modules that dynamically allocate merging ratios for large-scale video generation services.
Authors
Zhen Wang, Ning Lai, Qingtao Yu
Abstract
Token merging is an effective approach to reducing KV-cache memory consumption in autoregressive video generation by clustering temporally redundant tokens and retaining only a small set of representatives. Existing methods, however, typically apply a uniform merging ratio across all attention layers and heads. We observe that different layers and heads exhibit substantially different yet jointly correlated sensitivities to token merging, and that these sensitivity patterns remain consistent across prompts. Based on this observation, we introduce CALM, a Calibrated Allocation for Layer-Head Merging that uses offline calibration to improve online cache compression. Exhaustively evaluating all possible layer-head allocations is prohibitively expensive. To enable efficient calibration, we show that attention-output error strongly correlates with end-to-end video quality and can therefore serve as a lightweight proxy for allocation risk assessment. We evaluate different allocation strategies using this proxy and construct a fixed layer-head risk map offline. Given a cache budget at inference time, CALM uses dynamic programming to determine the best allocation of merging ratios. Experiments across three different autoregressive video models demonstrate that CALM consistently outperforms uniform token merging under the same memory budget, by reducing the LPIPS to original videos by around 21% with higher video-quality scores.