UMFANG: Optimierung der Schlüssel-Wert-Cache-Komprimierung bei der Generierung von Langkontexten

SCOPE: Optimizing Key-Value Cache Compression in Long-context Generation

December 18, 2024

Autoren: Jialong Wu, Zhenglin Wang, Linhai Zhang, Yilong Lai, Yulan He, Deyu Zhou

cs.AI

Zusammenfassung

Der Key-Value (KV) Cache ist zum Engpass von LLMs für die Generierung von Langkontexten geworden. Trotz zahlreicher Bemühungen in diesem Bereich wird die Optimierung für die Decodierungsphase im Allgemeinen vernachlässigt. Wir sind jedoch der Meinung, dass eine solche Optimierung entscheidend ist, insbesondere für Aufgaben zur Generierung langer Ausgaben, basierend auf den folgenden zwei Beobachtungen: (i) Übermäßige Kompression während der Vorausfüllphase, die einen spezifischen vollständigen Kontext erfordert, beeinträchtigt das Verständnis der Begründungsaufgabe; (ii) Abweichungen von schwerwiegenden Treffern treten bei Begründungsaufgaben mit langen Ausgaben auf. Daher wird SCOPE, ein einfaches, aber effizientes Framework, das die KV-Cache-Optimierung während der Vorausfüll- und Decodierungsphasen getrennt durchführt, vorgestellt. Speziell wird der KV-Cache während der Vorausfüllphase erhalten, um die wesentlichen Informationen zu bewahren, während eine neuartige Strategie auf Schiebebasis vorgeschlagen wird, um wesentliche schwerwiegende Treffer für die Decodierungsphase auszuwählen. Der Speicherverbrauch und der Speichertransfer werden weiter optimiert, indem adaptive und diskontinuierliche Strategien verwendet werden. Umfangreiche Experimente mit LongGenBench zeigen die Wirksamkeit und Verallgemeinerung von SCOPE sowie dessen Kompatibilität als Plug-In für andere Methoden zur KV-Kompression nur während der Vorausfüllung.

English

Key-Value (KV) cache has become a bottleneck of LLMs for long-context generation. Despite the numerous efforts in this area, the optimization for the decoding phase is generally ignored. However, we believe such optimization is crucial, especially for long-output generation tasks based on the following two observations: (i) Excessive compression during the prefill phase, which requires specific full context impairs the comprehension of the reasoning task; (ii) Deviation of heavy hitters occurs in the reasoning tasks with long outputs. Therefore, SCOPE, a simple yet efficient framework that separately performs KV cache optimization during the prefill and decoding phases, is introduced. Specifically, the KV cache during the prefill phase is preserved to maintain the essential information, while a novel strategy based on sliding is proposed to select essential heavy hitters for the decoding phase. Memory usage and memory transfer are further optimized using adaptive and discontinuous strategies. Extensive experiments on LongGenBench show the effectiveness and generalization of SCOPE and its compatibility as a plug-in to other prefill-only KV compression methods.