Unified multimodal models store cache more efficiently with task aware compression
UniCache: Task- and Type-Aware KV Cache Compression for Unified Multimodal Models
Machine Learning
Summary
Unified multimodal models combine different tasks like understanding images and generating text in one system, but storing temporary data called KV cache becomes costly as these tasks grow. The authors found that existing methods compress this data without considering different task needs or data types, leading to quality loss. They created UniCache, which smartly compresses cache by recognizing which parts matter most for each task without retraining the model. This approach keeps quality high while making storage up to five times smaller and speeding up processing for longer inputs.
What this means in practice
- •For ai infrastructure teams: Reduce memory use and increase speed when running unified multimodal models for diverse tasks by compressing KV cache in a task- and type-aware way.
- •For multimodal application developers: Enable longer context processing in models that integrate vision and language by using caching policies customized per task and cache type.
- •For cloud service providers: Offer multimodal AI models with lower hardware costs and faster response by leveraging UniCache’s cache compression under strict storage budgets.$Commercial implications: Makes scalable, cost-efficient multimodal AI services possible by improving memory compression and throughput, appealing to cloud-hosted AI platforms.
Authors
Wanqi Yang, Yuexiao Ma, Mei Xie, Xiawu Zheng, Shiwei Liu
Abstract
Unified multimodal models combine understanding, generation, and editing within a single network, offering a promising foundation for versatile multimodal applications. However, growing multimodal contexts make KV cache storage and access increasingly costly. Existing KV cache compression methods are typically tailored to specific tasks and single-modality caches, while overlooking changes in cache importance across tasks and timesteps. However, in unified multimodal models, each task involves multiple KV cache types, and both their composition and dynamics differ across tasks. As a result, a single compression policy overlooks task- and type-specific requirements, leading to the loss of critical information and degraded quality across tasks. Based on these findings, we propose UniCache, a training-free framework for task- and type-aware KV cache compression. UniCache identifies the cache segments activated by each task and assigns suitable compression policies through offline calibration. It coordinates their parallel execution under a shared storage budget through attention-guided allocation and task-aware temporal scheduling. Experiments show that UniCache achieves $5\times$ KV cache compression for understanding and editing and $2.5\times$ for generation with negligible quality loss, while increasing throughput by up to $1.78\times$ in long-context settings, significantly improving the practicality of scaling unified multimodal models to longer context.