Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LoCoCo: Dropping In Convolutions for Long Context Compression

Ruisi Cai, Yuandong Tian, Zhangyang Wang, Beidi Chen

Abstract

This paper tackles the memory hurdle of processing long context sequences in Large Language Models (LLMs), by presenting a novel approach, Dropping In Convolutions for Long Context Compression (LoCoCo). LoCoCo employs only a fixed-size Key-Value (KV) cache, and can enhance efficiency in both inference and fine-tuning stages. Diverging from prior methods that selectively drop KV pairs based on heuristics, LoCoCo leverages a data-driven adaptive fusion technique, blending previous KV pairs with incoming tokens to minimize the loss of contextual information and ensure accurate attention modeling. This token integration is achieved through injecting one-dimensional convolutional kernels that dynamically calculate mixing weights for each KV cache slot. Designed for broad compatibility with existing LLM frameworks, LoCoCo allows for straightforward "drop-in" integration without needing architectural modifications, while incurring minimal tuning overhead. Experiments demonstrate that LoCoCo maintains consistently outstanding performance across various context lengths and can achieve a high context compression rate during both inference and fine-tuning phases. During inference, we successfully compressed up to 3482 tokens into a 128-size KV cache, while retaining comparable performance to the full sequence - an accuracy improvement of up to 0.2791 compared to baselines at the same cache size. During post-training tuning, we also effectively extended the context length from 4K to 32K using a KV cache of fixed size 512, achieving performance similar to fine-tuning with entire sequences.

LoCoCo: Dropping In Convolutions for Long Context Compression

Abstract

This paper tackles the memory hurdle of processing long context sequences in Large Language Models (LLMs), by presenting a novel approach, Dropping In Convolutions for Long Context Compression (LoCoCo). LoCoCo employs only a fixed-size Key-Value (KV) cache, and can enhance efficiency in both inference and fine-tuning stages. Diverging from prior methods that selectively drop KV pairs based on heuristics, LoCoCo leverages a data-driven adaptive fusion technique, blending previous KV pairs with incoming tokens to minimize the loss of contextual information and ensure accurate attention modeling. This token integration is achieved through injecting one-dimensional convolutional kernels that dynamically calculate mixing weights for each KV cache slot. Designed for broad compatibility with existing LLM frameworks, LoCoCo allows for straightforward "drop-in" integration without needing architectural modifications, while incurring minimal tuning overhead. Experiments demonstrate that LoCoCo maintains consistently outstanding performance across various context lengths and can achieve a high context compression rate during both inference and fine-tuning phases. During inference, we successfully compressed up to 3482 tokens into a 128-size KV cache, while retaining comparable performance to the full sequence - an accuracy improvement of up to 0.2791 compared to baselines at the same cache size. During post-training tuning, we also effectively extended the context length from 4K to 32K using a KV cache of fixed size 512, achieving performance similar to fine-tuning with entire sequences.
Paper Structure (30 sections, 5 equations, 3 figures, 7 tables, 2 algorithms)

This paper contains 30 sections, 5 equations, 3 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Long-Context Inference
  4. Long-Context Fine-tuning
  5. Attention Approximation
  6. Language Model Design with Built-In Convolutions
  7. Methodology
  8. Segment-Level Attention with Long Sequences
  9. Convolution as a Context Compression Operator
  10. Convolutional Token Compressor
  11. Complexity Analysis
  12. Connection with Token Dropping
  13. Dropping-In Integration of LoCoCo
  14. Long-Context Efficient Inference
  15. Long-Context Extension
  16. ...and 15 more sections

Figures (3)

  • Figure 1: Overview of our pipeline. We process the long sequences block-wisely and maintain a fixed-size compressed memory.
  • Figure 2: Token merging via convolutional kernels as the drop-in" integration without modifying the original weights. Based on Llama-2-7B touvron2023llama, we inserted the convolutional heads on the top of self-attention, and tested the model performance on various few-shot downstream tasks. The input sequence typically consists of about 2000 tokens. We compare our method with zhang2023h, a token eviction strategy. We also provide the uncompressed case, where the model uses the full sequence.
  • Figure 3: Varying memory sizes during fine-tuning, evaluated on Proof-Pile-2 azerbayev2023llemma. Compared to zhang2023h, our method shows exceptional performance at large compression ratios, indicating the expressiveness of the merged token.