Токийская исследовательская лаборатория Sakana AI в декабре опубликовала техническую статью с описанием RePo. Метод позволяет языковым моделям присваивать токенам обученные позиции вместо фиксированных индексов в последовательности. Накладные расходы на вычисления составляют меньше 1%, а результаты стабильно улучшаются на задачах, где стандартное позиционное кодирование работает плохо.
Проблема плоской последовательности
Трансформеры читают входные данные как одну строку токенов. Позиция 0, позиция 1, позиция 2, и так далее до конца. Единственный структурный сигнал для модели — целочисленный индекс, который обычно кодируется через RoPE (Rotary Position Embedding). RoPE вращает векторы запросов и ключей в зависимости от расстояния между токенами.
Для многих задач это работает нормально. В естественном языке зависимости обычно локальные. Но современные сценарии промптинга всё чаще нарушают это допущение. RAG-системы засовывают релевантные фрагменты между нерелевантными. Таблицы линеаризуются в текст. Критически важная информация оказывается за тысячи токенов от того места, где модель должна её использовать.
Команда Sakana формулирует проблему через призму теории когнитивной нагрузки из педагогической психологии. Рабочая память ограничена. Когда тратишь ресурсы на мусор, на реальное рассуждение остаётся меньше. Аргумент такой: у LLM похожее узкое место, когда механизм внимания продирается через шум в поисках нужного.
Как работает RePo
Идея концептуально простая. Вместо того чтобы жёстко привязывать позицию к индексу токена в последовательности, дать модели выучить, какую позицию каждый токен должен занимать.
Для каждого токена небольшая сеть предсказывает вещественную позицию на основе скрытого состояния токена. Эти выученные позиции подставляются в стандартный механизм позиционного кодирования. Семантически связанные токены могут оказаться ближе друг к другу в пространстве позиций, даже если в исходном промпте они были далеко.
Реализация на GitHub показывает, что модуль применяется примерно с пятого слоя, то есть где-то на трети глубины сети. Каждая голова внимания учит собственные назначения позиций, хотя базовое представление позиций общее для голов одного слоя.
Обучение сквозное. Никакой специальной фазы предобучения. Параметры RePo оптимизируются вместе со всем остальным.
Что показывают бенчмарки
Команда дообучила OLMo-2 1B от Allen AI (полностью открытая языковая модель) на 50 миллиардах токенов с контекстным окном 4096. Сравнивали с несколькими базовыми линиями: стандартный RoPE, NoPE (вообще без позиционного кодирования) и гибридные подходы с чередованием слоёв RoPE и NoPE.
Три категории оценки, каждая нацелена на сценарий, где жёсткие позиции должны мешать:
Зашумлённый контекст. Бенчмарк RULER добавляет нерелевантный контент вокруг информации, которая нужна модели. RePo обгоняет RoPE в среднем на 11 пунктов. Конкретно на подзадаче «иголка в стоге сена»: 88,25 против 82,56.
Структурированные данные. Когда графы и таблицы сплющиваются в текст, исходная структура разрушается. На NLGraph и HybridQA RePo показывает улучшение на 1,94 пункта по сравнению с RoPE. Авторы отмечают, что NoPE на графовом датасете вообще работает лучше всех. Это намекает на то, что смещение в сторону локальности в позиционных кодировках может активно мешать при работе с графовыми структурами.
Длинный контекст. Тестирование на 8K и 16K токенов, то есть на длинах, которые модель никогда не видела при обучении. Разрыв растёт с увеличением контекста. На LongBench RePo набирает 28,31 против 21,07 у RoPE.
На бенчмарках общих способностей картина другая. ARC, HellaSwag, MMLU-Pro, TriviaQA: RePo не отстаёт от RoPE, разница в пределах пункта. Специализированные улучшения не ухудшают стандартную производительность.
Что модель на самом деле выучивает
Раздел анализа разбирает паттерны внимания. На задачах «иголка в стоге сена» RePo выделяет больше массы внимания на далёкие, но релевантные токены-иголки (2,013 против 1,754 у RoPE), одновременно снижая внимание к близким токенам запроса. Делает то, что хочется: концентрирует ресурсы на важном, а не на близком.
Выученные пространства позиций оказываются плотными и нелинейными. Разные головы внимания специализируются на разных диапазонах позиций. Одни занимаются локальной реорганизацией, другие глобальной. Визуализации в интерактивной демке показывают паттерны вроде плато (много токенов отображаются в похожие позиции) и резких скачков.
Оговорки
Тестировали на модели с 1 миллиардом параметров. Статья не заявляет, что подход масштабируется на большие размеры. Учитывая сложную динамику внимания, хотелось бы увидеть результаты на 7B и выше, прежде чем считать, что всё обобщается.
Цифра 1% накладных расходов — для их оптимизированной реализации. Если пытаться воспроизвести это на ванильном PyTorch без их инфраструктуры, цифры будут хуже.
И про теорию когнитивной нагрузки как обоснование. Красиво звучит, но проверить невозможно. Испытывают ли LLM что-то аналогичное ограничениям человеческой рабочей памяти — вопрос философский, бенчмарками на него не ответишь. Эмпирические результаты стоят сами по себе, теоретическая мотивация — скорее украшение.
Место в контексте
RePo присоединяется к растущему семейству подходов, которые пытаются сделать трансформеры более гибкими в отношении позиций. YaRN расширяет контекст через масштабирование углов. Варианты NoPE полностью убирают позиционное кодирование из части слоёв. Собственная работа Sakana по DroPE исследует отбрасывание позиционных эмбеддингов для расширения контекста.
Отличие RePo в том, что модель учится реорганизовывать, а не просто справляться с фиксированной организацией. Это шаг к моделям, которые могут структурировать собственную рабочую память, а не получать структуру, навязанную архитектурой.
Предобученная модель доступна на Hugging Face. Подхватят ли подход крупные лаборатории — зависит от того, сохранятся ли улучшения при масштабировании.
