SWARR: как reinforcement learning адаптирует скользящее внимание для математических рассуждений

Современные большие языковые модели (LLM) всё чаще используются для задач, требующих длинного контекста, однако механизм самовнимания (self-attention) имеет квадратичную вычислительную сложность. Одной из альтернатив является скользящее окно внимания (Sliding-Window Attention, SWA), но такие модели обычно уступают полному вниманию в точности рассуждений, особенно в математике.

В новой работе на arXiv исследователи предлагают рецепт под названием SWARR (Sliding-Window Attention with Reinforced Adaptation for Math Reasoning). Метод состоит из двух этапов: сначала предобученная модель с полным вниманием эффективно преобразуется в SWA с помощью контролируемой тонкой настройки (SFT), а затем на втором этапе применяется обучение с подкреплением (RL) для адаптации политики модели.

Авторы обнаружили, что после SFT модель SWA всё ещё заметно отстаёт от полного внимания. Причиной они называют несоответствие данных и архитектуры: большинство данных для SFT подготовлено для моделей с полным вниманием и может содержать дальние зависимости, которые SWA моделирует хуже. Поскольку обучение с подкреплением использует собственные генерации модели в условиях ограничения SWA, оно позволяет адаптировать траектории рассуждений под архитектуру.

Эксперименты на бенчмарках математических рассуждений показали, что SWARR существенно сокращает разрыв между SWA и полным вниманием. Потери точности, возникшие при конвертации, восстанавливаются при сохранении преимущества SWA в виде линейной сложности вычислений.

Ключевой эмпирический результат работы заключается в том, что применение RL меняет выводы о жизнеспособности SWA для математических рассуждений, которые можно было бы сделать только на основе конвертации и SFT. Это открывает путь к более эффективным моделям для длинного контекста.

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