Студенты НИУ ВШЭ и VK показали, как догнать трансформеры с линейным вниманием
Команда магистрантов Инженерно-математической школы НИУ ВШЭ и VK проверила, можно ли поднять качество моделей с линейным вниманием до уровня классического трансформера, не возвращаясь к softmax. Исследование сосредоточено на архитектурных настройках и схеме обучения, а не только на функции сходства.
Группа студентов мастерской по сервисам и платформам ИИ Инженерно-математической школы НИУ ВШЭ и VK — Павел Алексеев, Александр Рожкин, Андрей Аргаткин, Алексей Павлов и Матвей Каменев — представила результаты проекта, посвящённого качеству моделей с линейным вниманием. Работа опирается на архитектуру Dense Attention, ранее описанную руководителем направления Андреем Аргаткиным, и ставит вопрос о том, какие компоненты обеспечивают её высокие результаты. Авторы проверили, можно ли перенести эти решения на другие варианты линейного внимания, сохранив его главное преимущество — линейную сложность по длине последовательности.
Исходная мотивация исследования связана с хорошо известным ограничением стандартного трансформера: механизм self-attention требует квадратичных вычислений относительно длины входа, что резко удорожает обучение и применение моделей на длинном контексте. Линейное внимание обходит это за счёт замены softmax альтернативной функцией сходства и перегруппировки матричных умножений, снижая сложность до линейной при фиксированных размерностях признаков и значений. Однако на практике такие модели нередко уступают классическим трансформерам в задачах, где нужно моделировать сложные зависимости между удалёнными элементами последовательности. Целью команды было сократить этот разрыв на бенчмарке Long Range Arena и как минимум не проиграть квадратичному вниманию, не отказываясь от линейной вычислительной схемы.
Рабочая гипотеза студентов состояла в том, что линейные модели отстают не только из-за отказа от softmax, но и из-за архитектурных настроек, унаследованных от обычного трансформера. Число голов внимания, способ позиционного кодирования и чередование слоёв, подобранные для стандартной архитектуры, не обязательно оптимальны для линейного механизма. В экспериментах на Pathfinder32 из набора LRA команда варьировала позиционное кодирование, количество голов, схему локальных и глобальных слоёв, а также проверяла MLM-предобучение, гибридное и полиномиальное внимание. Отдельно рассматривался двунаправленный режим внимания, естественный для классификации, извлечения признаков и восстановления замаскированных токенов, тогда как для авторегрессионной генерации вводится каузальное ограничение.
Технически линейное внимание опирается на ассоциативность матричного умножения: вместо матрицы попарных оценок размера «длина на длину» сначала вычисляется матрица, зависящая только от размерности признаков, а затем она умножается на значения. В двунаправленном режиме запросы и ключи считаются один раз по всей последовательности, а агрегация и вычисление выходов требуют линейного числа операций. Именно эта схема объясняет вычислительную привлекательность подхода, но сама по себе не гарантирует качества: модели всё равно нужно передать информацию о порядке токенов и научить выделять полезные связи. Поэтому авторы сравнивали фиксированное синусоидальное кодирование, обучаемые позиционные эмбеддинги и Rotary Position Embedding, а также исследовали, как на результат влияет устройство слоёв и способ обучения.
Контекст работы задаёт общий тренд на экономичные архитектуры для длинного контекста. Пока одни команды пытаются масштабировать квадратичное внимание за счёт инженерных ухищрений и кэширования, другие ищут архитектурные альтернативы, снижающие сложность принципиально. В этом поле конкурируют разреженное внимание, рекуррентные подходы вроде Mamba и гибридные схемы, сочетающие локальные и глобальные слои. Отдельная линия связана с хранением состояния фиксированного размера вместо растущего KV-cache при авторегрессионной генерации, что напрямую влияет на стоимость инференса и объём памяти. Работа студентов ВШЭ и VK попадает именно в эту нишу, предлагая не новую функцию сходства, а систематическую настройку уже существующих линейных механизмов.
Для российского рынка и академического сообщества значение проекта двойное. С одной стороны, он демонстрирует, что прикладные исследования по эффективным архитектурам языковых моделей ведутся не только в крупных корпоративных лабораториях, но и в формате студенческих мастерских при участии индустриальных партнёров. С другой стороны, результаты важны для практических сценариев с длинными документами, где стоимость квадратичного внимания становится ограничивающим фактором. Если архитектурные настройки действительно сокращают разрыв в качестве, линейные модели могут стать более привлекательной основой для энкодерных задач и обработки больших массивов текста без пропорционального роста вычислительных затрат.
Сравнение с альтернативами пока остаётся неполным. Гибридные схемы, сочетающие локальное и глобальное внимание, и полиномиальные варианты предлагают компромисс между выразительностью и стоимостью, но требуют аккуратного подбора гиперпараметров. Рекуррентные архитектуры дают линейную сложность естественным образом, однако для двунаправленной обработки им приходится объединять прямой и обратный проходы, как это сделано, например, в Vision Mamba. Линейное внимание с перегруппировкой матричных умножений выигрывает в простоте реализации, но проигрывает в способности моделировать сложные зависимости, если не уделить внимание позиционной информации и схеме слоёв. Именно этот баланс и пытаются найти авторы. Открытым остаётся вопрос, насколько полученные на Pathfinder32 выводы переносятся на другие задачи LRA и на реальные языковые корпуса. Двунаправленный режим удобен для классификации и извлечения признаков, но для генерации текста нужен каузальный вариант, и там требования к качеству и памяти отличаются. Кроме того, предстоит понять, как архитектурные настройки линейного внимания ведут себя при масштабировании модели и на последовательностях существенно большей длины. Дальнейшие шаги команды, вероятно, будут связаны с расширением набора бенчмарков и проверкой гипотезы на более широком спектре линейных механизмов, чтобы отделить универсальные решения от частных случаев.