Использование языковых моделей для задач обучения с подкреплением долгое время казалось контринтуитивным из-за отсутствия явных пересечений между текстом и траекториями действий. Однако в новом обзоре научного исследования известный ИИ-блогер Янник Килчер (Yannic Kilcher) подробно разбирает работу, которая доказывает обратное. Предварительное обучение трансформеров на текстах Википедии неожиданно позволяет радикально ускорить сходимость и повысить эффективность алгоритмов офлайн-обучения с подкреплением (Offline RL).
🧠 Суть парадокса: Википедия против траекторий роботов 0:00
Янник Килчер начинает разбор со смелого заявления: результаты обсуждаемой научной статьи находятся «на грани абсурда». На первый взгляд, предобучение языковых моделей (Language Model pre-training) и обучение с подкреплением (Reinforcement Learning) — это две абсолютно разные области, у которых практически нет ничего общего. Тем не менее, авторы исследования — Майкл Рид (Michael Reid), Ютаро Ямада (Yutaro Yamada) и Шиксианг Шейн Гу (Shixiang Shane Gu) — продемонстрировали существование удивительного механизма переноса знаний из текстовой модальности в управление агентами.
Особый интерес вызывает то, что этот эффект не является универсальным для любого предобучения. Например, использование визуальных моделей, таких как Image GPT, не дало сопоставимых результатов. По мнению Килчера, в самом языковом моделировании скрывается нечто особенное, что делает его мощным априорным регуляризатором для последовательностей действий.
🤖 Переосмысление RL как задачи sequence-to-sequence 3:45
В основе рассматриваемого метода лежит концепция, трансформирующая традиционный взгляд на Offline RL. Вместо привычного итеративного поиска стратегий, этот подход фреймится как обычное обучение с учителем (supervised learning) для моделирования последовательностей в готовом датасете траекторий. Данная архитектурная парадигма заимствована из известной модели Decision Transformer.
В офлайн-обучении с подкреплением агент не может напрямую взаимодействовать со средой; вместо этого у него есть статический набор данных, собранный другими агентами. Каждая траектория представляет собой цепочку, состоящую из трех повторяющихся элементов:
Временная последовательность включает:
- Текущее состояние или наблюдение среды (State).
- Действие, совершенное агентом (Action).
- Полученная награда за шаг (Reward).
Поскольку эти данные естественным образом складываются в последовательность, их можно передать в крупную трансформерную модель для прогнозирования следующего токена. Однако у Decision Transformer есть классическое ограничение RL — фиксированное окно контекста (например, длиной в 4 токена), из-за чего модель не способна смотреть бесконечно далеко назад или вперед.
Ключевая модификация, используемая авторами, заключается в замене мгновенных наград на метрику «остаточной награды» (returns-to-go). Это сумма всех будущих наград от текущего шага до конца эпизода. Например, если агент ожидает суммарно получить 50 единиц награды и делает шаг стоимостью в 1 единицу, на следующем шаге значение возврата составит 49. Во время инференса (тестирования) инженеры могут искусственно задать максимально высокое значение остаточной награды, чтобы побудить модель выбирать наилучшие действия.
🛠️ Архитектура, выравнивание эмбеддингов и совместное обучение 10:15
В экспериментах исследователи задействовали несколько базовых архитектур, включая стандартную модель GPT-2 Small. Чтобы обеспечить честное сравнение с оригинальным Decision Transformer, авторы обучили собственную языковую модель под названием ChiBT, обладающую точно таким же количеством параметров. Также тестировались текстовый кодировщик из модели CLIP и авторегрессионная модель Image GPT, разбивающая изображения на последовательности пикселей или патчей.
Для успешной адаптации языкового трансформера к траекториям роботов авторам пришлось решить проблему несовпадения модальностей, внедрив два дополнительных механизма:
Дополнительные методы оптимизации:
- Выравнивание представлений (Alignment): линейные проекции отображают эмбеддинги наград, состояний и действий так, чтобы максимизировать их косинусное сходство с исходными текстовыми эмбеддингами из словаря предобученной модели.
- Совместное обучение (Co-training): в процессе настройки на траектории модель продолжает параллельно обучаться на текстовом корпусе, что заставляет её удерживать способность работать с обеими модальностями одновременно.
Хотя Янник Килчер отмечает, что необходимость постепенного снижения весов (decay) этих дополнительных функций потерь до нуля после 5000 шагов указывает на некоторую хрупкость настройки, общие результаты оказались устойчивыми.
📈 Результаты тестов: скорость сходимости и превосходство над бейзлайнами 16:39
Тестирование проводилось на бенчмарках D4RL и наборе данных 1% DQN Replay для игр Atari. Модели, прошедшие языковое предобучение, стабильно превосходили классический Decision Transformer во множестве задач. При этом модель Image GPT полностью провалила тесты, продемонстрировав худшие результаты.
Главным практическим преимуществом подхода стало радикальное ускорение сходимости (convergence speed). Янник Килчер приводит впечатляющие цифры:
Сравнительные тайминги обучения:
- Оригинальный Decision Transformer требует около 3 часов для обучения с нуля.
- Предобученная модель ChiBT аналогичного размера сходится всего за 43 минуты.
- Полноразмерная модель GPT-2, которая в 144 раза больше ChiBT по числу параметров, тратит на обучение всего 1 час 27 минут.
Ведущий, впрочем, призывает относиться к точным значениям наград с осторожностью. Из-за высокого стандартного отклонения (шума) на трех использованных случайных сидах (seeds) результаты в этой области традиционно нестабильны. Тем не менее, общая тенденция превосходства текстового предобучения прослеживается четко.
🎯 Секреты внимания: анализ паттернов Attention 20:40
Самой захватывающей частью работы Янник считает визуализацию матриц внимания (attention patterns) в различных моделях. У GPT-2 отчетливо видны диагональные полосы с шагом, кратным трем. Это обусловлено структурой траектории, где токены идут тройками (возврат, состояние, действие). Модель четко связывает текущее действие с предыдущими действиями, а состояние — с прошлыми состояниями.
В средних слоях GPT-2 начинает выделять конкретные ключевые состояния, удерживая на них фокус. Более того, в крайнем левом столбце матрицы видно, что GPT-2 постоянно обращается к самому первому токену остаточной награды всего эпизода. По словам Килчера, это критически важно, так как именно стартовая целевая награда определяет глобальную траекторию и стратегию поведения агента.
В отличие от неё, случайно инициализированная модель без предобучения демонстрирует чрезмерно жесткую, «переобученную» структуру внимания с шагом 3. Она игнорирует контекст между токенами, полагаясь только на самые последние шаги. По гипотезе Килчера, языковое предобучение работает как идеальный регуляризатор (prior), который заставляет модель удерживать в поле зрения широкий контекст, не скатываясь в слепое копирование последних действий. Image GPT, в свою очередь, страдает от фатального смещения в сторону новизны (recency bias), обращая внимание только на 2-3 токена позади себя, что делает её бесполезной в RL.
🧊 Масштабирование и провал «заморозки» параметров 34:17
Исследователи также проверили, как модели реагируют на увеличение масштаба. Выяснилось, что чистый Decision Transformer при наращивании миллионов параметров начинает быстро переобучаться и ухудшать результаты. Напротив, модели с текстовым предобучением отлично масштабируются: с ростом объема данных и параметров их средняя награда только растет.
Попытка зафиксировать (заморозить) основные веса трансформера и обучать только слои линейной проекции эмбеддингов привела к катастрофическому падению эффективности. По мнению авторов статьи, это связано с тем, что генеративное моделирование (генерация последовательности шагов) является принципиально более сложной задачей, чем дискриминативная классификация, где «заморозка» весов обычно работает успешно.
Янник Килчер резюмирует, что успех этой работы — яркое знамение времени. Современная индустрия ИИ движется к конвергенции, создавая единые унифицированные архитектурные решения для любых модальностей — будь то текст, код или управление робототехническими системами.