Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Полное состояние линейного внимания научились не читать на каждом шаге генерации. В нерецензированном препринте, где все числа получены самими авторами, SketchSSM сократил трафик доступа к состоянию примерно в 10 раз при близком среднем качестве. Метод может снять узкое место при обслуживании гибридных моделей большими пакетами запросов, не меняя их веса.
Почему неизменное состояние приходится читать снова
Гибридные модели заменяют большую часть обычного внимания линейным. Вместо растущего KV-кэша такой слой хранит рекуррентное состояние — матрицу фиксированного размера, в которой накапливается информация о предыдущих токенах.
Это позволяет увеличить пакет запросов, но переносит нагрузку на память GPU. На каждом шаге слой читает всю матрицу состояния, умножает её на новый запрос и записывает обновление. При больших пакетах именно обмен с памятью, а не арифметика, начинает определять время генерации.
ReplaySSM уже сокращает число записей: ключи и значения временно лежат в небольшом буфере, а накопленные изменения применяются разом. В использованном авторами окне из 16 шагов полное состояние обновляется только на границе окна. Однако каждый новый запрос всё равно читает его целиком, хотя между обновлениями матрица не меняется.
Сжатие самой матрицы решает эту проблему ценой накопления ошибок. Если удалить элементы состояния или снизить их точность, неточность попадает в следующее обновление, а затем влияет на все последующие шаги. SketchSSM разделяет чтение и запись: обновляет полное состояние без приближения, а упрощает только вычисление выходного вектора.
Как полный снимок превращается в компактную проекцию
Будущие запросы заранее неизвестны, но их можно представить через небольшой фиксированный набор базисных векторов. Этот базис SketchSSM подбирает заранее на калибровочных данных отдельно для каждой головы состояния. Переобучать модель для этого не нужно.
Когда буфер сбрасывается, ядро обновляет полную матрицу и один раз применяет её ко всем векторам базиса. Получившиеся выходные векторы образуют компактный эскиз состояния. До следующего обновления слой читает только этот эскиз и собирает ответ для текущего запроса как линейную комбинацию сохранённых векторов.
Коэффициенты комбинации зависят не только от запроса, но и от текущего полного состояния. Их карту пересчитывают на границе окна вместе с эскизом. Такой расчёт дороже фиксированной проекции, зато в экспериментах заметно лучше сохранял качество.
Одинаковый размер эскиза для всех голов тоже оказался неудачным. Одни головы хуже переносят приближение, а ошибки других сильнее влияют на итоговую функцию потерь. Поэтому SketchSSM заранее распределяет доступный ранг между головами: чувствительным оставляет больше векторов, устойчивым — меньше. Само полное состояние при этом остаётся доступно для точного обновления.
Когда SketchSSM меняет планы по инфраструктуре
Метод проверили на четырёх гибридных моделях с механизмами Mamba-2, GDN и KDA. Качество оценивали на MATH-500, AIME25, GPQA-D и LCB, а извлечение информации — на задачах RULER с контекстом 16K. Средняя точность и полнота в основном сохранились лучше, чем при квантовании или прореживании полного состояния.
На одном NVIDIA B300 специализированные ядра линейного внимания работали в 5,20–7,78 раза быстрее стандартной реализации vLLM в зависимости от механизма. Для Nemotron 3 Super максимальная пропускная способность генерации выросла до 2,64 раза. После оптимизации линейное внимание перестало занимать основную часть времени декодирования, поэтому дальнейший выигрыш уже зависит от остальных слоёв модели.
Работа меняет план внедрения, если команда уже выбрала гибридную архитектуру и рассчитывает обслуживать крупные пакеты запросов. В таком режиме стоит профилировать чтение рекуррентного состояния отдельно от матричных умножений: уменьшение KV-кэша само по себе не гарантирует быструю генерацию.
SketchSSM нельзя свести к настройке сервера. Нужны калибровка базиса, распределение рангов между головами и ядро, которое одновременно обновляет состояние и строит эскиз. Это отдельная ветка оптимизации среды выполнения, но не обучения модели.
Границы замеров тоже важны для планирования. Авторы использовали NVIDIA B300, vLLM и собственные ядра для конкретных семейств линейного внимания. Поэтому работу разумно рассматривать как подтверждение архитектурного приёма, а заявленное ускорение — как результат для целевой реализации, который придётся заново проверять на другом GPU и в другом стеке.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
