Ученые ВШЭ оптимизировали обучение генеративных потоковых нейросетей

© iStock

Исследователи факультета компьютерных наук НИУ ВШЭ улучшили метод обучения генеративных потоковых нейросетей для работы с неструктурированными задачами. Это поможет искать новые лекарства эффективнее. Результаты работы были представлены на одной из ведущих конференций по машинному обучению — ICLR 2025. Текст работы доступен в репозитории Arxiv.org.

Генеративные потоковые сети (GFlowNets) — особый тип алгоритмов машинного обучения, который строит сложные объекты шаг за шагом. Ученые применяют их для поиска белков, лекарств и оптимизации транспортных систем. 

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

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

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

Ученые НИУ ВШЭ нашли способ оптимизировать обратную модель с помощью метода Trajectory Likelihood Maximization (TLM). Они доработали алгоритмы обратной модели таким образом, чтобы она могла постоянно сверяться с действиями прямой модели.

Тимофей Грицаев

«Мы сделали так, что поиск оптимального решения стал похож на переговоры, в которых обе стороны готовы менять свою позицию. В  задачах с большой степенью неизвестности обратная модель — лишь вспомогательный инструмент, улучшающий результаты прямой модели. Мы искали способ сделать работу обратной модели гибкой и наконец смогли его получить», — поясняет один из авторов работы, стажер-исследователь Центра глубинного обучения и байесовских методов Института искусственного интеллекта и цифровых наук ФКН НИУ ВШЭ Тимофей Грицаев.

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

Никита Морозов

«Наш метод заметно быстрее просматривает пространство возможных решений и находит больше качественных вариантов. Этот подход в целом сближает генеративные модели с методами обучения с подкреплением», — говорит Никита Морозов, младший научный сотрудник Центра глубинного обучения и байесовских методов Института искусственного интеллекта и цифровых наук ФКН НИУ ВШЭ.

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