LinearKV: Кэширование для гибридных LLM без переобучения

Проблема гибридных архитектур

Современные методы ускорения LLM, такие как позиционно-независимое кэширование (PIC), опираются на полные механизмы внимания (full-attention), где ключевые и значения (KV-cache) хранятся по индексам токенов. Однако гибридные модели, сочетающие слои внимания с линейными рекуррентными слоями (например, Mamba), ломают эту логику: линейные слои выдают только состояние фиксированного размера, не позволяя конкатенировать или локально восстанавливать контекст.

Решение: Decoupled Initialization

Команда авторов (Yirui Liu, Xuelong Li и др.) предложила фреймворк LinearKV, который не требует дообучения. Ключевая идея — раздельная инициализация: каждый линейный слой маппит $K$ совпавших локальных состояний в одно начальное состояние, в то время как слои внимания работают по старой схеме. Это позволяет использовать существующие селекторы токенов (EPIC и др.) без изменений.

Ключевое открытие: «Одного состояния достаточно»

Авторы опровергли гипотезу о необходимости точного алгебраического сложения всех $K$ состояний (как в методе HYPIC). Оказалось, что использование одного сохраненного блока (single cached block initializer) не только проще, но и эффективнее на некоторых архитектурах.

Результаты бенчмарков

Сравнение проводилось на моделях GDN и Mamba-2 с использованием тестов LongBench QA и RULER (длины 8K–32K). Результаты наглядно демонстрируют преимущество простого подхода:

Использование одного состояния сократило время до первого токена (TTFT) до 0.46x от полного префилла, тогда как точное композиция состояний добавляла еще 5–17% накладных расходов. LinearKV совместим с существующими методами PIC и готов к внедрению в гибридные архитектуры.

Источник: arXiv cs.AI ↗