Zyphra перевела механизмы внимания на дерево: ускорение до 417 раз на миллиарде токенов

Токен, это кусочек слова, которым нейросеть меряет текст. Методы разреженного внимания, например DeepSeek Sparse Attention, читают не весь текст, а выборку, но выборку составляет индексатор, и он по-прежнему просматривает каждый токен. ZHA делает разреженным и сам поиск: так метод описывает техническая статья Zyphra на arXiv (номер 2610.42112), её пересказал Marktechpost.

Что измеряли?

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

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

ZHA строит над текстом оглавление. Ключи, то есть метки, по которым модель находит прочитанные фрагменты, режут на куски по 4 штуки и сжимают каждый кусок в сводку, а сводки снова группируют и сжимают, уровень за уровнем. При таких кусках 6 уровней накрывают 4 096 токенов одной верхней вершиной.

Сводки считают 16 обученных «контрольных запросов», которые прикидывают, сколько внимания любой будущий запрос уделит куску. Дальше действует правило top-p: модель берёт ветки, на которые приходится 95% ожидаемого внимания, и спускается только в них до исходных записей. Вопрос про один факт затрагивает мало веток, вопрос про разбросанные факты раскрывает больше.

Полная память модели о прочитанном, KV-кэш (хранилище ключей и значений всех прошлых токенов), остаётся целиком. Дерево лишь решает, что из неё читать.

Саму модель Zyphra не переучивает, она заморожена, учатся только контрольные запросы. Это дистилляция: маленькая надстройка учится повторять поведение плотного внимания на текстах длиной 32 тысячи токенов.

Для сравнения команда проверила два дерева:

• глубокое, куски по 4 ключа, 6 уровней;
• широкое, куски по 256 ключей, 1 уровень;
• соперник MSA (MiniMax Sparse Attention, разреженное внимание MiniMax), его индексатор учили на 40 млрд токенов, в 10 раз больше, чем ZHA.

Что обнаружили?

• Поиск фактов. На тесте RULER-13 при 256 тысячах токенов широкое дерево набрало 96,2% против 96,0% у плотного внимания, оба дерева держатся выше 93%, а MSA падает до 77,1%.
• Длинные тексты. На LongBench v2 глубокое дерево дало 63,4% против 58,0% у плотного, читая около 1,1% ключей. MSA читал 1% ключей, и глубокое дерево обошло его на 14,1 пункта.
• Общий уровень. На 7 тестах рассуждения, программирования и поиска ZHA отстаёт от плотного внимания на несколько пунктов, не больше.
• Устойчивость. Из 28 вариантов формы дерева точность на RULER держалась ровно при чтении от 10 до 30% ключей. Глубокие деревья сохраняли точность при большей разреженности.
• Скорость по расчёту. ZHA обгоняет плотное внимание около 48 тысяч токенов, а MSA около 480 тысяч. На 10 млн токенов блог Zyphra даёт выигрыш от 32 до 35 раз над плотным вниманием.
• Размер модели. Модели меньше 1 млрд параметров перенастраивались плохо, с 4 млрд параметров результат заметно рос.

Что это значит для вас?

Готового сервиса с ZHA нет, это исследование. Но оно показывает, как устроено чтение длинных документов в нейросетях и где оно ошибается.

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

Маркетологу. Длинный контекст дорог, и способы его удешевить важны для цены ИИ-сервисов. Но на реальных 256 тысячах токенов расчётный выигрыш от 1,93 до 2,28 раза, а сотни раз относятся к миллиарду токенов и остаются прогнозом.

Предпринимателю РФ и СНГ. О доступности метода в России источник ничего не сообщает, кода для ускорения нет. Если выбираете сервис под длинные договоры, проверяйте его на своём файле с заранее известными фактами.

Из доступных в РФ аналогов, YandexGPT и GigaChat.

По данным технической статьи Zyphra (arXiv 2610.42112) и блога Zyphra в пересказе Marktechpost