HiLS-Attention: иерархическое разреженное внимание для сверхдлинного контекста в языковых моделях

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

Авторы (публикацию на HuggingFace Papers представил исследователь Сян Ху, Xiang Hu) предлагают Hierarchical Landmark Sparse Attention (HiLS-Attention), механизм, в котором отбор релевантных чанков не задан эвристикой, а обучается сквозным образом (end-to-end) вместе с самой моделью, по функции потерь языкового моделирования (LM loss). Внимание строится иерархически: каждый запрос сначала независимо «просматривает» каждый отобранный чанк и извлекает специфичную для него информацию, а затем результаты объединяются с весами, равными оценкам релевантности (retrieval scores) этих чанков. Поскольку такие оценки релевантности встроены прямо в прямой проход внимания, они оптимизируются напрямую вместе с LM loss, это даёт сквозное обучение отбора и «нативную» разреженную тренировку без отдельной эвристической стадии поиска.

По экспериментам авторов, на длинах контекста, знакомых модели по обучению, HiLS-Attention показывает качество на уровне полного внимания, а местами превосходит его. За пределами тренировочной длины HiLS-Attention экстраполируется более чем в 64 раза при точности отбора нужных чанков 90%, по утверждению авторов, значительно дальше, чем способно полное внимание. Кроме того, уже существующие модели с полным вниманием можно перевести на HiLS-Attention лёгким продолженным дообучением (continued pretraining): это сохраняет исходное качество на привычных длинах контекста и добавляет способность работать со сверхдлинным контекстом. Вместе с разреженным доступом к KV-кэшу это, как утверждают авторы, снимает привычный компромисс между эффективностью и качеством, получаются модели с длинным контекстом, которые одновременно дешевле в вычислениях и точнее на задачах с длинным контекстом, чем аналоги с полным вниманием.

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

Полное (dense) внимание в LLM стоит квадратично от длины контекста и плохо экстраполируется за пределы длины обучения, это ограничивает масштабирование длинного контекста.
Предложен новый механизм, Hierarchical Landmark Sparse Attention (HiLS-Attention): чанковое разреженное внимание, где отбор релевантных чанков обучается end-to-end вместе с моделью по LM loss, а не задаётся эвристикой.
На длинах контекста, знакомых модели по обучению, HiLS-Attention не уступает полному вниманию, а местами превосходит его.
HiLS-Attention экстраполируется более чем в 64 раза за пределы длины обучения при точности отбора чанков 90%, заметно дальше, чем полное внимание.
Существующие модели с полным вниманием можно перевести на HiLS-Attention лёгким продолженным дообучением, сохранив качество и добавив поддержку сверхдлинного контекста при меньших вычислительных затратах за счёт разреженного доступа к KV-кэшу.

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

Длина контекста, одно из главных узких мест современных LLM: полное внимание требует вычислений, растущих квадратично с длиной текста, а модели теряют точность, как только контекст выходит за пределы того, что они видели при обучении. Из-за этого работа с длинными документами, большими кодовыми базами или долгими диалогами ИИ-агентов обходится дорого либо ненадёжна. HiLS-Attention предлагает разреженное внимание, в котором сам отбор нужных фрагментов текста обучается вместе с моделью, а не задаётся правилами извне, и, по данным авторов, это позволяет одновременно снизить стоимость вычислений, не потерять (а местами повысить) качество и сильно нарастить длину контекста, на которую модель способна экстраполироваться.

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

В первую очередь, исследователям и инженерам, которые проектируют и обучают LLM с длинным контекстом: тем, кто занимается обработкой больших документов, кода, длинных диалогов и памятью ИИ-агентов. Также это релевантно инженерам инфраструктуры, которые оптимизируют стоимость инференса за счёт разреженного внимания и работы с KV-кэшем, и командам, которые хотят продлить контекст уже готовых моделей без полного переобучения с нуля.

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

По описанию авторов, HiLS-Attention, это архитектурный механизм внимания, который встраивается в модель, а не отдельный инструмент. Ключевая практическая деталь: существующие модели с полным вниманием не обязательно обучать заново, их можно перевести на HiLS-Attention лёгким продолженным дообучением (continued pretraining), сохранив исходное качество на привычных длинах контекста и получив взамен экстраполяцию на сверхдлинный контекст. Это делает метод применимым как способ обновить уже готовые модели, а не только как схему для обучения новых моделей с нуля.

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

Источник, препринт научной статьи, выложенный на HuggingFace Papers; в доступном тексте, только аннотация (abstract), без полного текста статьи, деталей методологии, используемых датасетов и явного сравнения с конкретными базовыми методами. Все цифры (сопоставимое или более высокое качество, экстраполяция более чем в 64 раза, точность отбора 90%), это самоотчёт авторов из аннотации; независимого воспроизведения или стороннего анализа в тексте нет. На момент подготовки дайджеста публикация на HuggingFace Papers набрала 44 балла и 1 комментарий, то есть это ранний и пока не очень широкий отклик сообщества.

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

Аннотация не раскрывает детали экспериментов, какие датасеты и бенчмарки использовались, с какими именно методами сравнивали результаты и какого масштаба были модели, поэтому судить о полноте сравнения и воспроизводимости по одному тексту нельзя. Заявленная точность отбора чанков 90% при 64-кратной экстраполяции означает, что около 10% релевантных фрагментов модель всё же не находит, для задач, где важна каждая деталь, это может иметь значение. Перевод существующих моделей на HiLS-Attention всё равно требует дополнительного дообучения (пусть и лёгкого), то есть не является полностью бесплатным по вычислениям. Кроме того, в тексте нет ни критики метода, ни альтернативной точки зрения, только собственные результаты авторов.

«HiLS-Attention снимает привычный компромисс между эффективностью и качеством, позволяя создавать модели с длинным контекстом, которые одновременно эффективнее и результативнее на задачах с длинным контекстом, чем аналоги с полным вниманием.»

— из аннотации к статье о HiLS-Attention, HuggingFace Papers