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

Тренировки сложных моделей машинного обучения в облаке остаются чрезвычайно дорогим занятием. Согласно прошлогоднему докладу исследователей из Университета Вашингтон, стоимость обучения составляет в среднем $25 000 за две недели. OpenAI тратила на языковую модель GPT-2 $256 в час, а Google — $6912 на работу с моделью BERT.

Архитектура SEED RL основана на фреймворке Google TensorFlow 2.0. Она использует графические карты и тензорные процессоры (TPU), централизуя логические выводы модели, пишет Venture Beat. Для предотвращения ограничения передачи данных ИИ выполняет логические выводы централизованно вместе с элементом обучения, который тренирует модель, используя данные из распределенных логических выводов. Переменные и информация о состоянии целевой модели хранятся локально, тогда как наблюдения отсылаются ученику на каждом этапе. Задержка остается на минимальном уровне благодаря сетевой библиотеке, привязанной к открытому фреймфорку RPC.

Элемент обучения SEED RL может быть масштабирован на тысячи ядер, а число агентов — на тысячи машин. Один алгоритм — V-trace — предсказывает распределение действий, второй — R2D2 — выбирает действие на основе его предсказанной ценности.

Для оценки SEED RL разработчики использовали несколько популярных вариантов: среду Arcade Learning Environment, системы DeepMind Lab и Google Research Football. Им удалось решить ранее нерешенную задачу Google Research Football и добиться 2,4 млн кадров в секунду с 64 ядрами Cloud TPU. Это новый рекорд по сравнению с прошлым достижением распределенного агента.

Результаты говорят о значительном ускорении обучения, а поскольку этот подход значительно дешевле, чем использование графических процессоров, то и стоимость экспериментов существенно снижается. Авторы исследования полагают, что благодаря SEED RL обучение с подкреплением получило возможность использовать потенциал акселераторов наравне с другими методами глубокого обучения.