Сочетание следующего токена и видеодиффузии в компьютерном зрении и робототехнике

Новый метод позволяет обучать нейросети сортировать повреждённые данные, предсказывая следующие шаги. Он может эффективно планировать действия роботов, генерировать качественное видео и помогать ИИ-агентам ориентироваться в цифровых средах.

В текущем состоянии искусственного интеллекта модели предсказания последовательностей переживают настоящий бум благодаря своей способности анализировать данные. 

Простой пример — модели предсказания следующего токена, такие как ChatGPT, которые на основе контекста предсказывают каждое следующее слово (или токен) в последовательности, формируя ответы на запросы пользователей. Есть также модели диффузии для целых последовательностей, например, Sora, которые преобразуют текст в потрясающие и реалистичные изображения, поочередно «удаляя шум» из видеопоследовательностей.

Исследователи из лаборатории компьютерных наук и искусственного интеллекта (CSAIL) Массачусетского технологического института предложили несложное изменение в схеме обучения, которое значительно увеличивает гибкость этой технологии при работе с последовательностями.

В области компьютерного зрения и робототехники модели, работающие с предсказанием следующего токена и полными последовательностями, сталкиваются с определёнными компромиссами. Модели следующего токена могут генерировать последовательности разной длины. Однако, создавая такие последовательности, они не обладают возможностью учитывать долгосрочные цели, например, направлять генерацию последовательности к определённой цели, которая находится на 10 токенов впереди. Для этого необходимы дополнительные механизмы долгосрочного планирования. Модели диффузии могут выполнять такие выборки с учётом будущих состояний, но не обладают гибкостью моделей предсказания следующего токена в части создания последовательностей переменной длины.

Исследователи из CSAIL решили объединить преимущества обеих моделей и разработали методику обучения последовательных моделей под названием "Diffusion Forcing". Название отсылает к технике "Teacher Forcing", классической схеме обучения, которая разбивает задачу генерации целой последовательности на более простые этапы предсказания следующего токена (что напоминает, как хороший преподаватель помогает ученику понять сложную концепцию, разбивая её на более простые части).

Этот подход позволяет улучшить как гибкость планирования в робототехнике, так и качество генерации мультимедийных данных, что открывает новые перспективы для разработки более интеллектуальных и адаптивных систем ИИ.

Метод Diffusion Forcing представляет собой интересную и мощную комбинацию двух концепций машинного обучения: моделей диффузии и техники teacher forcing. Обе эти техники используют подходы к обучению, при которых предсказываются замаскированные (зашумленные) токены на основе немаскированных. В случае с моделями диффузии добавление шума к данным происходит поэтапно, что можно рассматривать как фрагментарное маскирование. 

Исследователи из MIT разработали метод Diffusion Forcing, который обучает нейросети очищать последовательности токенов, постепенно удаляя разное количество шума из каждого токена, одновременно предсказывая следующие элементы последовательности. Результат — гибкая и надежная модель последовательности, обеспечивающая более качественные искусственные видео и улучшенную точность принятия решений для роботов и агентов ИИ.

Суть метода заключается в том, что Diffusion Forcing эффективно справляется с обработкой зашумленных данных, предсказывая следующие шаги в задаче, что позволяет игнорировать отвлекающие визуальные элементы при выполнении манипуляций роботами. Метод не только генерирует стабильные и последовательные видео, но и может направлять ИИ-агентов через цифровые лабиринты. Это открывает перспективы для использования метода в быту и на производстве, где роботы могли бы адаптироваться к новым задачам, а также улучшить качество контента, создаваемого ИИ.

Как отмечает ведущий автор исследования, студент аспирантуры факультета электротехники и компьютерных наук MIT, а также член CSAIL Боюан Чен: «Модели последовательностей направлены на использование информации о прошлом для предсказания будущего, что является своего рода бинарным маскированием. Однако маскирование не обязательно должно быть бинарным». В Diffusion Forcing добавляется разный уровень шума к каждому токену, что эффективно действует как фракционное маскирование. На этапе тестирования система может «размаскировать» коллекцию токенов и уменьшить уровень шума, прогнозируя последовательность. Это позволяет модели «понимать», на какие данные можно опереться для преодоления ошибок, вызванных неизвестными или нестандартными входными данными.

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

Для генерации видео исследователи обучали Diffusion Forcing на игровом процессе в Minecraft и в ярких цифровых средах, созданных в симуляторе DeepMind Lab от Google. При подаче одного кадра видео метод создавал более стабильные и качественные ролики по сравнению с другими подходами, такими как модель диффузии полного цикла вроде Sora или модели, основанные на предсказании следующего токена, как у ChatGPT. Эти модели часто создавали видео с нестабильной картинкой, а в некоторых случаях не могли воспроизвести работающее видео уже после 72 кадров.

Кроме того, Diffusion Forcing может выступать в роли планировщика движений, который нацелен на достижение желаемых результатов или вознаграждений. Благодаря своей гибкости метод способен генерировать планы с разной временной перспективой, проводить поиск по дереву и учитывать интуицию. Учитывать, что будущее в более отдалённой перспективе всегда менее определённое, чем ближайшее. В задаче решения 2D-лабиринта Diffusion Forcing превзошёл шесть базовых моделей, генерируя более быстрые планы, которые приводили к цели. Это подтверждает, что метод может стать эффективным инструментом для планирования действий роботов в будущем.

В ряде своих экспериментов команда исследователей использовала метод Diffusion Forcing, который проявляет удивительную гибкость, действуя как модель последовательности, модель предсказания следующего токена или даже сочетая обе стратегии. По словам одного из авторов исследования, Чена, этот подход может стать мощной основой для создания так называемой мировой модели — искусственного интеллекта, который обучается на миллиардах видео с интернета и способен моделировать динамику окружающего мира. В результате такой системы роботы смогут выполнять новые задачи, "представляя" себе действия, которые им нужно выполнить, основываясь на информации о том, что происходит вокруг.

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

Сейчас исследовательская группа работает над расширением своего метода, используя более крупные наборы данных и новейшие трансформерные модели для повышения эффективности. Их конечная цель — создать нечто похожее на мозг робота, напоминающий ChatGPT, который позволит роботам выполнять задачи в новых, незнакомых им средах без необходимости в человеческой демонстрации.

Как отметил ведущий автор работы, Винсент Ситцман (ассистент профессора в MIT и член CSAIL, где он руководит группой Scene Representation), "С помощью Diffusion Forcing мы приближаем генерацию видео и робототехнику друг к другу". В долгосрочной перспективе ученые надеются использовать знания, хранящиеся в видео на просторах интернета, чтобы роботы могли выполнять задачи и помогать в повседневной жизни. Однако впереди еще множество захватывающих научных вызовов, например, научить роботов имитировать людей, наблюдая за ними, даже когда их собственные тела сильно отличаются от человеческих.

В состав авторов работы входят также исследователь MIT и недавний приглашенный ученый, Диего Марти Монсо, а также члены CSAIL: Илун Ду (аспирант кафедры EECS), Макс Симховиц (бывший постдок и будущий ассистент-профессор в Карнеги-Меллон) и Расс Тедрейк (профессор MIT, директор лаборатории робототехники Toyota Research Institute). Исследование поддержано Национальным научным фондом США, Сингапурским агентством оборонных наук и технологий, Институтом передовых исследований в области разведки и научным хабом Amazon.