Tensor Cache улучшает кэширование Transformer моделей, сочетая точность и экономию памяти

Кэш ключей и значений (KV cache) в авторегрессионных Transformer моделях растет линейно с длиной контекста, что создает проблемы при обработке длинных последовательностей. Традиционные методы скользящего окна ограничивают объем памяти, но полностью отбрасывают вытесненные токены, делая недоступной информацию за пределами окна.

Группа исследователей представила Tensor Cache — двухуровневую архитектуру, которая решает эту проблему. Первый уровень (L1) использует точное внимание softmax со скользящим окном для недавних токенов. Второй уровень (L2) представляет собой внешнюю память на основе быстрых весов фиксированного размера, куда попадают KV-пары, вытесненные из окна.

Вытесненные пары сжимаются в матрицу A на каждый слой модели. При чтении будущие запросы обращаются к этой памяти через единственное матричное умножение, используя тождество линейного внимания: q_t(k_i ? v_i) = ?q_t,k_i?v_i. Это позволяет эффективно получать доступ к информации из вытесненных токенов.

Архитектура включает обучаемый скалярный затвор для слияния выходов L1 и L2, а также per-head параметры затухания и скорости записи, которые тренируются сквозным образом. Авторы отмечают, что хотя внешняя память и тождество чтения известны, их новизна в использовании такого хранилища как второго уровня, питаемого исключительно вытеснениями из скользящего окна.

Также исследователи выявили проблему в распространенном обучающем приеме chunked-mean: при обновлении матрицы A через A??A+?(k??v?) возникают ложные кросс-токенные внешние произведения порядка C??C. Они предложили решение — параллельный weighted-sum scan, эквивалентный per-token записям с точностью до float32 epsilon.

Эксперименты проводились на задачах масштабирования систем, контролируемого ассоциативного припоминания, языкового моделирования с длинным контекстом и диагностики емкости памяти. Tensor Cache продемонстрировал улучшение границы между используемой памятью и качеством модели по сравнению с bounded-state baseline.

Новый метод позволяет Transformer моделям работать с длинными контекстами при фиксированном бюджете памяти, что актуально для приложений вроде анализа документов, диалоговых систем и генерации кода. Препринт доступен на arXiv.