CRISP ускорил внимание в языковых моделях до 5,3 раза на длинном контексте

Этап предзаполнения (prefilling) при работе языковых моделей с длинным контекстом требует, чтобы механизм внимания (self-attention) сравнил каждый токен с каждым, вычислительная стоимость растёт квадратично с длиной контекста, и это становится узким местом при инференсе. Существующие методы разреженного внимания решают проблему по-разному: статичные подходы задают фиксированные шаблоны или используют офлайн-профилирование, но не подстраиваются под структуру конкретного входа; более новые динамические методы направляют (маршрутизируют) каждую голову внимания на подходящий разреженный шаблон в реальном времени, но опираются на косвенные и затратные по вычислениям критерии маршрутизации, а их механизмы распределения бюджета внимания не учитывают иерархию распределения массы после softmax.

Авторы представляют CRISP (Cliff-awaRe Input-adaptive Sparse Prefilling) и решают две структурные проблемы такой динамической маршрутизации. Во-первых, они показывают, что решение о маршрутизации можно считать напрямую по структуре приближённой (proxy) карты внимания, и заменяют маршрутизацию на основе дивергенции Йенсена, Шеннона (JSD) на C_struct, структурный критерий, который измеряет массу внимания в позициях, совместимых с Vertical-Slash-паттерном. C_struct воспроизводит те же решения о маршрутизации, что и JSD, но убирает и затратное матричное умножение с пулингом, и последующий расчёт KL-дивергенции. Во-вторых, авторы формализуют понятие «обрыва массы» (post-softmax mass cliff) после softmax и теоретически доказывают, что строго кумулятивные пороги покрытия накапливают фоновый шум порядка O(n) при длинных контекстах. CRISP обходит эту проблему за счёт порога, учитывающего токены-«приёмники» (sink-aware) и опирающегося на уровень фонового шума.

В экспериментах на бенчмарках InfiniteBench, RULER и LongBench, проведённых на двух семействах моделей (какие именно, в тексте не названо), CRISP оказался сильнейшим среди разреженных методов в целом и на задачах, требующих точного поиска по контексту, сравнялся или превзошёл точное («плотное») внимание. По сравнению с базовыми методами CRISP отыграл до 28,0 процентных пунктов точности на задачах поиска и обеспечил ускорение внимания до 5,30 раза на контексте в 512 тысяч токенов, прирост скорости авторы объясняют прежде всего устранением накопления шума порядка O(n) при отборе токенов, при сохранении структурной целостности внимания.

Ключевые факты

CRISP, метод разреженного внимания для этапа предзаполнения (prefilling) длинного контекста в языковых моделях, где стандартное внимание масштабируется квадратично
Маршрутизацию на основе дивергенции Йенсена, Шеннона (JSD) заменяет структурный критерий C_struct: воспроизводит те же решения, но без затратного матричного умножения и расчёта KL-дивергенции
Строго кумулятивные пороги покрытия, по доказательству авторов, накапливают фоновый шум порядка O(n) на длинных контекстах, CRISP заменяет их на sink-aware порог по уровню шума
На бенчмарках InfiniteBench, RULER и LongBench (два семейства моделей, не названы) CRISP, сильнейший среди разреженных методов и сравнивается или превосходит точное внимание на задачах поиска по тексту
Прирост: до 28,0 процентных пунктов точности на задачах поиска относительно базовых методов и до 5,30 раза ускорения внимания на контексте в 512 тысяч токенов

Почему это важно

Механизм внимания в языковых моделях при обработке длинного контекста, квадратично растущее по стоимости место: чем длиннее вход, тем дороже становится этап предзаполнения. Прежние решения, фиксированные разреженные шаблоны или офлайн-профилирование, не подстраиваются под конкретный вход, а более новые динамические методы маршрутизации используют косвенные и затратные критерии и распределяют бюджет внимания, не учитывая, как на самом деле распределяется масса внимания после softmax. CRISP предлагает решение сразу для обеих проблем: дешёвый и точный критерий маршрутизации и порог отбора токенов, устраняющий накопление фонового шума на длинных контекстах.

Кому это важно

Работа адресована тем, кто занимается инференсом языковых моделей с длинным контекстом, разработчикам инфраструктуры для обслуживания моделей и исследователям методов разреженного внимания. Практическая ценность выше там, где важны и скорость обработки контекста в сотни тысяч токенов, и точность на задачах, требующих найти конкретный фрагмент в длинном тексте (поиск по документу, работа с длинными базами знаний).

Как это применить

CRISP встраивается в этап предзаполнения как замена существующей схемы динамической маршрутизации голов внимания к разреженным шаблонам: вместо расчёта JSD между полным и приближённым распределением внимания используется структурный критерий C_struct, измеряющий массу внимания в позициях, совместимых с Vertical-Slash-паттерном, это устраняет матричное умножение с пулингом и расчёт KL-дивергенции. Порог отбора токенов при этом строится не как строго кумулятивное покрытие, а как sink-aware порог, привязанный к уровню фонового шума. В тексте источника не сказано, опубликован ли код или веса модели.

Можно ли доверять

Источник, научная публикация (препринт), доступная в виде страницы на Hugging Face; в исходном тексте нет ни имён авторов и их принадлежности к организациям, ни названия площадки и даты публикации, эти данные в открытом тексте не указаны. Результаты опираются на три стандартных для области длинного контекста бенчмарка (InfiniteBench, RULER, LongBench), но конкретные модели из «двух семейств», на которых проводились эксперименты, в тексте не названы, как и разбивка результатов по отдельным бенчмаркам, приведены только итоговые (агрегированные) цифры. Публикация пока набрала немного откликов (8 баллов и 2 комментария на странице обсуждения).

Риски и подводные камни

Заявленные цифры, «до 28,0 процентных пунктов» и «до 5,30 раза», это верхние границы прироста, а не средние показатели, и без разбивки по конкретным бенчмаркам и моделям трудно оценить, насколько устойчив эффект в общем случае. Не названные семейства моделей и отсутствие данных о выпуске кода или весов ограничивают возможность независимой проверки и воспроизведения результатов.